LangChainのAgent機能を試してみる②：自然言語でSQLクエリを自動生成

pip install "langchain>=0.3.0" "langchain-community>=0.0.20" "langchain-openai>=0.0.5" "langchain-core>=0.3.0" python-dotenv openai requests

import os
import requests
import zipfile
from dotenv import load_dotenv
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import Tool
from langchain_openai import ChatOpenAI
from langchain_community.utilities.sql_database import SQLDatabase
from langchain.agents import AgentExecutor, create_openai_tools_agent

# 環境変数の読み込み
load_dotenv()

# Chinookデータベースのダウンロード（初回のみ）
def setup_chinook_db():
    db_path = "chinook.db"
    if not os.path.exists(db_path):
        print("Chinookデータベースをダウンロードしています...")
        zip_path = "chinook.zip"
        url = "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip"
        
        # ダウンロード
        response = requests.get(url)
        with open(zip_path, 'wb') as f:
            f.write(response.content)
        
        # 解凍
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(".")
        
        # クリーンアップ
        os.remove(zip_path)
        print("データベースのセットアップが完了しました")
    return db_path

# データベース接続
db_path = setup_chinook_db()
db = SQLDatabase.from_uri(f"sqlite:///{db_path}")

# データベース構造の確認
print("データベースのテーブル一覧:")
print(db.get_table_names())

# SQLエージェントの作成
def create_sql_agent(db_path: str, verbose: bool = True):
    # データベース接続の設定
    db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
    
    # LLMの初期化
    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    
    # 直接SQLクエリを実行する関数（ツールとして使用）
    def run_query(query: str) -> str:
        try:
            return db.run(query)
        except Exception as e:
            return f"Error: {str(e)}"
    
    # SQLツールを作成
    tools = [
        Tool(
            name="get_schema",
            func=lambda _: db.get_table_info(),
            description="データベース内のテーブル構造（スキーマ）情報を取得します。"
        ),
        Tool(
            name="query_sql",
            func=run_query,
            description="SQLクエリを実行してデータベースから情報を取得します。"
        )
    ]
    
    # エージェントプロンプトの設定
    agent_prompt = ChatPromptTemplate.from_messages([
        ("system", """あなたはデータベース分析の専門家です。SQLiteデータベースに対してクエリを実行し、
ユーザーの質問に答えるために必要な情報を取得できます。

使用可能なツール：
1. get_schema: データベーススキーマ情報を取得します。
2. query_sql: SQLクエリを実行して結果を取得します。

まず、get_schemaツールを使ってデータベース構造を理解し、
次にquery_sqlツールで適切なSQLクエリを実行してください。
複雑な分析では、段階的にクエリを実行し、結果に基づいて次のクエリを調整することも重要です。

最終的な回答は日本語で、データの分析結果や洞察を含めて提供してください。

このデータベースはChinookという音楽ストアのデータベースです。"""),
        ("human", "{input}"),
        ("assistant", "{agent_scratchpad}")  // 注意: v0.3では "agent" ではなく "assistant" を使用
    ])
    
    # エージェントの作成
    agent = create_openai_tools_agent(llm, tools, agent_prompt)
    
    # エージェントエグゼキューターの作成
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=verbose)
    
    return agent_executor

# エージェントを使った問い合わせ
def ask_database(question):
    print(f"質問: {question}")
    try:
        # SQLエージェントの作成
        agent_executor = create_sql_agent(db_path, verbose=True)
        
        # エージェントを実行
        response = agent_executor.invoke({"input": question})
        
        print("\n最終回答:")
        return response["output"]
    except Exception as e:
        print(f"エラーが発生しました: {e}")
        return str(e)

# 実行例
question = "最も売上の多い顧客トップ5を教えてください"
result = ask_database(question)
print(result)

あなたの質問: 売上の多い顧客トップ5

I'll analyze this database to find the top 5 customers with the highest sales.

First, I need to understand the database schema.

Action: get_schema
Action Input: 

