Text-to-SQLでオンデマンドなデータ分析を実現する

def create_text_to_sql_prompt(question: str, schema_info: str) -> str:
    prompt = f"""
あなたは優秀なSQLアナリストです。以下のデータベーススキーマを参考に、
ユーザーの質問に対するSQLクエリを生成してください。

データベーススキーマ:
{schema_info}

ユーザーの質問: {question}

重要なルール:
1. SELECT文のみを生成してください
2. 結果にはLIMIT 1000を付けてください
3. 日本語のカラム名には適切なエイリアスを付けてください
4. SQLクエリのみを返し、説明は不要です
5. 危険なキーワード（DROP、DELETE、UPDATE）は使用禁止

出力形式: SQLクエリのみ
"""
    return prompt

def get_database_schema(db_connection) -> str:
    """データベースから詳細なスキーマ情報を取得"""

    schema_query = """
    SELECT
        table_name,
        column_name,
        data_type,
        is_nullable,
        column_default
    FROM information_schema.columns
    WHERE table_schema = 'public'
    ORDER BY table_name, ordinal_position
    """

    schema_df = pd.read_sql(schema_query, db_connection)

    # スキーマ情報を構造化してLLM用のテキストに変換
    schema_text = ""
    for table in schema_df['table_name'].unique():
        table_cols = schema_df[schema_df['table_name'] == table]
        schema_text += f"\n{table}テーブル:\n"

        for _, col in table_cols.iterrows():
            schema_text += f"  - {col['column_name']}: {col['data_type']}"
            if col['is_nullable'] == 'NO':
                schema_text += " (NOT NULL)"
            schema_text += "\n"

    return schema_text

# 使用例
schema_info = """
salesテーブル (売上データ):
  - id: INTEGER (主キー)
  - customer_id: INTEGER (顧客ID)
  - product_id: INTEGER (商品ID)
  - amount: DECIMAL (売上金額)
  - quantity: INTEGER (数量)
  - sale_date: DATE (売上日)

customersテーブル (顧客データ):
  - id: INTEGER (主キー)
  - name: VARCHAR (顧客名)
  - region: VARCHAR (地域)
  - registration_date: DATE (登録日)

productsテーブル (商品データ):
  - id: INTEGER (主キー)
  - name: VARCHAR (商品名)
  - category: VARCHAR (カテゴリ)
  - price: DECIMAL (単価)
"""

# 必要なライブラリのインストール
uv add streamlit pandas matplotlib numpy openai

# Streamlitアプリの実行
uv run streamlit run app.py

# または通常のpipで
# pip install streamlit pandas matplotlib numpy openai
# streamlit run app.py

import streamlit as st
import pandas as pd
import sqlite3
import matplotlib.pyplot as plt
import numpy as np
from datetime import datetime, timedelta
import openai
import json
import re
from typing import Optional

def main():
    st.set_page_config(
        page_title="Text-to-SQL データ分析ダッシュボード",
        layout="wide"
    )

    st.title("Text-to-SQL データ分析ダッシュボード")
    st.markdown("自然言語でデータベースを分析し、SQL知識なしでインサイトを得ることができます")

    # サイドバー設定
    with st.sidebar:
        st.header("設定")

        # デモモード設定
        demo_mode = st.checkbox("デモモード", value=True, help="APIキー不要でサンプルデータを使用")

        if demo_mode:
            st.success("デモモードが有効です")
            st.info("サンプルデータベース（顧客・商品・売上）を使用します")
        else:
            api_key = st.text_input("OpenAI API Key", type="password")
            if not api_key:
                st.warning("本格モードにはAPIキーが必要です")

        st.markdown("---")

        # データベース情報表示
        if demo_mode:
            st.subheader("データベース構成")
            st.markdown("""
            **customers** (顧客)
            - id, name, region, created_at

            **products** (商品)
            - id, product_name, category, price

            **sales** (売上)
            - id, customer_id, product_id, amount, quantity, created_at
            """)

    # メイン分析エリア
    if demo_mode or ('api_key' in locals() and api_key):
        render_analysis_interface(demo_mode, locals().get('api_key'))
    else:
        st.warning("APIキーを入力するか、デモモードを有効にしてください")

