Autenticação para agentes de IA

Os agentes de IA geralmente precisam se autenticar em outros recursos para concluir tarefas. Por exemplo, um agente implantado pode precisar acessar um índice de Pesquisa Vetorial para consultar dados não estruturados ou o Registro de Prompt para carregar prompts dinâmicos.

Esta página aborda os métodos de autenticação disponíveis ao desenvolver e implantar agentes usando o Mosaic AI Agent Framework.

Métodos de autenticação

A tabela a seguir compara os métodos de autenticação disponíveis. Você pode misturar e combinar qualquer uma destas abordagens:

Autenticar em sistemas externos e servidores MCP

Para obter orientação sobre como autenticar APIs externas e servidores MCP do seu agente, consulte Conectar ferramentas de agente de IA a serviços externos. Esses recursos também podem ser consultados em nome do agente ou usuário, conforme descrito em Métodos de autenticação.

Escolha o método de autenticação certo para o seu recurso

Use este fluxograma para escolher o método de autenticação correto para cada recurso. Você pode combinar métodos conforme necessário, e um agente pode usar um método diferente para cada recurso, dependendo de seu caso de uso.

1. É necessário controle de acesso por usuário ou auditoria atribuída pelo usuário?

   • Sim → Usar autenticação em nome do utilizador
   • Não → Continue para a etapa 2

2. Todos os recursos suportam autenticação automática?

   • Sim → Usar passagem de autenticação automática (recomendado)
   • Sem → Usar autenticação manual

Passagem de autenticação automática

A passagem de autenticação automática é o método mais simples para acessar recursos gerenciados pelo Databricks. Declare dependências de recursos ao registar o agente, e o Databricks provisiona, roda e gere automaticamente credenciais temporárias quando o agente é implementado.

Esse comportamento de autenticação é semelhante ao comportamento "Executar como proprietário" nos painéis do Databricks. Recursos downstream, como tabelas do Unity Catalog, são acessados usando as credenciais de uma entidade de serviço com acesso de privilégios mínimos apenas aos recursos de que o agente precisa.

Como funciona a passagem de autenticação automática

Quando um agente é atendido atrás de um ponto de extremidade usando passagem de autenticação automática, o Databricks executa as seguintes etapas:

1. Verificação de permissão: O Databricks verifica se o criador do endpoint pode acessar todas as dependências especificadas durante o registo do agente.

2. Criação e concessões do principal de serviço: Um principal de serviço é criado para a versão do modelo do agente e recebe automaticamente acesso de leitura aos recursos do agente.

   Observação

   A entidade de serviço gerada pelo sistema não aparece nas listas da API nem na UI. Se a versão do modelo do agente for removida do endpoint, a entidade de serviço também será excluída.

3. Provisionamento e rotação de credenciais: credenciais de curta duração (um token M2M OAuth ) são injetadas no endpoint para um service principal, permitindo que o código do agente aceda aos recursos do Databricks. O Databricks também alterna as credenciais, garantindo que seu agente tenha acesso contínuo e seguro aos recursos dependentes.

Recursos suportados para passagem de autenticação automática

A tabela a seguir lista os recursos do Databricks que oferecem suporte à passagem de autenticação automática e as permissões que o criador do ponto de extremidade deve ter ao implantar o agente.

Implementar passagem de autenticação automática

Para habilitar a passagem de autenticação automática, especifique os recursos dependentes ao registrar o agente. Use o resources parâmetro da log_model() API:

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"),
    ]
  )

Autenticação em nome do utilizador

A autenticação em nome do usuário (OBO) permite que um agente atue como o usuário do Databricks que executa a consulta. Isto prevê:

• Acesso por usuário a dados confidenciais
• Controles de dados refinados impostos pelo Unity Catalog
• Os tokens de segurança são restritos ("downscoped") apenas às APIs que o seu agente declara, reduzindo assim o risco de uso indevido.

Requerimentos

• A autenticação em nome do usuário requer o MLflow 2.22.1 e superior.
• A autenticação em nome do usuário é desabilitada por padrão e deve ser habilitada por um administrador de espaço de trabalho. Analise as considerações de segurança antes de ativar esse recurso.

Recursos suportados pelo OBO

Os agentes com autenticação OBO podem acessar os seguintes recursos do Databricks:

Implementar autenticação OBO

Para habilitar a autenticação em nome do usuário, conclua as seguintes etapas:

1. Atualize as chamadas do SDK para especificar que os recursos são acessados em nome do usuário final.
2. Atualize o código do agente para inicializar o acesso OBO dentro da predict função, não no __init__, porque a identidade do usuário só é conhecida em tempo de execução.
3. Ao registrar o agente para implantação, declare os escopos da API REST do Databricks que o agente exige.

Os trechos a seguir demonstram como configurar o acesso em nome do usuário a diferentes recursos do Databricks. Ao inicializar ferramentas, trate os erros de permissão com elegância encapsulando a inicialização num bloco try-except.

Ferramenta Vetor Search Retriever

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")

Cliente de pesquisa vetorial

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
  )

Ponto de extremidade de serviço de modelo

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")

Conexões 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"},
)

