MLflow 2.16.0の新機能 : LangGraphサポートを試す

# 導入

気づけば8月末に[MLFlow](https://mlflow.org/) 2.16.0がリリースされていました。

https://mlflow.org/releases/2.16.0

このバージョンアップでの主な変更点は以下の通り。

> **LlamaIndexの強化**🦙 - LlamaIndex統合にさらなる柔軟性を提供するために、コードからのモデル機能のログ記録、拡張されたエンジンベースのログ記録、および外部ベクトルストアのサポートが追加されました。
>
> **LangGraphサポート** - LangChain統合を拡張し、エージェントフレームワークLangGraphをサポートしました。トレースとコードからのモデル機能を使用したログ記録のサポートにより、エージェントアプリケーションの作成と保存がこれまでになく簡単になりました！
>
> **AutoGenトレース** - MicrosoftのAutoGenフレームワークで構築されたマルチターンエージェントアプリケーションの完全自動トレースサポートがMLflowで利用可能になりました。mlflow.autogen.autolog()を有効にすることで、AutoGenで構築されたエージェントが自動的に計測されます。
>
> **AIゲートウェイのプラグインサポート** - MLflowのAIゲートウェイ（MLflowデプロイメントサーバーとも呼ばれる）で動作する独自のプロバイダーインターフェースを定義できるようになりました。インストール可能なプロバイダー定義を作成することで、任意のGenAIサービスにゲートウェイサーバーを接続できます。

上記のように、今回のアップデートで個人的に来てほしかったLangGraphサポートがきました。
LangChainのフレーバーを拡張する形で実装されているようです。

https://mlflow.org/docs/latest/llms/langchain/index.html#how-can-i-log-an-agent-built-with-langgraph-to-mlflow


また、気づいてなかったのですがModels From Codeという新？機能で実現されている模様。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/413ac634-1d5a-9f69-f236-8622d4cd28ad.png)

Models From Codeのガイドはこちら。

https://mlflow.org/docs/latest/model/models-from-code.html

というわけで、ごくごく単純なサンプルを作って挙動を見てみたいと思います。

検証はDatabricks on AWS上で行い、ノートブッククラスタはサーバレスを利用しました。
※　サーバレスは本当に便利だ。。。

# Step1. パッケージのインストール

ノートブックを作成し、必要なパッケージをインストール。
MLFlowは2.16.0以降をインストールします。

```python
%pip install langchain_core langgraph
%pip install -U "mlflow-skinny[databricks]>=2.16.0"
```

# Step2. 単純なLangGraphグラフの作成

LangGraphで単純なグラフを作成します。
(簡易化のために、固定文字列```This is sample message.```をメッセージとして返すだけのグラフにしました)

この処理では、以下のようなポイントがあります。
1. ```%%writefile```マジックコマンドで、セル内をファイルとして出力しています。これはMLflowでのロギングの際に、このコードをモデル定義として保持させるためです。
1. 処理の最後に```set_model```関数を呼び出しています。この関数に与えるインスタンスがMLFlowでモデルをロードする際に取得できるインスタンスとなります。

このセルを実行すると、ノートブックと同じフォルダ内に```langgraph_sample.py```というファイルが出来上がるはずです。

```python
%%writefile "./langgraph_sample.py"

from mlflow.models import set_model
from typing import Annotated

from typing_extensions import TypedDict

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_core.messages import AIMessage

class State(TypedDict):
    messages: Annotated[list, add_messages]

def chatbot(state: State):
    return {"messages": [AIMessage(content="This is sample message.")]}

graph_builder = StateGraph(State)
# Nodes
graph_builder.add_node("chatbot", chatbot)
# Edges
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)

graph = graph_builder.compile()

# MLflowでのロギングを有効にするために必要
set_model(graph)
```

# Step3. MLflowでのロギング

Step2で定義したLangGraph処理をMLFlowでロギングします。
ロギング時のパラメータ```lc_model```には、Step2で出力したpyファイルを指定。
また、任意ですが```registered_model_name```を指定してUnity Catalog上にモデルを保存しました。

※ 今回のモデルだと```output_example```で指定されるような出力にならないため、ここのスキーマ指定は厳密には誤っています。signatureの指定がないとUnity Catalogに登録できないので一旦適当に設定していますが、ちゃんと正しいスキーマを登録するようにしましょう。

```python
import mlflow
from mlflow.models import infer_signature

input_example = {
    "messages": [{"role": "user", "content": "what is the weather in seattle today?"}]
}
output_example = {
    "messages": [
        {"role": "user", "content": "what is the weather in seattle today?"},
        {"role": "user", "content": "This is sample message."},
    ]
}
signature = infer_signature(input_example, output_example)

with mlflow.start_run():
    model_info = mlflow.langchain.log_model(
        lc_model="./langgraph_sample.py",
        artifact_path="langgraph",
        input_example=input_example,
        signature=signature,
        registered_model_name="training.llm.langgraph_sample",
    )
```

# Step4. MLflowモデルのロードと実行

ロギングしたモデルを読み込んで実行します。
通常のLangGraphグラフとして利用できます。

```python
# MLflow Tracingの有効化
mlflow.langchain.autolog(log_traces=True)

uri = "models:/training.llm.langgraph_sample/1"
agent = mlflow.langchain.load_model(uri)
query = {
    "messages": [
        {
            "role": "user",
            "content": "Should I bring an umbrella today when I go to work in San Francisco?",
        }
    ]
}
agent.invoke(query)
```

```:出力
{'messages': [HumanMessage(content='Should I bring an umbrella today when I go to work in San Francisco?', id='322e556d-8a79-4527-9c90-d6bbfe7723c3'),
  AIMessage(content='This is sample message.', id='c9706a70-9a0f-44d1-b80e-c740b54f2e6b')]}
```

MLflow Tracingにも対応しているので、ノードごとの処理結果がわかりやすく可視化されます。ありがたい！

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/b894b3f5-5282-bfac-8862-47d741133d5e.png)

ちなみに、Mosaic AI Model Servingへの登録も当然可能です。
複雑なRAGパイプラインをLangChain&LangGraphで構築し、それをModel Serving Endpointとして公開し、フロントエンドからAPI呼び出しで実行、というアーキテクチャが簡単に（かつ信頼性高く）実現できますね。


# まとめ

MLflowもバージョンを重ねるごとにLLM関連の対応が進んできていますね。
Models from Codeも柔軟性が高く、非常にありがたいです。

LlamaIndexの拡張回りも気になってます。引き続きいろいろ確認してみたいと思います。