Authentification pour les agents IA

Les agents IA doivent souvent s’authentifier auprès d’autres ressources pour effectuer des tâches. Par exemple, un agent déployé peut avoir besoin d’accéder à un index Recherche vectorielle pour interroger des données non structurées ou le Registre d’invites pour charger des invites dynamiques.

Cette page décrit les méthodes d’authentification disponibles lors du développement et du déploiement d’agents à l’aide de Mosaic AI Agent Framework.

Méthodes d’authentification

Le tableau suivant compare les méthodes d’authentification disponibles. Vous pouvez combiner et mettre en correspondance l’une de ces approches :

S’authentifier auprès de systèmes externes et de serveurs MCP

Pour obtenir des conseils sur la façon de s’authentifier auprès d’API externes et de serveurs MCP à partir de votre agent, consultez Connecter les outils de l’agent IA aux services externes. Ces ressources peuvent également être interrogées au nom de l’agent ou de l’utilisateur, comme décrit dans les méthodes d’authentification.

Choisir la méthode d’authentification appropriée pour votre ressource

Utilisez cet organigramme pour choisir la méthode d’authentification appropriée pour chaque ressource. Vous pouvez combiner des méthodes en fonction des besoins, et un agent peut utiliser une méthode différente pour chaque ressource en fonction de son cas d’usage.

1. Le contrôle d’accès par utilisateur ou l’audit attribué par l’utilisateur est-il nécessaire ?

   • Oui → utiliser l’authentification au nom de l’utilisateur
   • Aucune → Passer à l’étape 2

2. Toutes les ressources prennent-elles en charge l’authentification automatique ?

   • Oui → Utiliser la passthrough d’authentification automatique (recommandé)
   • Aucune → Utiliser l’authentification manuelle

Passthrough d’authentification automatique

La passthrough d’authentification automatique est la méthode la plus simple pour accéder aux ressources gérées par Databricks. Déclarez les dépendances des ressources de l'agent lors de la journalisation, et Databricks approvisionne, renouvelle et gère automatiquement les informations d'identification éphémères lorsque l'agent est déployé.

Ce comportement d’authentification est similaire au comportement « Exécuter en tant que propriétaire » pour les tableaux de bord Databricks. Les ressources en aval telles que les tables du catalogue Unity sont accessibles à l’aide des informations d’identification d’un principal de service disposant d’un accès avec privilège minimum aux ressources dont l’agent a besoin.

Fonctionnement de la passthrough d’authentification automatique

Lorsqu’un agent est servi derrière un point de terminaison à l’aide d’une transmission automatique des informations d'identification, Databricks effectue les étapes suivantes :

1. Vérification des autorisations : Databricks vérifie que le créateur du point de terminaison peut accéder à toutes les dépendances spécifiées lors de la journalisation de l’agent.

2. Création de principal de service et octrois: un principal de service est créé pour la version du modèle de l'agent et reçoit automatiquement l'accès en lecture aux ressources de l'agent.

   Note

   Le principal de service généré par le système n’apparaît pas dans les listes d’API ou d’interface utilisateur. Si la version du modèle d’agent est supprimée du point de terminaison, le principal de service est également supprimé.

3. Provisionnement et rotation des informations d’identification : informations d’identification de courte durée (un jeton OAuth M2M) pour le principal de service sont injectées dans le point de terminaison, ce qui permet au code de l’agent d’accéder aux ressources Databricks. Databricks fait également pivoter les informations d’identification, ce qui garantit que votre agent a continué d’accéder en toute sécurité aux ressources dépendantes.

Ressources prises en charge pour la transmission automatique d'authentification

Le tableau suivant répertorie les ressources Databricks qui prennent en charge la passe d’authentification automatique et les autorisations dont le créateur de point de terminaison doit disposer lors du déploiement de l’agent.

Implémenter la transmission automatique d'authentification

