MLflow 2.20.1時点のChatAgentインターフェースについて調べて試してみる

# 2025/2/14追記

MLflow 2.20.2がリリースされました。
ChatAgentのインターフェース変更およびドキュメント記載が行われています。

https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?highlight=chatagent#mlflow.pyfunc.ChatAgent

以下の記事は2.20.2リリース前に書かれたものであることに注意ください。

# 導入

MLflow 2.20.0より、```ChatAgent``` base classが導入されました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/30fda89a-8eed-242f-6c68-b1869f5bae1a.png)

https://mlflow.org/releases/2.20.0

さらに関連する内容として、[MLflow 2.20.1](https://mlflow.org/releases/2.20.1)から上記```ChatAgent```を基にしたLangcChain/LangGraphのヘルパーコネクターが追加されています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/33fac510-5690-f9fd-cc10-0c2f21a5af20.png)


・・・と、リリースノート見ると書いてあるのですが、いったいこれが何かわかっていませんでした。
(Ver 2.20.1時点の公式ドキュメントには説明が記載されていない）

気になって[ソースコード](https://github.com/mlflow/mlflow/blob/master/mlflow/pyfunc/model.py#L397)を追ってみたら、最近のコミットで```ChatAgent```に対するDocStringが追加されていました。これでなんとなく内容がわかったので、テストコード等も含めて```ChatAgent```の動作を試してみます。

:::note warn

```ChatAgent```はまだ実験機能のようです。(```@experimental```タグが付いているため)
したがって将来的に仕様変更・廃止される可能性があります。

また、2月11日時点のリポジトリを見るに、次の公開バージョンからインターフェースの仕様変更が入りそうです。
次セクションの内容は現時点の```master```ブランチの内容から記載していますが、その後の試用は現在の最新公開バージョンであるver.2.20.1時点のパッケージを利用していることに注意してください。
:::

# MLflow ChatAgentインターフェースとは

[こちらのソース](https://github.com/mlflow/mlflow/blob/master/mlflow/pyfunc/model.py)における```ChatAgent```クラスのDocStringを邦訳して引用します。

:::note warn
2月11日時点の```master```ブランチにおける内容を記載しています。

ver.2.20.1のブランチには記載されていないことに注意してください。
また、将来的にこの内容は変更される可能性があります。
:::


> 
> **ChatAgentインターフェースとは？**
> 
> ChatAgentインターフェースは、会話型エージェントを作成するために設計されたチャットスキーマ仕様です。ChatAgentを使用すると、エージェントは以下のことができます：
> 
> - 複数のメッセージを返す
> - ツール呼び出しエージェントの中間ステップを返す
> - ツール呼び出しを確認する
> - 複数エージェントのシナリオをサポートする
> 
> エージェントを作成する際には、常に``ChatAgent``を使用する必要があります。また、将来的により多くのエージェント機能をサポートする柔軟性を持たせるために、シンプルなチャットモデル（例：プロンプトエンジニアリングされたLLM）であっても、```ChatModel <mlflow.pyfunc.ChatModel>```の代わりに```ChatAgent```を使用することをお勧めします。
> 
> ```ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>```スキーマは、OpenAI ChatCompletionスキーマと類似していますが、完全には互換性がありません。ChatAgentは追加の機能を提供し、OpenAI ```ChatCompletionRequest <mlflow.types.llm.ChatCompletionRequest>```とは以下の点で異なります：
> 
> - ツールや内部エージェント呼び出しのために、すべての入力/出力メッセージにオプションの``attachments``属性を追加し、視覚化や進行状況インジケータなどの追加出力を返すことができます
> - ``conversation_id``と``user_id``属性を持つ``context``属性を追加し、エージェントの動作をユーザーに応じて変更できるようにします
> - エージェントの動作を変更するための追加情報を渡すための任意の``dict[str, Any]``である``custom_inputs``属性を追加します
> 
> ```ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>```スキーマは、```ChatCompletionResponse <mlflow.types.llm.ChatCompletionResponse>```スキーマとは以下の点で異なります：
> 
> - 任意の追加情報を返すための``custom_outputs``キーを追加します
> - 出力に複数のメッセージを許可し、内部ツール呼び出しや最終的な回答に至るまでのエージェント間の通信の表示と評価を改善します
> 
> ツール呼び出しの詳細を示す```ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>```の例を以下に示します：
> 
> ```python
>     {
>         "messages": [
>             {
>                 "role": "assistant",
>                 "content": "",
>                 "id": "run-04b46401-c569-4a4a-933e-62e38d8f9647-0",
>                 "tool_calls": [
>                     {
>                         "id": "call_15ca4fcc-ffa1-419a-8748-3bea34b9c043",
>                         "type": "function",
>                         "function": {
>                             "name": "generate_random_ints",
>                             "arguments": '{"min": 1, "max": 100, "size": 5}',
>                         },
>                     }
>                 ],
>             },
>             {
>                 "role": "tool",
>                 "content": '{"content": "Generated array of 2 random ints in [1, 100]."',
>                 "name": "generate_random_ints",
>                 "id": "call_15ca4fcc-ffa1-419a-8748-3bea34b9c043",
>                 "tool_call_id": "call_15ca4fcc-ffa1-419a-8748-3bea34b9c043",
>             },
>             {
>                 "role": "assistant",
>                 "content": "The new set of generated random numbers are: 93, 51, 12, 7, and 25",
>                 "name": "llm",
>                 "id": "run-70c7c738-739f-4ecd-ad18-0ae232df24e8-0",
>             },
>         ],
>         "custom_outputs": {"random_nums": [93, 51, 12, 7, 25]},
>     }
> ```
> 
> **ChatAgentによるストリーミングエージェント出力**
> 
> エージェントの出力をストリーミングする方法の詳細については、```ChatAgent.predict_stream <mlflow.pyfunc.ChatAgent.predict_stream>```のドックストリングを参照してください。
> 
> **ChatAgentの作成**
> 
> ChatAgentインターフェースを使用してエージェントを作成することは、標準化されたインターフェースを持つモデルを作成するためのフレームワークに依存しない方法であり、MLflow pyfuncフレーバーでログ可能で、クライアント間で再利用可能で、サービングワークロードに対応しています。
> 
> 独自のエージェントを作成するには、``ChatAgent``をサブクラス化し、非ストリーミングおよびストリーミングの動作を定義するために``predict``およびオプションで``predict_stream``メソッドを実装します。任意のエージェント作成フレームワークを使用できますが、唯一の厳格な要件は``predict``インターフェースを実装することです。
> 
> ```python
> 
>     def predict(
>         self,
>         messages: list[ChatAgentMessage],
>         context: Optional[ChatContext] = None,
>         custom_inputs: Optional[dict[str, Any]] = None,
>     ) -> ChatAgentResponse: ...
> ```
> タイプヒントに一致する入力を使用してpredictおよびpredict_streamメソッドを呼び出すことに加えて、テストの容易さのために```ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>```スキーマに一致する単一の入力辞書を渡すこともできます。
> 
> ```python
> 
>     chat_agent = MyChatAgent()
>     chat_agent.predict(
>         {
>             "messages": [{"role": "user", "content": "10 + 10は何ですか？"}],
>             "context": {"conversation_id": "123", "user_id": "456"},
>         }
>     )
> ```
> 
> ``predict``および``predict_stream``の例の実装については、`LangChainChatAgent <https://github.com/mlflow/mlflow/blob/master/mlflow/langchain/chat_agent_langchain.py>`または`LangGraphChatAgent <https://github.com/mlflow/mlflow/blob/master/mlflow/langchain/chat_agent_langgraph.py>`クラスを参照してください。
> 
> **ChatAgentのログ記録**
> 
> LLMフレームワークの状況は常に進化しており、すべてのフレーバーがMLflowによってネイティブにサポートされるわけではないため、`Models-from-Code <https://mlflow.org/docs/latest/model/models-from-code.html>`ログ記録アプローチをお勧めします。
> 
> ```python
>     with mlflow.start_run():
>         logged_agent_info = mlflow.pyfunc.log_model(
>             artifact_path="agent",
>             python_model=os.path.join(os.getcwd(), "agent"),
>             # サービングエンドポイント、ツール、ベクトル検索インデックスをここに追加
>             resources=[],
>         )
> ```
> 
> モデルをログ記録した後、```ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>```スキーマを持つ単一の辞書でモデルをクエリできます。内部的には、``predict``および``predict_stream``メソッドが期待するPythonオブジェクトに変換されます。
> 
> ```python
>     loaded_model = mlflow.pyfunc.load_model(tmp_path)
>     loaded_model.predict(
>         {
>             "messages": [{"role": "user", "content": "10 + 10は何ですか？"}],
>             "context": {"conversation_id": "123", "user_id": "456"},
>         }
>     )
> ```
> 
> ChatAgentモデルのログ記録をできるだけ簡単にするために、MLflowには以下の機能が組み込まれています：
> 
> - 自動モデル署名推論
>     - ChatAgentをログ記録する際に署名を設定する必要はありません
>     - 入力および出力署名は、```ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>```および```ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>```スキーマに準拠するように自動的に設定されます
> - メタデータ
>     - モデルをログ記録する際に渡すメタデータに``{“task”: “agent/v2/chat”}``が自動的に追加されます
> - 入力例
>     - 入力例の提供はオプションであり、``mlflow.types.agent.CHAT_AGENT_INPUT_EXAMPLE``がデフォルトで提供されます
>     - 入力例を提供する場合は、```ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>```スキーマに一致する辞書であることを確認してください
> 
>     - ```python
>         input_example = {
>             "messages": [{"role": "user", "content": "MLflowとは何ですか？"}],
>             "context": {"conversation_id": "123", "user_id": "456"},
>         }
>         ```
> 
> **ChatModelからChatAgentへの移行**
> 
> ```List[ChatMessage] <mlflow.types.llm.ChatMessage>```および```ChatParams <mlflow.types.llm.ChatParams>```を受け取り、```ChatCompletionResponse <mlflow.types.llm.ChatCompletionResponse>```を出力する既存のChatModelを変換するには、以下の手順を実行します：
> 
> - ``ChatModel``の代わりに``ChatAgent``をサブクラス化します
> - ``ChatModel``の``load_context``実装からの機能を新しい``ChatAgent``の``__init__``メソッドに移動します
> - モデルの入力を辞書に変換する際に``.to_dict()``の代わりに``.model_dump_compat()``を使用します。例：``[msg.model_dump_compat() for msg in messages]``の代わりに``[msg.to_dict() for msg in messages]``
> - ```ChatCompletionResponse <mlflow.types.llm.ChatCompletionResponse>```の代わりに```ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>```を返します
> 
> 例えば、`Chat Model Intro <https://mlflow.org/docs/latest/llms/chat-model-intro/index.html#building-your-first-chatmodel>`のChatModelをChatAgentに変換することができます：
> 
> ChatModel実装:
> ```python
>     class SimpleOllamaModel(ChatModel):
>         def __init__(self):
>             self.model_name = "llama3.2:1b"
>             self.client = None
> 
>         def load_context(self, context):
>             self.client = ollama.Client()
> 
>         def predict(
>             self, context, messages: list[ChatMessage], params: ChatParams = None
>         ) -> ChatCompletionResponse:
>             ollama_messages = [msg.to_dict() for msg in messages]
>             response = self.client.chat(model=self.model_name, messages=ollama_messages)
>             return ChatCompletionResponse(
>                 choices=[{"index": 0, "message": response["message"]}],
>                 model=self.model_name,
>             )
> ```
> 
> ChatAgent実装:
> ```python
>     class SimpleOllamaModel(ChatAgent):
>         def __init__(self):
>             self.model_name = "llama3.2:1b"
>             self.client = None
>             self.client = ollama.Client()
> 
>         def predict(
>             self,
>             messages: list[ChatAgentMessage],
>             context: Optional[ChatContext] = None,
>             custom_inputs: Optional[dict[str, Any]] = None,
>         ) -> ChatAgentResponse:
>             ollama_messages = self._convert_messages_to_dict(messages)
>             response = self.client.chat(model=self.model_name, messages=ollama_messages)
>             return ChatAgentResponse(**{"messages": [response["message"]]})
> ```
> 
> **ChatAgentコネクタ**
> 
> MLflowは、人気のある作成フレームワークで作成されたエージェントをChatAgentでラップするための便利なAPIを提供しています。以下の例を参照してください：
> 
> - `LangGraphChatAgent <mlflow.langchain.chat_agent_langgraph.LangGraphChatAgent>`
>

というわけで、ChatAgentは```mlflow.pyfunc```モジュールにおけるPythonModelやChatModelに並ぶ新たなカスタムクラス実装のためのベースクラス・インターフェースとなります。
用途としては名前の通り(生成AI等利用における)Chat型エージェントの実装を想定したクラスだと思いますが、シンプルな用途であっても(従来使用が推奨されていた)ChatModelからChatAgentへ置き換えることが推奨されています。

※ ChatModel、思ったより短い命でした。。。技術進展が早い領域なのでしょうがない。。。

変化点はいろいろあるのですが、専用の入出力クラスの利用や入力の変換方法、
また```load_context```メソッドから```__init__```を利用するようにするなどChatModelから多くの変更が入っています。
採用する利点としては、AIエージェントによく用いられるツールの利用サポートや、複数メッセージ（メッセージ履歴）出力があるようです。最近のAIエージェント構築ユースケースに即した拡張ですね。

ChatModelからの移行方法も記載されているので、実際に簡単なChatAgentを基にしたカスタムクラスを定義・実行してみます。

# ChatAgentを使ったカスタムクラスの実装

:::note warn
これからの試用はMLflow ver.2.20.1を利用しています。

ver.2.20.1のChatAgentは```predict/predict_stream```メソッドが前章の解説と異なるインターフェースです。
次の公開バージョンで仕様が変わるであろうことにご注意ください。
:::

簡単なChatAgentを基にしたカスタムクラスを定義し、モデルのロギングから試験実行まで試してみます。
検証環境はDatabricks on AWS、クラスタはサーバレスです。

まず、ノートブックを作成して必要なパッケージをインストール。
```mlflow```のバージョンは```2.20.1```で固定しています。

```python
%pip install langgraph databricks-langchain
%pip install "mlflow-skinny[databricks] == 2.20.1"

%restart_python
```

それではChatAgentを継承した単純なカスタムクラスを作成します。
LLMへクエリを投げてその結果を返す単純なものです。

```python
from mlflow.pyfunc.model import ChatAgent
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatAgentParams
from typing import Optional, Any
from databricks_langchain import ChatDatabricks

class SimpleChatAgentModel(ChatAgent):
    def __init__(self):
        self.model_name = "databricks-meta-llama-3-3-70b-instruct"
        self.llm = ChatDatabricks(endpoint=self.model_name)

    def predict(
        self, messages: list[ChatAgentMessage], params: Optional[ChatAgentParams] = None
    ) -> ChatAgentResponse:
        # list[ChatAgentMessage]のメッセージ入力を辞書型に変換
        llm_messages = self._convert_messages_to_dict(messages)

        # デバッグ用に変換後入力内容を表示
        print(llm_messages)

        # LLMの実行。成功すればLangChainのAIMessageオブジェクトが返ってくる
        response = self.llm.invoke(llm_messages)

        # 結果をChatAgentResponseに変換して返す
        return ChatAgentResponse(**{"messages": [{"role": "assistant", "content": response.content}]})

model = SimpleChatAgentModel()
```

```predict```メソッドを実行して、推論を行います。
ChatAgentにおける```predict```メソッドはPythonModelやChatModelとはインターフェースが異なり、第1引数に```list[ChatAgentMessage]```の内容を指定して実行します。

```python
model.predict([ChatAgentMessage(role="user", content="MLflowとは何ですか？")])
```

また、```list[ChatAgentMessage]```のかわりに、特定要素を含んだ辞書型リストを指定することもできます。

```python
model.predict([{"role":"user", "content":"MLflowとは何ですか？"}])
```

```text:出力結果
[{'role': 'user', 'content': 'MLflowとは何ですか？', 'id': '41600400-eb37-4d00-b1ee-0b0ea8c4e034'}]
ChatAgentResponse(messages=[ChatAgentMessage(role='assistant', content='MLflowは、機械学習モデルの開発、管理、デプロイを支援するためのオープンソースのプラットフォームです。MLflowは、機械学習のライフサイクル全体を管理するためのツールを提供し、データサイエンティスト、エンジニア、DevOpsエンジニアが協力してモデルを開発、テスト、デプロイできるようにします。\n\nMLflowには、以下のような機能があります。\n\n1. **実験管理**：MLflowでは、実験を管理するための機能を提供します。実験とは、モデルをトレーニングするためのパラメータ、データ、アルゴリズムなどの設定を含む、モデル開発のプロセスです。MLflowでは、実験を実行し、結果を追跡し、比較することができます。\n2. **モデル管理**：MLflowでは、トレーニング済みのモデルを管理するための機能を提供します。モデルを保存し、バージョン管理し、デプロイすることができます。\n3. **モデルサービング**：MLflowでは、トレーニング済みのモデルをデプロイし、APIとして提供するための機能を提供します。モデルをREST APIとして公開し、クライアントからのリクエストに応答することができます。\n4. **モニタリング**：MLflowでは、デプロイされたモデルのパフォーマンスをモニタリングするための機能を提供します。モデルが予想どおりに動作しているかどうかを確認し、問題が発生した場合に通知を受けることができます。\n\nMLflowは、Python、R、Juliaなどのプログラミング言語で使用できます。また、Apache Spark、TensorFlow、PyTorchなどの機械学習フレームワークと統合することもできます。', name=None, id='a59ef3e2-eff2-493a-a782-07a9c10a6307', tool_calls=None, tool_call_id=None, attachments=None, finish_reason=None)], custom_outputs=None, usage=None)
```

では、作ったモデルを保存(ロギング)します。
従来のモデル保管と同じ方法になりますが、```input_schema```などのスキーマ指定は不要です。
また```input_example```の指定を行わなかった場合、デフォルトのexampleが自動設定されます。

```python
import mlflow

saved_model_path = "./model"

# モデルの保存
mlflow.pyfunc.save_model(
    python_model=model,
    path=saved_model_path,
)
```

実行後のモデル保管フォルダは以下のようなファイルが作られました。
```input_example.json```などもちゃんと作られています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/3e470f44-719d-b871-602e-2abe45c364c3.png)

```MLmodel```の中身は下記画像のようになっていました。
ポイントはmetadataに```task: agent/v2/chat```が設定されたり、signatureが自動設定されるところでしょうか。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/45c7abbb-fc11-47b1-9b7c-40564152f5e4.png)

では、保管したモデルをpyfuncモジュールでロードします。これは従来通りですね。

```python
# 保存モデルのロード
loaded_model = mlflow.pyfunc.load_model(saved_model_path)
```

では、再度```predict```メソッドを呼び出して実行します。

注意点は、Signatureに合わせたパラメータ指定が必要なことです。
つまり、保管前のモデルのように```list[ChatAgentMessage]```を直接渡すのではなく、```messages```キーを含んだ辞書型データを第一パラメータに設定する必要があります。

```python
# 保存モデルのロード
loaded_model.predict({"messages":[{"role":"user", "content":"MLflowとは何ですか？"}]})
```

```txt:出力結果
[{'role': 'user', 'content': 'MLflowとは何ですか？', 'id': '4acd273c-3636-4ee9-9c8e-4102b576f61e'}]
{'messages': [{'role': 'assistant',
   'content': 'MLflowは、機械学習モデルの開発を管理するためのオープンソースのプラットフォームです。機械学習のライフサイクル全体を管理することを目的として設計されており、モデル開発、学習、テスト、デプロイ、管理を一元的に行えるように支援します。\n\nMLflowには、以下のような主な機能があります。\n\n1. **実験管理 (Experiment Management)**: モデルの学習実験を管理し、パラメータ、メトリクス、成果物を追跡する機能。\n2. **モデル管理 (Model Management)**: 学習したモデルをバージョン管理し、モデルをデプロイし、モデルサービングを提供する機能。\n3. **モデルサービング (Model Serving)**: 学習したモデルをREST APIなどを通じて提供し、予測や推論を実行する機能。\n\nMLflowは、Pythonを中心に開発されており、多くの機械学習フレームワーク（TensorFlow、PyTorch、Scikit-learnなど）との統合が容易です。また、クラウドサービス（AWS、GCP、Azureなど）やオンプレミス環境でも利用できます。\n\nMLflowの主な利点は、機械学習モデルの開発を効率化し、再現性を高め、コラボレーションを促進することです。開発者は、MLflowを使用してモデル開発のプロセスを自動化し、モデルを迅速にデプロイし、モデルのパフォーマンスを継続的に監視することができます。',
   'id': 'f640c382-c1e4-4c33-87f4-dba74b62bee0'}]}
```

というわけで、保管＆読込・実行まで一連試すことが出来ました。

# まとめ

MLflowのChatAgentについて、現時点の最新公開バージョンであるv2.20.1を使って試してみました。
とはいえ、将来バージョンにてインターフェースの仕様変更が起こると思うので、ご利用は計画的に、かな。

まだわらかないところ、例えば従来のartifactsの扱いなど、もう少し調べたり試してみたりしてみたいと思います。ChatModelカスタムクラスからの置き換えはおそらく必要になってくると思うので。
また、今回は取り上げていませんがLangChain/LangGraphとのヘルパーコネクターがこのクラスをベースに公開されています。個人的にLangGraphと親和性高そうに思うので組み合わせていろいろやってみたいと思います。