2025/2/14追記

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

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

導入

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

さらに関連する内容として、MLflow 2.20.1から上記ChatAgentを基にしたLangcChain/LangGraphのヘルパーコネクターが追加されています。

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

気になってソースコードを追ってみたら、最近のコミットでChatAgentに対するDocStringが追加されていました。これでなんとなく内容がわかったので、テストコード等も含めてChatAgentの動作を試してみます。

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

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

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

こちらのソースにおけるChatAgentクラスのDocStringを邦訳して引用します。

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>の例を以下に示します：
    {
        "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インターフェースを実装することです。
    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>スキーマに一致する単一の入力辞書を渡すこともできます。
    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>ログ記録アプローチをお勧めします。
    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オブジェクトに変換されます。
    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>スキーマに一致する辞書であることを確認してください
  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実装:
    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実装:
    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を使ったカスタムクラスの実装

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

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

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

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

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

%restart_python

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

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]の内容を指定して実行します。

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

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

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

出力結果

[{'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が自動設定されます。

import mlflow

saved_model_path = "./model"

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

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

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

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

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

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

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

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

出力結果

[{'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と親和性高そうに思うので組み合わせていろいろやってみたいと思います。