Pydantic AIで作る！実践Text-to-SQLシステム構築ガイド 〜自然言語によるデータ抽出の自動化で分析業務を効率化〜

from dataclasses import dataclass
from typing import Any, Dict, Union
import pandas as pd 

@dataclass
class AnalysisContext:
    """分析コンテキスト: ワークフロー全体で共有される情報を保持するクラス"""
    query: str                  
    model_type: str             
    tables_info: str = None     
    sql_query: str = None       
    sql_check_result: Dict[str, Any] = None 
    execution_result: Union[pd.DataFrame, str] = None 
    final_answer: str = None    
    workflow_execution_count: int = 0 

import streamlit as st
from google.cloud import bigquery
from google.oauth2.credentials import Credentials
from dataclasses import dataclass, field
from typing import Union, List, TypedDict
import pandas as pd

from pydantic_ai.messages import ModelMessage
from pydantic_graph import BaseNode, End, GraphRunContext


from utils.logger import get_logger
from domain.ai.workflows.sql_generator.nodes.analysis_context import AnalysisContext
from utils.constants import DEFAULT_BILLING_PROJECT_ID

import domain.ai.workflows.sql_generator.nodes.logic.sql_generation_logic as sql_generation_logic
import domain.ai.workflows.sql_generator.nodes.logic.sql_check_node_logic as sql_check_node_logic
import domain.ai.workflows.sql_generator.nodes.logic.answer_formatting_logic as answer_formatting_logic

logger = get_logger()

MAX_EXECUTION_COUNT = 4 



@dataclass
class SQLGenerationNode(BaseNode[AnalysisContext]):
    """SQL生成ノード: ユーザーの質問からSQLを生成・修正する"""
    feedback: str = None 
    agent_messages: List[ModelMessage] = field(default_factory=list) 

    async def run(self, ctx: GraphRunContext) -> Union[End, "SQLCheckNode"]:
        self.context = ctx.state
        logger.info("SQLクエリ生成を開始します")

        
        if self.context.workflow_execution_count >= MAX_EXECUTION_COUNT:
            logger.warning(f"最大試行回数({MAX_EXECUTION_COUNT})に達しました。")
            return End(self.context) 

        try:
            
            result = await sql_generation_logic.run_logic(
                self.context, self.feedback, self.agent_messages
            )
            
            self.context.tables_info = result.data.get("tables_info")
            self.context.sql_query = result.data.get("sql_query")
            self.agent_messages += result.all_messages() 
            logger.debug(f"生成されたSQL: {self.context.sql_query}")
            
            return SQLCheckNode(agent_messages=self.agent_messages)
        except Exception as e:
            logger.error(f"SQLクエリの生成中にエラー: {e}")
            self.context.workflow_execution_count += 1
            
            return SQLGenerationNode(feedback=str(e), agent_messages=self.agent_messages)

@dataclass
class SQLCheckNode(BaseNode[AnalysisContext]):
    """SQLチェックノード: 生成されたSQLをレビューする"""
    agent_messages: List[ModelMessage] = field(default_factory=list) 

    async def run(self, ctx: GraphRunContext) -> Union[End, "SQLGenerationNode", "SQLExecutionNode"]:
        self.context = ctx.state
        logger.info("SQLクエリのチェックを開始します")
        try:
            
            result = await sql_check_node_logic.run_logic(self.context)
            check_result_data = result.data 
            self.context.sql_check_result = check_result_data 

            if check_result_data.get("is_valid"):
                logger.info("SQLクエリは有効です。実行ノードへ進みます。")
                
                return SQLExecutionNode()
            else:
                logger.warning("SQLクエリに問題あり。SQL生成ノードへ戻ります。")
                self.context.workflow_execution_count += 1
                feedback = check_result_data.get("overall_feedback", "SQLに問題があります。修正してください。")
                
                return SQLGenerationNode(feedback=feedback, agent_messages=self.agent_messages)
        except Exception as e:
            logger.error(f"SQLチェック中にエラー: {e}")
            self.context.workflow_execution_count += 1
            
            return SQLGenerationNode(feedback=f"SQLチェック中にエラーが発生しました: {e}", agent_messages=self.agent_messages)

