【langchain】Agentの実装をざっくり理解する

# はじめに
langchainのAgentは言語モデルに使用する関数（tool)を決定させるためのクラスです。Agentはtoolを決定するだけで実行はしません。タスクを完了するためにはtoolを実行し、その実行結果を言語モデルに渡す必要があり、その処理はAgentではなくAgentExecutorというクラスで記述されます。この辺りの関係性がややこしいので、整理してみました。

# Agent
langchainのAgentは言語モデルに使用する関数（tool)を決定させるためのクラスです。この処理は主にAgent.planというメソッドに記述されています。
よって、Agentの実装がどうなっているか知りたいときは、Agent.planを見るとよいです。
## コード
OpenAIFunctionsAgentを例に説明します。ソースコードを以下に示します。
参考: https://github.com/langchain-ai/langchain/blob/ddd07001f354cd09a76a61e1f5c678bf885506d2/libs/langchain/langchain/agents/openai_functions_agent/base.py#L142

```langchain/libs/langchain/langchain/agents/openai_functions_agent/base.py
    def plan(
        self,
        intermediate_steps: List[Tuple[AgentAction, str]],
        callbacks: Callbacks = None,
        with_functions: bool = True,
        **kwargs: Any,
    ) -> Union[AgentAction, AgentFinish]:
        """Given input, decided what to do.

        Args:
            intermediate_steps: Steps the LLM has taken to date, along with observations
            **kwargs: User inputs.

        Returns:
            Action specifying what tool to use.
        """
        agent_scratchpad = _format_intermediate_steps(intermediate_steps)
        selected_inputs = {
            k: kwargs[k] for k in self.prompt.input_variables if k != "agent_scratchpad"
        }
        full_inputs = dict(**selected_inputs, agent_scratchpad=agent_scratchpad)
        prompt = self.prompt.format_prompt(**full_inputs)
        messages = prompt.to_messages()
        if with_functions:
            predicted_message = self.llm.predict_messages(
                messages,
                functions=self.functions,
                callbacks=callbacks,
            )
        else:
            predicted_message = self.llm.predict_messages(
                messages,
                callbacks=callbacks,
            )
        agent_decision = _parse_ai_message(predicted_message)
        return agent_decision
```

入力に対して、toolを決定する処理が書かれています。
入力は、変数名がわかりにくいですが、\**kwargsです。デフォルトでは

```Python
Agent.plan(input="input")
```

のような感じで、inputというキーワード引数で指定します。
また、intermediate_stepsという位置引数もあります。これは過去の入力履歴のようなものです。
intermediate_stepsとs\**kwargsからfull_inputsを作り、それをllmに渡しています。OpenAIFunctionsAgentなのでfunction callingを使っていることがわかります。
returnの直前に_parse_ai_message関数を使って、llmの出力をagent_decisionに変換しています。agent_decisionはllmの出力がfunction_callの場合はAgentAction、そうでない場合はAgentFinishとなります。
planの内部ではtoolの実行はされていないことも見てわかります。


## 実験
OpenAIFunctionsAgent.planを実行してみます。

```Python
from langchain.agents import load_tools
from langchain.agents import OpenAIFunctionsAgent
from langchain.chat_models import ChatOpenAI

# LLM ラッパーを初期化
llm = ChatOpenAI(temperature=0.7)
# ツールの一覧を作成
tools = load_tools(["llm-math"], llm=llm)

agent = OpenAIFunctionsAgent.from_llm_and_tools(llm, tools)
print(agent.plan([], input="128と345の積は?"))
```

```
_FunctionsAgentAction(tool='Calculator', tool_input='128 * 345', log='\nInvoking: `Calculator` with `128 * 345`\n\n\n', message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'Calculator', 'arguments': '{\n  "__arg1": "128 * 345"\n}'}}, example=False)])
```
toolとtool_inputが適切に決定されていることがわかります。
intermediate_steps(planの第一引数)はとりあえず空にしています。
message_logにfunction callingの出力も含まれています。
planの段階では、toolの実行は行われていません。

toolを使う必要がない入力の場合、どうなるか試してみます。

```Python
from langchain.schema import AgentAction
agent = OpenAIFunctionsAgent.from_llm_and_tools(llm, tools)
print(agent.plan([], input="こんにちは!"))
```

```
AgentFinish(return_values={'output': 'こんにちは! お困りのことはありますか？私はAIアシスタントですので、お手伝いできることがあれば教えてください。'}, log='こんにちは! お困りのことはありますか？私はAIアシスタントですので、お手伝いできることがあれば教えてください。')
```

AgentActionではなくAgentFinishが返ってきました。

intermediate_stepsが空でない場合を試してみます。
planの戻り値と手動で作った実行結果（observation)を組み合わせて、first_stepを作ってみます。

```Python
from langchain.schema import AgentAction
first_step = (agent.plan([], input="128と345の積は？ "), "Answer: 44160")
agent = OpenAIFunctionsAgent.from_llm_and_tools(llm, tools)
print(agent.plan([first_step], input="128と345の積は？ "))
```