def create_sample_database():
    """サンプルデータベースの作成"""
    conn = sqlite3.connect(':memory:')

    # 再現可能なランダムシード
    np.random.seed(42)

    # 顧客データ
    regions = ['東京', '大阪', '名古屋', '福岡', '仙台']
    customers_data = {
        'id': range(1, 101),
        'name': [f'顧客{i:03d}' for i in range(1, 101)],
        'region': np.random.choice(regions, 100),
        'created_at': pd.date_range('2023-01-01', periods=100, freq='D')
    }
    customers_df = pd.DataFrame(customers_data)

    # 商品データ
    categories = ['電子機器', '本・雑誌', '衣類', '食品・飲料', 'スポーツ']
    products_data = {
        'id': range(1, 51),
        'product_name': [f'商品{chr(65 + (i % 26))}{i//26 + 1:02d}' for i in range(50)],
        'category': np.random.choice(categories, 50),
        'price': np.random.randint(500, 50000, 50)
    }
    products_df = pd.DataFrame(products_data)

    # 売上データ（より現実的な分布）
    base_date = datetime(2023, 1, 1)
    sales_data = []

    for i in range(1000):
        # 時期による季節変動を模擬
        days_offset = np.random.randint(0, 365)
        sale_date = base_date + timedelta(days=days_offset)

        # 商品価格に基づく販売量調整
        product_id = np.random.randint(1, 51)
        product_price = products_df[products_df['id'] == product_id]['price'].iloc[0]

        # 高価格商品ほど販売量少なめ
        if product_price > 20000:
            quantity = np.random.randint(1, 3)
        elif product_price > 10000:
            quantity = np.random.randint(1, 5)
        else:
            quantity = np.random.randint(1, 10)

        sales_data.append({
            'id': i + 1,
            'customer_id': np.random.randint(1, 101),
            'product_id': product_id,
            'amount': product_price * quantity,
            'quantity': quantity,
            'created_at': sale_date
        })

    sales_df = pd.DataFrame(sales_data)

    # データベースに保存
    customers_df.to_sql('customers', conn, index=False, if_exists='replace')
    products_df.to_sql('products', conn, index=False, if_exists='replace')
    sales_df.to_sql('sales', conn, index=False, if_exists='replace')

    return conn

def generate_mock_sql(question: str) -> str:
    """デモモード用のモックSQL生成"""

    # 質問を小文字に変換してキーワードマッチング
    q = question.lower()

    # パターンマッチングによるSQL生成
    if any(word in q for word in ['売上', '合計', '総売上', 'total']):
        if any(word in q for word in ['今月', '月']):
            return """
SELECT SUM(amount) as total_sales, COUNT(*) as transaction_count
FROM sales
WHERE strftime('%Y-%m', created_at) = '2023-06'
"""
        elif any(word in q for word in ['地域', 'region']):
            return """
SELECT c.region, SUM(s.amount) as total_sales, COUNT(s.id) as transaction_count
FROM sales s
JOIN customers c ON s.customer_id = c.id
GROUP BY c.region
ORDER BY total_sales DESC
"""
        else:
            return """
SELECT SUM(amount) as total_sales, COUNT(*) as total_transactions
FROM sales
"""

    elif any(word in q for word in ['トップ', 'top', '上位']):
        if any(word in q for word in ['商品', 'product']):
            return """
SELECT p.product_name, p.category, SUM(s.amount) as total_sales
FROM sales s
JOIN products p ON s.product_id = p.id
GROUP BY p.id, p.product_name, p.category
ORDER BY total_sales DESC
LIMIT 10
"""
        elif any(word in q for word in ['顧客', 'customer']):
            return """
SELECT c.name, c.region, SUM(s.amount) as total_purchase
FROM sales s
JOIN customers c ON s.customer_id = c.id
GROUP BY c.id, c.name, c.region
ORDER BY total_purchase DESC
LIMIT 10
"""

    elif any(word in q for word in ['カテゴリ', 'category']):
        return """
SELECT p.category,
       COUNT(s.id) as transaction_count,
       SUM(s.amount) as total_sales,
       AVG(s.amount) as avg_sale_amount
FROM sales s
JOIN products p ON s.product_id = p.id
GROUP BY p.category
ORDER BY total_sales DESC
"""

    elif any(word in q for word in ['月別', '推移', 'trend']):
        return """
SELECT strftime('%Y-%m', created_at) as month,
       SUM(amount) as monthly_sales,
       COUNT(*) as monthly_transactions
FROM sales
GROUP BY strftime('%Y-%m', created_at)
ORDER BY month
"""

    elif any(word in q for word in ['平均', 'average', 'avg']):
        return """
SELECT p.category,
       AVG(s.amount) as avg_sales_amount,
       AVG(s.quantity) as avg_quantity
FROM sales s
JOIN products p ON s.product_id = p.id
GROUP BY p.category
ORDER BY avg_sales_amount DESC
"""

    # デフォルトのSQL
    return """
SELECT COUNT(*) as total_records,
       SUM(amount) as total_sales,
       AVG(amount) as average_sale
FROM sales
"""

