LangGraphを使ったマルチデータソース対応のRAGをstep by stepで作ってみる

本記事は  生成AIウィーク  2日目の記事です。
👨‍💻  1日目  ▶▶ 本記事 ▶▶  3日目  👩‍💻

はじめに

こんにちは堤です。

LLMでは外部データを活用してより精度の高い回答を生成するRetrieval-Augmented Generation（RAG）が広がりつつあります。RAGという言葉自体も一般的になりつつありますが、精度向上のための手法も色々考えられており、様々なアプローチが存在します。また一方で最近はAgent の活用が進んでおり、それに伴いLangGraphが注目されています。本記事では、LangGraphを用いたマルチデータソース対応のRAGをステップバイステップで実装してみたいと思います。

LangGraphとは

LangGraphは、LLMを活用したアプリケーションを構築するためのライブラリです。ステップや状態をグラフ構造で管理することで、複雑なワークフローやエージェントの設計を容易にします。特に、ループを含むフローや動的な状態管理が求められるアプリケーションにおいて、LangGraphを活用することで、柔軟かつ効率的なシステム構築が可能になります。

langchain-ai.github.io

構築の流れ

ここからはLangGraphを用いたRAGを構築していきます。 今回は

1. データソース1つのRAG
2. データソース2つのRAG
3. データソース2つ+Web検索のRAG

の3Stepで構築していきたいと思います。 今回はLangGraphの基本的なコンポーネントや文法については触れないので、公式ドキュメントや他のブログをご参照ください。

1. データソース1つのRAG

まずは1つのデータソースのみを参照するシンプルなRAGをLangGraphで構築します。 今回はWikipediaの「野球」ページを読み込ませてその内容を回答してもらいます。

コード例

from typing import Annotated

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_community.document_loaders import UnstructuredURLLoader
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict

llm = ChatOpenAI(model="gpt-4o-mini")
embedding_model = OpenAIEmbeddings()


vectorstore = Chroma(
    embedding_function=embedding_model, persist_directory="./chroma_db"
)


class State(TypedDict):
    messages: Annotated[list[str], add_messages]


def retrieve(state: State):
    question = state["messages"][-1].content
    retriever = vectorstore.similarity_search(question, k=3)
    context = "\n\n".join([doc.page_content for doc in retriever])
    return {
        "messages": [
            {
                "role": "user",
                "content": f"以下のコンテキスト情報を使用して回答してください。\n\n{context}",
            }
        ]
    }


def generate_answer(state: State):
    response = llm.invoke(state["messages"])
    return {"messages": [response]}

def load_data(urls):
    loader = UnstructuredURLLoader(urls)
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=50)
    splits = text_splitter.split_documents(docs)
    return splits

urls = ["https://ja.wikipedia.org/wiki/%E9%87%8E%E7%90%83"]


docs = load_data(urls)
vectorstore.add_documents(docs)

builder = StateGraph(State)

builder.add_node("retrieve", retrieve)
builder.add_node("generate_answer", generate_answer)

builder.add_edge(START, "retrieve")
builder.add_edge("retrieve", "generate_answer")
builder.add_edge("generate_answer", END)

graph = builder.compile()

グラフ構造

LangGraphはグラフ構造をMermaid記法で出力することができます。 可視化するとこのような構造になります。

このようなシンプルな構造であればLangGraphを使う意味はあまりなく、LangChainでも問題なさそうです。

実行例

構築したワークフローを実行してみます。

response = graph.invoke(
     {"messages": [{"role": "user", "content": "世界初のプロ球団は？"}]},
     debug=True,
)
print(response["messages"][-1].content)

世界初のプロ球団はシンシナティ・レッドストッキングス（Cincinnati Reds）です。1869年に設立されました。

Wikipediaから情報を読み取って回答していることが実行ログから確認できます。

実行ログ

[-1:checkpoint] State at the end of step -1:
{'messages': []}
[0:tasks] Starting 1 task for step 0:
- __start__ -> {'messages': [{'content': '世界初のプロ球団は？', 'role': 'user'}]}
[0:writes] Finished step 0 with writes to 1 channel:
- messages -> [{'content': '世界初のプロ球団は？', 'role': 'user'}]
[0:checkpoint] State at the end of step 0:
{'messages': [HumanMessage(content='世界初のプロ球団は？', additional_kwargs={}, response_metadata={}, id='accb2f91-4c53-436e-a246-175fd9864071')]}
[1:tasks] Starting 1 task for step 1:
- retrieve -> {'messages': [HumanMessage(content='世界初のプロ球団は？', additional_kwargs={}, response_metadata={}, id='accb2f91-4c53-436e-a246-175fd9864071')]}
[1:writes] Finished step 1 with writes to 1 channel:
- messages -> [{'content': '以下のコンテキスト情報を使用して回答してください。\n'
             '\n'
             '1869年には世界初のプロ球団であるシンシナティ・レッドストッキングスが設立され、1871年には世界初のプロ野球リーグであるナショナル・アソシエーションが設立された。このリーグ自体は5年で破綻したも\n'
             '\n'
             '1869年には世界初のプロ球団であるシンシナティ・レッドストッキングスが設立され、1871年には世界初のプロ野球リーグであるナショナル・アソシエーションが設立された。このリーグ自体は5年で破綻したも\n'
             '\n'
             '基礎がつくられたことから、一般に同国が野球発祥の地とされている。1869年には最初のプロチームが生まれ、アメリカで有数の人気スポーツとなり、国民的娯楽となった。元々は21点先取制、勝敗だけでなく両チー',
  'role': 'user'}]
