AIエージェントの記録と登録（Model Serving）
新しいユースケースの場合、Databricksは、エージェントコード、サーバー構成、およびデプロイワークフローを完全に制御するために、Databricks Appsにエージェントをデプロイすることをお勧めします。AIエージェントを作成してDatabricks Appsにデプロイするを参照してください。既存のエージェントを移行するには、Model ServingからDatabricks Appsにエージェントを移行するを参照してください。
Log AI agents using Custom Agents.エージェントのログ記録は、開発プロセスの基礎となります。ログ記録は、エージェントのコードと構成の「現時点」をキャプチャするため、品質を評価できます。
要件
ロギングの前に AI エージェントを作成します。
Databricks は、databricks-sdk の最新バージョンをインストールすることを推奨しています。
% pip install databricks-sdk
コードベースの記録
Databricks では、エージェントをログに記録するときに MLflow の Models from Code機能 を使用することをお勧めします。
このアプローチでは、エージェントのコードはPythonファイルとしてキャプチャされ、Python環境はパッケージのリストとしてキャプチャされます。エージェントがデプロイされると、Python環境が復元され、エージェントのコードが実行されて、エージェントがメモリにロードされ、Endpointが呼び出されたときに呼び出せるようになります。
このアプローチを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モデルにLogします:Pythonmlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)resourcesパラメーターは、エージェントにサービスを提供するために必要な Databricksマネージドリソース (AI Searchインデックス、または基盤モデルを提供する提供エンドポイントなど) を宣言します。詳細については、自動認証パススルーを実装するを参照してください。 -
モデルをデプロイします。「生成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モデルにLogsをします:Pythonmlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)resourcesパラメーターは、エージェントにサービスを提供するために必要な Databricksマネージドリソース (AI Searchインデックス、または基盤モデルを提供する提供エンドポイントなど) を宣言します。詳細については、自動認証パススルーを実装するを参照してください。 -
モデルをデプロイします。「生成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()を参照してください。