AI エージェントの認証

AI エージェントでは、タスクを完了するために他のリソースに対して認証を行う必要があることがよくあります。 たとえば、デプロイされたエージェントは、ベクター検索インデックスにアクセスして非構造化データを照会したり、プロンプト レジストリにアクセスして動的プロンプトを読み込んだりする必要がある場合があります。

このページでは、Mosaic AI Agent Framework を使用してエージェントを開発およびデプロイするときに使用できる認証方法について説明します。

認証方法

次の表は、使用可能な認証方法を比較しています。 次のいずれかの方法を組み合わせることができます。

外部システムと MCP サーバーに対する認証

エージェントから外部 API と MCP サーバーに対して認証する方法のガイダンスについては、「 AI エージェント ツールを外部サービスに接続する」を参照してください。 これらのリソースは、「 認証方法」で説明されているように、エージェントまたはユーザーに代わってクエリを実行することもできます。

リソースに適した認証方法を選択する

このフローチャートを使用して、各リソースに適した認証方法を選択します。 必要に応じてメソッドを組み合わせることができます。エージェントは、そのユース ケースに応じてリソースごとに異なる方法を使用できます。

1. ユーザーごとのアクセス制御またはユーザー属性監査は必要ですか?

   • はい →ユーザーの代理認証を使用する
   • →なし 手順 2 に進む

2. すべてのリソースが 自動認証をサポートしていますか?

   • はい → 自動認証パススルー を使用する (推奨)
   • 手動認証→使用しない

自動認証パススルー

自動認証パススルーは、Databricks マネージド リソースにアクセスするための最も簡単な方法です。 エージェントのログ記録時にリソースの依存関係を宣言し、Databricks はエージェントのデプロイ時に有効期間の短い資格情報を自動的にプロビジョニング、ローテーション、管理します。

この認証動作は、Databricks ダッシュボードの "所有者として実行" 動作に似ています。 Unity カタログ テーブルなどのダウンストリーム リソースには、エージェントが必要とするリソースのみに最小限の特権でアクセスできるサービス プリンシパルの資格情報を使用してアクセスされます。

自動認証パススルーのしくみ

エージェントが自動認証パススルーを使用してエンドポイントの背後で処理される場合、Databricks は次の手順を実行します。

1. アクセス許可の検証: Databricks は、エンドポイント作成者がエージェントのログ記録中に指定されたすべての依存関係にアクセスできることを確認します。

2. サービス プリンシパルの作成と許可: サービス プリンシパルがエージェント モデルのバージョンに対して作成されて、エージェント リソースへの読み取りアクセスを自動的に許可されます。

   注

   システムによって生成されたサービス プリンシパルは、API または UI の一覧には表示されません。 エージェント モデルのバージョンがエンドポイントから削除されると、サービス プリンシパルも削除されます。

3. 資格情報のプロビジョニングとローテーション: サービス プリンシパルの有効期間が短い資格情報 (M2M OAuth トークン) がエンドポイントに挿入されて、エージェントのコードが Databricks のリソースにアクセスできるようになります。 また、Databricks によって資格情報のローテーションが行われ、エージェントは依存リソースへのセキュリティ保護されたアクセスを確実に継続できます。

自動認証パススルーでサポートされているリソース

次の表は、自動認証パススルーをサポートしている Databricks リソースと、エージェントのデプロイ時にエンドポイント作成者が持っている必要のあるアクセス許可の一覧です。

自動認証パススルーを実装する

自動認証パススルーを有効にするには、 エージェントをログに記録するときに依存リソースを指定します。 resources API の log_model() パラメーターを使用します。

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

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model="agent.py",
    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"),
      DatabricksLakebase(database_instance_name="lakebase_instance_name"),
    ]
  )

ユーザーの代理認証

On-Behalf-of-User (OBO) 認証を使用すると、エージェントはクエリを実行する Databricks ユーザーとして機能できます。 これにより、次の機能が提供されます。

• 機密データへのユーザーごとのアクセス
• Unity カタログによって適用されるきめ細かいデータ コントロール
• セキュリティ トークンは、エージェントが宣言する API のみに制限 ("ダウンスコープ") され、誤用のリスクが軽減されます

Requirements

