wandb による LangChain のトレースを試す

2023年7月27日 11:01

「wandb」による「LangChain」のトレースを試したので、まとめました。

1. wandb (Weights & Biases)

「wandb」 (Weights & Biases)は、機械学習の実験の追跡および可視化プラットフォームです。「LangChain」では環境変数を1つ指定するだけで、「LangChain」の内部状態をトレースすることができます。

「wandb」のサイトでアカウントを作成してください。個人利用は無料になります。詳しくは、「Pricing」を参照してください。

2. トレースするコードの準備

Google Colab でトレースするコードを準備します。

(1) パッケージのインストール。

# パッケージのインストール
!pip install langchain
!pip install openai

(2) 環境変数の準備。
以下のコードの <OpenAI_APIのAPIキー> にはOpenAI APIのAPIキーを指定します。(有料)

import openai

# OpenAI APIキーの準備
openai.api_key = "<OpenAI_APIのAPIキー>"

(3) LangChainのコードの準備。
今回は、「llm-math」ツールで計算するエージェントを使います。

from langchain.agents import initialize_agent, load_tools
from langchain.agents import AgentType
from langchain.llms import OpenAI

# エージェントの準備
llm = OpenAI(temperature=0)
tools = load_tools(["llm-math"], llm=llm)
agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

# エージェントの実行
agent.run("123*4を計算機で計算してください")

> Entering new AgentExecutor chain...
 I need to use the calculator to solve this
Action: Calculator
Action Input: 123*4
Observation: Answer: 492
Thought: I now know the final answer
Final Answer: 492

> Finished chain.
492

2. トレースの実行

トレースの実行手順は、次のとおりです。

(1) wandbパッケージのインストール。

# パッケージのインストール
!pip install wandb

(2) 環境変数の準備。

import os
os.environ["LANGCHAIN_WANDB_TRACING"] = "true"

(3) エージェントの実行。

# エージェントの実行
agent.run("123*4を計算機で計算してください")

「Google Colab」のログで、wandbのAPIキーの入力を促されたら、「https://wandb.ai/authorize」をクリックしてAPIキーをコピーし、テキストボックスに入力します。

(4) wandbのダッシュボードでトレース内容を確認。
「Google Colab」のログのリンク (https://wandb.ai/<user>/uncategorized/runs/<id>) からも開くことができます。

「LangChain」内部すべての「Prompt→Completion」を確認できるため、レスポンス失敗時にどの処理が原因なのかがわかります。

【おまけ】特定ブロックのみをトレース

アプリケーション全体ではなく、特定ブロックのみをトレースすることもできます。環境変数 (LANGCHAIN_WANDB_TRACING) を使わず、「wandb_tracing_enabled()」を使います。

from langchain.callbacks import wandb_tracing_enabled

# ブロック内のトレースの有効化
with wandb_tracing_enabled():
    agent.run("123*4を計算機で計算してください")

参考

この記事が気に入ったらサポートをしてみませんか？