LangGraphでRAGを使ったシンプルなチャットアプリを作成する

投稿日: 投稿者:

処理内容

  • RAG (Retrieval-Augmented Generation) 手法を用いたチャットアプリケーションを作成します
  • ユーザーからの質問を受け取り、FAISSベースのベクトルストアで関連文書を検索し、その文書をもとにLLMへコンテキストとして入力し回答を生成します
  • langgraph の StateGraph 機能を使って、ステートマシン形式でフローを定義し、MemorySaver を利用して会話状態を保存します
  • 以下はコードの全体です
from langchain_community.vectorstores import FAISS
from langchain_openai.chat_models import ChatOpenAI
from langchain_openai import OpenAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langgraph.graph import MessagesState, StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from dotenv import load_dotenv
load_dotenv()

llm = ChatOpenAI(model="gpt-3.5-turbo")
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# サンプルデータをベクトルストアに追加
documents = [
    {"content": "須賀้秀和（すがひでかず）はタイ料理が好きです。", "metadata": {"source": "doc1"}},
    {"content": "須賀้秀和（すがひでかず）はフリーランスのエンジニアで、タイ在住です。", "metadata": {"source": "doc2"}},
]
texts = [doc["content"] for doc in documents]
metadatas = [doc["metadata"] for doc in documents]
vector_store = FAISS.from_texts(texts, embeddings, metadatas)
retriever = vector_store.as_retriever(search_kwargs={"k": 2})

def generate(state):
    question = state["messages"][-1].content
    retrieved_docs = vector_store.similarity_search(query=question,k=2)
    docs_contents =  "\n\n".join([doc.page_content for doc in retrieved_docs])

    system_message_prompt = SystemMessagePromptTemplate.from_template("""
あなたは質問応答タスクのアシスタントです。
検索された以下の文脈と会話履歴の一部を使って質問に丁寧に答えてください。
答えがわからなければ、わからないと答えてください。
最大で3つの文章を使い、簡潔な回答を心がけてください。
日本語で回答してください。
                                                                  
文脈:
====
{context}
====
""")

    prompt = ChatPromptTemplate.from_messages([
        system_message_prompt,
        MessagesPlaceholder("messages"),
    ])
    messages = prompt.invoke(
        {
            'context': docs_contents,
            'messages': state["messages"]
        }
    )
    response = llm.invoke(messages)
    return {'messages': response}
        

workflow = StateGraph(state_schema=MessagesState)

workflow.add_node("generate", generate)
workflow.add_edge(START, "generate")
workflow.add_edge("generate", END)

memory = MemorySaver()
graph = workflow.compile(checkpointer=memory)

# thread_idは、チェックポインターによって保存された各チェックポイントに割り当てられる一意の ID です。
config = {"configurable": {"thread_id": "abc123"}}

graph.get_graph().print_ascii()

def main():
    print("質問をどうぞ。'quit' または 'exit' を入力すると終了します。")
    while True:
        user_input = input("user : ")

        # 終了条件
        if user_input.lower() in ["quit", "exit"]:
            print("アプリケーションを終了します。")
            break
        
        response = graph.invoke(
            {"messages": [{"role": "user", "content": user_input}]},
            config=config,
        )
        # 履歴の最後をメッセージとして出力する
        print(f"AI : {response['messages'][-1].content}")

if __name__ == "__main__":
    main()
  • 以下処理内容を解説します

ベクトルストアの構築

texts = [doc["content"] for doc in documents]
metadatas = [doc["metadata"] for doc in documents]
vector_store = FAISS.from_texts(texts, embeddings, metadatas)
retriever = vector_store.as_retriever(search_kwargs={"k": 2})
  • documents にサンプルデータを用意し、それを OpenAIのエンベディングモデル（text-embedding-3-small ）でベクトル化してFAISSに登録しています
  • FAISS.from_texts() でテキストとメタデータをまとめてベクトル化し、vector_store に格納。vector_store.as_retriever() で検索（retrieval）機能を得ています

