Autenticación para agentes de IA

A menudo, los agentes de inteligencia artificial necesitan autenticarse en otros recursos para completar tareas. Por ejemplo, un agente implementado podría necesitar acceder a un índice de búsqueda vectorial para consultar datos no estructurados o el Registro de mensajes para cargar avisos dinámicos.

En esta página se tratan los métodos de autenticación disponibles al desarrollar e implementar agentes mediante Mosaic AI Agent Framework.

Métodos de autenticación

En la tabla siguiente se comparan los métodos de autenticación disponibles. Puede combinar y hacer coincidir cualquiera de estos enfoques:

Autenticación en sistemas externos y servidores MCP

Para obtener instrucciones sobre cómo autenticarse en API externas y servidores MCP desde el agente, consulte Conexión de herramientas del agente de IA a servicios externos. Estos recursos también se pueden consultar en nombre del agente o usuario, como se describe en Métodos de autenticación.

Elección del método de autenticación adecuado para el recurso

Use este diagrama de flujo para elegir el método de autenticación adecuado para cada recurso. Puede combinar métodos según sea necesario y un agente puede usar un método diferente para cada recurso en función de su caso de uso.

1. ¿Se requiere el control de acceso por usuario o la auditoría con atributos de usuario?

   • Sí → Usar autenticación por cuenta del usuario
   • No → Continuar con el paso 2

2. ¿Todos los recursos admiten la autenticación automática?

   • Sí → Usar el paso de autenticación automática (recomendado)
   • No → usar autenticación manual

Traspaso automático de autenticación

El paso directo de autenticación automática es el método más sencillo para acceder a los recursos administrados por Databricks. Declare las dependencias de recursos cuando se registre el agente, y Databricks aprovisiona, rota y administra automáticamente credenciales temporales cuando se implementa el agente.

Este comportamiento de autenticación es similar al comportamiento de "Ejecutar como propietario" para los paneles de Databricks. A los recursos downstream, como las tablas del Catálogo de Unity, se accede mediante las credenciales de una entidad de servicio con el acceso de privilegio mínimo solo a los recursos que necesita el agente.

Funcionamiento del paso a través de autenticación automática

Cuando se proporciona un agente detrás de un punto de conexión mediante el paso directo de autenticación automática, Databricks realiza los pasos siguientes:

1. Comprobación de permisos: Databricks comprueba que el creador del punto de conexión puede acceder a todas las dependencias especificadas durante el registro del agente.

2. Creación y concesión de permisos de entidad de servicio: se crea una entidad de servicio para la versión del modelo del agente y se concede automáticamente acceso de lectura a los recursos del agente.

   Nota:

   La entidad de servicio generada por el sistema no aparece en las listas de la API ni de la interfaz de usuario. Si la versión del modelo del agente se quita del punto de conexión, también se elimina la entidad de servicio.

3. Aprovisionamiento y rotación de credenciales: las credenciales de corta duración (un token de OAuth de M2M) para la entidad de servicio se insertan en el punto de conexión, lo que permite que el código del agente acceda a los recursos de Databricks. Databricks también rota las credenciales, asegurando que el agente tenga acceso continuo y seguro a los recursos dependientes.

Recursos admitidos para el paso directo de autenticación automática

En la tabla siguiente se enumeran los recursos de Databricks que admiten el paso directo de autenticación automática y los permisos que debe tener el creador del punto de conexión al implementar el agente.

Implementación del paso directo de autenticación automática

Para habilitar el paso directo de autenticación automática, especifique los recursos dependientes al registrar el agente. Use el resources parámetro de la 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"),
    ]
  )

Autenticación en representación del usuario

La autenticación en nombre del usuario (OBO) permite a un agente actuar como el usuario de Databricks que ejecuta la consulta. Esto proporciona:

• Acceso por usuario a datos confidenciales
• Controles de datos específicos aplicados por el catálogo de Unity
• Los tokens de seguridad están restringidos ("downscoped") solo a las API que declara el agente, lo que reduce el riesgo de uso indebido.

Requisitos

• La autenticación en nombre del usuario requiere MLflow 2.22.1 y versiones posteriores.
• La autenticación en nombre de usuario está deshabilitada de forma predeterminada y debe estar habilitada por un administrador del área de trabajo. Revise las consideraciones de seguridad antes de habilitar esta característica.

Recursos compatibles con OBO

Los agentes con autenticación de OBO pueden acceder a los siguientes recursos de Databricks:

Implementación de la autenticación de OBO

Para habilitar la autenticación en nombre del usuario, complete los pasos siguientes:

1. Actualice las llamadas del SDK para especificar que se accede a los recursos en nombre del usuario final.
2. Actualice el código del agente para inicializar el acceso de OBO dentro de la predict función, no en __init__, porque la identidad del usuario solo se conoce en tiempo de ejecución.
3. Al registrar el agente para la implementación, declare los ámbitos de la API REST de Databricks que requiere el agente.

Los fragmentos de código siguientes muestran cómo configurar el acceso en nombre del usuario a distintos recursos de Databricks. Al inicializar las herramientas, controle correctamente los errores de permisos envolviendo la inicialización en un bloque try-except.

Herramienta de búsqueda de vectores

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 búsqueda de vectores

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
  )

