AIエージェントの登録
Mosaic AI Agent Framework を使用して AI エージェント を記録します。 エージェントの記録は、開発プロセスの基本です。 ロギングでは、エージェントのコードと構成の「時点」をキャプチャするため、構成の品質を評価できます。
必要条件
ロギングの前に AI エージェントを作成します。
Databricks では、最新バージョンの databricks-sdkをインストールすることをお勧めします。
% pip install databricks-sdk
コードベースの記録
Databricks では、エージェントをログに記録するときに MLflow Models from Code機能 を使用することをお勧めします。
このアプローチでは、エージェントのコードはPythonファイルとしてキャプチャされ、Python環境はパッケージのリストとしてキャプチャされます。 エージェントがデプロイされると、Python環境が復元され、エージェントのコードが実行されてエージェントがメモリにロードされ、エンドポイントが呼び出されたときにエージェントを呼び出せるようになります。
このアプローチをAPIs mlflow.models.predict() などのデプロイ前検証 と組み合わせて使用することで、エージェントが配信のためにデプロイされたときに確実に実行されるようにすることができます。
コードベースのログ記録の例については、ChatAgent 作成ノートブックの例を参照してください。
ロギングにおけるモデルシグネチャの推論
Databricks では、 ChatAgent インターフェイスを使用してエージェントを作成することをお勧めします。ChatAgent を使用している場合は、このセクションをスキップできます。MLflow は、エージェントの有効なシグネチャを自動的に推論します。
ChatAgent インターフェイスを使用しない場合は、次のいずれかの方法を使用して、ログ時にエージェントの MLflow モデルシグネチャを指定する必要があります。
- シグネチャを手動で定義する
- MLflow のモデルシグネチャ推論機能を使用して、指定した入力例に基づいてエージェントのシグネチャを自動的に生成します。 このアプローチは、シグネチャを手動で定義するよりも便利です。
MLflow モデル シグネチャは、入力と出力を検証して、エージェントが AI Playground やレビュー アプリなどのダウンストリーム ツール。 また、エージェントの使用方法に関する他のアプリケーションのガイドも示します 効果的。
以下の LangChain と PyFunc の例では、モデルシグネチャの推論を使用しています。
ロギング時にモデルシグネチャを自分で明示的に定義する場合は、 MLflow ドキュメント - シグネチャを使用してモデルを記録する方法を参照してください。
LangChainによるコードベースのロギング
次の手順とコードサンプルは、LangChainを使用してエージェントをログに記録する方法を示しています。
-
コードを使用してノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は
agent.pyです。 ノートブックまたはファイルには、ここでは LangChain エージェント (lc_agentと呼びます) が含まれている必要があります。 -
ノートブックまたはファイルに mlflow.models.set_model(lc_agent) を含めます。
-
ドライバーノートブックとして機能する新しいノートブックを作成します（この例では
driver.pyと呼ばれます）。 -
ドライバー ノートブックで、次のコードを使用して ML
agent.pyを実行し、結果を MLflow モデルに記録します。Pythonmlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)resourcesパラメーターは、エージェントにサービスを提供するために必要な Databricks管理のリソース (ベクトル検索インデックスや、基盤モデルを提供する提供エンドポイントなど) を宣言します。詳細については、「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と呼ばれます）。 -
ドライバー ノートブックで、次のコードを使用して ML
agent.pyを実行し、結果を MLflow モデルに記録します。Pythonmlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)resourcesパラメーターは、エージェントにサービスを提供するために必要な Databricks管理のリソース (ベクトル検索インデックスや、基盤モデルを提供する提供エンドポイントなど) を宣言します。詳細については、「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)
Databricks リソースの認証
AIエージェントは、タスクを完了するために他のリソースに対して認証する必要があることがよくあります。 たとえば、エージェントは、非構造化データをクエリするためにベクトル検索インデックスにアクセスする必要がある場合があります。
依存リソースの認証で説明されているように、モデルサービングは、エージェントをデプロイするときに、Databricksマネージド リソースと外部リソースの両方に対する認証をサポートします。
モデルサービングは、 Databricks管理リソースに対して 2 種類の認証をサポートしています。
- システム認証: エージェント サービスプリンシパルが、エージェントのログ時に指定された任意の依存リソースにアクセスできるようにします。 これは、共有リソースや機密性の低いリソース(公開ドキュメントを含むベクトル検索インデックスなど)にアクセスする場合に便利です
- [ベータ版] ユーザー代理認証: エージェントがエンド ユーザーの資格情報を使用して Databricks リソースにアクセスできるようにします。これは、エージェントが機密データにアクセスしたり、リモート APIにクエリを実行してユーザーごとにアクションを実行したりする必要があるシナリオに役立ちます
自動認証パススルー (システム認証) のリソースを指定します
最も一般的な Databricks リソースの種類については、Databricks では、ログ記録中にエージェントのリソース依存関係を事前に宣言することをサポートしており、これを推奨しています。 これにより、エージェントをデプロイするときに 自動認証パススルー が有効になり、Databricks は、エージェント エンドポイント内からこれらのリソース依存関係に安全にアクセスするための有効期間の短い資格情報を自動的にプロビジョニング、ローテーション、管理します。
自動認証パススルーを有効にするには、次のコードに示すように、log_model() API の resources パラメーターを使用して依存リソースを指定します。
import mlflow
from mlflow.models.resources import (
DatabricksVectorSearchIndex,
DatabricksServingEndpoint,
DatabricksSQLWarehouse,
DatabricksFunction,
DatabricksGenieSpace,
DatabricksTable,
DatabricksUCConnection
)
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"),
DatabricksUCConnection(connection_name="your_connection_name"),
]
)
Databricks では、すべてのエージェント フレーバーの resources を手動で指定することをお勧めします。
mlflow.langchain.log_model(...)を使用して LangChain エージェントをログに記録するときにリソースを指定しない場合、MLflow はリソースのベストエフォート型自動推論を実行します。ただし、これによりすべての依存関係がキャプチャされるわけではないため、エージェントの提供時またはクエリ時に認証エラーが発生する可能性があります。
次の表に、自動認証パススルーをサポートする Databricks リソースと、リソースをログに記録するために必要な最小mlflowバージョンを示します。
リソースタイプ | リソースをログに記録するために必要な最小 |
|---|---|
ベクトル検索インデックス |
|
モデルサービングエンドポイント |
|
SQLウェアハウス |
|
Unity Catalogの関数 |
|
Genieスペース |
|
Unity Catalog テーブル |
|
Unity Catalog 接続 |
|
ユーザー代理認証
ベータ版
この機能は ベータ版です。
ユーザー代理認証を使用するエージェントをログに記録する場合は、エージェント コード内でエンド ユーザーとしてアクションを実行するために必要な Databricks API スコープの最小セットを指定します。これにより、エージェントはデプロイ時にエンドユーザーに代わってアクションを実行するための最小特権アクセス権を持つようになり、不正なアクションを防止し、トークンの誤用のリスクを最小限に抑えることでセキュリティを強化します。
以下は、いくつかの一般的な種類の Databricks リソースにアクセスするために必要なスコープの一覧です。
Databricks リソース | 必要な API スコープ |
|---|---|
ベクトル検索インデックス |
|
モデルサービングエンドポイント |
|
SQLウェアハウス |
|
UC接続 |
|
Genieスペース |
|
ユーザー代理認証を有効にするには、次の例に示すように、MLflow AuthPolicy を log_model()に渡します。MLflow AuthPolicy には、次の 2 つのコンポーネントがあります。
system_auth_policy: システム認証のリソースを指定します。通常、エージェントは、共有リソースのシステム認証(モデルサービングエンドポイントのクエリなど)を、機密性の高いリソースやリソースにアクセスするためのユーザー代理認証と組み合わせて使用 APIsuser_auth_policy: エージェントの代理ユーザー認証に必要な API スコープを指定します
from mlflow.models.resources import DatabricksServingEndpoint
from mlflow.models.auth_policy import SystemAuthPolicy, UserAuthPolicy, AuthPolicy
resources = [
DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct")
]
# Specify resources here for system authentication
system_auth_policy = SystemAuthPolicy(resources=resources)
# Specify the minimal set of API scopes needed for on-behalf-of-user authentication
# When deployed, the agent can access Databricks resources and APIs
# on behalf of the end user, but only via REST APIs that are covered by the list of
# scopes below
user_auth_policy = UserAuthPolicy(
api_scopes=[
"serving.serving-endpoints",
"vectorsearch.vector-search-endpoints",
"vectorsearch.vector-search-indexes",
]
)
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
...
# Instead of passing `resources` (which only supports system authentication),
# pass an auth_policy to log_model to enable both system authentication and
# on-behalf-of-user authentication
auth_policy=AuthPolicy(
system_auth_policy=system_auth_policy,
user_auth_policy=user_auth_policy
)
)
OpenAIクライアントの自動認証
エージェントが OpenAI クライアントを使用している場合は、Databricks SDK を使用してデプロイ中に自動的に認証します。 Databricks SDK には、承認が自動的に構成された OpenAI クライアントを構築するためのラッパーが用意されています。 ノートブックで次のコマンドを実行します。
% pip install databricks-sdk[openai]
from databricks.sdk import WorkspaceClient
def openai_client(self):
w = WorkspaceClient()
return w.serving_endpoints.get_open_ai_client()
次に、モデルサービングエンドポイントを resources の一部として指定して、デプロイ時に自動的に認証します。
エージェントを 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)