【LangChain】データベースを扱う④ LangGraphを用いたグラフ実装
前回の記事では、sql-query-system-promptを使って自然言語からのSQL生成と、SQLDatabaseを使ってpostgreSQLからのデータ抽出を検証しました。
今回はLangGraphを用いて、質問を投げかけるところからデータ抽出工程、解答生成までを実装しました。
LangGraphについて
- ノードの実態は関数。stateを引数にとって、変化させたstateのkey-valueを返す。
- Stateやstructured outputの型定義は、pydanticのBaseModelを使うのが結局一番わかりやすい
用いたデータ
よかったらコピペで使ってください。
（再生回数は適当です）
クリックで展開
Artistテーブル
アーティスト名 デビュー年
Mrs. GREEN APPLE 2015年
YOASOBI 2019年
back number 2011年
Ado 2020年
Vaundy 2019年
Official髭男dism 2012年
King Gnu 2017年
Creepy Nuts 2017年
優里 2019年
あいみょん 2015年
Musicテーブル
曲名 発売年月日 再生回数 アーティスト名
ケセラセラ 2024年4月5日 1億回 Mrs. GREEN APPLE
WanteD! WanteD! 2017年8月30日 9,000万回 Mrs. GREEN APPLE
青と夏 2018年7月18日 1.1億回 Mrs. GREEN APPLE
アイドル 2023年4月12日 1.5億回 YOASOBI
夜に駆ける 2019年12月15日 2億回 YOASOBI
群青 2020年9月1日 1.8億回 YOASOBI
水平線 2021年8月9日 2億回 back number
高嶺の花子さん 2013年6月26日 1.2億回 back number
瞬き 2017年12月20日 1.5億回 back number
新時代 2022年6月8日 1.8億回 Ado
うっせぇわ 2020年10月23日 3億回 Ado
踊 2021年4月27日 2.2億回 Ado
花占い 2023年3月1日 5,000万回 Vaundy
怪獣の花唄 2020年7月10日 1億回 Vaundy
東京フラッシュ 2019年11月6日 8,000万回 Vaundy
ミックスナッツ 2022年4月15日 1.2億回 Official髭男dism
Pretender 2019年5月15日 3.5億回 Official髭男dism
I LOVE... 2020年2月12日 2.5億回 Official髭男dism
一途 2021年12月29日 9,000万回 King Gnu
白日 2019年2月22日 3.2億回 King Gnu
三文小説 2020年12月2日 1.1億回 King Gnu
Bling-Bang-Bang-Born 2024年2月14日 3億回 Creepy Nuts
よふかしのうた 2018年9月4日 2億回 Creepy Nuts
Lazy Boy 2021年11月3日 7,000万回 Creepy Nuts
ドライフラワー 2020年10月25日 2.5億回 優里
ベテルギウス 2021年11月4日 1.7億回 優里
ミズキリ 2023年7月15日 8,000万回 優里
マリーゴールド 2018年8月8日 3億回 あいみょん
裸の心 2020年6月17日 2.8億回 あいみょん
愛を伝えたいだとか 2017年5月3日 1.9億回 あいみょん
コード全体
- LangSmithを使っています
クリックして展開
from langchain_community.utilities import SQLDatabase
# SQLDatabaseを使って、PostgreSQLに接続
db = SQLDatabase.from_uri("postgresql://username:@localhost:5432/postgres", schema="public20250109")
# print('DB名:', db.dialect)
# print('テーブル名:', db.get_usable_table_names())
# print('テーブル情報:', db.get_table_info())
# 環境変数の設定
import getpass
import os
import ast
from dotenv import load_dotenv
load_dotenv()
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "今回のプロジェクトの名前"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
if not os.environ.get("LANGCHAIN_API_KEY"):
os.environ["LANGCHAIN_API_KEY"] = getpass.getpass("Enter API key for Langchain: ")
if not os.environ.get("OPENAI_API_KEY"):
os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")
from langchain import hub
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
llm = ChatOpenAI(model="gpt-4o-mini")
query_prompt_template = hub.pull("langchain-ai/sql-query-system-prompt")
class State(BaseModel):
question: str = Field(..., description="ユーザーの質問")
query: str = Field("", description="LLMによって生成されたSQLクエリ")
result: str = Field("", description="SQLクエリの実行結果")
answer: str = Field("", description="ユーザーの質問に対する回答")
class QueryOutput(BaseModel):
"""SQLクエリを返すLLMの出力"""
query: str = Field(..., description="モデルによって生成されたSQLクエリ")
def write_query(state: State):
"""Generate SQL query to fetch information."""
prompt = query_prompt_template.invoke(
{
"dialect": db.dialect,
"top_k": 10,
"table_info": db.get_table_info(),
"input": state.question,
}
)
structured_llm = llm.with_structured_output(QueryOutput)
result = structured_llm.invoke(prompt)
return {"query": result.query}
def write_result(state: State):
"""SQLを実行し、結果をstateに書き込む"""
result = str(db.run(state.query)) # デフォルトだとinclude_columnsはFalseであり、結果は文字列で書かれたリストになる
return {"result": result}
def write_answer(state: State):
"""resultから回答を生成してstateに書き込む"""
template = """データ抽出結果から、ユーザーの質問に対する回答を生成してください。
データ抽出結果: {result}
ユーザーの質問: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
chain = prompt | llm
answer = chain.invoke({"result": state.result, "question": state.question})
return {"answer":answer}
# ワークフローの定義
from langgraph.graph import StateGraph
workflow = StateGraph(State)
workflow.add_node(write_query)
workflow.add_node(write_result)
workflow.add_node(write_answer)
workflow.set_entry_point("write_query")
workflow.add_edge("write_query", "write_result")
workflow.add_edge("write_result", "write_answer")
workflow.set_finish_point("write_answer")
compiled = workflow.compile()
# 実行
initial_state = State(question="viewsの回数が一番多いアーティストの再生回数は?")
result = compiled.invoke(initial_state)
print(result["answer"])
解説
write_query
def write_query(state: State):
"""ユーザーの質問からSQLクエリを生成してstateに書き込む"""
query_prompt_template = hub.pull("langchain-ai/sql-query-system-prompt")
prompt = query_prompt_template.invoke(
{
"dialect": db.dialect,
"top_k": 10,
"table_info": db.get_table_info(),
"input": state.question,
}
)
structured_llm = llm.with_structured_output(QueryOutput)
result = structured_llm.invoke(prompt)
return {"query": result.query}
この関数の解説は前回やったので不要かと思います。
hubにあるlangchain-ai/sql-query-system-promptを使えば、自然言語をいい感じのクエリに変換してくれます。
LangSmithでは、この部分の実行結果は以下の通りでした
write_result
def write_result(state: State):
"""SQLを実行し、結果をstateに書き込む"""
result = str(db.run(state.query)) # デフォルトだとinclude_columnsはFalseであり、結果は文字列で書かれたリストになる
return {"result": result}
SQLDatabaseのdbはrunにクエリを渡して実行すると、結果を返してくれます。
この結果はデフォルトでは文字列です。（リストではありません）
(method) def run(
command: str | Executable,
fetch: Literal['all', 'one', 'cursor'] = "all",
include_columns: bool = False,
*,
parameters: Dict[str, Any] | None = None,
execution_options: Dict[str, Any] | None = None
) -> (str | Sequence[Dict[str, Any]] | Result[Any])
なぜ文字列かというと、様々なデータベースシステムとの互換性を保つため、そして結果を簡単に表示したりログに記録したりするため、リストにするより文字列にした方が扱いやすいからだそうです。
もしリストの方がいい！という場合は、以下のようにast.literal_evalを使えばリストに変換できます
def write_result(state: State):
"""SQLを実行し、結果をstateに書き込む"""
result = str(db.run(state.query))
# もし文字列のresultをリストに変換したい場合は、ast.literal_evalを使う
result = ast.literal_eval(result)
return {"result": result}
LangSmithの結果は以下の通りです
write_answer
def write_answer(state: State):
"""resultから回答を生成してstateに書き込む"""
template = """データ抽出結果から、ユーザーの質問に対する回答を生成してください。
データ抽出結果: {result}
ユーザーの質問: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
chain = prompt | llm
answer = chain.invoke({"result": state.result, "question": state.question})
return {"answer":answer}
ここも解説は不要ですかね。
ごく普通のllmを使って、結果から回答文を生成しています。
LangSmithの結果は以下の通り
実行結果
以下の通り結果を得ることができました。
content='データ抽出結果によると、再生回数が一番多いアーティストは「あいみょん」で、再生回数は77,000です。' additional_kwargs={'refusal': None} response_metadata={'token_usage': {'completion_tokens': 40, 'prompt_tokens': 158, 'total_tokens': 198, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_f2cd28694a', 'finish_reason': 'stop', 'logprobs': None} id='run-5ea5ed35-9a29-482a-8dbf-aa12b399b16f-0' usage_metadata={'input_tokens': 158, 'output_tokens': 40, 'total_tokens': 198, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}
ついでにあいみょんの楽曲も貼っとこ
まとめ
たぶん一番シンプルな形で、自然言語をもとにクエリ作成→DBからのデータ抽出→解答生成を実現できたと思います。
LangGraphを用いているので、今後はAIエージェントを介して解答のファクトチェックや、最新データの取得等も拡大していけます
Discussion