def render_analysis_interface(demo_mode: bool, api_key: Optional[str] = None):
    """分析インターフェースの描画"""

    # 質問例の表示
    st.markdown("### 質問例")

    col1, col2, col3 = st.columns(3)

    with col1:
        example_questions_basic = [
            "今月の売上合計は？",
            "売上トップ10の商品を教えて",
            "地域別の売上を教えて"
        ]
        for q in example_questions_basic:
            if st.button(q, key=f"basic_{q}"):
                st.session_state.selected_question = q

    with col2:
        example_questions_intermediate = [
            "カテゴリ別の売上と取引数は？",
            "月別の売上推移を教えて",
            "平均売上が高いカテゴリは？"
        ]
        for q in example_questions_intermediate:
            if st.button(q, key=f"intermediate_{q}"):
                st.session_state.selected_question = q

    with col3:
        example_questions_advanced = [
            "売上トップ顧客の購買パターンは？",
            "地域×カテゴリのクロス分析"
        ]
        for q in example_questions_advanced:
            if st.button(q, key=f"advanced_{q}"):
                st.session_state.selected_question = q

    st.markdown("---")

    # 質問入力
    default_question = st.session_state.get('selected_question', '')
    user_question = st.text_area(
        "データについて自由に質問してください:",
        value=default_question,
        placeholder="例: 今月の地域別売上を前年同期と比較して、商品カテゴリ別の売上推移を教えて",
        height=100
    )

    col1, col2 = st.columns([4, 1])

    with col1:
        analyze_button = st.button("分析実行", type="primary", use_container_width=True)

    with col2:
        if st.button("クリア", use_container_width=True):
            st.session_state.selected_question = ''
            st.rerun()

    # 分析実行
    if analyze_button and user_question:
        analyze_question(user_question, demo_mode, api_key)

def analyze_question(question: str, demo_mode: bool, api_key: Optional[str] = None):
    """質問を分析してSQLを生成・実行"""

    with st.spinner("AIが分析中..."):
        # SQLクエリ生成
        if demo_mode:
            sql_query = generate_mock_sql(question)
        else:
            sql_query = generate_sql_with_api(question, api_key)

        if sql_query:
            # 生成されたSQLの表示
            st.markdown("### 生成されたSQLクエリ")
            st.code(sql_query, language="sql")

            # セキュリティチェック
            if is_safe_sql(sql_query):
                # SQLクエリ実行
                df = execute_sql_query(sql_query)

                if df is not None and len(df) > 0:
                    display_results(df, question)
                else:
                    st.warning("データが見つかりませんでした")
            else:
                st.error("セキュリティチェックに引っかかりました。安全でないSQLクエリです。")

