実験記録: LangGraphのグラフはmlflow.langchainフレーバーで保管できるのか？

先に結論

MLflow 2.14.1時点ではできませんでした。

実験背景

LangGraphのグラフをMLflowで実験管理したい。でもFlavorは(まだ)出ていない。
あれ、でもLangGraphのグラフってRunnable系の何かを継承していなかったっけ？ということは既にできるんじゃないか？というのが実験背景です。

mlflowのバージョンは2.14.1、langgraphは0.1.1で確認しました。

調査①

LangGraphのコードを見るに、最終的にコンパイルされたグラフはCompiledGraphのインスタンスになります。CompiledGraphはPregelというクラスを継承しています。

class CompiledGraph(Pregel):

そしてPregelはRunnableSerializableを継承しています。
RunnableSerializableはlangchain_core.runnablesに含まれるクラスです。

class Pregel(
    RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]
):

MLflowのlangchainフレーバーがRunnableSerializableをChainとして記録するようになっていれば、モデルロギングできるのでは？

調査②

というわけで検証実験。
こちらの記事のコードでLangGraphのグラフを作成し、以下のコードで保存。

# 検証コード
with mlflow.start_run():
    logged_chain_info = mlflow.langchain.log_model(
        lc_model=graph,
        artifact_path="chain",
    )

実行した結果、以下のようなエラーとなりました。

出力

MlflowException: MLflow langchain flavor only supports subclasses of langchain.chains.base.Chain, langchain.agents.agent.AgentExecutor, langchain.schema.BaseRetriever, langchain.schema.runnable.RunnableSequence, langchain.schema.runnable.RunnableLambda, langchain.schema.runnable.RunnableParallel, langchain.schema.runnable.RunnablePassthrough, langchain.schema.runnable.passthrough.RunnableAssign instances, found CompiledStateGraph

残念ながら対応しているrunnableモジュールのクラスの中に、RunnableSerializableがなく、チェックで弾かれました。

まとめ

というわけで、現時点ではダメでした。
そりゃそうだろ、という感じですがこういうのはやってみないとわからなかったりするので。。。

MLflowのコントリビュータの方々、LangGraphの対応よろしくお願いします！（お前がやれという声が聞こえる。。。)