```
AgentFinish(return_values={'output': '128と345の積は44160です。'}, log='128と345の積は44160です。')
```

toolの実行結果が存在する状態のため、タスクの完了を意味するAgentFinishが帰ってきました。

# AgentExecutor
AgentExecutorはtoolを実行してタスクを完了するためのクラスです。
実行すべきtoolの決定のためにAgent.planが使われます。
ソースを見るとわかるのですが、AgentExecutorはChainを継承したクラスです。したがって、Chainだと思って挙動を確認すると理解しやすいです。

少し余談ですが、langchainのチュートリアルによく出てくるinitialize_agent関数の戻り値はAgentではなくAgentExecutorです。またinitialize_agentにはagentという引数を与えることができるのですが、このagentはAgentクラスではなくAgentTypeという文字列です。initialize_agentはAgentTypeを受け取って、AgentTypeに応じたAgentを内部で生成し、さらにそのAgentを与えて初期化したAgentExecutorクラスを返す関数なのです。ややこしいですね。
## コード
AgentExecutorはChainであり、Chainの処理の実態はChain._callに書かれています。よって、AgentExecutor._callを見ると、AgentExecutorがどういう処理を実行するクラスなのかわかります。ソースコードを見てみます。

```Python:langchain/libs/langchain/langchain/agents/agent.py
    def _call(
        self,
        inputs: Dict[str, str],
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        """Run text through and get agent response."""
        # Construct a mapping of tool name to tool for easy lookup
        name_to_tool_map = {tool.name: tool for tool in self.tools}
        # We construct a mapping from each tool to a color, used for logging.
        color_mapping = get_color_mapping(
            [tool.name for tool in self.tools], excluded_colors=["green", "red"]
        )
        intermediate_steps: List[Tuple[AgentAction, str]] = []
        # Let's start tracking the number of iterations and time elapsed
        iterations = 0
        time_elapsed = 0.0
        start_time = time.time()
        # We now enter the agent loop (until it returns something).
        while self._should_continue(iterations, time_elapsed):
            next_step_output = self._take_next_step(
                name_to_tool_map,
                color_mapping,
                inputs,
                intermediate_steps,
                run_manager=run_manager,
            )
            if isinstance(next_step_output, AgentFinish):
                return self._return(
                    next_step_output, intermediate_steps, run_manager=run_manager
                )

            intermediate_steps.extend(next_step_output)
            if len(next_step_output) == 1:
                next_step_action = next_step_output[0]
                # See if tool should return directly
                tool_return = self._get_tool_return(next_step_action)
                if tool_return is not None:
                    return self._return(
                        tool_return, intermediate_steps, run_manager=run_manager
                    )
            iterations += 1
            time_elapsed = time.time() - start_time
        output = self.agent.return_stopped_response(
            self.early_stopping_method, intermediate_steps, **inputs
        )
        return self._return(output, intermediate_steps, run_manager=run_manager)
```

ざっくりとは、AgentFinishが返ってくるまで、self._take_next_stepを繰り返し実行する、という処理になっています（他の終了条件もありますが、一旦スルーします）

_take_next_stepを見てみます。

```Python:langchain/libs/langchain/langchain/agents/agent.py
    def _take_next_step(
        self,
        name_to_tool_map: Dict[str, BaseTool],
        color_mapping: Dict[str, str],
        inputs: Dict[str, str],
        intermediate_steps: List[Tuple[AgentAction, str]],
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Union[AgentFinish, List[Tuple[AgentAction, str]]]:
        """Take a single step in the thought-action-observation loop.

        Override this to take control of how the agent makes and acts on choices.
        """
        try:
            intermediate_steps = self._prepare_intermediate_steps(intermediate_steps)

            # Call the LLM to see what to do.
            output = self.agent.plan(
                intermediate_steps,
                callbacks=run_manager.get_child() if run_manager else None,
                **inputs,
            )
        except OutputParserException as e:
　　　　　　　　　　　　　　　　　　　　　　　　　"""例外処理の詳細は省略"""
            return [(output, observation)]
        # If the tool chosen is the finishing tool, then we end and return.
        if isinstance(output, AgentFinish):
            return output
        actions: List[AgentAction]
        if isinstance(output, AgentAction):
            actions = [output]
        else:
            actions = output
        result = []
        for agent_action in actions:
            if run_manager:
                run_manager.on_agent_action(agent_action, color="green")
            # Otherwise we lookup the tool
            if agent_action.tool in name_to_tool_map:
                tool = name_to_tool_map[agent_action.tool]
                return_direct = tool.return_direct
                color = color_mapping[agent_action.tool]
                tool_run_kwargs = self.agent.tool_run_logging_kwargs()
                if return_direct:
                    tool_run_kwargs["llm_prefix"] = ""
                # We then call the tool on the tool input to get an observation
                observation = tool.run(
                    agent_action.tool_input,
                    verbose=self.verbose,
                    color=color,
                    callbacks=run_manager.get_child() if run_manager else None,
                    **tool_run_kwargs,
                )
            else:
                tool_run_kwargs = self.agent.tool_run_logging_kwargs()
                observation = InvalidTool().run(
                    {
                        "requested_tool_name": agent_action.tool,
                        "available_tool_names": list(name_to_tool_map.keys()),
                    },
                    verbose=self.verbose,
                    color=None,
                    callbacks=run_manager.get_child() if run_manager else None,
                    **tool_run_kwargs,
                )
            result.append((agent_action, observation))
        return result
```