• ユーザーの代理認証には、MLflow 2.22.1 以降が必要です。
• ユーザーの代理認証は既定で無効になっており、ワークスペース管理者が有効にする必要があります。この機能を有効にする前に 、セキュリティに関する考慮事項 を確認してください。

OBO でサポートされているリソース

OBO 認証を使用するエージェントは、次の Databricks リソースにアクセスできます。

OBO 認証を実装する

ユーザーの代理認証を有効にするには、次の手順を実行します。

1. SDK 呼び出しを更新して、エンド ユーザーの代わりにリソースにアクセスすることを指定します。
2. ユーザー ID は実行時にのみ認識されるため、predictではなく、__init__関数内で OBO アクセスを初期化するようにエージェント コードを更新します。
3. エージェントをデプロイ用にログに記録するときは、エージェントに必要な Databricks REST API スコープを宣言します。

次のスニペットは、さまざまな Databricks リソースへのユーザーの代理アクセスを構成する方法を示しています。 ツールを初期化するときは、 try-except ブロックで初期化をラップすることで、アクセス許可エラーを適切に処理します。

ベクター検索リトリーバーツール

from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials
from databricks_langchain import VectorSearchRetrieverTool

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy = ModelServingUserCredentials())

vector_search_tools = []
# Exclude exception handling if the agent should fail
# when users lack access to all required Databricks resources
try:
  tool = VectorSearchRetrieverTool(
    index_name="<index_name>",
    description="...",
    tool_name="...",
    workspace_client=user_client # Specify the user authorized client
    )
    vector_search_tools.append(tool)
except Exception as e:
    _logger.debug("Skipping adding tool as user does not have permissions")

ベクター検索クライアント

from databricks.vector_search.client import VectorSearchClient
from databricks.vector_search.utils import CredentialStrategy

# Configure a VectorSearch Client to use on behalf of end
# user authentication
user_authenticated_vsc = VectorSearchClient(credential_strategy=CredentialStrategy.MODEL_SERVING_USER_CREDENTIALS)
# Exclude exception handling if the agent should fail when
# users lack access to all required Databricks resources
try:
  vs_index = user_authenticated_vsc.get_index(endpoint_name="endpoint_name", index_name="index_name")
  ...
except Exception as e:
  _logger.debug("Skipping Vector Index because user does not have permissions")

MCP

from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials
from databricks_mcp import DatabricksMCPClient

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

mcp_client = DatabricksMCPClient(
    server_url="<mcp_server_url>",
    workspace_client=user_client, # Specify the user client here
  )

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

from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

# Exclude exception handling if the agent should fail
# when users lack access to all required Databricks resources
try:
  user_client.serving_endpoints.query("endpoint_name", input="")
except Exception as e:
  _logger.debug("Skipping Model Serving Endpoint due to no permissions")

UC 接続

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
from databricks_ai_bridge import ModelServingUserCredentials

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

user_client.serving_endpoints.http_request(
  conn="connection_name",
  method=ExternalFunctionRequestHttpMethod.POST,
  path="/api/v1/resource",
  json={"key": "value"},
  headers={"extra_header_key": "extra_header_value"},
)

Genie Spaces (WorkspaceClient)

from databricks_langchain.genie import GenieAgent
from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials


# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())


genie_agent = GenieAgent(
    genie_space_id="space-id",
    genie_agent_name="Genie",
    description="This Genie space has access to sales data in Europe"
    client=user_client
)

# Use the Genie SDK methods available through WorkspaceClient
try:
    response = agent.invoke("Your query here")
except Exception as e:
    _logger.debug("Skipping Genie due to no permissions")

Genie Spaces (LangChain)

from databricks.sdk import WorkspaceClient
from databricks_ai_bridge import ModelServingUserCredentials
from databricks_langchain.genie import GenieAgent

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

genie_agent = GenieAgent(
    genie_space_id="<genie_space_id>",
    genie_agent_name="Genie",
    description="Genie_description",
    client=user_client, # Specify the user client here
  )

予測関数でエージェントを初期化する

ユーザーの ID はクエリ時にのみ認識されるため、エージェントのpredictメソッドではなく、predict_streamまたは__init__内の OBO リソースにアクセスする必要があります。 これにより、リソースが呼び出し間で分離されます。

