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 busca vetorial para consultar dados não estruturados ou o Registro de Prompts 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 combinar qualquer uma dessas abordagens:

Autenticar em sistemas externos e servidores MCP

Para obter diretrizes sobre como autenticar em 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 do usuário, conforme descrito nos métodos de Autenticação.

Escolha o método de autenticação correto para o recurso

Use esse 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. O controle de acesso por usuário ou a auditoria atribuída pelo usuário são necessários?

   • Sim → Usar Autenticação em Nome do Usuário
   • Sem → Continuar para a etapa 2

2. Todos os recursos dão suporte à autenticação automática?

   • Sim → usar a passagem de autenticação automática (recomendado)
   • Não → 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 logar o agente, e o Databricks provisiona, rotaciona e gerencia automaticamente credenciais de curta duração quando o agente é implantado.

Esse comportamento de autenticação é semelhante ao comportamento "Executar como proprietário" para dashboards do Databricks. Recursos downstream como tabelas do Catálogo do Unity são acessados usando as credenciais de uma entidade de serviço com acesso de privilégio mínimo 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 a passagem de autenticação automática, o Databricks executa as seguintes etapas:

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

2. Criação e concessão de entidade de serviço: uma entidade de serviço é criada para a versão do modelo do agente e é automaticamente recebe acesso de leitura aos recursos do agente.

   Observação

   A entidade de serviço gerada pelo sistema não aparece em listagens de API ou interface do usuário. Se a versão do modelo do agente for removida do ponto de extremidade, 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) para a entidade de serviço são injetadas no ponto de extremidade, permitindo que o código do agente acesse os recursos do Databricks. O Databricks também gira as credenciais, garantindo que o agente tenha acesso seguro e contínuo aos recursos dependentes.

Recursos com suporte para passagem de autenticação automática

A tabela a seguir lista os recursos do Databricks que dão 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 em log. 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 usuário

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

• Acesso por usuário a dados confidenciais
• Controles de dados refinados impostos pelo Catálogo do Unity
• Os tokens de segurança são limitados ("com escopo reduzido") apenas às APIs que seu agente declara, reduzindo o risco de uso indevido.

Requirements

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

Recursos suportados por OBO

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

Implementar a 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 dentro __init__, porque a identidade do usuário só é conhecida no runtime.
3. Ao registrar o agente para implantação, declare os escopos da API REST do Databricks que o agente requer.

Os snippets a seguir demonstram como configurar o acesso em nome do usuário para diferentes recursos do Databricks. Ao inicializar ferramentas, trate os erros de permissão de forma adequada, encapsulando a inicialização em um bloco try-except.

Ferramenta de Busca Vetorial

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 de Vetor

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 do 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 do 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 do Gênio (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 recursos OBO dentro predict ou predict_streamnão no método do __init__ agente. Isso garante que os recursos sejam isolados entre 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 em log

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

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

Para habilitar a autenticação em nome de outro usuário, 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 SDK do Databricks 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 autenticar automaticamente no momento da implantação.

Considerações sobre 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 endpoint de serviço em nome do usuário. No entanto, o endpoint de serviço pode acessar escopos adicionais de API que o agente original não está autorizado a usar.

Notebooks de exemplo OBO

O notebook a seguir mostra como criar um agente com o Vector Search usando autorização em nome do usuário.

Em nome da autorização do usuário com a pesquisa de vetor

Obter laptop

O notebook a seguir mostra como criar um agente que dá suporte à execução de SQL em um SQL Warehouse usando autorização em nome do usuário. Isso permite que o agente invoque com segurança as funções do Catálogo do Unity usando as credenciais do usuário. Observação: no momento, esta é a maneira recomendada de executar as funções UC com OBO, uma vez que a execução do Spark sem servidor com OBO ainda não é suportada.

Em nome de autorização do usuário para execução de comandos SQL

Obter laptop

Autenticação manual

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

• O recurso dependente não dá suporte à passagem de autenticação automática
• O agente precisa usar credenciais diferentes das do implementador 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 token para entidades de serviço com recursos automáticos de atualização de token:

1. Crie um principal de serviço e gere credenciais OAuth.

2. Conceda as permissões da entidade de serviço a qualquer recurso do Databricks que o agente tenha acesso a privilégios para acessar recursos do Databricks. Para acessar o registro de prompts, conceda CREATE FUNCTION, EXECUTE e MANAGE permissões no esquema do Unity Catalog para armazenar prompts.

3. Crie segredos do Databricks para as credenciais do OAuth.

4. Configure as credenciais do 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 as credenciais para se conectar à área 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 do 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 pat (Token de Acesso Pessoal) fornece uma configuração mais simples para ambientes de desenvolvimento e teste, embora exija um gerenciamento de credenciais mais manual:

1. Obtenha um PAT usando um principal de serviço ou uma conta pessoal:

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

   1. Crie uma entidade de serviço.
   2. Conceda as permissões da entidade de serviço a qualquer recurso do Databricks que o agente tenha acesso a privilégios para acessar recursos do Databricks. Para acessar o registro de prompts, conceda CREATE FUNCTION, EXECUTE e MANAGE permissões no esquema do Catálogo do Unity usado para armazenar prompts.
   3. Crie um PAT para a entidade de serviço.

   Conta pessoal:

   1. Crie um PAT para uma conta pessoal.

2. Armazene o PAT com segurança 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}}}}}"
       },
   )