Pour activer la passe directe d’authentification automatique, spécifiez les ressources dépendantes lorsque vous connectez l’agent. Utilisez le resources paramètre de l’API log_model() :

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

Authentification déléguée pour le compte de l’utilisateur

L’authentification OBO (On-behalf-of-user) permet à un agent d’agir en tant qu’utilisateur Databricks qui exécute la requête. Cela fournit les éléments suivants :

• Accès par utilisateur aux données sensibles
• Contrôles de données affinés appliqués par le catalogue Unity
• Les jetons de sécurité sont limités (« réduits ») aux API déclarées par votre agent, ce qui réduit le risque d’utilisation abusive

Spécifications

• L’authentification au nom de l’utilisateur nécessite MLflow 2.22.1 et versions ultérieures.
• L’authentification au nom de l’utilisateur est désactivée par défaut et doit être activée par un administrateur d’espace de travail. Passez en revue les considérations de sécurité avant d’activer cette fonctionnalité.

Ressources prises en charge par OBO

Les agents avec l’authentification OBO peuvent accéder aux ressources Databricks suivantes :

Implémenter l’authentification OBO

Pour activer l’authentification au nom de l’utilisateur, procédez comme suit :

1. Mettez à jour les appels du Kit de développement logiciel (SDK) pour spécifier que les ressources sont accessibles au nom de l’utilisateur final.
2. Mettez à jour le code de l’agent pour initialiser l’accès OBO à l’intérieur de la predict fonction, et non dans __init__, car l’identité de l’utilisateur est connue uniquement au moment de l’exécution.
3. Lors de l'enregistrement de l'agent pour le déploiement, déclarez les étendues de l’API REST Databricks requises par l’agent.

Les extraits de code suivants montrent comment configurer l’accès au nom de l’utilisateur à différentes ressources Databricks. Lors de l’initialisation des outils, gérez correctement les erreurs d’autorisation en encapsulant l’initialisation dans un try-except bloc.

Outil de récupération de recherche vectorielle

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

Client de recherche vectorielle

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
  )

Point de terminaison de mise en service de modèles

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

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

Initialiser l’agent dans la fonction de prédiction

Étant donné que l’identité de l’utilisateur est connue uniquement au moment de la requête, vous devez accéder aux ressources OBO à l’intérieur predict ou predict_streamnon dans la méthode de l’agent __init__ . Cela garantit que les ressources sont isolées entre les appels.

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

Déclarer des étendues d’API REST lors de la journalisation de l’agent

Lorsque vous connectez votre agent OBO pour le déploiement, vous devez répertorier les étendues de l’API REST Databricks que votre agent appelle pour le compte de l’utilisateur. Cela garantit que l’agent suit le principe du privilège minimum : les jetons sont limités aux API dont votre agent a besoin, ce qui réduit le risque d’actions non autorisées ou d’utilisation incorrecte des jetons.

Voici une liste des étendues requises pour accéder à plusieurs types courants de ressources Databricks :

Pour activer l’authentification au nom de l’utilisateur, transmettez un MLflow AuthPolicy à 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
        )
    )

Authentification OBO pour les clients OpenAI

Pour les agents qui utilisent le client OpenAI, utilisez le Kit de développement logiciel (SDK) Databricks pour s’authentifier automatiquement pendant le déploiement. Databricks SDK a un wrapper pour construire le client OpenAI avec l’authentification automatiquement configurée, 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()

Ensuite, spécifiez le point de terminaison Model Serving dans le cadre de resources pour l'authentification automatique lors du déploiement.

Considérations relatives à la sécurité OBO

Tenez compte des considérations de sécurité suivantes avant d’activer l’authentification au nom de l’utilisateur avec des agents.