CREATE TABLE "Album" (
	"AlbumId" INTEGER NOT NULL, 
	"Title" NVARCHAR(160) NOT NULL, 
	"ArtistId" INTEGER NOT NULL, 
	PRIMARY KEY ("AlbumId"), 
	FOREIGN KEY("ArtistId") REFERENCES "Artist" ("ArtistId")
)
... (スキーマ情報が続く) ...

Now I understand the database structure. To find the top 5 customers with the highest sales, I need to look at the Invoice and Customer tables.

Action: query_sql
Action Input: SELECT c.CustomerId, c.FirstName, c.LastName, SUM(i.Total) AS TotalSpent
FROM Customer c
JOIN Invoice i ON c.CustomerId = i.CustomerId
GROUP BY c.CustomerId
ORDER BY TotalSpent DESC
LIMIT 5;

CustomerId | FirstName | LastName | TotalSpent
59         | Helena    | Holý     | 49.62
29         | Robert    | Brown    | 47.62
25         | Victor    | Stevens  | 46.62
53         | Ladislav  | Kovács   | 45.62
52         | Hugh      | O'Reilly | 45.62

Based on the query results, I can now provide the answer.

最終回答:
最も売上の多い顧客トップ5は以下の通りです：

1. Helena Holý（顧客ID: 59）- 総支出額: $49.62
2. Robert Brown（顧客ID: 29）- 総支出額: $47.62
3. Victor Stevens（顧客ID: 25）- 総支出額: $46.62
4. Ladislav Kovács（顧客ID: 53）- 総支出額: $45.62
5. Hugh O'Reilly（顧客ID: 52）- 総支出額: $45.62

これらの顧客が店舗での支出額上位5名となります。特にHelena Holýさんが最も多くの金額を支出しています。

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_community.utilities.sql_database import SQLDatabase

# データベース接続
db = SQLDatabase.from_uri(f"sqlite:///chinook.db")

# データベーススキーマ情報を取得
db_schema = db.get_table_info()

# LLMの初期化
llm = ChatOpenAI(model="gpt-4o", temperature=0)

# SQLクエリ生成チェーン
sql_generation_prompt = ChatPromptTemplate.from_messages([
    ("system", """あなたはSQLの専門家です。以下のデータベーススキーマを元に、
ユーザーの質問に答えるための適切なSQLクエリを作成してください。
SQLiteの文法を使って、読みやすく効率的なクエリを書いてください。

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

回答は必ずSQLクエリのみを返してください。説明は含めないでください。"""),
    ("human", "{question}")
])

# 結果解釈チェーン
interpretation_prompt = ChatPromptTemplate.from_messages([
    ("system", """あなたはデータアナリストです。SQLクエリの結果を分析し、
ユーザーの質問に対する洞察を日本語で提供してください。
専門用語を避け、一般の人にもわかりやすく説明してください。"""),
    ("human", """元の質問: {question}
実行されたクエリ: {query}
クエリの結果: {result}

この結果から何がわかりますか？""")
])

# 完全なチェーンの構築
chain = (
    RunnablePassthrough.assign(
        schema=lambda _: db_schema,
        query=lambda x: (
            sql_generation_prompt 
            | llm 
            | StrOutputParser()
        ).invoke({"schema": db_schema, "question": x["question"]})
    )
    .assign(
        result=lambda x: db.run(x["query"])
    )
    | interpretation_prompt
    | llm
    | StrOutputParser()
)

# 実行例
result = chain.invoke({"question": "最も売れている音楽ジャンルは何ですか？"})
print(result)

あなたの質問: 学校で聞くのにオススメの音楽は？

質問: 学校で聞くのにオススメの音楽は？

...(about 5 lines omitted)...
3. **Samba De Uma Nota Só (One Note Samba)**
4. **Por Causa De Você**
5. **Ligia**
6. **Fotografia**
7. **Dindi (Dindi)**
8. **Se Todos Fossem Iguais A Você (Instrumental)**
9. **Falando De Amor**
10. **Angela**

これらの曲は、心地よいメロディーとリズムで、集中力を高めたり、リラックスしたりするのに役立つでしょう。ぜひ、試してみてください。

