Amazon Athena × Bedrock Converse APIで自然言語からSQLを自動生成・実行するエージェントを構築してみた

ユーザーの質問
  ↓
アプリケーションサーバー（FastAPI 等）
  ↓
bedrock_runtime.converse()  ← tool_use でSQLツールを定義
  ↓
Claude が SQL を生成（stop_reason: "tool_use"）
  ↓
アプリが SQL を受け取り、Athena に直接実行（Lambda 不要）
  ↓
結果を toolResult としてメッセージに追加
  ↓
bedrock_runtime.converse()  ← 結果をもとに回答生成
  ↓
（必要なら追加SQLを生成 → 再実行 → ループ）
  ↓
stop_reason: "end_turn" で回答完了

import pandas as pd

df = pd.read_excel("data.xlsx", sheet_name="Sheet1")
df.to_parquet("output/data.parquet", engine="pyarrow", compression="snappy")

s3://your-bucket/sql/
  ├── sales/
  │   └── sales.parquet
  ├── employees/
  │   └── employees.parquet
  └── products/
      └── products.parquet

s3://your-bucket/
  ├── sql/                  ← 元データ（Parquet）
  │   ├── sales/
  │   └── employees/
  └── athena-results/       ← Athena が自動生成（クエリ結果）
      ├── <query-id-1>.csv
      ├── <query-id-1>.csv.metadata
      ├── <query-id-2>.csv
      ├── <query-id-2>.csv.metadata
      └── ...（クエリ実行のたびに増える）

-- クエリ1: データベースを作成
CREATE DATABASE IF NOT EXISTS your_database
COMMENT 'Structured data for SQL agent';

-- クエリ2: テーブル例：売上データ
CREATE EXTERNAL TABLE your_database.sales (
  `region`         STRING   COMMENT '地域名',
  `fiscal_year`    INT      COMMENT '会計年度',
  `quarter`        STRING   COMMENT '四半期 (1Q, 2Q, 3Q, 4Q)',
  `revenue_plan`   BIGINT   COMMENT '計画売上高',
  `revenue_actual`  BIGINT  COMMENT '実績売上高'
)
STORED AS PARQUET
LOCATION 's3://your-bucket/sql/sales/'
TBLPROPERTIES ('parquet.compression'='SNAPPY');

SELECT * FROM your_database.sales LIMIT 10;

SELECT COUNT(*), MIN(fiscal_year), MAX(fiscal_year) FROM your_database.sales;

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "athena:StartQueryExecution",
                "athena:GetQueryExecution",
                "athena:GetQueryResults",
                "athena:StopQueryExecution",
                "athena:GetWorkGroup"
            ],
            "Resource": "*"
        },
        {
            "Effect": "Allow",
            "Action": [
                "glue:GetDatabase",
                "glue:GetTable",
                "glue:GetTables",
                "glue:GetPartition",
                "glue:GetPartitions"
            ],
            "Resource": "*"
        },
        {
            "Effect": "Allow",
            "Action": [
                "s3:GetBucketLocation",
                "s3:ListBucket"
            ],
            "Resource": "arn:aws:s3:::your-bucket"
        },
        {
            "Effect": "Allow",
            "Action": [
                "s3:PutObject",
                "s3:GetObject"
            ],
            "Resource": "arn:aws:s3:::your-bucket/*"
        }
    ]
}

"""athena.py — Athena クエリ実行モジュール"""

import asyncio
import re
import time
import boto3

MAX_ROWS = 200
POLL_INTERVAL = 0.5
MAX_WAIT = 30

ATHENA_DATABASE = "your_database"
ATHENA_WORKGROUP = "primary"
ATHENA_OUTPUT_S3 = "s3://your-bucket/athena-results/"

_SELECT_ONLY = re.compile(r"^\s*SELECT\b", re.IGNORECASE)

def _validate_sql(sql: str) -> None:
    """SELECT 文以外を拒否する。"""
    if not _SELECT_ONLY.match(sql):
        raise ValueError(
            f"Only SELECT statements are allowed. Received: {sql[:80]!r}"
        )