from mlflow.pyfunc import ResponsesAgent

class OBOResponsesAgent(ResponsesAgent):
  def initialize_agent():
    user_client = WorkspaceClient(
      credentials_strategy=ModelServingUserCredentials()
    )
    system_authorized_client = WorkspaceClient()
    ### Use the clients above to access resources with either system or user authentication

  def predict(
    self, request
  ) -> ResponsesAgentResponse:
    agent = initialize_agent() # Initialize the Agent in Predict

    agent.predict(request)
    ...

エージェントのログ記録時に REST API スコープを宣言する

OBO エージェントをデプロイ用にログに記録するときは、エージェントがユーザーの代わりに呼び出す Databricks REST API スコープを一覧表示する必要があります。 これにより、エージェントは最小限の特権の原則に従います。トークンは、エージェントが必要とする API のみに制限され、承認されていないアクションやトークンの誤用の可能性が減ります。

いくつかの一般的な種類の Databricks リソースにアクセスするために必要なスコープの一覧を次に示します。

ユーザーの代理認証を有効にするには、MLflow AuthPolicy を log_model()に渡します。

import mlflow
from mlflow.models.auth_policy import AuthPolicy, SystemAuthPolicy, UserAuthPolicy
from mlflow.models.resources import DatabricksServingEndpoint

# System policy: resources accessed with system credentials
system_policy = SystemAuthPolicy(
    resources=[DatabricksServingEndpoint(endpoint_name="my_endpoint")]
)

# User policy: API scopes for OBO access
user_policy = UserAuthPolicy(api_scopes=[
    "serving.serving-endpoints",
    "vectorsearch.vector-search-endpoints",
    "vectorsearch.vector-search-indexes"
])

# Log the agent with both policies
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        name="agent",
        python_model="agent.py",
        auth_policy=AuthPolicy(
            system_auth_policy=system_policy,
            user_auth_policy=user_policy
        )
    )

OpenAI クライアントの OBO 認証

OpenAI クライアントを使用するエージェントの場合は、デプロイ時に Databricks SDK を使用して自動的に認証します。 Databricks SDK には、認証が自動的に構成された OpenAI クライアントを構築するためのラッパー get_open_ai_client()。

% 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 の一部として Model Serving エンドポイントを指定します。

OBO のセキュリティに関する考慮事項

エージェントでユーザーの代理認証を有効にする前に、次のセキュリティに関する考慮事項を検討してください。

リソース アクセスの拡張: エージェントは、ユーザーに代わって機密性の高いリソースにアクセスできます。 スコープによって API が制限されますが、エンドポイントでは、エージェントが明示的に要求するよりも多くのアクションが許可される場合があります。 たとえば、 serving.serving-endpoints API スコープは、ユーザーに代わってサービス エンドポイントを実行するアクセス許可をエージェントに付与します。 ただし、サービス エンドポイントは、元のエージェントが使用を許可されていない追加の API スコープにアクセスできます。

OBO ノートブックの例

次のノートブックは、ユーザーの代理承認を使用して Vector Search を使用してエージェントを作成する方法を示しています。

ベクター検索を使用したユーザー承認の代理

ノートブックを入手

次のノートブックは、ユーザーの代理承認を使用して SQL Warehouse での SQL 実行をサポートするエージェントを作成する方法を示しています。 これにより、エージェントはユーザー資格情報を使用して Unity カタログ関数を安全に呼び出すことができます。 注: OBO を使用したサーバーレス Spark 実行はまだサポートされていないため、これは現在、OBO で UC 関数を実行するための推奨される方法です。

ユーザー承認を代理して行う SQL 実行

ノートブックを入手

手動認証

手動認証を使用すると、エージェントのデプロイ時に資格情報を明示的に指定できます。 この方法は柔軟性が最も高いですが、より多くのセットアップと継続的な資格情報管理が必要です。 次の場合は、このメソッドを使用します。

• 依存リソースは、自動認証パススルーをサポートしていません
• エージェントは、エージェント デプロイツール以外の資格情報を使用する必要があります
• エージェントが Databricks の外部の外部リソースまたは API にアクセスする
• デプロイされたエージェントがプロンプト レジストリにアクセスする

OAuth 認証 (推奨)