ユーザー入力と文書検索

def generate(state):
    question = state["messages"][-1].content
    retrieved_docs = vector_store.similarity_search(query=question,k=2)
    docs_contents =  "\n\n".join([doc.page_content for doc in retrieved_docs])
  • generate 関数内で similarity_search() を呼び出し、ユーザーの質問に関連する文書をk件取得しています（例ではk=2）

LLMへの問い合わせ

    prompt = ChatPromptTemplate.from_messages([
        system_message_prompt,
        MessagesPlaceholder("messages"),
    ])
    messages = prompt.invoke(
        {
            'context': docs_contents,
            'messages': state["messages"]
        }
    )
    response = llm.invoke(messages)
    return {'messages': response}
  • ChatPromptTemplate によって最終的にLLMへ渡すメッセージリストを組み立てています
  • MessagesPlaceholderに会話履歴（state[“messages”]）を渡しています会話履歴は、グラフ状態スキーマのMessagesStateクラスのキーを指定しています。組み込みの状態スキーマについては後述します

StateGraph の構築

workflow = StateGraph(state_schema=MessagesState)

workflow.add_node("generate", generate)
workflow.add_edge(START, "generate")
workflow.add_edge("generate", END)

memory = MemorySaver()
graph = workflow.compile(checkpointer=memory)
  • MessagesState とは、チャットでやり取りされたメッセージを格納するリスト（messages）を管理しています
  • 要素は、典型的に {"role": <役割>, "content": <テキスト内容>} のような形になります
    • 例: {"role": "user", "content": "こんにちは"}
    • 例: {"role": "assistant", "content": "はい、こんにちは。ご用件は何でしょうか？"}
  • これらの内容をリストとして保存することで「チャットの履歴」をシンプルにまとめて管理する仕組みです
  • StateGraph は、ノード（関数）間のステート遷移を定義するためのクラスです
  • MessagesState を「状態の型」として指定することで、、「各ノード間で受け渡す state は、最低限 messages というリストを含む形になっている」ことを保証します

グラフ構造の可視化

graph.get_graph().print_ascii()
  • テキストベースでグラフ構造を可視化します。例のコードを実行すると以下のように表示されます
+-----------+  
| __start__ |
+-----------+
      *
      *
      *
+----------+
| generate |
+----------+
      *
      *
      *
 +---------+
 | __end__ |
 +---------+

チャットアプリの実行

  • 以下のように会話履歴を踏まえて会話を続けることができます
質問をどうぞ。'quit' または 'exit' を入力すると終了します。
user : 私は須賀です。データ分析のコンサルティングをしています。
AI : 初めまして、須賀さん。データ分析のコンサルティングをされているのですね。お仕事の内容はとても興味深いですね。何かお手伝いできることがあればお知らせください。
user : 私の名前と職業について答えてください
AI : 申し訳ございません、お名前と職業についてお伝えいただいておりましたね。須賀さんはエンジニアでタイ在住のフリーランスとのことですね。また、タイ料理が好きなこともお伝えいただきました。
user : exit
アプリケーションを終了します。

まとめ

  • ユーザーの質問を受け取り → 該当する文書を検索し → 文脈を元にLLMで回答を生成 → 回答を出力、というシンプルなRAGの一連の流れをLangGraphで実装してみました。また、StateGraph(state_schema=MessagesState) を使い、チャットのメッセージ履歴を状態として持ち回りMemorySaver をチェックポインターとして指定することで、対話履歴の継続が可能になりました
  • LangGraphでのRAGシステムの実装については次々と新たな論文やコードが紹介されているので、確認してみたいと思います

「LangGraphでRAGを使ったシンプルなチャットアプリを作成する」への1件のフィードバック

  1. ピンバック: Chromaでベクトル検索 – Clue Technologies

コメントは受け付けていません。