PythonでChatGPT APIを使ってみる 第4回 LangChainと検索APIの連携

前回の記事ではLangChainでChatGPTを使う方法を記事にしました。

今回はLangChainとChatGPT、そして検索エンジンであるTavilyのAPIを連携した使い方を記事にしようと思います。

Tavilyとは

Tavilyとは、AIエージェント向けにカスタマイズされた検索エンジンです。

公式サイトはこちらです。

リアルタイムで正確な検索結果、インテリジェントなクエリ提案、詳細なリサーチ機能を提供するとのことです。

公式サイトにログインし、APIキーを発行することでAPIへのアクセスが可能になります。

TavilyをLangChainで使用する

TavilyをLangChainで使用するためには、以下のコマンドでパッケージをインストールする必要があります。

% pip install -qU "langchain-community>=0.2.11" tavily-python

発行したAPIキーをTAVILY_API_KEY環境変数に指定することで、利用が可能になります。

Tavily Search APIのリファレンスはこちら。

Tavilyで検索を行うためのインスタンスは以下のように作成します。

from langchain_community.tools import TavilySearchResults

tool = TavilySearchResults(
    max_results=5,
    search_depth="advanced",
    include_answer=True,
    include_raw_content=True,
    include_images=False,
    # include_domains=[...],
    # exclude_domains=[...],
    # name="...", # overwrite default tool name
    # description="...", # overwrite default tool description
    # args_schema=..., # overwrite default args_schema: BaseModel
)

今回の例では以下の設定でインスタンスを作成しています。

• 検索結果を5件取得
• 詳細な検索を行う
• 検索結果に簡易回答を含む
• 元の検索結果の内容を含む
• 検索結果に画像を含めない

以前使用したBing Search APIと比べると、以下の機能を容易に使用することができます。

• APIにクエリを渡した際にフォローアップクエスチョンを提示
• HTMLコンテンツの内部を全て取得
• HTMLコンテンツの要約

これらの機能により、ChatGPTの入力クエリとしてより適切な検索結果を取得できると考えています。

そのため今回はBing Search APIではなく、Tavilyを使用してみようと思いました。

チェーンの作成

LangChainの要であるチェーンの作成は、前回と同様にパイプラインを用いて以下のように行います。

# プロンプトテンプレート
prompt = ChatPromptTemplate.from_template(
    "あなたは優秀アシスタントです。提供された検索結果に基づいて、ユーザーの問い合わせに答えてください。\n\n"
    "User query: {user_query}\n\nSearch results:\n{search_results}"
)

# ChatGPTの初期化
llm = ChatOpenAI(model="gpt-4o")

# パイプライン（RunnableSequence）の作成
llm_chain = prompt | llm

# チェーン全体
def search_and_answer(user_query: str):
    # Tavily Searchで検索
    search_results = tavily_tool.invoke({"query": user_query})
    
    # 検索結果を整形
    results_text = "\n".join(
        f"URL: {result['url']}\nContent: {result['content']}" for result in search_results
    )
    
    # LLMチェーンで応答を生成
    response = llm_chain.invoke({"user_query": user_query, "search_results": results_text})
    return response

このようにTavilyの検索結果をChatGPTの入力にすることで、検索結果からユーザーからの問い合わせに対する回答を生成することができます。

最終的な全体コードは以下の通りです。

from dotenv import load_dotenv
load_dotenv()

from langchain_community.tools import TavilySearchResults
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

# Tavily Search Toolの初期化
tavily_tool = TavilySearchResults(
    max_results=5,
    search_depth="advanced",
    include_answer=True,
    include_raw_content=True,
    include_images=False,
)

# プロンプトテンプレート
prompt = ChatPromptTemplate.from_template(
    "You are a helpful assistant. Answer the user's query based on the provided search results.\n\n"
    "User query: {user_query}\n\nSearch results:\n{search_results}"
)

# ChatGPTの初期化
llm = ChatOpenAI(model="gpt-4o")

# パイプライン（RunnableSequence）の作成
llm_chain = prompt | llm

# チェーン全体
def search_and_answer(user_query: str):
    # Tavily Searchで検索
    search_results = tavily_tool.invoke({"query": user_query})
    
    # 検索結果を整形
    results_text = "\n".join(
        f"URL: {result['url']}\nContent: {result['content']}" for result in search_results
    )
    
    # LLMチェーンで応答を生成
    response = llm_chain.invoke({"user_query": user_query, "search_results": results_text})
    return response

# 整形して表示する関数
def format_raw_answer(raw_answer):
    # AIMessageオブジェクトのプロパティに直接アクセス
    content = raw_answer.content if hasattr(raw_answer, "content") else "No content available."
    metadata = raw_answer.response_metadata if hasattr(raw_answer, "response_metadata") else {}
    token_usage = metadata.get("token_usage", {}) if metadata else {}
    model_name = metadata.get("model_name", "Unknown Model") if metadata else "Unknown Model"
    
    formatted_output = (
        f"Answer: {content}\n"
        f"\n---\n"
        f"Model: {model_name}\n"
        f"Token Usage:\n"
        f"  - Input Tokens: {token_usage.get('prompt_tokens', 'N/A')}\n"
        f"  - Output Tokens: {token_usage.get('completion_tokens', 'N/A')}\n"
        f"  - Total Tokens: {token_usage.get('total_tokens', 'N/A')}\n"
    )
    return formatted_output

# テスト実行
if __name__ == "__main__":
    user_query = input("Enter your query: ")
    raw_answer = search_and_answer(user_query)

    # 整形して表示
    print("\nChatGPT's Answer:")
    print(format_raw_answer(raw_answer))

以下は実行結果になります。

Enter your query: 2024年のM-1グランプリの優勝者は誰ですか？

ChatGPT's Answer:
Answer: 2024年のM-1グランプリの優勝者は「令和ロマン」です。彼らは史上初の二連覇を達成しました。

---
Model: gpt-4o-2024-08-06
Token Usage:
  - Input Tokens: 5527
  - Output Tokens: 39
  - Total Tokens: 5566

おわりに

4回に分けて記事を書きましたが、LangChainはまだまだ数多くの機能が存在し、全貌を把握することは難しく感じます。

AIを使用するための学習をAIに任せられるようになったら、もっと習得が容易になるだろうか…と思いつつ、記事を締めくくりたいと思います。