Espaços Genie (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")

Espaços Genie (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
  )

Inicializar o agente na função de previsão

Como a identidade do usuário só é conhecida no momento da consulta, você deve acessar os recursos do OBO dentro predict ou predict_stream, não no método do __init__ agente. Isso garante que os recursos sejam isolados entre as invocações.

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)
    ...

Declarar escopos da API REST ao registrar o agente

Ao registrar seu agente OBO para implantação, você deve listar os escopos da API REST do Databricks que o agente chama em nome do usuário. Isso garante que o agente siga o princípio do menor privilégio: os tokens são restritos apenas às APIs que seu agente exige, reduzindo a chance de ações não autorizadas ou uso indevido do token.

Abaixo está uma lista de escopos necessários para acessar vários tipos comuns de recursos do Databricks:

Para habilitar a autenticação em nome de um utilizador, passe um MLflow AuthPolicy para 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
        )
    )

Autenticação OBO para clientes OpenAI

Para agentes que usam o cliente OpenAI, use o SDK do Databricks para autenticar automaticamente durante a implantação. O Databricks SDK tem um wrapper para construir o cliente OpenAI com autenticação configurada automaticamente, 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()

Em seguida, especifique o ponto de extremidade do Serviço de Modelo como parte de resources para que a autenticação ocorra automaticamente no momento da implementação.

Considerações de segurança do OBO

Considere as seguintes considerações de segurança antes de habilitar a autenticação em nome do usuário com agentes.

Acesso expandido a recursos: os agentes podem acessar recursos confidenciais em nome dos usuários. Embora os escopos restrinjam APIs, os pontos de extremidade podem permitir mais ações do que o agente solicita explicitamente. Por exemplo, o escopo da serving.serving-endpoints API concede a um agente permissão para executar um ponto de extremidade de serviço em nome do usuário. No entanto, o endpoint de serviço pode aceder a escopos adicionais de API que o agente original não está autorizado a usar.

Exemplos de notebooks OBO

O bloco de anotações a seguir mostra como criar um agente com a Pesquisa Vetorial usando a autorização em nome do usuário.

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

Obter caderno

O bloco de anotações a seguir mostra como criar um agente que ofereça suporte à execução SQL em um SQL Warehouse usando a autorização em nome do usuário. Isso permite que o agente invoque com segurança as funções do Unity Catalog usando credenciais de usuário. Nota: Esta é atualmente a maneira recomendada de executar funções UC com OBO, pois a execução de faísca sem servidor com OBO ainda não é suportada.

Em nome da autorização do usuário com execução SQL

Obter caderno

Autenticação manual

A autenticação manual permite especificar explicitamente as credenciais durante a implantação do agente. Esse método tem a maior flexibilidade, mas requer mais configuração e gerenciamento contínuo de credenciais. Use este método quando:

• O recurso dependente não suporta passagem de autenticação automática
• O agente precisa usar credenciais diferentes daquelas do despachante do agente
• O agente acessa recursos externos ou APIs fora do Databricks
• O agente implantado acessa o registro de prompt

Autenticação OAuth (recomendado)

O OAuth é a abordagem recomendada para autenticação manual porque tem autenticação segura baseada em tokens para entidades de serviço com recursos de atualização automática de token:

1. Crie uma principal de serviço e gere as credenciais do OAuth.

2. Conceda permissões à entidade de serviço a qualquer recurso do Databricks ao qual o agente tenha acesso privilégios para acessar os recursos do Databricks. Para aceder ao registro de prompt, conceda as permissões CREATE FUNCTION, EXECUTE e MANAGE no esquema do Catálogo Unity para armazenar prompts.

3. Crie segredos do Databricks para as credenciais OAuth.

4. Configure as credenciais OAuth no código do agente.

   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. Use os segredos para se conectar ao espaço de trabalho:

   w = WorkspaceClient(
     host=os.environ["DATABRICKS_HOST"],
     client_id=os.environ["DATABRICKS_CLIENT_ID"],
     client_secret = os.environ["DATABRICKS_CLIENT_SECRET"]
   )
   

6. Ao implantar com agents.deploy(), inclua as credenciais OAuth como variáveis de ambiente:

   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}}}}}"
       },
   )
   

Autenticação PAT

A autenticação de Token de Acesso Pessoal (PAT) fornece uma configuração mais simples para ambientes de desenvolvimento e teste, embora exija mais gerenciamento manual de credenciais:

1. Obtenha uma PAT utilizando uma entidade de serviço ou uma conta pessoal:

   Entidade de serviço (recomendada para segurança):

   1. Crie uma entidade de serviço.
   2. Conceda permissões à entidade de serviço a qualquer recurso do Databricks ao qual o agente tenha acesso privilégios para acessar os recursos do Databricks. Para aceder ao registo de prompts, conceda CREATE FUNCTION, EXECUTE e MANAGE permissões no esquema do Unity Catalog usado para armazenar prompts.
   3. Crie uma PAT para a entidade de serviço.

   Conta pessoal:

   1. Crie um PAT para uma conta pessoal.

2. Armazene o PAT de forma segura, criando um segredo do Databricks para o PAT.

3. Configure a autenticação PAT no código do agente:

   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. Ao implantar o agente usando agents.deploy(), inclua o PAT como uma variável de ambiente.

   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}}}}}"
       },
   )