def _run_query_sync(sql: str) -> dict:
    """Athena でクエリを実行し、結果を返す（同期）。"""
    athena = boto3.client("athena")

    response = athena.start_query_execution(
        QueryString=sql,
        QueryExecutionContext={"Database": ATHENA_DATABASE},
        WorkGroup=ATHENA_WORKGROUP,
        ResultConfiguration={"OutputLocation": ATHENA_OUTPUT_S3},
    )
    execution_id = response["QueryExecutionId"]

    # ポーリングで完了を待つ
    deadline = time.time() + MAX_WAIT
    while time.time() < deadline:
        status = athena.get_query_execution(QueryExecutionId=execution_id)
        state = status["QueryExecution"]["Status"]["State"]
        if state == "SUCCEEDED":
            break
        if state in ("FAILED", "CANCELLED"):
            reason = status["QueryExecution"]["Status"].get(
                "StateChangeReason", "unknown"
            )
            raise RuntimeError(f"Athena query {state}: {reason}")
        time.sleep(POLL_INTERVAL)
    else:
        raise TimeoutError(
            f"Athena query did not complete within {MAX_WAIT}s"
        )

    # 結果を取得
    paginator = athena.get_paginator("get_query_results")
    pages = paginator.paginate(QueryExecutionId=execution_id)

    rows: list[list] = []
    columns: list[str] = []
    for page in pages:
        result = page["ResultSet"]
        if not columns:
            columns = [
                c["Label"]
                for c in result["ResultSetMetadata"]["ColumnInfo"]
            ]
        for row in result["Rows"][1 if not rows else 0:]:
            values = [d.get("VarCharValue") for d in row["Data"]]
            rows.append(values)
            if len(rows) >= MAX_ROWS:
                return {"columns": columns, "rows": rows, "row_count": len(rows)}

    return {"columns": columns, "rows": rows, "row_count": len(rows)}

async def execute_query(sql: str) -> dict:
    """非同期ラッパー。スレッドプールで同期処理を実行する。"""
    _validate_sql(sql)
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, _run_query_sync, sql)

EXECUTE_SQL_TOOL = {
    "toolSpec": {
        "name": "execute_sql_query",
        "description": (
            "Presto SQL SELECTクエリをAthenaで実行し、結果をJSONで返します。"
        ),
        "inputSchema": {
            "json": {
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "実行するPresto SQL SELECT文",
                    }
                },
                "required": ["sql"],
            }
        },
    }
}

"""agent_loop.py — Converse API を使った SQL エージェントループ"""

import json
import boto3
from athena import execute_query

MODEL_ID = "ap-northeast-1.anthropic.claude-sonnet-4-6"
MAX_ITERATIONS = 4  # SQL実行の最大反復回数

bedrock = boto3.client("bedrock-runtime", region_name="ap-northeast-1")

# ツール定義（Step 4 で作成したもの）
TOOL_CONFIG = {"tools": [EXECUTE_SQL_TOOL]}

def load_system_prompt() -> str:
    with open("system_prompt.md", encoding="utf-8") as f:
        return f.read()

