AI エージェントをログに記録して登録する
重要
この機能はパブリック プレビュー段階にあります。
Mosaic AI エージェント フレームワークを使用して AI エージェントをログに記録します。 エージェントのログは、開発プロセスの基礎となります。 ログにより、"特定の時点" でのエージェントのコードと構成がキャプチャされるため、構成の品質を評価できます。
要件
ログに記録する前に AI エージェントを作成します。
コードベースのログ記録
Databricks では、エージェントのログ記録時に、コード機能 から MLflow の
この方法では、エージェントのコードが Python ファイルとしてキャプチャされ、Python 環境がパッケージの一覧としてキャプチャされます。 エージェントがデプロイされると、Python 環境が復元され、エージェントのコードが実行され、エージェントがメモリに読み込まれるため、エンドポイントが呼び出されたときに呼び出すことができます。
この方法と、mlflow.models.predict() などのデプロイ前検証 API を使用して、サービスを提供するためにデプロイされたときにエージェントが確実に実行されるようにすることができます。
エージェントまたはエージェントをログするコードは、エージェント コードとは別のノートブックに含める必要があります。 このノートブックは、ドライバー ノートブックと呼ばれます。 ノートブックの例については、「サンプル ノートブック」を参照してください。
ログ記録中にモデル署名を推論する
ログ記録中に、エージェントの入力スキーマと出力スキーマを指定する MLflow Model Signatureを定義する必要があります。 署名は入力と出力を検証して、エージェントが AI Playground やレビュー アプリなどのダウンストリーム ツールと正しく対話することを確認します。 また、エージェントを効果的に使用する方法に関する他のアプリケーションのガイドも示します。
Databricks では、MLflow の Model Signature 推論機能を使用して、指定した入力例に基づいてエージェントの署名を自動的に生成することをお勧めします。 この方法は、署名を手動で定義するよりも便利です。
以下の LangChain と PyFunc の例では、Model Signature 推論を使用します。
ログ記録時にモデル署名を自分で明示的に定義する場合は、「MLflow のドキュメント - 署名を使用してモデルをログに記録する方法」を参照してください。
LangChain を使用したコードベースのログ
次の手順とコード サンプルでは、LangChain を使用してエージェントをログに記録する方法を示します。
コードを含むノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は
agent.py
です。 ノートブックまたはファイルには、ここでlc_agent
と呼ばれる LangChain エージェントが含まれている必要があります。ノートブックまたはファイルに mlflow.models.set_model(lc_agent) を含めます。
ドライバー ノートブックとして機能する新しいノートブックを作成します (この例では
driver.py
と呼ばれます)。ドライバー ノートブックで、次のコードを使用して
agent.py
を実行し、MLflow モデルに結果をログに記録します。mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
resources
パラメーターは、ベクター検索インデックスや基盤モデルを提供するサービス エンドポイントなど、エージェントにサービスを提供するために必要な Databricks マネージド リソースを宣言します。 詳細については、「自動認証パススルーのリソースを指定する」を参照してください。モデルをデプロイします。 「生成 AI アプリケーション用にエージェントをデプロイする」を参照してください。
サービス環境が読み込まれると、
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
を実行し、MLflow モデルに結果をログに記録します。mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
resources
パラメーターは、ベクター検索インデックスや基盤モデルを提供するサービス エンドポイントなど、エージェントにサービスを提供するために必要な Databricks マネージド リソースを宣言します。 詳細については、「自動認証パススルーのリソースを指定する」を参照してください。モデルをデプロイします。 「生成 AI アプリケーション用にエージェントをデプロイする」を参照してください。
サービス環境が読み込まれると、
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-mixtral-8x7b-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)
自動認証パススルーのリソースを指定する
多くの場合、AI エージェントは、タスクを完了するために他のリソースに対して認証を行う必要があります。 たとえば、エージェントは、非構造化データのクエリを実行するために Vector Search インデックスにアクセスする必要がある場合があります。
「依存リソースの
最も一般的な Databricks リソースの種類については、Databricks では、ログ記録中にエージェントのリソース依存関係を事前に宣言することをサポートし、推奨しています。 これにより、エージェントのデプロイ時 自動認証パススルー が可能になります。Databricks は、有効期間の短い資格情報を自動的にプロビジョニング、ローテーション、管理して、エージェント エンドポイント内からこれらのリソースの依存関係に安全にアクセスします。
自動認証パススルーを有効にするには、次のコードに示すように、log_model() API の resources
パラメーターを使用して依存リソースを指定します。
import mlflow
from mlflow.models.resources import (
DatabricksVectorSearchIndex,
DatabricksServingEndpoint,
DatabricksSQLWarehouse,
DatabricksFunction,
DatabricksGenieSpace,
DatabricksTable,
)
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
python_model=agent_notebook_path,
artifact_path="agent",
input_example=input_example,
example_no_conversion=True,
# Specify resources for automatic authentication passthrough
resources=[
DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
DatabricksFunction(function_name="ml.tools.python_exec"),
DatabricksGenieSpace(genie_space_id="your_genie_space_id"),
DatabricksTable(table_name="your_table_name"),
]
)
Databricks では、すべてのエージェント フレーバーに resources
を手動で指定することをお勧めします。
Note
mlflow.langchain.log_model(...)
を使用して LangChain エージェントをログに記録するときにリソースを指定しない場合、MLflow はリソースのベスト エフォート自動推論を実行します。 ただし、これによりすべての依存関係がキャプチャされるとはかからず、エージェントにサービスを提供したりクエリを実行したりするときに承認エラーが発生する可能性があります。
次の表に、自動認証パススルーをサポートする Databricks リソースと、リソースのログ記録に必要な最小 mlflow
バージョンを示します。
リソースの種類 | リソースをログに記録するために必要な最小 mlflow バージョン |
---|---|
ベクター検索インデックス | mlflow 2.13.1 以降が必要です |
モデル サービング エンドポイント | mlflow 2.13.1 以降が必要です |
SQL Warehouse | mlflow 2.16.1 以降が必要です |
Unity Catalog 関数 | mlflow 2.16.1 以降が必要です |
ジーニー空間 | mlflow 2.17.1 以降が必要です |
Unity カタログ テーブル | mlflow 2.18.0 以降が必要です |
エージェントを Unity カタログに登録する
エージェントをデプロイする前に、エージェントを Unity カタログに登録する必要があります。 エージェントパッケージをモデルとして Unity カタログに登録する。 その結果、エージェント内のリソースの承認に Unity カタログのアクセス許可を使用できます。
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)