@dataclass
class SQLExecutionNode(BaseNode[AnalysisContext]):
    """SQL実行ノード: 検証済みSQLをBigQueryで実行する"""
    async def run(self, ctx: GraphRunContext) -> Union[End, "SQLGenerationNode", "AnswerFormattingNode"]:
        self.context = ctx.state
        logger.info("SQLクエリの実行を開始します")
        try:
            token = st.session_state.get("token")
            credentials = Credentials(token) if token else None
            client = bigquery.Client(project=DEFAULT_BILLING_PROJECT_ID, credentials=credentials)

            logger.debug(f"Executing SQL: {self.context.sql_query}")
            query_job = client.query(self.context.sql_query)
            df = query_job.to_dataframe()

            self.context.execution_result = df
            logger.info(f"SQLクエリを実行しました: 行数={len(df)}")
            
            return AnswerFormattingNode(agent_messages=[]) 
        except Exception as e:
            logger.error(f"SQLクエリの実行中にエラー: {e}")
            self.context.workflow_execution_count += 1
            
            return SQLGenerationNode(feedback=f"SQLクエリの実行中にエラーが発生しました: {str(e)}", agent_messages=[]) 

@dataclass
class AnswerFormattingNode(BaseNode[AnalysisContext]):
    """回答フォーマットノード: SQL実行結果を自然言語で解説する"""
    agent_messages: List[ModelMessage] = field(default_factory=list) 

    async def run(self, ctx: GraphRunContext) -> End[AnalysisContext]:
        self.context = ctx.state
        logger.info("回答のフォーマットを開始します")
        try:
            
            result = await answer_formatting_logic.run_logic(self.context, self.agent_messages)
            self.context.final_answer = result.data 
            
            
            logger.info("回答のフォーマットが完了しました")
        except Exception as e:
            logger.error(f"回答のフォーマット中にエラー: {e}")
            self.context.final_answer = f"回答生成中にエラーが発生しました: {e}"
        
        return End(self.context)

import re
import json

from google.cloud import bigquery
from google.oauth2.credentials import Credentials
from .utils import to_json 

class BigQueryTool:
    def __init__(self, billing_project_id: str, access_token: str | None = None):
        """BigQueryクライアントを初期化"""
        if access_token:
            credentials = Credentials(access_token)
        else:
            credentials = None 
        self.client = bigquery.Client(
            project=billing_project_id, credentials=credentials
        )

    def get_schema(self, table_id: str) -> str:
        """指定されたテーブルのスキーマ情報(カラム名, 型, モード)をJSON文字列で返す"""
        if not self._validate_table_id(table_id):
             raise ValueError(f"無効なテーブルID形式です: {table_id}")
        try:
            table = self.client.get_table(table_id)
            schema_info = [
                {"name": field.name, "type": field.field_type, "mode": field.mode}
                for field in table.schema
            ]
            return to_json(schema_info)
        except Exception as e:
            raise ValueError(f"テーブル '{table_id}' のスキーマ取得中にエラーが発生しました: {e}") from e

    def sample_data(self, table_id: str, num_rows: int = 5) -> str:
        """指定されたテーブルから指定行数のサンプルデータをJSON文字列で返す"""
        if not self._validate_table_id(table_id):
            raise ValueError(f"無効なテーブルID形式です: {table_id}")
        query = f"SELECT * FROM `{table_id}` LIMIT {num_rows}" 
        try:
            df = self.client.query(query).to_dataframe()
            if df.empty:
                return to_json([]) 
            return to_json(df) 
        except Exception as e:
            raise ValueError(f"テーブル '{table_id}' のサンプルデータ取得中にエラーが発生しました: {e}") from e

    def query(self, query: str, limit: int = -1) -> str:
        """与えられたSQLクエリを実行し、結果をJSON文字列で返す (limit指定可能)"""
        try:
            df = self.client.query(query).to_dataframe()
            if df.empty:
                return to_json([])
            if limit  0: 
                return to_json(df)
            else:
                return to_json(df.head(limit))
        except Exception as e:
            error_query_snippet = query[:100] + '...' if len(query) > 100 else query
            raise ValueError(f"クエリ実行中にエラーが発生しました (Query: '{error_query_snippet}'): {e}") from e

    def _validate_table_id(self, table_id: str) -> bool:
        """テーブルIDが 'project.dataset.table' 形式か簡易チェック"""
        
        pass