async def run_agent(user_message: str, history: list[dict]) -> str:
    """
    ユーザーの質問に対してエージェントループを実行し、最終回答を返す。

    フロー:
    1. Converse API にメッセージ + ツール定義を送信
    2. stop_reason が "tool_use" なら → SQL抽出 → Athena実行 → 結果をメッセージに追加 → 再送信
    3. stop_reason が "end_turn" なら → テキスト回答を返す
    4. 最大反復回数に達したら強制終了
    """
    system_prompt = load_system_prompt()
    messages = _build_messages(history, user_message)

    for iteration in range(MAX_ITERATIONS):
        # ── Converse API 呼び出し ──
        response = bedrock.converse(
            modelId=MODEL_ID,
            system=[{"text": system_prompt}],
            messages=messages,
            toolConfig=TOOL_CONFIG,
            inferenceConfig={"maxTokens": 4096, "temperature": 0},
        )

        assistant_message = response["output"]["message"]
        stop_reason = response["stopReason"]

        # ── Claude がテキスト回答を返した場合 → 完了 ──
        if stop_reason != "tool_use":
            for block in assistant_message.get("content", []):
                if "text" in block:
                    return block["text"]
            return ""

        # ── Claude がツール呼び出しを返した場合 ──
        tool_use_block = _extract_tool_use(assistant_message)
        if not tool_use_block:
            return "SQLの生成に失敗しました。"

        sql = tool_use_block["input"]["sql"]
        tool_use_id = tool_use_block["toolUseId"]

        print(f"[iteration {iteration + 1}] SQL: {sql}")

        # ── Athena で SQL を実行 ──
        try:
            result = await execute_query(sql)
        except (ValueError, RuntimeError, TimeoutError) as exc:
            result = {"error": str(exc)}

        # ── 結果を toolResult として会話に追加 ──
        messages = messages + [
            assistant_message,  # Claude のアシスタントメッセージ（tool_use を含む）
            {
                "role": "user",
                "content": [
                    {
                        "toolResult": {
                            "toolUseId": tool_use_id,
                            "content": [
                                {"text": json.dumps(result, ensure_ascii=False)}
                            ],
                        }
                    }
                ],
            },
        ]
        # ループの次の反復で再度 converse() を呼ぶ

    return "最大反復回数に達しました。"

def _extract_tool_use(message: dict) -> dict | None:
    """アシスタントメッセージから execute_sql_query の toolUse ブロックを抽出する。"""
    for block in message.get("content", []):
        if "toolUse" in block and block["toolUse"]["name"] == "execute_sql_query":
            return block["toolUse"]
    return None

def _build_messages(history: list[dict], current: str) -> list[dict]:
    """会話履歴を Converse API のメッセージ形式に変換する。"""
    messages = []
    for turn in history:
        role = turn.get("role")
        content = turn.get("content", "")
        if role in ("user", "assistant") and content:
            messages.append({"role": role, "content": [{"text": content}]})
    messages.append({"role": "user", "content": [{"text": current}]})
    return messages

messages = [user_message]

┌─── ループ開始（最大 4 回）────────────────────────┐
│                                                    │
│  response = converse(messages, toolConfig)          │
│                                                    │
│  stop_reason == "end_turn"?  ──→ テキスト回答を返す │
│        │ No                                        │
│        ↓                                           │
│  stop_reason == "tool_use"                         │
│        │                                           │
│        ↓                                           │
│  SQL を抽出 → Athena で実行                         │
│        │                                           │
│        ↓                                           │
│  messages += [assistant_msg, toolResult_msg]        │
│        │                                           │
│        └──→ ループ先頭に戻る                        │
│                                                    │
└────────────────────────────────────────────────────┘

あなたは[ドメイン]のデータアナリストです。

## 回答ルール
1. 検索結果から正確な数値を引用して回答してください
2. 数値には必ず単位を付けてください
3. データが見つからない場合は「該当データが見つかりませんでした」と回答してください
4. 推測や概算は行わないでください

## SQLクエリ記述ルール
- 構文: Presto SQL（Athena互換）
- 日本語カラム名はダブルクォートで囲む（例: "売上高"）
- 文字列値はシングルクォート（例: '1Q'）

## テーブル定義

### テーブル名: sales
| カラム | 型 | 説明 | 値の例 |
|--------|-----|------|--------|
| `region` | STRING | 地域名 | 東京, 大阪 |
| `fiscal_year` | INT | 会計年度 | 2021〜2026 |
| `quarter` | STRING | 四半期 | 1Q, 2Q, 3Q, 4Q, 通期 |
| `revenue_plan` | BIGINT | 計画売上高（百万円） | — |
| `revenue_actual` | BIGINT | 実績売上高（百万円） | — |

## データの注意点
- [集計時に二重計上を避けるフィルタ条件]
- [特定のカラム値の組み合わせ制約]
- [欠損データの年度・期間]
- [曖昧な用語の判定ルール]

2025年度1Qの全地域の売上実績合計を教えてください。

[iteration 1] SQL: SELECT SUM("revenue_actual") AS total FROM sales WHERE "fiscal_year" = 2025 AND "quarter" = '1Q'