OAuth は、自動トークン更新機能を備えたサービス プリンシパルに対してセキュリティで保護されたトークンベースの認証を備えるため、手動認証に推奨される方法です。

1. サービス プリンシパルを作成 し、 OAuth 資格情報を生成します。

2. サービス プリンシパルに、エージェントがアクセスできる任意の Databricks リソースへのアクセス許可を付与します。 プロンプト レジストリにアクセスするには、プロンプトを格納するために Unity カタログ スキーマに対する CREATE FUNCTION、 EXECUTE、および MANAGE のアクセス許可を付与します。

3. OAuth 資格情報を使用してDatabricks シークレットを作成します。

4. エージェント コードで OAuth 資格情報を構成します。

   import os
   
   # Configure OAuth authentication for Prompt Registry access
   # Replace with actual secret scope and key names
   secret_scope_name = "your-secret-scope"
   client_id_key = "oauth-client-id"
   client_secret_key = "oauth-client-secret"
   
   os.environ["DATABRICKS_HOST"] = "https://<your-workspace-url>"
   os.environ["DATABRICKS_CLIENT_ID"] = dbutils.secrets.get(scope=secret_scope_name, key=client_id_key)
   os.environ["DATABRICKS_CLIENT_SECRET"] = dbutils.secrets.get(scope=secret_scope_name, key=client_secret_key)
   

5. シークレットを使用してワークスペースに接続します。

   w = WorkspaceClient(
     host=os.environ["DATABRICKS_HOST"],
     client_id=os.environ["DATABRICKS_CLIENT_ID"],
     client_secret = os.environ["DATABRICKS_CLIENT_SECRET"]
   )
   

6. agents.deploy()を使用してデプロイする場合は、環境変数として OAuth 資格情報を含めます。

   agents.deploy(
       UC_MODEL_NAME,
       uc_registered_model_info.version,
       environment_vars={
           "DATABRICKS_HOST": "https://<your-workspace-url>",
           "DATABRICKS_CLIENT_ID": f"{{{{secrets/{secret_scope_name}/{client_id_key}}}}}",
           "DATABRICKS_CLIENT_SECRET": f"{{{{secrets/{secret_scope_name}/{client_secret_key}}}}}"
       },
   )
   

PAT 認証

個人用アクセス トークン (PAT) 認証を使用すると、開発環境とテスト環境のセットアップが簡単になりますが、より多くの手動資格情報管理が必要になります。

1. サービス プリンシパルまたは個人アカウントを使用して PAT を取得します。

   サービス プリンシパル (セキュリティ対策として推奨):

   1. サービス プリンシパルを作成します。
   2. サービス プリンシパルに、エージェントがアクセスできる任意の Databricks リソースへのアクセス許可を付与します。 プロンプト レジストリにアクセスするには、プロンプトの格納に使用する Unity カタログ スキーマに対する CREATE FUNCTION、 EXECUTE、および MANAGE のアクセス許可を付与します。
   3. サービス プリンシパルの PAT を作成します。

   個人用アカウント:

   1. 個人用アカウントの PAT を作成します。

2. PAT の Databricks シークレットを作成して、 PAT を安全に格納します。

3. エージェント コードで PAT 認証を構成します。

   import os
   
   # Configure PAT authentication for Prompt Registry access
   # Replace with your actual secret scope and key names
   secret_scope_name = "your-secret-scope"
   secret_key_name = "your-pat-key"
   
   os.environ["DATABRICKS_HOST"] = "https://<your-workspace-url>"
   os.environ["DATABRICKS_TOKEN"] = dbutils.secrets.get(scope=secret_scope_name, key=secret_key_name)
   
   # Validate configuration
   assert os.environ["DATABRICKS_HOST"], "DATABRICKS_HOST must be set"
   assert os.environ["DATABRICKS_TOKEN"], "DATABRICKS_TOKEN must be set"
   

4. agents.deploy()を使用してエージェントをデプロイする場合は、環境変数として PAT を含めます。

   agents.deploy(
       UC_MODEL_NAME,
       uc_registered_model_info.version,
       environment_vars={
           "DATABRICKS_HOST": "https://<your-workspace-url>",
           "DATABRICKS_TOKEN": f"{{{{secrets/{secret_scope_name}/{secret_key_name}}}}}"
       },
   )