MLflowによるLangChainオートロギングをDatabricksで動かしてみる

こちらのページを見つけました。

https://mlflow.org/docs/latest/llms/langchain/autologging.html

> MLflowのLangChainフレーバーでは、明示的なロギング文の必要なしに、LangChainのモデルと処理に関する重要な詳細情報を記録できるようにするためのパワフルな機能であるオートロギングをサポートしています。MLflow LangChainオートロギングは、トレース、モデル、シグネチャなどを含むさまざまなモデルの側面をカバーします。

なんと。トレースののみならずモデルも記録してくれるとは。

Databricksで試してみます。デフォルトでは[ワークスペースモデルレジストリ](https://docs.databricks.com/ja/machine-learning/manage-model-lifecycle/workspace-model-registry.html)にモデルが記録されるので、Unity Catalog(UC)に登録するようにしていきます。

# LangChainのオートロギング

ライブラリのバージョンに注意してインストールします。langchainとlangchain_communityのバージョンに注意しないと、モデルサービングエンドポイントへのデプロイでエラーになります。

```py
%pip install langchain_community==0.3.9 langchain
%pip install -U openai langchain mlflow[langchain]
%pip install langgraph
%pip install textstat spacy
%restart_python
```

以下のバージョンで動作確認できました。

![Screenshot 2025-01-21 at 17.48.16.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/fdf9bfa9-fede-3fd8-f1d8-c122055de543.png)

事前に定義した[シークレット](https://docs.databricks.com/ja/security/secrets/index.html)からOpenAIのAPIキーを設定します。

```py
import os
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")
```

```py
import os
from operator import itemgetter

from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnableLambda

import mlflow

# 以下のコメントアウトを解除して、langchainの自動ロギングの全機能を使用する
# %pip install `langchain_community>=0.0.16`
# これらの2つのライブラリは、テキスト分析関連のアーティファクトを自動的にログに記録するために必要
# %pip install textstat spacy

# UCにモデルを登録するように設定
mlflow.set_registry_uri("databricks-uc")
```

```py
assert "OPENAI_API_KEY" in os.environ, "OPENAI_API_KEY環境変数を設定してください。"

# mlflowのlangchainオートロギングを有効にする
# 注意: レトリーバーを含まないモデルのみ自動ロギングをサポート
mlflow.langchain.autolog(
    log_input_examples=True,
    log_model_signatures=True,
    log_models=True,
    registered_model_name="users.takaaki_yayoi.lc_model",
)

prompt_with_history_str = """
以下はあなたと人間との間の履歴です: {chat_history}

さて、この質問に答えてください: {question}
"""
prompt_with_history = PromptTemplate(
    input_variables=["chat_history", "question"], template=prompt_with_history_str
)


def extract_question(input):
    return input[-1]["content"]


def extract_history(input):
    return input[:-1]


llm = OpenAI(temperature=0.9)

# LCELを使用してチェーンを構築
chain_with_history = (
    {
        "question": itemgetter("messages") | RunnableLambda(extract_question),
        "chat_history": itemgetter("messages") | RunnableLambda(extract_history),
    }
    | prompt_with_history
    | llm
    | StrOutputParser()
)

inputs = {"messages": [{"role": "user", "content": "MLflowの所有者は誰ですか?"}]}

print(chain_with_history.invoke(inputs))
# サンプル出力:
# "1. Databricks\n2. Microsoft\n3. Google\n4. Amazon\n\n答えを入力してください: 1\n\n
# 正解！MLflowはDatabricksによって開発されたオープンソースプロジェクトです。 ...
```

MLflow Tracingが表示されることに加え、オートロギングによってモデルがUnity Catalogに登録されます。

![Screenshot 2025-01-21 at 17.51.27.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/5c17cd01-936c-3eba-b09e-f105da71cb42.png)

![Screenshot 2025-01-21 at 17.52.24.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/8c4e10b6-f37b-b870-d602-ba0059888e6b.png)

```py
# モデルと関連するアーティファクトを自動的にログに記録
# `lc_model`という名前のモデルが登録され、PyFuncモデルとして再ロード可能
model_name = "users.takaaki_yayoi.lc_model"
model_version = 3
loaded_model = mlflow.pyfunc.load_model(f"models:/{model_name}/{model_version}")
print(loaded_model.predict(inputs))
```

![Screenshot 2025-01-21 at 17.53.07.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0fbabaec-1d18-775b-ba34-a956e3bb34e3.png)

UCに登録されたモデルをモデルサービングエンドポイントにデプロイします。この際に環境変数に上で指定したシークレットを指定します。

![Screenshot 2025-01-21 at 17.53.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/7b71b44e-b61f-6a6d-f43d-fc9711b0ccf8.png)

デプロイが完了すると、問い合わせを行なって動作確認を行うことができます。

![Screenshot 2025-01-21 at 17.47.30.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/1a140055-f9ec-ae26-3913-9eed1686329e.png)

これは思った以上にお手軽です。

# LangGraphのトレース

このオートロギングはLangGraphもサポートしています。ただし、トレースのみとなります。

```py
from typing import Literal

import mlflow

from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

# LangGraph (LangChain) のトレースを有効にする
mlflow.langchain.autolog()

# オプション: トラッキングURIと実験を設定する
#mlflow.set_tracking_uri("http://localhost:5000")
#mlflow.set_experiment("LangGraph")


@tool
def get_weather(city: Literal["nyc", "sf"]):
    """天気情報を取得するために使用します。"""
    if city == "nyc":
        return "nycでは曇りかもしれません"
    elif city == "sf":
        return "sfはいつも晴れています"


llm = ChatOpenAI(model="gpt-4o-mini")
tools = [get_weather]
graph = create_react_agent(llm, tools)

# グラフを呼び出す
result = graph.invoke(
    {"messages": [{"role": "user", "content": "sfの天気はどうですか?"}]}
)
```
![Screenshot 2025-01-21 at 17.59.02.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/287c2d70-e141-5de5-2826-11b54b3fc5c2.png)


### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)