Modelo de puntos de conexión de servicio

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

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

Genie Spaces (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")

Genie Spaces (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 el agente en la función predict

Dado que la identidad del usuario solo se conoce en el momento de la consulta, debe acceder a los recursos de OBO dentro predict de o predict_stream, no en el método del __init__ agente. Esto garantiza que los recursos estén aislados entre invocaciones.

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

Declaración de ámbitos de la API REST al registrar el agente

Al registrar el agente de OBO para la implementación, debe enumerar los ámbitos de la API REST de Databricks a los que llama el agente en nombre del usuario. Esto garantiza que el agente siga el principio de privilegios mínimos: los tokens están restringidos solo a las API que requiere el agente, lo que reduce la posibilidad de acciones no autorizadas o uso indebido de tokens.

A continuación se muestra una lista de ámbitos necesarios para acceder a varios tipos comunes de recursos de Databricks:

Para habilitar la autenticación en nombre de un usuario, pase un MLflow AuthPolicy a 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
        )
    )

Autenticación de OBO para clientes OpenAI

Para los agentes que usan el cliente de OpenAI, use el SDK de Databricks para autenticarse automáticamente durante la implementación. El SDK de Databricks tiene un contenedor para construir el cliente openAI con la autenticación configurada automáticamente, 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()

A continuación, especifique el punto de conexión De servicio de modelos como parte de para autenticarse automáticamente en el momento de resources la implementación.

Consideraciones de seguridad de OBO

Tenga en cuenta las siguientes consideraciones de seguridad antes de habilitar la autenticación en nombre del usuario con agentes.

Acceso expandido a los recursos: los agentes pueden acceder a recursos confidenciales en nombre de los usuarios. Aunque los ámbitos restringen las API, los puntos de conexión pueden permitir más acciones que las solicitudes explícitas del agente. Por ejemplo, el ámbito de la serving.serving-endpoints API concede a un agente permiso para ejecutar un punto de conexión de servicio en nombre del usuario. Sin embargo, el punto de conexión de servicio puede acceder a ámbitos de API adicionales que el agente original no está autorizado a usar.

Cuadernos de ejemplo de OBO

En el cuaderno siguiente se muestra cómo crear un agente con la búsqueda vectorial mediante la autorización en nombre del usuario.

En nombre de la autorización de usuario con la búsqueda de vectores

Obtener el cuaderno

En el cuaderno siguiente se muestra cómo crear un agente que admita la ejecución de SQL en un almacén de SQL mediante la autorización en nombre del usuario. Esto permite al agente invocar de forma segura funciones de Catálogo de Unity mediante credenciales de usuario. Nota: Actualmente, esta es la manera recomendada de ejecutar funciones UC con OBO, ya que la ejecución de Spark sin servidor con OBO todavía no se admite.

En nombre de la autorización de usuario con la ejecución de SQL

Obtener el cuaderno

Autenticación manual

La autenticación manual permite especificar explícitamente las credenciales durante la implementación del agente. Este método tiene la mayor flexibilidad, pero requiere más configuración y administración de credenciales en curso. Use este método cuando:

• El recurso dependiente no admite la transferencia automática de autenticación.
• El agente debe usar credenciales distintas de las del implementador del agente.
• El agente accede a recursos externos o API fuera de Databricks.
• El agente implementado accede al registro de comandos.

Autenticación de OAuth (recomendado)

OAuth es el enfoque recomendado para la autenticación manual porque tiene autenticación segura basada en tokens para entidades de servicio con funcionalidades de actualización automática de tokens:

1. Cree una entidad de servicio y genere credenciales de OAuth.

2. Conceda a la entidad de servicio permisos para cualquier recurso de Databricks al que el agente tenga privilegios de acceso permisos para acceder a recursos de Databricks. Para acceder al registro de mensajes, conceda permisos en CREATE FUNCTION, EXECUTE y MANAGE en el esquema del Catálogo de Unity para almacenar mensajes.

3. Cree secretos de databricks para las credenciales de OAuth.

4. Configure las credenciales de OAuth en el código del 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 los secretos para conectarse al área de trabajo:

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

6. Al implementar con agents.deploy(), incluya las credenciales de OAuth como variables de entorno:

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

Autenticación PAT

La autenticación de token de acceso personal (PAT) proporciona una configuración más sencilla para entornos de desarrollo y pruebas, aunque requiere una administración de credenciales más manual:

1. Obtenga un PAT mediante una entidad de servicio o una cuenta personal:

   Entidad de servicio (recomendada para la seguridad):

   1. Cree una entidad de servicio.
   2. Conceda a la entidad de servicio permisos para cualquier recurso de Databricks al que el agente tenga privilegios de acceso permisos para acceder a recursos de Databricks. Para acceder al registro de mensajes, conceda los permisos en CREATE FUNCTION, EXECUTE y MANAGE en el esquema del catálogo de Unity que se usa para almacenar mensajes.
   3. Cree un PAT para la entidad de servicio.

   Cuenta personal:

   1. Cree un PAT para una cuenta personal.

2. Almacene el PAT de forma segura mediante la creación de un secreto de Databricks para el PAT.

3. Configure la autenticación PAT en el código del 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. Al implementar el agente mediante agents.deploy(), incluya el PAT como una variable de entorno:

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