[1:checkpoint] State at the end of step 1:
{'messages': [HumanMessage(content='世界初のプロ球団は？', additional_kwargs={}, response_metadata={}, id='accb2f91-4c53-436e-a246-175fd9864071'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\n1869年には世界初のプロ球団であるシンシナティ・レッドストッキングスが設立され、1871年には世界初のプロ野球リーグであるナショナル・アソシエーションが設立された。このリーグ自体は5年で破綻したも\n\n1869年には世界初のプロ球団であるシンシナティ・レッドストッキングスが設立され、1871年には世界初のプロ野球リーグであるナショナル・アソシエーションが設立された。このリーグ自体は5年で破綻したも\n\n基礎がつくられたことから、一般に同国が野球発祥の地とされている。1869年には最初のプロチームが生まれ、アメリカで有数の人気スポーツとなり、国民的娯楽となった。元々は21点先取制、勝敗だけでなく両チー', additional_kwargs={}, response_metadata={}, id='6633d4cb-8275-41b0-9f36-a747667fb083')]}
[2:tasks] Starting 1 task for step 2:
- generate_answer -> {'messages': [HumanMessage(content='世界初のプロ球団は？', additional_kwargs={}, response_metadata={}, id='accb2f91-4c53-436e-a246-175fd9864071'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\n1869年には世界初のプロ球団であるシンシナティ・レッドストッキングスが設立され、1871年には世界初のプロ野球リーグであるナショナル・アソシエーションが設立された。このリーグ自体は5年で破綻したも\n\n1869年には世界初のプロ球団であるシンシナティ・レッドストッキングスが設立され、1871年には世界初のプロ野球リーグであるナショナル・アソシエーションが設立された。このリーグ自体は5年で破綻したも\n\n基礎がつくられたことから、一般に同国が野球発祥の地とされている。1869年には最初のプロチームが生まれ、アメリカで有数の人気スポーツとなり、国民的娯楽となった。元々は21点先取制、勝敗だけでなく両チー', additional_kwargs={}, response_metadata={}, id='6633d4cb-8275-41b0-9f36-a747667fb083')]}
[2:writes] Finished step 2 with writes to 1 channel:
- messages -> [AIMessage(content='世界初のプロ球団はシンシナティ・レッドストッキングス（Cincinnati Reds）です。1869年に設立されました。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 36, 'prompt_tokens': 256, 'total_tokens': 292, '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_00428b782a', 'finish_reason': 'stop', 'logprobs': None}, id='run-d7bf082a-acab-4987-9d79-e2319ecb9dd1-0', usage_metadata={'input_tokens': 256, 'output_tokens': 36, 'total_tokens': 292, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]
[2:checkpoint] State at the end of step 2:
{'messages': [HumanMessage(content='世界初のプロ球団は？', additional_kwargs={}, response_metadata={}, id='accb2f91-4c53-436e-a246-175fd9864071'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\n1869年には世界初のプロ球団であるシンシナティ・レッドストッキングスが設立され、1871年には世界初のプロ野球リーグであるナショナル・アソシエーションが設立された。このリーグ自体は5年で破綻したも\n\n1869年には世界初のプロ球団であるシンシナティ・レッドストッキングスが設立され、1871年には世界初のプロ野球リーグであるナショナル・アソシエーションが設立された。このリーグ自体は5年で破綻したも\n\n基礎がつくられたことから、一般に同国が野球発祥の地とされている。1869年には最初のプロチームが生まれ、アメリカで有数の人気スポーツとなり、国民的娯楽となった。元々は21点先取制、勝敗だけでなく両チー', additional_kwargs={}, response_metadata={}, id='6633d4cb-8275-41b0-9f36-a747667fb083'),
              AIMessage(content='世界初のプロ球団はシンシナティ・レッドストッキングス（Cincinnati Reds）です。1869年に設立されました。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 36, 'prompt_tokens': 256, 'total_tokens': 292, '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_00428b782a', 'finish_reason': 'stop', 'logprobs': None}, id='run-d7bf082a-acab-4987-9d79-e2319ecb9dd1-0', usage_metadata={'input_tokens': 256, 'output_tokens': 36, 'total_tokens': 292, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}

2. データソース2つのRAG

次は「サッカー」に関するWikipediaのデータソースを1つ増やして2つのデータソースから情報を取得するようにします。

コード例

重複が多いので、新しく作成したRouterノードのみ記載します。

class RouterResponse(TypedDict):
    goto: Literal["node1", "node2"]


def router_node(state: State) -> Command[Literal["node1", "node2"]]:
    llm = ChatOpenAI(model="gpt-4o-mini")
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "質問が野球に関するものであればnode1に、サッカーの質問であればnode2にルーティングしてください。",
            ),
            ("user", "{message}"),
        ]
    )
    chain = prompt | llm.with_structured_output(RouterResponse)
    response = chain.invoke({"message": state["messages"][-1]})
    return Command(goto=response["goto"])

どちらのデータソースを使用するかは、Routerノードが判断して適切なデータソースにルーティングするようにします。 Command機能を使うことで動的に次のノードを決定するので簡潔に書くことができます。

blog.langchain.dev

グラフ構造

実行例

response = graph.invoke(
     {"messages": [{"role": "user", "content": "野球の起源となった球技は何とされている？"}]},
     debug=True,
)
print(response["messages"][-1].content)

野球の起源となった球技としては、イギリスの「タウンボール」が挙げられます。タウンボールはイギリス系移民によってアメリカに持ち込まれ、そこから変化していったと考えられています。多くの研究者が、この過程を経て野球が形成されたと見ています。

実行ログ

[-1:checkpoint] State at the end of step -1:
{'messages': []}
[0:tasks] Starting 1 task for step 0:
- __start__ -> {'messages': [{'content': '野球の起源となった球技は何とされている？', 'role': 'user'}]}
[0:writes] Finished step 0 with writes to 1 channel:
- messages -> [{'content': '野球の起源となった球技は何とされている？', 'role': 'user'}]
[0:checkpoint] State at the end of step 0:
{'messages': [HumanMessage(content='野球の起源となった球技は何とされている？', additional_kwargs={}, response_metadata={}, id='c74dc0c3-5eea-4f43-866d-f8416e49f44e')]}
[1:tasks] Starting 1 task for step 1:
- router_node -> {'messages': [HumanMessage(content='野球の起源となった球技は何とされている？', additional_kwargs={}, response_metadata={}, id='c74dc0c3-5eea-4f43-866d-f8416e49f44e')]}
{'goto': 'node1'}
[1:writes] Finished step 1 with writes to 0 channels:

[1:checkpoint] State at the end of step 1:
{'messages': [HumanMessage(content='野球の起源となった球技は何とされている？', additional_kwargs={}, response_metadata={}, id='c74dc0c3-5eea-4f43-866d-f8416e49f44e')]}
[2:tasks] Starting 1 task for step 2:
- node1 -> {'messages': [HumanMessage(content='野球の起源となった球技は何とされている？', additional_kwargs={}, response_metadata={}, id='c74dc0c3-5eea-4f43-866d-f8416e49f44e')]}
[2:writes] Finished step 2 with writes to 1 channel:
- messages -> [{'content': '以下のコンテキスト情報を使用して回答してください。\n'
             '\n'
             '野球の起源は明確にはされていないが、イギリスの球技である「タウンボール」がイギリス系移民によってアメリカに持ち込まれた後に変化し、野球として形成されたと考える研究者が多い。1830年代から1840年\n'
             '\n'
             '類似競技\n'
             '\n'
             '[編集]\n'
             '\n'
             '野球とは異なる起源を持つ競技\n'
             '\n'
             '[編集]\n'
             '\n'
             'クリケット\n'
             '\n'
             'ラプター\n'
             '\n'
             'ラウンダーズ\n'
             '\n'
             'スティックボール\n'
             '\n'
             'シュラークバル\n'
             '\n'
             '野球から派生した競技\n'
             '\n'
             '[編集]\n'
             '\n'
             'ソフトボール\n'
             '\n'
             'スティックボール\n'
             '\n'
             'シュラークバル\n'
             '\n'
             '野球から派生した競技\n'
             '\n'
             '[編集]\n'
             '\n'
             'ソフトボール\n'
             '\n'
             'ペサパッロ\n'
             '\n'
             'ティーボール\n'
             '\n'
             'キックベースボール\n'
             '\n'
             'ハンドベースボール\n'
             '\n'
             'ラケットベースボール',
  'role': 'user'}]
[2:checkpoint] State at the end of step 2:
{'messages': [HumanMessage(content='野球の起源となった球技は何とされている？', additional_kwargs={}, response_metadata={}, id='c74dc0c3-5eea-4f43-866d-f8416e49f44e'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\n野球の起源は明確にはされていないが、イギリスの球技である「タウンボール」がイギリス系移民によってアメリカに持ち込まれた後に変化し、野球として形成されたと考える研究者が多い。1830年代から1840年\n\n類似競技\n\n[編集]\n\n野球とは異なる起源を持つ競技\n\n[編集]\n\nクリケット\n\nラプター\n\nラウンダーズ\n\nスティックボール\n\nシュラークバル\n\n野球から派生した競技\n\n[編集]\n\nソフトボール\n\nスティックボール\n\nシュラークバル\n\n野球から派生した競技\n\n[編集]\n\nソフトボール\n\nペサパッロ\n\nティーボール\n\nキックベースボール\n\nハンドベースボール\n\nラケットベースボール', additional_kwargs={}, response_metadata={}, id='7a7ce564-7ec9-45c0-a45d-965ba4399cf7')]}
[3:tasks] Starting 1 task for step 3:
- generate_answer -> {'messages': [HumanMessage(content='野球の起源となった球技は何とされている？', additional_kwargs={}, response_metadata={}, id='c74dc0c3-5eea-4f43-866d-f8416e49f44e'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\n野球の起源は明確にはされていないが、イギリスの球技である「タウンボール」がイギリス系移民によってアメリカに持ち込まれた後に変化し、野球として形成されたと考える研究者が多い。1830年代から1840年\n\n類似競技\n\n[編集]\n\n野球とは異なる起源を持つ競技\n\n[編集]\n\nクリケット\n\nラプター\n\nラウンダーズ\n\nスティックボール\n\nシュラークバル\n\n野球から派生した競技\n\n[編集]\n\nソフトボール\n\nスティックボール\n\nシュラークバル\n\n野球から派生した競技\n\n[編集]\n\nソフトボール\n\nペサパッロ\n\nティーボール\n\nキックベースボール\n\nハンドベースボール\n\nラケットベースボール', additional_kwargs={}, response_metadata={}, id='7a7ce564-7ec9-45c0-a45d-965ba4399cf7')]}
[3:writes] Finished step 3 with writes to 1 channel:
- messages -> [AIMessage(content='野球の起源となった球技としては、イギリスの「タウンボール」が挙げられます。タウンボールはイギリス系移民によってアメリカに持ち込まれ、そこから変化していったと考えられています。多くの研究者が、この過程を経て野球が形成されたと見ています。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 90, 'prompt_tokens': 242, 'total_tokens': 332, '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_00428b782a', 'finish_reason': 'stop', 'logprobs': None}, id='run-8a1dd045-cf16-4f8a-a377-58385e3c461c-0', usage_metadata={'input_tokens': 242, 'output_tokens': 90, 'total_tokens': 332, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]
[3:checkpoint] State at the end of step 3:
{'messages': [HumanMessage(content='野球の起源となった球技は何とされている？', additional_kwargs={}, response_metadata={}, id='c74dc0c3-5eea-4f43-866d-f8416e49f44e'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\n野球の起源は明確にはされていないが、イギリスの球技である「タウンボール」がイギリス系移民によってアメリカに持ち込まれた後に変化し、野球として形成されたと考える研究者が多い。1830年代から1840年\n\n類似競技\n\n[編集]\n\n野球とは異なる起源を持つ競技\n\n[編集]\n\nクリケット\n\nラプター\n\nラウンダーズ\n\nスティックボール\n\nシュラークバル\n\n野球から派生した競技\n\n[編集]\n\nソフトボール\n\nスティックボール\n\nシュラークバル\n\n野球から派生した競技\n\n[編集]\n\nソフトボール\n\nペサパッロ\n\nティーボール\n\nキックベースボール\n\nハンドベースボール\n\nラケットベースボール', additional_kwargs={}, response_metadata={}, id='7a7ce564-7ec9-45c0-a45d-965ba4399cf7'),
              AIMessage(content='野球の起源となった球技としては、イギリスの「タウンボール」が挙げられます。タウンボールはイギリス系移民によってアメリカに持ち込まれ、そこから変化していったと考えられています。多くの研究者が、この過程を経て野球が形成されたと見ています。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 90, 'prompt_tokens': 242, 'total_tokens': 332, '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_00428b782a', 'finish_reason': 'stop', 'logprobs': None}, id='run-8a1dd045-cf16-4f8a-a377-58385e3c461c-0', usage_metadata={'input_tokens': 242, 'output_tokens': 90, 'total_tokens': 332, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}

response = graph.invoke(
     {"messages": [{"role": "user", "content": "サッカーの世界初の国際試合はどこ対どこ？"}]},
     debug=True,
)
print(response["messages"][-1].content)

サッカーの世界初の公式国際試合は、1867年に行われたイングランドとスコットランドの試合で、スコアは0-0の引き分けでした。

実行ログ

[-1:checkpoint] State at the end of step -1:
{'messages': []}
[0:tasks] Starting 1 task for step 0:
- __start__ -> {'messages': [{'content': 'サッカーの世界初の国際試合はどこ対どこ？', 'role': 'user'}]}
[0:writes] Finished step 0 with writes to 1 channel:
- messages -> [{'content': 'サッカーの世界初の国際試合はどこ対どこ？', 'role': 'user'}]
[0:checkpoint] State at the end of step 0:
{'messages': [HumanMessage(content='サッカーの世界初の国際試合はどこ対どこ？', additional_kwargs={}, response_metadata={}, id='29d8ca6c-a4e3-4253-a284-2d7eda634fad')]}
[1:tasks] Starting 1 task for step 1:
- router_node -> {'messages': [HumanMessage(content='サッカーの世界初の国際試合はどこ対どこ？', additional_kwargs={}, response_metadata={}, id='29d8ca6c-a4e3-4253-a284-2d7eda634fad')]}
{'goto': 'node2'}
[1:writes] Finished step 1 with writes to 0 channels:

[1:checkpoint] State at the end of step 1:
{'messages': [HumanMessage(content='サッカーの世界初の国際試合はどこ対どこ？', additional_kwargs={}, response_metadata={}, id='29d8ca6c-a4e3-4253-a284-2d7eda634fad')]}
[2:tasks] Starting 1 task for step 2:
- node2 -> {'messages': [HumanMessage(content='サッカーの世界初の国際試合はどこ対どこ？', additional_kwargs={}, response_metadata={}, id='29d8ca6c-a4e3-4253-a284-2d7eda634fad')]}
[2:writes] Finished step 2 with writes to 1 channel:
- messages -> [{'content': '以下のコンテキスト情報を使用して回答してください。\n'
             '\n'
             'で、世界で最初の“公式”国際試合が、イングランドとスコットランドの間で実施された。スコアは0－0の引き分けだった[25]。その後1880年代までに、スコットランド、ウェールズ、アイルランドではサッカー\n'
             '\n'
             'で、世界で最初の“公式”国際試合が、イングランドとスコットランドの間で実施された。スコアは0－0の引き分けだった[25]。その後1880年代までに、スコットランド、ウェールズ、アイルランドではサッカー\n'
             '\n'
             'このFAルールでの初の試合、つまり世界初の「サッカー」の試合は、1863年12月19日にイングランドで行われたリッチモンド対バーンズ戦で、0-0の引き分けだった[25]。',
  'role': 'user'}]
[2:checkpoint] State at the end of step 2:
{'messages': [HumanMessage(content='サッカーの世界初の国際試合はどこ対どこ？', additional_kwargs={}, response_metadata={}, id='29d8ca6c-a4e3-4253-a284-2d7eda634fad'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\nで、世界で最初の“公式”国際試合が、イングランドとスコットランドの間で実施された。スコアは0－0の引き分けだった[25]。その後1880年代までに、スコットランド、ウェールズ、アイルランドではサッカー\n\nで、世界で最初の“公式”国際試合が、イングランドとスコットランドの間で実施された。スコアは0－0の引き分けだった[25]。その後1880年代までに、スコットランド、ウェールズ、アイルランドではサッカー\n\nこのFAルールでの初の試合、つまり世界初の「サッカー」の試合は、1863年12月19日にイングランドで行われたリッチモンド対バーンズ戦で、0-0の引き分けだった[25]。', additional_kwargs={}, response_metadata={}, id='f081b9fe-7441-4c27-a77d-2d3938b4ce70')]}
[3:tasks] Starting 1 task for step 3:
- generate_answer -> {'messages': [HumanMessage(content='サッカーの世界初の国際試合はどこ対どこ？', additional_kwargs={}, response_metadata={}, id='29d8ca6c-a4e3-4253-a284-2d7eda634fad'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\nで、世界で最初の“公式”国際試合が、イングランドとスコットランドの間で実施された。スコアは0－0の引き分けだった[25]。その後1880年代までに、スコットランド、ウェールズ、アイルランドではサッカー\n\nで、世界で最初の“公式”国際試合が、イングランドとスコットランドの間で実施された。スコアは0－0の引き分けだった[25]。その後1880年代までに、スコットランド、ウェールズ、アイルランドではサッカー\n\nこのFAルールでの初の試合、つまり世界初の「サッカー」の試合は、1863年12月19日にイングランドで行われたリッチモンド対バーンズ戦で、0-0の引き分けだった[25]。', additional_kwargs={}, response_metadata={}, id='f081b9fe-7441-4c27-a77d-2d3938b4ce70')]}
[3:writes] Finished step 3 with writes to 1 channel:
- messages -> [AIMessage(content='サッカーの世界初の公式国際試合は、1867年に行われたイングランドとスコットランドの試合で、スコアは0-0の引き分けでした。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 49, 'prompt_tokens': 249, 'total_tokens': 298, '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_13eed4fce1', 'finish_reason': 'stop', 'logprobs': None}, id='run-70ee4d89-00fa-40f1-8a3e-ff79a17722ca-0', usage_metadata={'input_tokens': 249, 'output_tokens': 49, 'total_tokens': 298, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]
[3:checkpoint] State at the end of step 3:
{'messages': [HumanMessage(content='サッカーの世界初の国際試合はどこ対どこ？', additional_kwargs={}, response_metadata={}, id='29d8ca6c-a4e3-4253-a284-2d7eda634fad'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\nで、世界で最初の“公式”国際試合が、イングランドとスコットランドの間で実施された。スコアは0－0の引き分けだった[25]。その後1880年代までに、スコットランド、ウェールズ、アイルランドではサッカー\n\nで、世界で最初の“公式”国際試合が、イングランドとスコットランドの間で実施された。スコアは0－0の引き分けだった[25]。その後1880年代までに、スコットランド、ウェールズ、アイルランドではサッカー\n\nこのFAルールでの初の試合、つまり世界初の「サッカー」の試合は、1863年12月19日にイングランドで行われたリッチモンド対バーンズ戦で、0-0の引き分けだった[25]。', additional_kwargs={}, response_metadata={}, id='f081b9fe-7441-4c27-a77d-2d3938b4ce70'),
              AIMessage(content='サッカーの世界初の公式国際試合は、1867年に行われたイングランドとスコットランドの試合で、スコアは0-0の引き分けでした。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 49, 'prompt_tokens': 249, 'total_tokens': 298, '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_13eed4fce1', 'finish_reason': 'stop', 'logprobs': None}, id='run-70ee4d89-00fa-40f1-8a3e-ff79a17722ca-0', usage_metadata={'input_tokens': 249, 'output_tokens': 49, 'total

それぞれの質問にあったデータソースを参照し、適切な回答をしていることが確認できます。

3. データソース2つ+Web検索のRAG

今度はデータソースにない質問が来た場合、Web検索を行って回答を行うRAGを作成します。 Web検索にはTavily Search APIを使用します。

コード例

コード全文

from typing import Annotated

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_community.document_loaders import UnstructuredURLLoader
from langchain_community.tools import TavilySearchResults
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.types import Command
from typing_extensions import Literal, TypedDict


llm = ChatOpenAI(model="gpt-4o-mini")
embedding_model = OpenAIEmbeddings()


vectorstore_a = Chroma(
    embedding_function=embedding_model, persist_directory="./chroma_db_a"
)

vectorstore_b = Chroma(
    embedding_function=embedding_model, persist_directory="./chroma_db_b"
)


class State(TypedDict):
    messages: Annotated[list[str], add_messages]
    query: str


class RouterResponse(TypedDict):
    goto: Literal["node1", "node2", "generate_query"]


class QueryResponse(TypedDict):
    query: str


def router_node(state: State) -> Command[Literal["node1", "node2", "generate_query"]]:
    llm = ChatOpenAI(model="gpt-4o-mini")
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "質問が野球に関するものであればnode1に、サッカーの質問であればnode2, それ以外の質問はgenerate_queryにルーティングしてください。",
            ),
            ("user", "{message}"),
        ]
    )
    chain = prompt | llm.with_structured_output(RouterResponse)
    response = chain.invoke({"message": state["messages"][-1]})
    print(response)
    return Command(goto=response["goto"])


def retrieve(state: State, vectorstore: Chroma):
    question = state["messages"][-1].content
    retriever = vectorstore.similarity_search(question, k=3)
    context = "\n\n".join([doc.page_content for doc in retriever])
    return {
        "messages": [
            {
                "role": "user",
                "content": f"以下のコンテキスト情報を使用して回答してください。\n\n{context}",
            }
        ]
    }


def generate_query(state: State):
    llm = ChatOpenAI(model="gpt-4o-mini")
    message = state["messages"][-1].content
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "Web検索用のクエリを生成してください。"),
            ("user", "{message}"),
        ]
    )
    chain = prompt | llm.with_structured_output(QueryResponse)
    response = chain.invoke({"message": state["messages"][-1]})
    print(response)
    return {"query": response["query"]}


def node3(state: State):
    search = TavilySearchResults(max_results=3)
    result = search.invoke(state["query"])
    print(result)
    context = "\n\n".join([doc["content"] for doc in result])
    return {
        "messages": [
            {
                "role": "user",
                "content": f"以下のコンテキスト情報を使用して回答してください。\n\n{context}",
            }
        ]
    }


def generate_answer(state: State):
    response = llm.invoke(state["messages"])
    return {"messages": [response]}


def node1(state: State, vectorstore=vectorstore_a):
    return retrieve(state, vectorstore)


def node2(state: State, vectorstore=vectorstore_b):
    return retrieve(state, vectorstore)


urls_a = ["https://ja.wikipedia.org/wiki/%E9%87%8E%E7%90%83"]
urls_b = ["https://ja.wikipedia.org/wiki/%E3%82%B5%E3%83%83%E3%82%AB%E3%83%BC"]


def load_data(urls):
    loader = UnstructuredURLLoader(urls)
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splits = text_splitter.split_documents(docs)
    return splits


docs_a = load_data(urls_a)
docs_b = load_data(urls_b)
vectorstore_a.add_documents(docs_a)
vectorstore_b.add_documents(docs_b)

retriever_a = vectorstore_a.as_retriever(search_kwargs={"k": 3})
retriever_b = vectorstore_b.as_retriever(search_kwargs={"k": 3})

builder = StateGraph(State)

builder.add_edge(START, "router_node")

builder.add_node(router_node)
builder.add_node(node1)
builder.add_node(node2)
builder.add_node(generate_query)
builder.add_node(node3)
builder.add_node(generate_answer)

builder.add_edge("node1", "generate_answer")
builder.add_edge("node2", "generate_answer")
builder.add_edge("generate_query", "node3")
builder.add_edge("node3", "generate_answer")
builder.add_edge("generate_answer", END)

graph = builder.compile()

Web検索するためのノードをnode3として追加します。

def node3(state: State):
    search = TavilySearchResults(max_results=3)
    result = search.invoke(state["query"])
    print(result)
    context = "\n\n".join([doc["content"] for doc in result])
    return {
        "messages": [
            {
                "role": "user",
                "content": f"以下のコンテキスト情報を使用して回答してください。\n\n{context}",
            }
        ]
    }

また、Web検索する前は検索に適したクエリに変換するためのノードを追加します。

def generate_query(state: State):
    llm = ChatOpenAI(model="gpt-4o-mini")
    message = state["messages"][-1].content
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "Web検索用のクエリを生成してください。"),
            ("user", "{message}"),
        ]
    )
    chain = prompt | llm.with_structured_output(QueryResponse)
    response = chain.invoke({"message": state["messages"][-1]})
    print(response)
    return {"query": response["query"]}

グラフ構造

実行例

response = graph.invoke(
     {"messages": [{"role": "user", "content": "パリオリンピックの男子バスケットボールの優勝国はどこ？"}]},
     debug=True,
)
print(response["messages"][-1].content)

パリオリンピックの男子バスケットボールで優勝した国はアメリカ合衆国です。決勝戦ではフランスを98対87で破り、5大会連続17度目の金メダルを獲得しました。

実行ログ

[-1:checkpoint] State at the end of step -1:
{'messages': []}
[0:tasks] Starting 1 task for step 0:
- __start__ -> {'messages': [{'content': 'パリオリンピックの男子バスケットボールの優勝国はどこ？', 'role': 'user'}]}
[0:writes] Finished step 0 with writes to 1 channel:
- messages -> [{'content': 'パリオリンピックの男子バスケットボールの優勝国はどこ？', 'role': 'user'}]
[0:checkpoint] State at the end of step 0:
{'messages': [HumanMessage(content='パリオリンピックの男子バスケットボールの優勝国はどこ？', additional_kwargs={}, response_metadata={}, id='a1b7b778-d44c-439e-84a1-52e58ca0b5a0')]}
[1:tasks] Starting 1 task for step 1:
- router_node -> {'messages': [HumanMessage(content='パリオリンピックの男子バスケットボールの優勝国はどこ？', additional_kwargs={}, response_metadata={}, id='a1b7b778-d44c-439e-84a1-52e58ca0b5a0')]}
{'goto': 'generate_query'}
[1:writes] Finished step 1 with writes to 0 channels:

[1:checkpoint] State at the end of step 1:
{'messages': [HumanMessage(content='パリオリンピックの男子バスケットボールの優勝国はどこ？', additional_kwargs={}, response_metadata={}, id='a1b7b778-d44c-439e-84a1-52e58ca0b5a0')]}
[2:tasks] Starting 1 task for step 2:
- generate_query -> {'messages': [HumanMessage(content='パリオリンピックの男子バスケットボールの優勝国はどこ？', additional_kwargs={}, response_metadata={}, id='a1b7b778-d44c-439e-84a1-52e58ca0b5a0')]}
{'query': '2024 パリオリンピック 男子バスケットボール 優勝国はどこ'}
[2:writes] Finished step 2 with writes to 1 channel:
- query -> '2024 パリオリンピック 男子バスケットボール 優勝国はどこ'
[2:checkpoint] State at the end of step 2:
{'messages': [HumanMessage(content='パリオリンピックの男子バスケットボールの優勝国はどこ？', additional_kwargs={}, response_metadata={}, id='a1b7b778-d44c-439e-84a1-52e58ca0b5a0')],
 'query': '2024 パリオリンピック 男子バスケットボール 優勝国はどこ'}
[3:tasks] Starting 1 task for step 3:
- node3 -> {'messages': [HumanMessage(content='パリオリンピックの男子バスケットボールの優勝国はどこ？', additional_kwargs={}, response_metadata={}, id='a1b7b778-d44c-439e-84a1-52e58ca0b5a0')],
 'query': '2024 パリオリンピック 男子バスケットボール 優勝国はどこ'}
[{'url': 'https://ja.wikipedia.org/wiki/2024%E5%B9%B4%E3%83%91%E3%83%AA%E3%82%AA%E3%83%AA%E3%83%B3%E3%83%94%E3%83%83%E3%82%AF%E3%81%AE%E3%83%90%E3%82%B9%E3%82%B1%E3%83%83%E3%83%88%E3%83%9C%E3%83%BC%E3%83%AB%E7%AB%B6%E6%8A%80', 'content': '... 優勝国… アメリカ合衆国の旗 · アメリカ合衆国; オリンピック開催国… フランスの旗 · フランス. 世界最終予選の組み合わせ抽選は2023年10月5日に行われた。16か国は直近'}, {'url': 'https://www.olympics.com/ja/news/paris2024-basketball-men-france-united-states-of-america', 'content': 'パリ2024バスケットボール男子決勝が現地時間8月10日にベルシー・アリーナで行われた。決勝に進出したのは開催国・フランス代表と、オリンピック4連覇'}, {'url': 'https://www.yomiuri.co.jp/olympic/2024/20240811-OYT1T50061/', 'content': 'パリオリンピック・男子バスケットボールの決勝が10日行われ、アメリカが98―87でフランスを破り、5大会連続17度目となる金メダルを獲得した。'}]
[3:writes] Finished step 3 with writes to 1 channel:
- messages -> [{'content': '以下のコンテキスト情報を使用して回答してください。\n'
             '\n'
             '... 優勝国… アメリカ合衆国の旗 · アメリカ合衆国; オリンピック開催国… フランスの旗 · フランス. '
             '世界最終予選の組み合わせ抽選は2023年10月5日に行われた。16か国は直近\n'
             '\n'
             'パリ2024バスケットボール男子決勝が現地時間8月10日にベルシー・アリーナで行われた。決勝に進出したのは開催国・フランス代表と、オリンピック4連覇\n'
             '\n'
             'パリオリンピック・男子バスケットボールの決勝が10日行われ、アメリカが98―87でフランスを破り、5大会連続17度目となる金メダルを獲得した。',
  'role': 'user'}]
[3:checkpoint] State at the end of step 3:
{'messages': [HumanMessage(content='パリオリンピックの男子バスケットボールの優勝国はどこ？', additional_kwargs={}, response_metadata={}, id='a1b7b778-d44c-439e-84a1-52e58ca0b5a0'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\n... 優勝国… アメリカ合衆国の旗 · アメリカ合衆国; オリンピック開催国… フランスの旗 · フランス. 世界最終予選の組み合わせ抽選は2023年10月5日に行われた。16か国は直近\n\nパリ2024バスケットボール男子決勝が現地時間8月10日にベルシー・アリーナで行われた。決勝に進出したのは開催国・フランス代表と、オリンピック4連覇\n\nパリオリンピック・男子バスケットボールの決勝が10日行われ、アメリカが98―87でフランスを破り、5大会連続17度目となる金メダルを獲得した。', additional_kwargs={}, response_metadata={}, id='e3b281ba-b89d-47c8-923d-99839b256499')],
 'query': '2024 パリオリンピック 男子バスケットボール 優勝国はどこ'}
[4:tasks] Starting 1 task for step 4:
- generate_answer -> {'messages': [HumanMessage(content='パリオリンピックの男子バスケットボールの優勝国はどこ？', additional_kwargs={}, response_metadata={}, id='a1b7b778-d44c-439e-84a1-52e58ca0b5a0'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\n... 優勝国… アメリカ合衆国の旗 · アメリカ合衆国; オリンピック開催国… フランスの旗 · フランス. 世界最終予選の組み合わせ抽選は2023年10月5日に行われた。16か国は直近\n\nパリ2024バスケットボール男子決勝が現地時間8月10日にベルシー・アリーナで行われた。決勝に進出したのは開催国・フランス代表と、オリンピック4連覇\n\nパリオリンピック・男子バスケットボールの決勝が10日行われ、アメリカが98―87でフランスを破り、5大会連続17度目となる金メダルを獲得した。', additional_kwargs={}, response_metadata={}, id='e3b281ba-b89d-47c8-923d-99839b256499')],
 'query': '2024 パリオリンピック 男子バスケットボール 優勝国はどこ'}
[4:writes] Finished step 4 with writes to 1 channel:
- messages -> [AIMessage(content='パリオリンピックの男子バスケットボールで優勝した国はアメリカ合衆国です。決勝戦ではフランスを98対87で破り、5大会連続17度目の金メダルを獲得しました。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 62, 'prompt_tokens': 230, 'total_tokens': 292, '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_00428b782a', 'finish_reason': 'stop', 'logprobs': None}, id='run-9e37f769-7c99-46d1-aaae-de1c2b3a0021-0', usage_metadata={'input_tokens': 230, 'output_tokens': 62, 'total_tokens': 292, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]
[4:checkpoint] State at the end of step 4:
{'messages': [HumanMessage(content='パリオリンピックの男子バスケットボールの優勝国はどこ？', additional_kwargs={}, response_metadata={}, id='a1b7b778-d44c-439e-84a1-52e58ca0b5a0'),
              HumanMessage(content='以下のコンテキスト情報を使用して回答してください。\n\n... 優勝国… アメリカ合衆国の旗 · アメリカ合衆国; オリンピック開催国… フランスの旗 · フランス. 世界最終予選の組み合わせ抽選は2023年10月5日に行われた。16か国は直近\n\nパリ2024バスケットボール男子決勝が現地時間8月10日にベルシー・アリーナで行われた。決勝に進出したのは開催国・フランス代表と、オリンピック4連覇\n\nパリオリンピック・男子バスケットボールの決勝が10日行われ、アメリカが98―87でフランスを破り、5大会連続17度目となる金メダルを獲得した。', additional_kwargs={}, response_metadata={}, id='e3b281ba-b89d-47c8-923d-99839b256499'),
              AIMessage(content='パリオリンピックの男子バスケットボールで優勝した国はアメリカ合衆国です。決勝戦ではフランスを98対87で破り、5大会連続17度目の金メダルを獲得しました。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 62, 'prompt_tokens': 230, 'total_tokens': 292, '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_00428b782a', 'finish_reason': 'stop', 'logprobs': None}, id='run-9e37f769-7c99-46d1-aaae-de1c2b3a0021-0', usage_metadata={'input_tokens': 230, 'output_tokens': 62, 'total_tokens': 292, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})],
 'query': '2024 パリオリンピック 男子バスケットボール 優勝国はどこ'}

サッカーと野球以外の質問なので、データソースは参照せず、Web検索を行い、正しい回答を導き出していることが確認できます。

まとめ

今回はLangGraphを用いたRAGを構築してみました。LangGraphは記述量が多くなりがちですが、ブラックボックスなく処理が書けるのがいいなと思っています。今回は単純なルーティングのワークフローのRAGでしたが、次はもっと自律的に質問に対処してくれるAgenticなRAGを構築してみたいです。

執筆者堤 拓哉

インフラエンジニア。データ分析やデータ基盤周りの話に興味があります。