def generate_sql_with_api(question: str, api_key: str) -> Optional[str]:
    """実際のAPI呼び出しでSQL生成"""

    schema_info = """
    データベーススキーマ:

    1. customers (顧客テーブル)
       - id: 顧客ID (INTEGER, PRIMARY KEY)
       - name: 顧客名 (TEXT)
       - region: 地域 (TEXT) - 値例: '東京', '大阪', '名古屋', '福岡', '仙台'
       - created_at: 顧客登録日 (DATE)

    2. products (商品テーブル)
       - id: 商品ID (INTEGER, PRIMARY KEY)
       - product_name: 商品名 (TEXT)
       - category: カテゴリ (TEXT) - 値例: '電子機器', '本・雑誌', '衣類', '食品・飲料', 'スポーツ'
       - price: 価格 (INTEGER)

    3. sales (売上テーブル)
       - id: 売上ID (INTEGER, PRIMARY KEY)
       - customer_id: 顧客ID (INTEGER, FOREIGN KEY)
       - product_id: 商品ID (INTEGER, FOREIGN KEY)
       - amount: 売上金額 (INTEGER)
       - quantity: 販売数量 (INTEGER)
       - created_at: 売上日時 (DATE)
    """

    client = openai.OpenAI(api_key=api_key)

    prompt = f"""
あなたは優秀なSQLアナリストです。以下のデータベーススキーマを参考に、
ユーザーの質問に対するSQLクエリを生成してください。

{schema_info}

ユーザーの質問: {question}

重要なルール:
1. SELECT文のみを生成してください
2. 結果にはLIMIT 100を付けてください
3. 日本語のカラム名には適切なエイリアスを付けてください
4. SQLクエリのみを返し、説明は不要です
5. 危険なキーワード（DROP、DELETE、UPDATE、INSERT、CREATE、ALTER）は使用禁止
6. SQLiteの構文を使用してください
7. 日付操作にはstrftime関数を使用してください

出力形式: SQLクエリのみ
"""

    try:
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            max_tokens=500
        )

        sql_query = response.choices[0].message.content.strip()

        # SQLクエリの清浄化
        if sql_query.startswith("```sql"):
            sql_query = sql_query[6:]
        if sql_query.startswith("```"):
            sql_query = sql_query[3:]
        if sql_query.endswith("```"):
            sql_query = sql_query[:-3]

        return sql_query.strip()

    except Exception as e:
        st.error(f"SQLクエリ生成エラー: {e}")
        return None

def is_safe_sql(sql: str) -> bool:
    """SQLクエリの安全性をチェック"""

    # 危険なキーワードのチェック
    dangerous_keywords = [
        'DROP', 'DELETE', 'INSERT', 'UPDATE', 'CREATE', 'ALTER',
        'TRUNCATE', 'EXEC', 'EXECUTE', 'UNION', '--', '/*', '*/',
        'SCRIPT', 'SHUTDOWN', 'SP_', 'XP_'
    ]

    sql_upper = sql.upper()

    for keyword in dangerous_keywords:
        if keyword in sql_upper:
            return False

    # SELECT文のみ許可
    if not sql_upper.strip().startswith('SELECT'):
        return False

    return True

def execute_sql_query(sql_query: str) -> Optional[pd.DataFrame]:
    """安全なSQLクエリ実行"""

    try:
        conn = create_sample_database()
        df = pd.read_sql_query(sql_query, conn)
        conn.close()
        return df

    except Exception as e:
        st.error(f"クエリ実行エラー: {e}")
        return None

def display_results(df: pd.DataFrame, question: str):
    """分析結果の表示"""

    st.markdown("### 分析結果")

    # メトリクス表示
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("レコード数", len(df))
    with col2:
        st.metric("カラム数", len(df.columns))
    with col3:
        # 数値カラムがある場合、合計を表示
        numeric_cols = df.select_dtypes(include=[np.number]).columns
        if len(numeric_cols) > 0:
            total_value = df[numeric_cols[0]].sum()
            st.metric("合計値", f"{total_value:,.0f}")

    # データテーブル表示
    st.dataframe(df, use_container_width=True)

if __name__ == "__main__":
    # セッション状態の初期化
    if 'selected_question' not in st.session_state:
        st.session_state.selected_question = ''

    main()