メインコンテンツまでスキップ

AIエージェントの登録

Mosaic AI Agent Framework を使用して AI エージェントをログに記録します。エージェントのログ記録は、開発プロセスの基礎です。ログは、エージェントのコードと構成の「特定の時点」をキャプチャするため、構成の品質を評価できます。

必要条件​

ロギングの前に AI エージェントを作成します。

Databricks では、最新バージョンの databricks-sdkをインストールすることをお勧めします。

Python
% pip install databricks-sdk

コードベースの記録​

Databricks では、エージェントをログに記録するときに MLflow の Code からのモデル機能 を使用することをお勧めします。

このアプローチでは、エージェントのコードは Python ファイルとしてキャプチャされ、Python 環境はパッケージのリストとしてキャプチャされます。エージェントがデプロイされると、Python 環境が復元され、エージェントのコードが実行されてエージェントがメモリにロードされ、エンドポイントが呼び出されたときに呼び出せるようにします。

このアプローチをAPI mlflow.models.predict() などのデプロイ前検証 と組み合わせて使用することで、エージェントが配信のためにデプロイされたときに確実に実行されるようにすることができます。

コードベースのログの例については、「 ResponsesAgent オーサリングのノートブックの例」を参照してください。

ロギングにおけるモデルシグネチャの推論​

注記

Databricks では、 ResponsesAgent インターフェイスを使用してエージェントを作成することをお勧めします。ResponsesAgent を使用している場合は、このセクションをスキップできます。MLflow は、エージェントの有効な署名を自動的に推測します。

ResponsesAgent インターフェイスを使用しない場合は、次のいずれかの方法を使用して、ログ時にエージェントの MLflow モデルシグネチャを指定する必要があります。

  1. シグネチャを手動で定義する
  2. MLflow のモデル署名推論機能を使用して、指定した入力例に基づいてエージェントの署名を自動的に生成します。このアプローチは、署名を手動で定義するよりも便利です。

MLflow モデル シグネチャは、入力と出力を検証して、エージェントが AI Playground やレビュー アプリなどのダウンストリーム ツール。 また、エージェントの使用方法に関する他のアプリケーションのガイドも示します 効果的。

以下の LangChain と PyFunc の例では、モデルシグネチャの推論を使用しています。

ロギング時にモデルシグネチャを自分で明示的に定義する場合は、 MLflow ドキュメント - シグネチャを使用してモデルを記録する方法を参照してください。

LangChainによるコードベースのロギング​

次の手順とコードサンプルは、LangChainを使用してエージェントをログに記録する方法を示しています。

  1. コードを使用してノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は agent.pyです。 ノートブックまたはファイルには、ここでは LangChain エージェント ( lc_agentと呼びます) が含まれている必要があります。

  2. ノートブックまたはファイルに mlflow.models.set_model(lc_agent) を含めます。

  3. ドライバーノートブックとして機能する新しいノートブックを作成します（この例ではdriver.pyと呼ばれます）。

  4. ドライバー ノートブックで、次のコードを使用して ML agent.py を実行し、結果を MLflow モデルに記録します。

    Python
    mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)

    resources パラメーターは、エージェントにサービスを提供するために必要な Databricks管理のリソース (ベクトル検索インデックスや、基盤モデルを提供する提供エンドポイントなど) を宣言します。詳細については、「Databricks リソースの認証」を参照してください。

  5. モデルをデプロイします。生成AIアプリケーションのためのエージェントのデプロイを参照してください。

  6. サービング環境が読み込まれると、 agent.py が実行されます。

  7. サービングリクエストが来るとlc_agent.invoke(...)が呼び出されます。

Python

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 を使用してエージェントをログに記録する方法を示しています。

  1. コードを使用してノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は agent.pyです。 ノートブックまたはファイルには、 PyFuncClassという名前の PyFunc クラスが含まれている必要があります。

  2. ノートブックまたはファイルに mlflow.models.set_model(PyFuncClass) を含めます。

  3. ドライバーノートブックとして機能する新しいノートブックを作成します（この例ではdriver.pyと呼ばれます）。

  4. ドライバー ノートブックで、次のコードを使用して agent.py を実行し、log_model() を使用して結果を MLflow モデルに記録します。

    Python
    mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)

    resources パラメーターは、エージェントにサービスを提供するために必要な Databricks管理のリソース (ベクトル検索インデックスや、基盤モデルを提供する提供エンドポイントなど) を宣言します。詳細については、「Databricks リソースの認証」を参照してください。

  5. モデルをデプロイします。生成AIアプリケーションのためのエージェントのデプロイを参照してください。

  6. サービング環境が読み込まれると、 agent.py が実行されます。

  7. サービングリクエストが来るとPyFuncClass.predict(...)が呼び出されます。

Python
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エージェントは、タスクを完了するために他のリソースに対して認証する必要があることがよくあります。 たとえば、エージェントは、非構造化データをクエリするためにベクトル検索インデックスにアクセスする必要がある場合があります。

依存リソースの認証で説明されているように、モデルサービングは、エージェントをデプロイするときに、Databricksマネージド リソースと外部リソースの両方に対する認証をサポートします。