async def run_logic(context: AnalysisContext, feedback: Optional[str] = None, agent_messages: Optional[List[ModelMessage]] = None):
    
    bq_tool = BigQueryTool(...) 

    
    system_prompt = """
あなたはECサイト分析に特化したBigQueryのSQLエキスパートです。与えられた情報をもとに、ユーザーの質問に答えるための正確で効率的なSQLクエリを生成してください。

あなたは以下の2ステップで処理を行います。
step1: テーブル情報の取得(table_info)
step2: SQLの生成(sql_query)

# step1: テーブル情報の取得 (table_info)
このステップでは、ユーザーの質問を分析し、SQLクエリを生成するために必要な情報を整理します。

1.  **ユーザーの質問を分析し、必要なデータを特定する。**
    *   例：「過去3ヶ月間のアクティブユーザー数を知りたい」「特定の商品カテゴリで最もレビュー評価の高い商品は何か？」「キャンセル率が高い支払い方法は？」など。
    *   注意：ユーザーの質問が「売上分析」なのか「顧客行動分析」なのか「商品分析」なのか、大まかな分析の目的を特定してください。
    
2.  **適切なBigQueryテーブルを選択する。**
    *   下記の「利用可能なBigQueryテーブル」セクションを参照し、質問に答えるために最も適切なテーブルを1つ以上選択してください。
    *   選択した各テーブルについて、なぜそのテーブルが必要なのか理由を簡潔に説明してください。
    
3.  **テーブル同士の結合条件を特定する（必要な場合）。**
    *   複数のテーブルを使用する場合、それらをどのように結合するか（例: `orders.user_id = users.user_id`）を具体的に特定してください。
    *   **重要：不必要にテーブルの結合を提案しないでください。** 単一テーブルで回答できる場合は、その旨を明確にしてください。
    
4.  **テーブルのカラムを特定する。**
    *   選択したテーブルから、質問に答えるために必要なカラムを特定してください。
    *   **最重要：必ず `get_schema` ツールを使ってテーブルに存在するカラム名を確認し、指定してください。** テーブルに存在しないカラム名を指定することは厳禁です。
    *   特定した各カラムについて、なぜそのカラムが必要なのか、集計においてどのような役割を果たすのか（例: `total_amount` は売上集計のため、`prefecture` は地域別集計のグルーピングキーとして利用）を説明してください。
    
5.  **フィルター条件を特定する。**
    *   ユーザーの質問に含まれる条件（例: 「過去3ヶ月間」「特定の商品カテゴリ」「評価が4以上」など）を抽出し、どのテーブルのどのカラムに適用するかを具体的に特定してください。
    *   日付範囲のフィルターについては、`order_timestamp` や `created_at` などのタイムスタンプ型カラムを利用することを検討してください。

## 利用可能なBigQueryテーブル 

### 売上・注文関連
-   `your-project-id.ec_data.orders`
    -   ECサイトの注文ヘッダー情報が格納されたテーブル。注文ごとの基本的なデータ（注文ID、顧客ID、注文日時、合計金額、注文ステータスなど）が含まれます。
    -   `status` カラムの値：`pending` (処理中), `shipped` (発送済), `delivered` (配送完了), `cancelled` (キャンセル済)。
    -   注文日時は `order_timestamp` (TIMESTAMP型) を使用してください。
-   `your-project-id.ec_data.order_items`
    -   注文明細情報テーブル。どの注文でどの商品がいくつ、いくらで売れたかの詳細情報。
    -   `orders` テーブルとは `order_id` で結合可能です。

### ユーザー関連
-   `your-project-id.ec_data.users`
    -   顧客マスタテーブル。顧客ID、氏名、メールアドレス、登録日時、居住都道府県などの情報が含まれます。
    -   登録日時は `created_at` (TIMESTAMP型) を使用してください。
    -   `orders` テーブルとは `user_id` で結合可能です。

### 商品関連
-   `your-project-id.ec_data.products`
    -   商品マスタテーブル。商品ID、商品名、カテゴリ、現在の単価などの情報が含まれます。
    -   `order_items` テーブルとは `product_id` で結合可能です。
-   `your-project-id.ec_data.reviews`
    -   商品に対する顧客からのレビュー情報が格納されたテーブル。商品ID、顧客ID、評価点（1～5の整数）、コメント本文、レビュー投稿日時などが含まれます。
    -   `products` テーブルとは `product_id` で、`users` テーブルとは `user_id` で結合可能です。
    -   レビュー評価は `rating` カラムを参照してください。

### 行動ログ関連 (高度な分析用)
-   `your-project-id.ec_data.page_views`
    -   ウェブサイト内のページ閲覧ログ。セッションID、顧客ID（ログイン時）、閲覧ページパス、閲覧日時、参照元URLなどの情報が含まれます。
    -   コンバージョン分析やユーザー行動追跡に使用します。

### table_infoの出力形式の説明
以下の形式で、マークダウンで記述してください。
1.  **利用するテーブル**
    *   テーブル名1: 選定理由
    *   テーブル名2: 選定理由
2.  **テーブルの結合条件** (必要な場合のみ)
    *   `テーブルA.カラムX = テーブルB.カラムY` : 結合理由
3.  **テーブルのフィルター条件**
    *   テーブル名1.カラムA: 条件 (例: `> '2023-01-01'`) : 理由
4.  **利用するテーブル項目**
    *   テーブル名1.カラムA: 項目の役割、選定理由
    *   テーブル名1.カラムB: 項目の役割、選定理由
    *   集計関数(カラムC) AS 集計後カラム名: 項目の役割、選定理由

	**※ステップ2以降のSQL生成のために、この`table_info`は非常に重要です。正確かつ網羅的に記述してください。**
	**※上記の手順2（テーブル選択）から手順5（フィルター条件特定）までは、必要に応じて `get_schema` や `sample_data` ツールを積極的に利用して、テーブルやカラムの存在、データ型、具体的な値を確認しながら進めてください。**
	

# step2: SQLの生成 (sql_query)
このステップでは、step1で作成した`table_info`に基づいて、BigQueryで実行可能な標準SQLクエリを生成します。

以下の条件に厳密に従ってください：
※重要事項：step1で定義した`table_info`の内容と完全に一致するSQLを生成してください。`table_info`に記載のないテーブルやカラムを使用しないでください。

## 作業手順
1.  **BigQuery の標準SQL**を使用する。
2.  **CTE (Common Table Expression) を使用して、クエリを構造化する。**
    *   読みやすく、メンテナンスしやすいSQLを目指してください。
3.  **`import_テーブル名` プレフィックスのCTEでテーブルをインポートする。**
    *   例: `import_orders AS (SELECT ... FROM \`your-project-id.ec_data.orders\`)`
    *   この段階で、必要なカラム選択と基本的なフィルター（日付範囲など）を行ってください。
4.  **`logic_処理内容` プレフィックスのCTEでデータ加工や集計を行う。**
    *   例: `logic_user_aggregation AS (SELECT user_id, COUNT(DISTINCT order_id) AS order_count FROM import_orders GROUP BY user_id)`
    *   JOIN処理もこの段階で行うことが多いです。
5.  **各CTEの目的、主要な処理内容について、SQL内に詳細な日本語コメントを追加する。**
    *   なぜこの処理をしているのか、第三者にも理解できるように記述してください。
6.  **可能な限り効率的なクエリを作成する。**
    *   不要なJOINや複雑なサブクエリ、過度なウィンドウ関数の使用は避けてください。
    *   `SELECT *` は最終的な出力以外では極力避け、必要なカラムのみを選択してください。
7.  **ARRAY型やSTRUCT型のカラムの扱いには注意してください。**
    *   適切に`UNNEST`して処理してください（今回のECサイト例では該当するものは明示していませんが、一般的な注意点として）。
8.  **集計クエリで結果行数が少ないと明確に判断できる場合を除き、最終的なSELECT文に `LIMIT 100` などを適切に設定することを検討してください。**
    *   ただし、ユーザーが全件データを期待している場合は不要です。

## SQLのベストプラクティス
SQLを記述する際は、以下の書き方を強く推奨します。

### 書き方①: CTE（Common Table Expression）を強く推奨
(前回の回答で示した詳細なCTEの書き方説明)
...
### 書き方⑤: なぜこういう処理をしてるかという理由をクエリ内のコメントとして日本語で書こう

### sql_queryの出力形式の説明
生成するSQLクエリは、以下のテンプレートに従ってください。
```sql
-- ユーザーの質問: {ユーザーの質問}
-- 生成日時: {現在の日時}

-- Step 1: 必要なテーブルからのデータ抽出と前処理
WITH
import_orders AS (
    -- 注文テーブルから必要な情報を抽出
    SELECT
        order_id,
        user_id,
        order_timestamp,
        total_amount,
        status
    FROM
        `your-project-id.ec_data.orders`
    -- 必要に応じてWHERE句で初期フィルター
    -- WHERE order_timestamp >= 'YYYY-MM-DD'
),

import_users AS (
    -- 顧客テーブルから必要な情報を抽出
    SELECT
        user_id,
        prefecture
    FROM
        `your-project-id.ec_data.users`
),

-- Step 2: データ加工や集計ロジック
logic_calculate_sales_by_prefecture AS (
    -- 注文と顧客情報を結合し、都道府県別の売上を集計
    SELECT
        u.prefecture,
        SUM(o.total_amount) AS total_sales_amount,
        COUNT(DISTINCT o.order_id) AS total_order_count
    FROM
        import_orders AS o
    INNER JOIN
        import_users AS u ON o.user_id = u.user_id
    WHERE
        o.status="delivered" -- 例: 配送完了したもののみを集計
    GROUP BY
        u.prefecture
),

-- Step 3: 最終結果の整形 (この例では不要なので省略されることも多い)
final_result AS (
    SELECT
        prefecture,
        total_sales_amount,
        total_order_count
    FROM
        logic_calculate_sales_by_prefecture
    ORDER BY
        total_sales_amount DESC
)

-- 最終的なクエリ実行
SELECT
    *
FROM
    final_result
-- LIMIT 100 -- 必要に応じて
;

# step3：
table_infoとsql_queryをキーとしたJSON形式で、ステップ1とステップ2の結果をそれぞれ出力してください。
{
    "table_info": "...",
    "sql_query": "..."
}

"""