import os
import requests
import zipfile
from typing import List, Dict, Any, TypedDict, Literal, Optional, Union
from dotenv import load_dotenv
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langchain_openai import ChatOpenAI
from langchain_community.utilities.sql_database import SQLDatabase
from langchain_core.tools import tool
from langgraph.graph import END, START, StateGraph
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import ToolNode

# 環境変数の読み込み
load_dotenv()

# Chinookデータベースのダウンロード（初回のみ）
def setup_chinook_db():
    db_path = "chinook.db"
    if not os.path.exists(db_path):
        print("Chinookデータベースをダウンロードしています...")
        zip_path = "chinook.zip"
        url = "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip"
        
        # ダウンロード
        response = requests.get(url)
        with open(zip_path, 'wb') as f:
            f.write(response.content)
        
        # 解凍
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(".")
        
        # クリーンアップ
        os.remove(zip_path)
        print("データベースのセットアップが完了しました")
    return db_path

# 状態の型定義
class AgentState(TypedDict):
    messages: List[Union[HumanMessage, AIMessage, ToolMessage]]

# SQLツールの作成
def setup_tools(db_path: str):
    # データベース接続の設定
    db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
    
    # スキーマ取得ツール
    @tool
    def get_schema() -> str:
        """データベース内のテーブル構造（スキーマ）情報を取得します。"""
        return db.get_table_info()
    
    # クエリ実行ツール
    @tool
    def query_sql(query: str) -> str:
        """SQLクエリを実行してデータベースから情報を取得します。"""
        try:
            return db.run(query)
        except Exception as e:
            return f"エラー: {str(e)}"
    
    return [get_schema, query_sql]

# エージェント（LLM）ノード
def agent(state: AgentState, llm, tools):
    # メッセージの取得
    messages = state["messages"]
    
    # システムメッセージを追加（入力されていない場合）
    system_message_exists = any(
        isinstance(m, HumanMessage) and "あなたはデータベース分析の専門家です" in m.content
        for m in messages
    )
    
    if not system_message_exists:
        messages.insert(0, HumanMessage(content="""あなたはデータベース分析の専門家です。SQLiteデータベースに対してクエリを実行し、ユーザーの質問に答えるために必要な情報を取得できます。

使用可能なツール：
1. get_schema: データベーススキーマ情報を取得します。
2. query_sql: SQLクエリを実行して結果を取得します。

まず、get_schemaツールを使ってデータベース構造を理解し、次にquery_sqlツールで適切なSQLクエリを実行してください。
複雑な分析では、段階的にクエリを実行し、結果に基づいて次のクエリを調整することも重要です。

以下のガイドラインに従ってください：
- クエリを実行する前に、必ずスキーマを確認してください
- SQLiteの文法を使用してください
- 不明なテーブルやカラムに関しては、スキーマ情報を再確認してください
- 最終的な回答は日本語で、データの分析結果や洞察を含めて提供してください
- クエリと結果だけでなく、あなたの思考プロセスも説明してください

このデータベースはChinookという音楽ストアのデータベースで、アーティスト、アルバム、トラック、顧客、請求書などの情報が含まれています。"""))
    
    # LLMによる応答を取得
    response = llm.bind_tools(tools).invoke(messages)
    
    # メッセージリストに追加
    messages.append(response)
    
    return {"messages": messages}