モデルサービングは、 Databricks管理リソースに対して認証するための2つの方法をサポートしています。

  • システム認証: エージェントは、 Databricks サービスプリンシパルを使用して、ログ時に宣言した依存リソースにアクセスします。 この方法は、すべてのエージェントユーザーが同じアクセス権を持つ必要がある共有組織レベルのリソースに使用します。たとえば、組織が承認したドキュメントの共有ベクトル検索インデックス、共有モデルサービングエンドポイント、またはアプリケーションで使用される読み取り専用の SQLウェアハウスなどです。

  • [パブリックプレビュー] ユーザーに代わって承認: エージェントは、呼び出し元のユーザー ID を使用して Databricks リソースにアクセスします。これは、ユーザーごとのガバナンスを適用する必要がある場合、またはユーザースコープのデータにアクセスする必要がある場合に使用します。たとえば、行/列レベルのポリシー、個人用Genie spaces、またはユーザー所有の接続を持つテーブルをUnity Catalogしたり、ユーザー属性の監査が必要な場合などです。

監査の目的で、システム認証はサービスプリンシパルへのアクセスを属性化し、ユーザー代理認証はそれを個々のユーザーに属性付けします。

自動認証パススルー (システム認証) のリソースを指定します​

最も一般的な Databricks リソースの種類については、Databricks では、ログ記録中にエージェントのリソース依存関係を事前に宣言することをサポートしており、これを推奨しています。 これにより、エージェントをデプロイするときに 自動認証パススルー が有効になり、Databricks は、エージェント エンドポイント内からこれらのリソース依存関係に安全にアクセスするための有効期間の短い資格情報を自動的にプロビジョニング、ローテーション、管理します。

自動認証パススルーを有効にするには、次のコードに示すように、log_model() API の resources パラメーターを使用して依存リソースを指定します。

Python
import mlflow
from mlflow.models.resources import (
  DatabricksVectorSearchIndex,
  DatabricksServingEndpoint,
  DatabricksSQLWarehouse,
  DatabricksFunction,
  DatabricksGenieSpace,
  DatabricksTable,
  DatabricksUCConnection,
  DatabricksApp
)

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-meta-llama-3-3-70b-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"),
      DatabricksApp(app_name="app_name"),
    ]
  )

Databricks では、すべてのエージェント フレーバーの resources を手動で指定することをお勧めします。

注記

mlflow.langchain.log_model(...)を使用して LangChain エージェントをログに記録するときにリソースを指定しない場合、MLflow はリソースのベストエフォート型自動推論を実行します。ただし、これによりすべての依存関係がキャプチャされるわけではないため、エージェントの提供時またはクエリ時に認証エラーが発生する可能性があります。

次の表に、自動認証パススルーをサポートする Databricks リソースと、リソースをログに記録するために必要な最小mlflowバージョンを示します。

リソースタイプ

リソースをログに記録するために必要な最小 mlflow バージョン

ベクトル検索インデックス

mlflow 2.13.1 以降が必要です

モデルサービングエンドポイント

mlflow 2.13.1 以降が必要です

SQLウェアハウス

mlflow 2.16.1以降が必要です

Unity Catalogの関数

mlflow 2.16.1以降が必要です

Genieスペース

mlflow 2.17.1 以降が必要です

Unity Catalog テーブル

mlflow 2.18.0 以降が必要です

Unity Catalog 接続

mlflow 2.17.1 以降が必要です

ユーザーに代わって承認​

備考

プレビュー

この機能は パブリック プレビュー段階です。

ユーザー代理承認を使用するエージェントをログに記録する場合は、エージェントがエンド ユーザーに代わって動作するために必要な Databricks REST API スコープの最小セットを宣言する必要があります。

これにより、エージェントは最小権限の原則に従うことができます: トークン エージェントが必要とする APIs のみに制限され、不正なアクションやトークンの悪用の可能性が減ります。

以下は、いくつかの一般的な種類の Databricks リソースにアクセスするために必要なスコープの一覧です。

リソースタイプ

必要な API スコープ

モデルサービングエンドポイント

serving.serving-endpoints

ベクトル検索エンドポイント

vectorsearch.vector-search-endpoints

ベクトル検索インデックス

vectorsearch.vector-search-indexes

SQLウェアハウス

sql.warehouses

UC 接続

catalog.connections そして serving.serving-endpoints

MCP Genie spaces

mcp.genie

MCP UC 関数

mcp.functions

MCP ベクトル検索

mcp.vectorsearch

MCPの外部機能

mcp.external

「 ユーザー代理認証を使用したエージェントの展開」を参照してください。

OpenAIクライアントの自動認証​

エージェントが OpenAI クライアントを使用している場合は、Databricks SDK を使用してデプロイ中に自動的に認証します。Databricks SDK には、承認が自動的に構成された OpenAI クライアントを構築するためのラッパーが用意されています ( get_open_ai_client())。ノートブックで次のコマンドを実行します。

Python
% pip install databricks-sdk[openai]
Python
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 のアクセス許可を使用できます。

Python
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()を参照してください。

次のステップ​

に最終更新