共用方式為


AI 代理程式的驗證

AI 代理程式通常需要向其他資源進行驗證,才能完成工作。 例如,已部署的代理程式可能需要存取向量搜尋索引來查詢非結構化資料,或存取提示登錄來載入動態提示。

本頁介紹了使用 Mosaic AI 代理框架開發和部署代理時可用的身份驗證方法。

身份驗證方法

下表比較了可用的驗證方法。 您可以混合搭配下列任何一種方法:

方法 Description 安全性狀態 設定複雜度
自動驗證通過 代理程式會以部署代理程式之使用者的許可權執行
Databricks 會自動管理宣告資源的短期認證
短期憑證、自動輪換 低 - 在記錄時指定相依項目
以他人身份進行的使用者驗證 (OBO) 代理程式會以提出請求的終端使用者權限來執行 使用具有受限範圍的一般使用者認證 中 - 需要範圍宣告和執行階段初始化
手動驗證 使用環境變數明確提供認證 長期憑證需要輪替管理 高 - 需要手動認證管理

向外部系統和 MCP 伺服器進行驗證

如需如何從代理程式向外部 API 和 MCP 伺服器進行驗證的指引,請參閱將 AI 代理程式工具連線到外部服務。 也可以代表代理程式或使用者查詢這些資源,如 驗證方法中所述。

為您的資源選擇正確的驗證方法

使用此流程圖可為每一個資源選擇正確的鑑別方法。 您可以視需要組合方法,代理程式可以根據其使用案例,針對每個資源使用不同的方法。

  1. 是否需要每個使用者存取控制或使用者屬性稽核?

  2. 所有資源都 支援自動驗證嗎?

自動驗證通過

自動驗證傳遞是存取 Databricks 管理資源的最簡單方法。 在記錄代理程式時宣告資源相依性,Databricks 會在部署代理程式時自動佈建、輪替和管理短期認證。

此驗證行為類似於 Databricks 儀錶板的「以擁有者身分執行」行為。 Unity 目錄資料表等下游資源是使用服務主體的認證來存取,其最低許可權存取權只適用於代理程式所需的資源。

自動驗證傳遞的運作方式

當使用自動驗證傳遞在端點後方執行代理時,Databricks 會執行下列步驟:

  1. 許可權驗證:Databricks 會驗證端點建立者是否可以存取代理程式記錄期間指定的所有相依性。

  2. 服務主體的建立和授權:會為代理模型版本建立一個服務主體,並自動授予對代理資源的讀取許可權。

    備註

    系統產生的服務主體不會出現在 API 或 UI 清單中。 如果代理程式模型版本已從端點移除,服務主體也會一併刪除。

  3. 認證布建和輪替:會將短期認證(M2M OAuth 令牌)注入端點,以允許服務主體的代理代碼存取 Databricks 資源。 Databricks 也會輪替憑證,確保代理程式能夠持續且安全地存取相依資源。

支援自動驗證通過的資源

下表列出支援自動驗證傳遞的 Databricks 資源,以及部署代理程式時端點建立者必須具備的許可權。

備註

Unity Catalog 資源在父架構上需要 USE SCHEMA,在父目錄上需要 USE CATALOG

資源類型 權限 最低 MLflow 版本
SQL 資料倉儲 Use Endpoint 2.16.1 或以上
模型服務端點 Can Query 2.13.1 或以上
Unity 目錄功能 EXECUTE 2.16.1 或以上
精靈空間 Can Run 2.17.1 或以上
向量搜尋索引 Can Use 2.13.1 或以上
Unity 目錄數據表 SELECT 2.18.0 或以上
Unity 目錄連線 Use Connection 2.17.1 或以上
Lakebase databricks_superuser 3.3.2或以上

實作自動驗證傳遞

若要啟用自動驗證傳遞,請在 記錄代理程式時指定相依資源。 使用 resources API 的 log_model() 參數:

備註

請記得也記錄所有下游相依資源。 例如,如果您記錄 Genie 空間,也必須記錄其資料表、SQL 倉儲和 Unity 目錄函式。

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 資源:

