はじめに
Langchainを用いてエージェントを実行する際、トークン数を記録するためにOpenAICallbackHandler及びget_openai_callbackを使用することが一般的です。しかし、特定のバージョン（例：langchain==0.1.11以前）では、トークン数が正しく記録されない問題が発生していました。本記事では、その問題の原因と解決方法について解説します。
問題の背景
特定の状況下で、AgentExecutorを使用してエージェントを実行すると、トークン数が記録されない問題が発生していました。これは特にRunnableAgent.planでrunnable.streamが実行される場合に顕著でした。
参考：問題の詳細
-
AzureChatOpenAIはlangchain_core.language_models.chat_model.BaseChatModelを継承しています。 -
BaseChatModelはRunnableを継承しており、streamメソッド内でrunmanager.on_llm_endが呼び出されますが、このときllm_outputが引数として渡されないため、OpenAICallbackHandlerが正しく機能しません。 -
OpenAICallbackHandlerはcallback処理において、on_llm_endが呼び出された際に、引数にLLMResult.llm_outputを読み取ることでトークン数の利用状況を記録します。llm_outputはstream実行時に情報が取得できないため、設定されません。 - 一方、
BaseChatModelのgenerateメソッドではllm_outputが設定されています。通常のAzureChatOpenAIでinvokeからLLMの処理を実行した際はトークン消費状況が記録されるということです。
このため、ストリーミングの場合にはトークン利用量が記録されないという問題が発生していました。
解決方法
Langchainのバージョン0.1.12以降では、この問題に対する解決策が提供されています。具体的には、AgentExecutorが継承するRunnableAgentにstream_runnableというメンバ変数が追加され、plan実行時にストリーミングを使用するかバッチ処理を使用するかを選択できるようになりました。
実装例
以下のコード例では、stream_runnableをFalseに設定することで、ストリーミングを使用せずにバッチ処理を行い、トークン数の情報を正しく取得できるようにしています。
from langchain import AgentExecutor
from langchain.callbacks.openai_info import OpenAICallbackHandler
from langchain_community.callbacks import get_openai_callback
from langchain_core.runnables.config import RunnableConfig
# stream_runnableをFalseに設定
# 細かい設定はこちら：https://qiita.com/kash203/items/e264ea57d43fa9a43a3d
agent_executor = AgentExecutor(
agent=agent,
tools=tools,
verbose=True,
stream_runnable=False,
)
# エージェントを実行（パターン１：ハンドラをagent_executorに設定する方法）
openai_callback_handler = OpenAICallbackHandler()
config = RunnableConfig(callbacks=[openai_callback_handler])
pattern1 = agent_executor.invoke({"input": "what is LangChain?."}, config=config)
print(openai_callback_handler) # トークン数を表示
# エージェントを実行（パターン２：get_openai_callbackを用いる方法）
with get_openai_callback() as cb:
pattern2 = agent_executor.invoke({"input": "what is LangChain?."})
print(cb) # トークン数を表示
# 表示例
"""
Tokens Used: 1000
Prompt Tokens: 800
Completion Tokens: 200
Successful Requests: 2
Total Cost (USD): $0.0075
"""
このようにすることで、エージェント実行時にトークン数が正しく記録されるようになります。
ただし、トークンの料金については最新モデルのものがLangchainに設定されていないケースがありますので、必要に応じて自身で計算する必要がある可能性がある点に注意です。
詳細は以下の一覧にモデルの情報があるかどうかをチェックしてください。
https://github.com/langchain-ai/langchain/blob/873d06c009409560d52cde3f1a2b39f268319112/libs/community/langchain_community/callbacks/openai_info.py#L8-L100
おわりに
Langchainを用いたエージェント実行時に発生していたトークン数記録の問題は、バージョン0.1.12以降で提供されたstream_runnableオプションを利用することで解決できます。エージェントのストリーミングが必要な場合は要注意ですが、基本的にはこの設定を使用することでトークン数の正確な記録が可能となります。
この情報が、Langchainを利用する開発者の皆様の問題解決の一助となれば幸いです。今後も最新のアップデート情報をチェックし、効率的な開発を行ってください。
謝辞
本記事の作成にあたり、ChatGPTの協力を得ました。高度な自然言語処理能力を持つChatGPTのおかげで、複雑な問題の解決策を明確に示すことができました。心より感謝申し上げます。
謝辞の通り、自身のメモ書きから、おおよその内容をChatGPTに記事化してもらいました。謝辞自体もChatGPTが考えています。