Implantar um agente usando a autenticação no local do usuário

info

Beta

Esse recurso está na versão beta.

atenção

Embora a autenticação em nome do usuário seja uma ferramenta poderosa para garantir o acesso seguro a dados confidenciais. Ele permite que os usuários do workspace criem agentes que atuam em nome de outros usuários no Databricks. Durante a versão beta, ele é desativado por default e deve ser ativado por um administrador de workspace. Analise as considerações de segurança da autenticação no local do usuário antes de ativar esse recurso.

Com a autenticação no local do usuário, os agentes implantados via Mosaic AI servindo modelo podem acessar Databricks recurso usando a identidade do Databricks usuário final que consultou o agente. Isso permite o acesso a informações confidenciais por usuário, com aplicação refinada do controle de acesso a dados em Unity Catalog.

A autenticação no local do usuário restringe ainda mais os tokens de usuário recebidos por meio de downscoping, garantindo que os tokens expostos ao código do agente sejam limitados a acessar apenas o APIs específico definido pelo autor do agente. Isso aumenta a segurança, evitando ações não autorizadas e reduzindo o risco de uso indevido de tokens.

Ao criar seu agente, o senhor pode continuar a usar os SDKs existentes para acessar Databricks recurso, como índices de pesquisa vetorial. Para habilitar o acesso do usuário ao recurso:

1. No código do agente, atualize as chamadas SDK para indicar que o recurso deve ser acessado em nome do usuário final do agente
2. No momento do registro do agente (antes da implementação do agente), especifique os escopos da API REST do usuário final exigidos pelo seu agente. Para obter mais detalhes, consulte Autenticação em nome do usuário

Para ver um exemplo de ponta a ponta de autenticação em nome do usuário, consulte Exemplo de ponta a ponta.

Os recursos a seguir são compatíveis com a autenticação no local do usuário com agentes.

Databricks recurso	Clientes compatíveis
Índice de pesquisa vetorial	databricks_langchain.VectorSearchRetrieverTool, databricks_openai.VectorSearchRetrieverTool ou VectorSearchClient
Endpoint do serviço de modelos	databricks.sdk.WorkspaceClient
SQL Warehouse	databricks.sdk.WorkspaceClient
Conexões UC	databricks.sdk.WorkspaceClient
Tabelas e funções de UC	Atualmente, a Databricks não oferece suporte a clientes diretos para acessar tabelas UC ou funções UC com autenticação no local do usuário. Em vez disso, incentivamos os usuários a usar o Genie para acessar dados estruturados com autenticação no local do usuário
Espaço Genie	databricks_langchain.GenieAgent ou databricks_openai.GenieAgent

Ao inicializar ferramentas com o cliente em nome do usuário, você pode agrupar a inicialização da ferramenta em um bloco try-except ou permitir que os erros sejam expostos aos usuários. Ao lidar com os erros, o agente ainda pode dar a melhor resposta, mesmo que o usuário final não tenha acesso a todas as ferramentas necessárias. No entanto, se o senhor optar por não tratar os erros durante a inicialização da ferramenta, o agente emitirá um erro se o usuário não tiver algum recurso necessário.

Configurar SDKs

Os snippets abaixo demonstram como configurar o acesso do usuário a diferentes Databricks recurso usando vários SDKs

Vector Search Retriever Tool

Python

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)

Vector Search Client

Python

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)

Model Serving Endpoint

Python

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 Connections

Python

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 Spaces

Python

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
  )

Inicializando o agente

a autenticação em nome do usuário é compatível com a interface do ChatAgent. Ao usar a autenticação no local do usuário, a identidade do usuário final só é conhecida quando o agente implantado é consultado, ou seja, nas funções predict e predict_stream da interface do ChatAgent. Como resultado, o senhor deve executar qualquer acesso do usuário ao recurso em nome do usuário (por exemplo liste os índices de pesquisa vetorial aos quais o usuário final tem acesso) a partir desses métodos, em vez do método __init__ da sua implementação do ChatAgent. Isso garante que os recursos sejam isolados entre as invocações

Python

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)

Considerações de segurança

Há algumas implicações de segurança a serem consideradas antes de ativar a autenticação em nome do usuário com agentes:

1. Acesso a recursos confidenciais Databricks: a ativação da autenticação na metade do usuário permite que os agentes acessem recursos confidenciais Databricks. Embora tenhamos implementado os escopos do API para restringir os recursos que os desenvolvedores podem acessar e reduzir o risco de uso indevido de tokens, alguns riscos ainda permanecem. Por exemplo, o escopo da API serving.serving-endpoints concede permissão a um agente para executar um endpoint de serviço em nome do usuário. No entanto, o próprio endpoint de atendimento pode ter acesso a escopos de API adicionais que o agente original não está autorizado a usar.
2. Não há suporte para o consentimento do usuário final: Durante a fase beta atual, os usuários do agente não podem view ou consentir com os escopos Databricks REST API exigidos por um agente. Os usuários são responsáveis por garantir que confiam nas pessoas com permissões "Can gerenciar" no servidor endpoint para realizar ações em Databricks em seu nome.

Exemplo de ponta a ponta

O Notebook a seguir mostra como criar um agente com pesquisa vetorial usando a autenticação no local do usuário.

Em nome da autenticação do usuário com pesquisa vetorial

Open notebook in new tab