AIエージェントのログ記録と登録する (Model Serving)
新しいユースケースの場合、Databricksは、エージェントのコード、サーバー設定、およびデプロイワークフローを完全に制御するために、Databricks Appsへのエージェントのデプロイを推奨しています。AIエージェントを作成してDatabricks Appsにデプロイするを参照してください。既存のエージェントを移行するには、Model ServingからDatabricks Appsにエージェントを移行するを参照してください。
カスタムエージェントを使用してAIエージェントを記録します。エージェントのログ記録は、開発プロセスの基礎となります。ログ記録により、エージェントのコードと構成の「時点」がキャプチャされ、その品質を評価できます。
要件
ロギングの前に AI エージェントを作成します。
Databricksでは、最新バージョンのdatabricks-sdkをインストールすることをお勧めします。
% pip install databricks-sdk
コードベースの記録
Databricks は、エージェントをログに記録する際に MLflow の Models from Code機能を使用することをお勧めします。
このアプローチでは、エージェントのコードはPythonファイルとしてキャプチャされ、Python環境はパッケージのリストとしてキャプチャされます。エージェントがデプロイされると、Python環境が復元され、エージェントのコードが実行されてエージェントがメモリにロードされるため、エンドポイントが呼び出されたときに呼び出すことができます。
このアプローチをAPI mlflow.models.predict() などのデプロイ前検証 と組み合わせて使用することで、エージェントが配信のためにデプロイされたときに確実に実行されるようにすることができます。
コードベースのログ記録の例については、ResponsesAgent オーサリング ノートブックの例を参照してください。
ロギングにおけるモデルシグネチャの推論
Databricks では、 ResponsesAgent インターフェイスを使用してエージェントを作成することをお勧めします。ResponsesAgent を使用している場合は、このセクションをスキップできます。MLflow は、エージェントの有効なシグネチャを自動的に推論します。
ResponsesAgent インターフェイスを使用しない場合は、次のいずれかの方法を使用して、ログ時にエージェントの MLflow モデルシグネチャを指定する必要があります。
- シグネチャを手動で定義する
- MLflow のモデルシグネチャ推論機能を使用して、指定した入力例に基づいてエージェントのシグネチャを自動的に生成します。このアプローチは、シグネチャを手動で定義するよりも便利です。
MLflow モデルシグネチャは、AI Playground やレビューアプリなどのダウンストリームツールとエージェントが正しくやり取りできるように、入力と出力を検証します。また、他のアプリケーションがエージェントを効果的に使用する方法を案内します。
以下の LangChain と PyFunc の例では、モデルシグネチャの推論を使用しています。
ロギング時にモデルシグネチャを自分で明示的に定義する場合は、 MLflow ドキュメント - シグネチャを使用してモデルを記録する方法を参照してください。
LangChainによるコードベースのログ記録
次の手順とコードサンプルでは、LangChain を使用してエージェントをログに記録する方法を説明します。
-
コードを使用してノートブックまたはPythonファイルを作成します。この例では、ノートブックまたはファイルの名前は
agent.pyです。ノートブックまたはファイルには、ここではlc_agentと呼ばれるLangChainエージェントが含まれている必要があります。 -
ノートブックまたはファイルにmlflow.models.set_model(lc_agent)を含めます。
-
ドライバーノートブックとして機能する新しいノートブックを作成します（この例では
driver.pyと呼ばれます）。 -
ドライバーノートブックで、次のコードを使用して
agent.pyを実行し、結果を MLflow モデルに記録します:Pythonmlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)resourcesパラメーターは、AI Searchインデックスや基盤モデルを提供するサービングエンドポイントなど、エージェントにサービスを提供するために必要なDatabricks管理のリソースを宣言します。情報については、「自動認証パススルーの実装」を参照してください。 -
モデルをデプロイします。生成AIアプリケーション (Model Serving) のためのエージェントのデプロイをご覧ください。
-
サービング環境がロードされると、
agent.pyが実行されます。 -
サービングリクエストが来ると
lc_agent.invoke(...)が呼び出されます。
import mlflow
code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"
# Input example used by MLflow to infer Model Signature
input_example = {
"messages": [
{
"role": "user",
"content": "What is Retrieval-augmented Generation?",
}
]
}
# example using langchain
with mlflow.start_run():
logged_agent_info = mlflow.langchain.log_model(
lc_model=code_path,
model_config=config_path, # If you specify this parameter, this configuration is used by agent code. The development_config is overwritten.
artifact_path="agent", # This string is used as the path inside the MLflow model where artifacts are stored
input_example=input_example, # Must be a valid input to the agent
example_no_conversion=True, # Required
)
print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")
# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.langchain.load_model(logged_agent_info.model_uri)
model.invoke(example)
PyFuncによるコードベースのログ記録
以下の手順とコードサンプルは、PyFunc を使用してエージェントをログに記録する方法を示します。
-
コードを使ってノートブックまたは Python ファイルを作成します。この例では、ノートブックまたはファイルの名前は
agent.pyです。ノートブックまたはファイルには、PyFuncClassという名前のPyFuncクラスが含まれている必要があります。 -
mlflow.models.set_model(PyFuncClass)をノートブックまたはファイルに含めます。 -
ドライバーノートブックとして機能する新しいノートブックを作成します（この例では
driver.pyと呼ばれます）。 -
ドライバーノートブックで、以下のコードを使用して
agent.pyを実行し、log_model()を使用して結果をMLflowモデルに記録します:Pythonmlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)resourcesパラメーターは、AI Searchインデックスや基盤モデルを提供するサービングエンドポイントなど、エージェントにサービスを提供するために必要なDatabricks管理のリソースを宣言します。情報については、「自動認証パススルーの実装」を参照してください。 -
モデルをデプロイします。生成AIアプリケーション (Model Serving) のためのエージェントのデプロイをご覧ください。
-
サービング環境がロードされると、
agent.pyが実行されます。 -
サービングリクエストが来ると
PyFuncClass.predict(...)が呼び出されます。
import mlflow
from mlflow.models.resources import (
DatabricksServingEndpoint,
DatabricksVectorSearchIndex,
)
code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"
# Input example used by MLflow to infer Model Signature
input_example = {
"messages": [
{
"role": "user",
"content": "What is Retrieval-augmented Generation?",
}
]
}
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
python_model=agent_notebook_path,
artifact_path="agent",
input_example=input_example,
resources=resources_path,
example_no_conversion=True,
resources=[
DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct"),
DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
]
)
print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")
# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.pyfunc.load_model(logged_agent_info.model_uri)
model.invoke(example)
Databricksリソースの認証
AIエージェントは多くの場合、タスクを完了するために他のリソースに対して認証する必要があります。例えば、デプロイされたエージェントは、非構造化データをクエリするためにAI Searchインデックスにアクセスしたり、動的プロンプトをロードするためにPrompt Registryにアクセスしたりする必要がある場合があります。
自動認証パススルーと代理認証には、エージェントのログ記録中に構成が必要です。
エージェントを Unity Catalogに登録する
エージェントをデプロイする前に、エージェントを Unity Catalog に登録する必要があります。エージェントを登録すると、Unity Catalog にモデルとしてパッケージ化されます。その結果、エージェント内のリソースの認可に Unity Catalog の権限を使用できます。
import mlflow
mlflow.set_registry_uri("databricks-uc")
catalog_name = "test_catalog"
schema_name = "schema"
model_name = "agent_name"
model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=model_name)
詳しくは、mlflow.register_model()を参照してください。