こちらのページを見つけました。
MLflowのLangChainフレーバーでは、明示的なロギング文の必要なしに、LangChainのモデルと処理に関する重要な詳細情報を記録できるようにするためのパワフルな機能であるオートロギングをサポートしています。MLflow LangChainオートロギングは、トレース、モデル、シグネチャなどを含むさまざまなモデルの側面をカバーします。
なんと。トレースののみならずモデルも記録してくれるとは。
Databricksで試してみます。デフォルトではワークスペースモデルレジストリにモデルが記録されるので、Unity Catalog(UC)に登録するようにしていきます。
LangChainのオートロギング
ライブラリのバージョンに注意してインストールします。langchainとlangchain_communityのバージョンに注意しないと、モデルサービングエンドポイントへのデプロイでエラーになります。
%pip install langchain_community==0.3.9 langchain
%pip install -U openai langchain mlflow[langchain]
%pip install langgraph
%pip install textstat spacy
%restart_python
以下のバージョンで動作確認できました。
事前に定義したシークレットからOpenAIのAPIキーを設定します。
import os
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")
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")
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に登録されます。
# モデルと関連するアーティファクトを自動的にログに記録
# `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))
UCに登録されたモデルをモデルサービングエンドポイントにデプロイします。この際に環境変数に上で指定したシークレットを指定します。
デプロイが完了すると、問い合わせを行なって動作確認を行うことができます。
これは思った以上にお手軽です。
LangGraphのトレース
このオートロギングはLangGraphもサポートしています。ただし、トレースのみとなります。
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の天気はどうですか?"}]}
)
