AI 代理身份验证

AI 代理通常需要向其他资源进行身份验证才能完成任务。 例如，部署的代理可能需要访问矢量搜索索引来查询非结构化数据或提示注册表以加载动态提示。

本页介绍使用马赛克 AI 代理框架开发和部署代理时可用的身份验证方法。

身份验证方法

下表比较了可用的身份验证方法。 可以混合和匹配以下任一方法：

为资源选择正确的身份验证方法

使用此流程图为每个资源选择正确的身份验证方法。 可以根据需要组合方法，代理可以为每个资源使用不同的方法，具体取决于其用例。

1. 是否需要按用户访问控制或用户属性审核？

   • 是→使用 代表用户的身份验证
   • 没有→继续执行步骤 2

2. 所有资源是否 支持自动身份验证？

   • 是→请使用 自动身份验证透传 （建议）
   • 无→使用 手动身份验证

自动身份验证直通

自动认证透传是访问 Databricks 托管资源的最简单方法。 记录代理时声明资源依赖项，在部署代理时，Databricks 会自动预配、轮换和管理生存期较短的凭据。

此身份验证行为类似于 Databricks 仪表板的“以所有者身份运行”行为。 使用服务主体的凭据访问下游资源（如 Unity Catalog 表），仅赋予代理所需的资源最小权限访问。

自动身份验证透传的工作原理

使用自动身份验证透传在终结点后部署代理时，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"),
    ]
  )

代表用户身份验证

代理用户 （OBO） 身份验证允许代理充当运行查询的 Databricks 用户。 这提供：

• 对敏感数据的按用户访问
• Unity 目录强制执行的细粒度数据控件
• 安全令牌仅限代理声明的 API（“缩小范围”），从而减少滥用风险

要求

• 代表用户身份验证需要 MLflow 2.22.1 及更高版本。
• 默认情况下禁用代表用户身份验证，并且必须由工作区管理员启用。启用此功能之前，请查看 安全注意事项 。

OBO 支持的资源

使用 OBO 身份验证的代理可以访问以下 Databricks 资源：

实现 OBO 身份验证

若要启用代表用户身份验证，请完成以下步骤：

1. 更新 SDK 调用以指定代表最终用户访问资源。
2. 更新代理代码以初始化函数内的 OBO 访问，而不是在函数中predict__init__，因为用户标识仅在运行时才已知。
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
  )

在预测函数中初始化代理

由于用户的标识仅在查询时已知，因此必须在代理的方法中访问 OBO 资源， 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 的一部分，以便在部署时自动进行身份验证。

OBO 安全注意事项

在向代理启用代表用户身份验证之前，请考虑以下安全注意事项。

扩展的资源访问：代理可以代表用户访问敏感资源。 虽然作用域限制 API，但终结点可能允许比代理显式请求更多的作。 例如， serving.serving-endpoints API 范围授予代理权限，以代表用户运行服务终结点。 但是，服务终结点可以访问原始代理无权使用的其他 API 范围。

OBO 示例笔记本

以下笔记本演示如何使用代表用户授权创建具有矢量搜索的代理。

代表使用矢量搜索的用户授权

获取笔记本

以下笔记本演示如何使用代表用户授权在 SQL Warehouse 上创建支持 SQL 执行的代理。 这使代理能够使用用户凭据安全地调用 Unity 目录函数。 注意：目前建议使用 OBO 执行 UC 函数，因为尚不支持使用 OBO 执行无服务器 Spark。

代表用户授权的 SQL 执行

获取笔记本

手动身份验证

手动身份验证允许在代理部署期间显式指定凭据。 此方法具有最高的灵活性，但需要更多的设置和持续进行的凭据管理。 在以下情况下使用此方法：

• 依赖资源不支持自动身份验证传递
• 代理需要使用代理部署程序以外的其他凭据
• 代理访问 Databricks 之外的资源或 API
• 已部署的代理程序访问 提示注册表

OAuth 身份验证（建议）

OAuth 是手动身份验证的推荐方法，因为它对具有自动令牌刷新功能的服务主体具有安全的基于令牌的身份验证：

1. 创建服务主体 并 生成 OAuth 凭据。

2. 向服务主体授予访问任何代理有权访问的 Databricks 资源的权限。 若要访问提示注册表，请在 Unity 目录架构上授予CREATE FUNCTIONEXECUTE和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 FUNCTIONEXECUTE和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}}}}}"
       },
   )