sql_agent = Agent(
    model=context.model_type,
    result_type=SQLGenerationResult, 
    tools=[ 
        Tool(wrap_error(bq_tool.get_schema)),
        Tool(wrap_error(bq_tool.sample_data)),
        Tool(wrap_error(bq_tool.query)),
    ],
    system_prompt=system_prompt 
)


user_input = f"""ユーザーの質問: {context.query}"""


result = await sql_agent.run(user_input, message_history=agent_messages)

return result

from pydantic_ai import Agent, Tool
from ..tools.bigquery_tool import BigQueryTool
from .utils import wrap_error


async def run_logic(context: AnalysisContext):
    
    bq_tool = BigQueryTool(...) 

    
    system_prompt = """あなたはECサイト分析に詳しいBigQueryのSQLレビュアーです。...

与えられたSQLクエリを分析し、以下の観点でレビューしてください：

1. 最初にSQLをLIMIT10件で実行して構文エラーがないかをチェックする。 （注意！：必ず最初にチェックしてください。ここでエラーになったら後続のチェックは不要です。）エラーなったらSQLを修正してください。
2. 抽出仕様とSQLが整合しているかをチェックする。
  - 抽出条件が満たされているか? 
  - 抽出項目が適切か? 
  - 仕様に記載されていないロジックが含まれていないか?
3. パフォーマンス上の問題がないか
4. セキュリティ上の問題がないか
5. ベストプラクティスからの逸脱がないか

レビューにおいては、ThinkToolを頻繁に利用してください。

以下の形式でJSONを出力してください（必ず日本語で）：
json
{
  "is_valid": true or false,
  "issues": [
    {
      "severity": "error|warning|info",
      "description": "問題の詳細な説明",
      "line": "問題のある行番号または行の内容",
      "suggestion": "修正案",
      
    }
  ],
  "overall_feedback": "全体的なフィードバック"
}
"""



sql_check_agent = Agent(
    model=context.model_type,
    result_type=SQLCheckResult, 
    tools=[ 
        Tool(wrap_error(bq_tool.query)),
    ],
    system_prompt=system_prompt 
)


result = await sql_check_agent.run(user_input, model_settings={'temperature': 0.0})

return result