Databricks 資源 相容的用戶端
向量搜尋索引 databricks_langchain.VectorSearchRetrieverTooldatabricks_openai.VectorSearchRetrieverToolVectorSearchClient
模型服務端點 databricks.sdk.WorkspaceClient
SQL 資料倉儲 databricks.sdk.WorkspaceClient
UC 連線 databricks.sdk.WorkspaceClient
UC 表格和函數 databricks.sdk.WorkspaceClient (若要存取 UC 表,您必須使用 SQL 陳述式執行 API 來使用 SQL 查詢)
精靈空間 databricks.sdk.WorkspaceClient (建議)、 databricks_langchain.GenieAgentdatabricks_ai_bridge.GenieAgent
模型內容通訊協定 (MCP) databricks_mcp.DatabricksMCPClient

實作 OBO 驗證

若要啟用代表使用者驗證,請完成下列步驟:

  1. 更新 SDK 呼叫,以指定代表一般使用者存取資源。
  2. 更新代理程式程式碼以在函數內 predict 初始化 OBO 存取,而不是在 __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"},
)

精靈空間 (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")

精靈空間(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
  )

在 predict 函數中初始化代理程式

因為使用者的身分只有在查詢時才知道,所以您必須在 或 predictpredict_stream存取 OBO 資源,而不是在代理程式的方法__init__中。 這可確保資源在調用之間隔離。

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 資源所需的範圍清單:

資源類型 必要的 API 範圍
模型服務端點 serving.serving-endpoints
向量搜尋端點 vectorsearch.vector-search-endpoints
向量搜尋索引 vectorsearch.vector-search-indexes
SQL 數據倉庫 sql.warehousessql.statement-execution
精靈空間 dashboards.genie
UC 連線 catalog.connectionsserving.serving-endpoints
Databricks 應用程式 apps.apps
MCP 精靈空間 mcp.genie
MCP UC 函數 mcp.functions
MCP 向量搜尋 mcp.vectorsearch
MCP DBSQL mcp.sqlsql.warehousessql.statement-execution
MCP 外部功能 mcp.external

若要啟用受使用者委託的驗證,請將 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,但端點可能允許比代理程式明確要求的更多動作。 例如,API 範圍會 serving.serving-endpoints 授與代理程序許可權,以代表使用者執行服務端點。 不過,服務端點可以存取原始代理程式未獲授權使用的其他 API 範圍。

OBO 範例筆記本

以下筆記本說明如何使用代替使用者授權,透過向量搜索創建代理。

拿筆記本

下列筆記本示範如何使用代表使用者授權,建立支援在 SQL 倉儲上執行 SQL 的代理程式。 這可讓代理程式使用使用者認證安全地叫用 Unity 目錄函式。 附註: 這是目前使用 OBO 執行 UC 函數的建議方式,因為尚不支援使用 OBO 執行無伺服器 Spark。

代表使用者進行授權及執行 SQL

拿筆記本

手動驗證

手動驗證允許在代理程式部署期間明確指定認證。 此方法具有最大的靈活性,但需要更多的設定和持續的憑證管理。 在下列情況下使用此方法:

  • 相依資源不支援自動身份驗證通過
  • 代理程式需要使用非代理程式部署者的其他認證資料。
  • 代理程式會存取 Databricks 外部的外部資源或 API
  • 已部署的代理程式會存取 提示註冊表

這很重要

覆寫安全環境變數會停用代理程式所依賴的其他資源的自動傳遞。

OAuth 是手動驗證的推薦方法,因為它提供了安全、以權杖為基礎的驗證,適用於服務主體並具有自動權杖重新整理功能:

  1. 建立服務主體產生 OAuth 認證

  2. 將服務主體的許可權授予代理程式能存取的任何 Databricks 資源,以便代理程式存取 Databricks 資源。 若要存取提示註冊表,請授予 CREATE FUNCTIONEXECUTEMANAGE Unity Catalog 結構的許可權來儲存提示。

  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 資源,以便代理程式存取 Databricks 資源。 若要存取提示登錄,請授與 CREATE FUNCTIONEXECUTEMANAGE 用來儲存提示之 Unity 目錄架構的許可權。
    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}}}}}"
        },
    )