Accès aux ressources étendues : les agents peuvent accéder aux ressources sensibles pour le compte des utilisateurs. Bien que les étendues limitent les API, les points de terminaison peuvent autoriser plus d’actions que votre agent demande explicitement. Par exemple, l’étendue de l’API serving.serving-endpoints accorde à un agent l’autorisation d’exécuter un point de terminaison de service pour le compte de l’utilisateur. Toutefois, le point de terminaison de service peut accéder à des étendues d’API supplémentaires que l’agent d’origine n’est pas autorisé à utiliser.

Exemples de notebooks OBO

Le notebook suivant vous montre comment créer un agent avec recherche vectorielle à l'aide de l'autorisation au nom de l'utilisateur.

Au nom de l’autorisation de l’utilisateur avec recherche vectorielle

Obtenir un ordinateur portable

Le notebook suivant vous montre comment créer un agent qui prend en charge l’exécution SQL sur un SQL Warehouse à l’aide de l’autorisation de la part de l’utilisateur. Cela permet à l’agent d’appeler en toute sécurité les fonctions du catalogue Unity à l’aide des informations d’identification de l’utilisateur. Remarque : Il s’agit actuellement de la méthode recommandée pour exécuter des fonctions UC avec OBO, car l’exécution serverless de Spark avec OBO n’est pas encore prise en charge.

Sous l'autorisation de l'utilisateur pour l'exécution de SQL

Obtenir un ordinateur portable

Authentification manuelle

L’authentification manuelle permet de spécifier explicitement les informations d’identification pendant le déploiement de l’agent. Cette méthode a la plus grande flexibilité, mais nécessite davantage de configuration et de gestion continue des informations d’identification. Utilisez cette méthode quand :

• La ressource dépendante ne prend pas en charge la transmission automatique d’authentification
• L’agent doit utiliser des informations d’identification autres que celles de l’agent deployer
• L’agent accède aux ressources externes ou aux API en dehors de Databricks
• L’agent déployé accède au prompt registry

Authentification OAuth (recommandé)

OAuth est l’approche recommandée pour l’authentification manuelle, car elle dispose d’une authentification sécurisée basée sur des jetons pour les principaux de service avec des fonctionnalités d’actualisation automatique des jetons :

1. Créez un principal de service et générez des informations d’identification OAuth.

2. Accordez les autorisations du principal de service à toutes les ressources Databricks auxquelles l’agent a accès aux privilèges pour accéder aux ressources Databricks. Pour accéder au registre d’invites, accordez les autorisations CREATE FUNCTION, EXECUTE et MANAGE sur le schéma du catalogue Unity pour stocker des invites.

3. Créez des secrets Databricks pour les identifiants OAuth.

4. Configurez les informations d’identification OAuth dans le code de l’agent :

   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. Utilisez les secrets pour vous connecter à l’espace de travail :

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

6. Lors du déploiement avec agents.deploy(), incluez les informations d’identification OAuth en tant que variables d’environnement :

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

Authentification PAT

L’authentification par jeton d’accès personnel (PAT) offre une configuration plus simple pour les environnements de développement et de test, bien qu’elle nécessite une gestion des informations d’identification plus manuelle :

1. Obtenez un jeton d'accès personnel à l’aide d’un service principal ou d’un compte personnel :

   Entité de service (recommandée pour la sécurité) :

   1. Créez un principal de service.
   2. Accordez les autorisations du principal de service à toutes les ressources Databricks auxquelles l’agent a accès aux privilèges pour accéder aux ressources Databricks. Pour accéder au registre d'instructions, accordez les autorisations CREATE FUNCTION, EXECUTE, et MANAGE sur le schéma du catalogue Unity utilisé pour stocker les instructions.
   3. Créez un pater pour le principal de service.

   Compte personnel :

   1. Créez un PAT pour un compte personnel.

2. Stockez le PAT en toute sécurité en créant un secret Databricks pour le PAT.

3. Configurez l’authentification PAT dans le code de l’agent :

   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. Lors du déploiement de l’agent à l’aide de agents.deploy(), incluez le PAT comme variable d’environnement :

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