self.agent.planでtool(action)を決定し、tool.runでその実行結果(observation)を得て、actionとobservationをペアで返す、という処理になっています。

よって、AgentExecutor.run (_call) は、内部でAgent.planを使いながらタスクを完了するためのクラス、と理解できます。

## 実験
run, \_\_call__, _call, _take_next_stepと実行する関数を徐々に深くして、出力がどのように変化するか追ってみます。

AgentExecutor.runを実行してみます。

```Python
from langchain.agents import load_tools
from langchain.agents import OpenAIFunctionsAgent
from langchain.chat_models import ChatOpenAI
from langchain.agents import AgentExecutor

# LLM ラッパーを初期化
llm = ChatOpenAI(temperature=0.7)
# ツールの一覧を作成
tools = load_tools(["llm-math"], llm=llm)

agent = OpenAIFunctionsAgent.from_llm_and_tools(llm, tools)

agent_executor = AgentExecutor.from_agent_and_tools(agent, tools, verbose=True)
print(agent_executor.run("123と345をかけてください"))
```

```
> Entering new AgentExecutor chain...

Invoking: `Calculator` with `123 * 345`


Answer: 42435123と345をかけると、答えは42435です。

> Finished chain.
123と345をかけると、答えは42435です。
```

outputとして、toolの実行結果そのものではなく、toolの実行結果を踏まえてllmで生成したと思われる文字列が得られました。

__call__を実行してみます。

```Python
print(agent_executor({"input":"123と345をかけてください"}))
```

```
> Entering new AgentExecutor chain...

Invoking: `Calculator` with `123 * 345`


Answer: 42435123と345をかけると、42435になります。

> Finished chain.
{'input': '123と345をかけてください', 'output': '123と345をかけると、42435になります。'}
```

辞書が得られました。outputはまだ、llmで生成した結果です。

_callを実行してみます。

```Python
print(agent_executor._call({"input":"123と345をかけてください"}))
```

```
Answer: 42435{'output': '123と345をかけると、42435になります。'}
```
まだllmで生成した結果が得られています。

_take_next_stepを実行してみます。
name_to_tool_map, color_mappingが必須引数なので、_callの実装を真似して作ります。
intermediate_stepsも必須引数なので、空にしておきます。

```Python
from langchain.utils.input import get_color_mapping
name_to_tool_map = {tool.name: tool for tool in tools}
color_mapping = get_color_mapping(
            [tool.name for tool in tools], excluded_colors=["green", "red"]
        )
inputs = {"input":"123と345をかけてください"}
intermediate_steps=[]
print(agent_executor._take_next_step(name_to_tool_map, color_mapping, inputs, intermediate_steps))
```

```
Answer: 42435[(_FunctionsAgentAction(tool='Calculator', tool_input='123 * 345', log='\nInvoking: `Calculator` with `123 * 345`\n\n\n', message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'Calculator', 'arguments': '{\n  "__arg1": "123 * 345"\n}'}}, example=False)]), 'Answer: 42435')]
```

ようやくtoolの実行結果の生データ（"Answer: 42435")が得られました。
戻り値がAgentActionなので、ループが終了せず、_call内で２回目の_take_next_stepが実行されるときに、"Answer: 42435"を踏まえた文章が生成される、という流れだと想像できます。

# おわりに
AgentとAgentExecutorの関係を改めて整理しておきます。

Agentは使うtoolと引数を決定するためのクラスであり、Agent.planというメソッドが主に使われます。Agent.planは、toolを使う必要がある場合はtoolとその引数（AgentAction）、使う必要がない場合はただの文章(AgentFinish)を返します。

AgentExecutorはAgent.planを使いながらタスクを完了するためのクラスです。AgentExecutorはChainを継承しているため、Chainだと思うと、理解が進みやすいです。内部では、Agent.planで決定されたtoolの実行や、戻り値がAgentFinishになるまでAgent.planを繰り返し呼びます。

このため、独自のagent機能を作りたい場合、tool決定のロジックを調整したい場合はAgentのplanを、tool決定後の処理を調整した場合はAgentExecutorの_callや_take_next_stepを書き換えればよいとざっくり理解できます。

一方で、今回でまだわかっていないこととして、Agentに記憶を持たせる方法や、Agentのバリエーションごとの違いなどがあります。これらの点についても余裕があれば深掘りしてみたいです。