Zenn
🦁

RAG From Scratch をやってみた (2/6) : Routing

2025/03/08に公開
LangChain
RAG
tech

概要

今回はRAGの構成のなかのRoutingに関する部分を読み解いてみた内容です。

こちらのnotebookを実行してみながらまとめました。
https://github.com/langchain-ai/rag-from-scratch/blob/main/rag_from_scratch_10_and_11.ipynb

Routingはその名の通り、処理の流れもイメージしやすいのでコードと実行結果の例の紹介が中心となります。

実行環境

必要なライブラリをインストール

.ipynbファイル
! pip install langchain_community tiktoken langchain-openai langchainhub chromadb langchain

環境変数に、OpenAI APIの認証情報等をセットします。（今回はAzure OpenAIのAPIを利用しました。）

環境変数の読み込み
.ipynbファイル
from IPython import display # 結果を見やすくするライブラリインポート
import os
from dotenv import load_dotenv
load_dotenv() # 環境変数を読み込み
os.environ['LANGCHAIN_TRACING_V2'] = 'true'
os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'
.env.template
# Azure OpenAI Service のAPI情報をセットする
AZURE_OPENAI_ENDPOINT=
AZURE_OPENAI_API_KEY=

# Azureのデプロイメント名をセット
DEPLOYMENT_NAME=
# APIのバージョンをセット
API_VERSION=

# Azureのembedding modelのデプロイメント名をセット
EMBE_DEPLOYMENT_NAME=
# embedding modelのAPIバージョンをセット
EMBE_API_VERSION=

# LangSmithのAPIKEYをセット（ない場合も実行可）
LANGCHAIN_API_KEY=

LLMにはgpt-4o-mini, 埋め込みモデルはtext-embedding-3-largeを使用しました。

Logical and Semantic routing

ユーザの質問文から意図を読み取り、適切なアクションを判別する仕組みです。
cookbook内では、以下２パターンが紹介されています。

  • 検索対象のデータソースの選択
  • 回答生成を促すシステムプロンプトの選択

検索対象のデータソースの選択


Function Callingを使用して、検索対象のデータソースを判別しています。
※ コードは判定部分のみで、質問文受け取り～回答生成までの流れではありません。

検索対象のデータソースを判別するコード
from typing import Literal

from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from langchain_openai import AzureChatOpenAI

# データモデル
class RouteQuery(BaseModel):
    """ユーザーの質問を最も関連性のあるデータソースにルーティングする。"""

    datasource: Literal["python_docs", "js_docs", "golang_docs"] = Field(
        ...,
        description="ユーザーの質問に基づいて、最も関連性のあるデータソースを選択する",
    )

# 関数呼び出しを持つLLM
llm = AzureChatOpenAI(
    openai_api_version=os.getenv("API_VERSION"),
    azure_deployment=os.getenv("DEPLOYMENT_NAME"),
    temperature=0
)
structured_llm = llm.with_structured_output(RouteQuery)

# プロンプト
system = """あなたはユーザーの質問を適切なデータソースにルーティングする専門家です。

質問が指しているプログラミング言語に基づいて、それを関連するデータソースにルーティングしてください。"""

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        ("human", "{question}"),
    ]
)

# ルーターを定義
router = prompt | structured_llm

作成したルータ―を使用して、検証するコードと実行結果は以下です：

python
question = """次のコードが動作しない理由は何ですか：

from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_messages(["human", "speak in {language}"])
prompt.invoke("french")
"""
result = router.invoke({"question": question})
result
output
RouteQuery(datasource='python_docs')
python
result.datasource
output
'python_docs'

質問文から、正しくpython_docs(pythonコードに関するドキュメント)が選択されました。

routerとこのresult.datasourceを使って、データソースの分岐を行う処理の例は以下です。

コード例
pyton
def choose_route(result):
    if "python_docs" in result.datasource.lower():
        ### Logic here 
        return "chain for python_docs"
    elif "js_docs" in result.datasource.lower():
        ### Logic here 
        return "chain for js_docs"
    else:
        ### Logic here 
        return "golang_docs"

from langchain_core.runnables import RunnableLambda

full_chain = router | RunnableLambda(choose_route)
full_chain.invoke({"question": question})
output
'chain for python_docs'

回答生成を促すシステムプロンプトの選択


もう１つのrouterは、ユーザの質問文から関連するシステムプロンプトを選択するものです。
質問文とプロンプトが似ているか？は同じembedding modelを使用してベクトル化し、コサイン類似度で判断しているようです。

例では、物理学の教授のロールのプロンプト（PHYSICS）、もしくは数学者のロールのプロンプト（MATH）のどちらかを選択してユーザの質問「ブラックホールとは何ですか」に回答しています。
※ 関連文書は取得しません。

プロンプトの切替コード
python
from langchain.utils.math import cosine_similarity
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import AzureOpenAIEmbeddings

# 2つのプロンプト
physics_template = """あなたは非常に賢い物理学の教授です。\
物理学に関する質問に簡潔で理解しやすい方法で答えるのが得意です。\
もし質問に答えられない場合は、それを認めます。

質問は次の通りです：
{query}"""

math_template = """あなたは非常に優れた数学者です。数学の質問に答えるのが得意です。\
なぜなら、難しい問題をその構成要素に分解し、各部分を解決した後に、全体を組み立てて広範な質問に答えることができるからです。

質問は次の通りです：
{query}"""

# プロンプトを埋め込む
embeddings = AzureOpenAIEmbeddings(
    azure_deployment=os.environ.get("EMBE_DEPLOYMENT_NAME"),
    openai_api_version=os.environ.get("EMBE_API_VERSION"),
)
prompt_templates = [physics_template, math_template]
prompt_embeddings = embeddings.embed_documents(prompt_templates)

# 質問をプロンプトにルーティング
def prompt_router(input):
    # 質問を埋め込む
    query_embedding = embeddings.embed_query(input["query"])
    # 類似度を計算
    similarity = cosine_similarity([query_embedding], prompt_embeddings)[0]
    most_similar = prompt_templates[similarity.argmax()]
    # 選ばれたプロンプト
    print("MATH を使用中" if most_similar == math_template else "PHYSICS を使用中")
    return PromptTemplate.from_template(most_similar)


chain = (
    {"query": RunnablePassthrough()}
    | RunnableLambda(prompt_router)
    | llm
    | StrOutputParser()
)

answer = chain.invoke("ブラックホールとは何ですか")
display.Markdown(answer)

このコードの実行結果は以下となります。

output
PHYSICS を使用中

ブラックホールとは、非常に強い重力を持つ天体のことです。これは、星がその生涯の終わりに重力崩壊を起こし、非常に小さな体積に圧縮されることで形成されます。ブラックホールの重力は、光さえも脱出できないほど強いため、直接見ることはできませんが、その周囲の物質や光の挙動から存在を推測することができます。

ブラックホールには主に3つのタイプがあります：

恒星ブラックホール：大きな星が超新星爆発を経て形成される。
超大質量ブラックホール：銀河の中心に存在し、数百万から数十億倍の太陽質量を持つ。
中間質量ブラックホール：恒星ブラックホールと超大質量ブラックホールの中間の質量を持つが、存在はまだ確認されていない。
ブラックホールの境界には「事象の地平線」と呼ばれる境界があり、これを越えると何も戻ってこれません。

ユーザの質問文により関連する、物理学の教授のロールのプロンプト（PHYSICS）の方が選択されて回答生成されていることがわかります。

Discussion