# ツール実行ノード
def tools_executor(state: AgentState, tools):
    messages = state["messages"]
    last_message = messages[-1]
    
    # ツール呼び出しがある場合
    if isinstance(last_message, AIMessage) and hasattr(last_message, 'tool_calls') and last_message.tool_calls:
        tool_calls = last_message.tool_calls
        
        for tool_call in tool_calls:
            # 修正：ツールコール形式のパターンを追加
            if hasattr(tool_call, 'name'):
                # LangChainのオブジェクト形式
                action = tool_call.name
                args = tool_call.args
                tool_call_id = getattr(tool_call, 'id', None)
            elif isinstance(tool_call, dict) and 'function' in tool_call:
                # OpenAI標準形式
                action = tool_call['function']['name']
                args = tool_call['function']['arguments']
                tool_call_id = tool_call.get('id', None)
            elif isinstance(tool_call, dict) and 'name' in tool_call:
                # 新しく追加する形式: nameが直接dictのトップレベルにある
                action = tool_call['name']
                args = tool_call.get('args', {})
                tool_call_id = tool_call.get('id', None)
            else:
                print(f"未知のツールコール形式: {tool_call}")
                continue
            
            # 引数がJSON文字列の場合はパース
            if isinstance(args, str):
                import json
                try:
                    args = json.loads(args)
                except:
                    args = {}
            
            # ツールを検索
            for tool in tools:
                if tool.name == action:
                    # ツールを実行
                    try:
                        if isinstance(args, dict):
                            # 辞書の場合は引数として渡す
                            result = tool.invoke(args)
                        else:
                            # それ以外の場合は直接渡す
                            result = tool.invoke(args)
                    except Exception as e:
                        result = f"エラー: {str(e)}"
                    
                    # 結果をツールメッセージとして追加
                    messages.append(ToolMessage(
                        content=str(result), 
                        tool_call_id=tool_call_id if tool_call_id else "unknown_id"
                    ))
                    break  # ツールが見つかったらループを抜ける
    
    return {"messages": messages}

# 次のステップを決定する関数を修正
def decide_next_step(state: AgentState) -> Literal["agent", "tools", END]:
    messages = state["messages"]
    
    if not messages:
        return "agent"
    
    last_message = messages[-1]
    
    # ツール呼び出しの検出をより堅牢に
    has_tool_calls = False
    if isinstance(last_message, AIMessage):
        if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
            has_tool_calls = True
        elif hasattr(last_message, 'additional_kwargs'):
            if 'tool_calls' in last_message.additional_kwargs and last_message.additional_kwargs['tool_calls']:
                has_tool_calls = True
    
    # ツール呼び出しがある場合はツール実行へ
    if has_tool_calls:
        return "tools"
    # ツールからの応答がある場合はエージェントへ
    elif isinstance(last_message, ToolMessage):
        return "agent"
    # 普通のAIメッセージで終了
    elif isinstance(last_message, AIMessage):
        return END
    # その他の場合はエージェントへ
    else:
        return "agent"

# SQLエージェントグラフの作成
def create_sql_agent_graph(db_path: str):
    # LLMの初期化
    llm = ChatOpenAI(model="gpt-4o", temperature=0)
    
    # ツールの作成
    tools = setup_tools(db_path)
    
    # グラフの作成
    workflow = StateGraph(AgentState)
    
    # ノードの追加
    workflow.add_node("agent", lambda state: agent(state, llm, tools))
    workflow.add_node("tools", lambda state: tools_executor(state, tools))
    
    # エッジの設定
    workflow.add_edge(START, "agent")
    workflow.add_conditional_edges("agent", decide_next_step)
    workflow.add_conditional_edges("tools", decide_next_step)
    
    # グラフのコンパイル
    return workflow.compile()

# エージェントを使った問い合わせ
def ask_database(question: str):
    print(f"質問: {question}")
    try:
        # データベースのパスを取得
        db_path = setup_chinook_db()
        
        # SQLエージェントグラフの作成
        agent_executor = create_sql_agent_graph(db_path)
        
        # 初期状態の設定
        initial_state = AgentState(
            messages=[HumanMessage(content=question)]
        )
        
        # エージェントを実行
        result = agent_executor.invoke(initial_state)
        
        print("\n最終回答:")
        # 最後のAIメッセージを取得
        for message in reversed(result["messages"]):
            if isinstance(message, AIMessage) and not (hasattr(message, 'tool_calls') and message.tool_calls):
                return message.content
        
        return "回答が見つかりませんでした。"
    except Exception as e:
        print(f"エラーが発生しました: {e}")
        return str(e)

# メイン実行部分
if __name__ == "__main__":
    print("Chinookデータベースチャットボット（LangGraph版）へようこそ！")
    print("質問を入力してください（終了するには 'exit' と入力）")
    print("例: '最も売上の多い顧客トップ5を教えてください'")
    
    while True:
        user_input = input("\nあなたの質問: ")
        if user_input.lower() in ["exit", "quit", "終了"]:
            print("チャットボットを終了します。お役に立てて嬉しいです！")
            break
            
        result = ask_database(user_input)
        print(result)