使用代表用户身份验证部署代理

通过代表用户身份认证，通过 Mosaic AI 模型服务部署的代理可以使用查询该代理的 Databricks 最终用户的身份访问 Databricks 资源。 这样就可以按用户访问敏感信息，并在 Unity Catalog 中精细实现数据访问控制。

代表用户身份验证通过缩小范围进一步限制传入的用户令牌，确保公开给代理代码的令牌仅限于访问代理创作者定义的特定 API。 这通过防止未经授权的作并减少令牌滥用的风险来提高安全性。

创作代理时，可以继续使用现有的 SDK 来访问 Databricks 资源，例如矢量搜索索引。 若要对资源启用代表用户访问，请执行以下步骤：

1. 在代理代码中，更新 SDK 调用以指示应代表代理最终用户访问资源
2. 在代理日志记录时间（在代理部署之前），指定代理所需的最终用户 REST API 范围。 有关详细信息，请参阅代表用户身份验证

有关代表用户身份验证的端到端示例，请参阅 端到端示例。

以下资源与代理的代表用户身份验证兼容。

使用以用户身份客户端初始化工具时，可以将工具初始化包装在 try-except 块中，或者让错误直接暴露给用户。 通过处理错误，即使最终用户无法访问所需的所有工具，代理仍然可以做出最佳响应。 但是，如果选择在工具初始化期间不处理错误，如果用户缺少任何所需的资源，代理将引发错误。

配置 SDK

以下代码片段演示如何使用各种 SDK 配置代表用户访问不同的 Databricks 资源

矢量搜索检索器工具

from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider 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 you want the agent to 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 authenticated 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 you want the agent to 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)

模型服务终结点

from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider 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 you want the agent to 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

# 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 空间

from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider 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
  )

初始化代理

代表用户身份验证与 ChatAgent 接口兼容。 使用代表用户身份验证时，仅当查询已部署的代理时，最终用户的标识才已知，即在 ChatAgent 接口的 predict 和 predict_stream 函数中。 因此，必须从这些方法中代表用户执行对资源的任何访问（例如，列出最终用户有权访问的矢量搜索索引），而不是在 ChatAgent 实现的 __init__ 方法中执行。 这可以确保资源在调用之间隔离

from mlflow.pyfunc import ChatAgent


class LangGraphChatAgent(ChatAgent):
  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 authorization

  def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
  ) -> ChatAgentResponse:
    agent = initialize_agent() # Initialize the Agent in Predict
    request = {"messages": self._convert_messages_to_dict(messages)}

    messages = []
    for event in self.agent.stream(request, stream_mode="updates"):
      for node_data in event.values():
        messages.extend(
          ChatAgentMessage(**msg) for msg in node_data.get("messages", [])
        )
    return ChatAgentResponse(messages=messages)

安全注意事项

在启用代理代表用户身份验证之前，需要考虑一些安全问题：

1. 访问敏感的 Databricks 资源：启用代表用户身份验证会允许代理访问敏感的 Databricks 资源。 虽然我们已经实现了 API 范围来限制开发人员可以访问和降低令牌滥用风险的资源，但仍存在一些风险。 例如， serving.serving-endpoints API 范围授予代理权限，以代表用户执行服务终结点。 但是，服务终结点本身可能有权访问原始代理无权使用的其他 API 范围。
2. 不支持最终用户同意：在当前 beta 阶段，代理用户无法查看或同意代理所需的 Databricks REST API 范围。 用户负责确保他们信任服务端点上具有“可管理”权限的人员，在 Databricks 中代表他们采取行动。

端到端示例

以下笔记本演示如何通过代替用户身份验证创建具备矢量搜索功能的代理程序。

使用矢量搜索的代表用户身份验证

获取笔记本