LangChainには色んなイベント発生時にログを取ったりすることができるCallbacksという機能がありますが、その使い方で四苦八苦したので記録しようと思います。

イベントの種類に関してはこちらを参考にしました。

前提

使用OS : Windows10
使用言語 : Python 3.11.5
AzureOpenAI : gpt4-turbo 2024-02-15-preview
ライブラリバージョン:
langchain==0.1.5
openai==1.11.0

やりたかったこと

agentとfunction callingを組み合わせたときの各イベント(agentの動作、chainの動作、toolの動作、Modelの動作等)のログを取りたかった。

勘違いしてたこと

agentにコールバックを設定してれば、agentの中でchainやらModelやToolやらが動いてるんだからログを取れるだろうと思ってました。

しかし、実際にagent実行時に動いたイベントは以下でした。

イベント名	説明
on_chain_start	チェーンの動作開始時
on_chain_end	チェーンの動作終了時
on_chain_error	チェーンでエラーが発生した時
on_agent_action	エージェントアクション時
on_agent_finish	エージェントの終了時

もしかして、ModelとかToolにも設定しないと全部とれないのかなと思い試してみたところ、とれました！

実装まとめ

シンプルにまとめるために、一部コード省略して書いてます。

・Model

from typing import Dict,List,Union,Any
from langchain_core.messages.base import BaseMessage
from langchain.schema.output import LLMResult
from langchain_openai import AzureChatOpenAI
from langchain.callbacks.base import BaseCallbackHandler

class ModelCallbackHandler(BaseCallbackHandler):
    """Base callback handler that can be used to handle callbacks from langchain."""

    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> Any:
        """Run when LLM starts running."""

    def on_chat_model_start(
        self, serialized: Dict[str, Any], messages: List[List[BaseMessage]], **kwargs: Any
    ) -> Any:
        """Run when Chat Model starts running."""

    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
        """Run on new LLM token. Only available when streaming is enabled."""

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
        """Run when LLM ends running."""

    def on_llm_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> Any:
        """Run when LLM errors."""


llm = AzureChatOpenAI(
    openai_api_version="AOAIのバージョン",
    azure_deployment="AOAIのデプロイ名",
    callbacks=[ModelCallbackHandler()]
)

・Tool

from typing import Dict,Union,Any
from pydantic import BaseModel , Field
from langchain.agents.tools import BaseTool
from langchain.callbacks.base import BaseCallbackHandler
from langchain.agents import load_tools

class ToolCallbackHandler(BaseCallbackHandler):
    """Base callback handler that can be used to handle callbacks from langchain."""

    def on_tool_start(
        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
    ) -> Any:
        """Run when tool starts running."""

    def on_tool_end(self, output: str, **kwargs: Any) -> Any:
        """Run when tool ends running."""

    def on_tool_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> Any:
        """Run when tool errors."""

    
# langchainに備わっているツールを使用する場合
tools = load_tools(
    tool_names=["ツールの名前"],callbacks=[ToolCallbackHandler()]
    )

# 作成したカスタムツール使用する場合
class CustomToolInput(BaseModel):
    query: str = Field(..., description="引数の説明書き")

class CustomTool(BaseTool):
    name = "カスタムツールの名前"
    description = "カスタムツールの説明書き"
    args_schema = CustomToolInput

    def _run(self, query: str) -> str:
        # 同期時の処理を書く
        return query
    
    async def _arun(self, query: str) -> str:
        # 非同期時の処理を書く
        return self._run(query)

custom_tool = CustomTool()
custom_tool.callbacks = [ToolCallbackHandler()]
tools = [custom_tool]

・Agent

from typing import Dict,Union,Any
from langchain_core.agents import AgentAction, AgentFinish
from langchain.callbacks.base import BaseCallbackHandler
from langchain_openai import AzureChatOpenAI
from langchain.schema import SystemMessage
from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
from langchain.agents.agent import AgentExecutor

class AgentCallbackHandler(BaseCallbackHandler):
    """Base callback handler that can be used to handle callbacks from langchain."""

    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> Any:
        """Run when chain starts running."""

    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
        """Run when chain ends running."""

    def on_chain_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> Any:
        """Run when chain errors."""
    
    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
        """Run on agent action."""

    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
        """Run on agent end."""

    
llm = AzureChatOpenAI(...)
tools = [...]
prompt = [SystemMessage("あなたは優秀なAIアシスタントです。")]
agent = OpenAIFunctionsAgent(
    llm=llm,
    tools = tools,
    prompt = prompt
)
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    callbacks=[AgentCallbackHandler()]
)


agent_executor.invoke(input={"input":"こんにちわ"})

わからなかったこと

以下の2つだけどういう実装で発火してくれるのかわかりませんでした。

イベント名	説明
on_llm_start	LLMの動作開始時
on_text	任意のテキストが出力された時

on_llm_startは普通にLLMで発火してくれると思ったら、on_chat_model_startしか発火しなかった...なぜ...？

on_textはドキュメントみてもよくわかりませんでした。
任意のテキストって...何...。

ご存じの方がいらっしゃれば教えて頂けると幸いです。

追記

LnagChainでは最新バージョンを使用すると、OpenAIFunctionsAgent使用時に0.2以降のバージョンで使用できなくなるという警告が出ます。
推奨ではcreate_openai_functions_agentを使用することとされていますが、当記事執筆時はcreate_openai_functions_agentを使用するとなぜか強制的にストリーミングとなってしまいトークン数が取得できないというバグが発生しておりました。
そのため、OpenAIFunctionsAgentにて検証しておりますことをお断りさせて頂きます。