LangChainのCallbacks

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

イベントの種類に関しては[こちら](https://zenn.dev/umi_mori/books/prompt-engineer/viewer/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
```python
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
```python
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
```python
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にて検証しておりますことをお断りさせて頂きます。