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

info

Beta

Esse recurso está na versão beta.

atenção

A autenticação em nome do usuário está desativada por default e deve ser ativada por um administrador de workspace. Analise as considerações de segurança para autenticação em nome do usuário antes de habilitar este 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 de “ SDK ” para indicar que os recursos são acessados 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.

Requisitos

dica

A Databricks recomenda instalar a versão mais recente do cliente MLflow Python ao desenvolver agentes.

A autorização do usuário requer mlflow 2.22.1 ou superior.

Recurso compatível

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 optar por não tratar erros durante a inicialização da ferramenta, o agente apresentará 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_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 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_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 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
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

Python

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
  )

Inicializando o agente

A autenticação em nome do usuário é compatível com a interface ResponsesAgent.

Ao utilizar a autenticação em nome do usuário, a identidade do usuário final só é conhecida quando o agente implantado é consultado, ou seja, dentro das funções predict e predict_stream da interface ResponsesAgent. Como resultado, é necessário realizar qualquer acesso em nome do usuário ao recurso (por exemplo, listar índices de pesquisa vetorial aos quais o usuário final tem acesso) a partir desses métodos, em vez de no método __init__ da implementação do ResponsesAgent. Isso garante que os recursos sejam isolados entre invocações.

Python

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 authorization

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

    agent.predict(request)
    ...

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: ativar a autenticação em nome do usuário permite que os agentes acessem recursos confidenciais Databricks. Embora tenhamos implementado escopos de API para restringir o recurso ao qual os desenvolvedores podem acessar e mitigar o risco de uso indevido de tokens, alguns riscos ainda permanecem. Por exemplo, o escopo serving.serving-endpoints API concede a um agente permissão para executar um endpoint em nome do usuário. No entanto, o próprio endpoint que está sendo servido pode ter acesso a escopos API adicionais que o agente original não está autorizado a utilizar.
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