Journaliser et enregistrer des agents d’IA

2025-04-16

Enregistrer les agents IA avec Mosaic AI Agent Framework. La journalisation d’un assistant est la base du processus de développement. La journalisation capture un « point dans le temps » du code et de la configuration de l’agent afin de pouvoir évaluer la qualité de la configuration.

Spécifications

Créez un agent IAavant de le journaliser.

Databricks recommande d’installer la dernière version du databricks-sdk.

% pip install databricks-sdk

Journalisation basée sur le code

Databricks recommande d'utiliser la fonctionnalité Models from Code de MLflow lors de l'enregistrement des agents.

Dans cette approche, le code de l’agent est capturé en tant que fichier Python et l’environnement Python est capturé sous la forme d’une liste de packages. Lorsque l’agent est déployé, l’environnement Python est restauré et le code de l’agent est exécuté pour charger l’agent en mémoire afin qu’il puisse être appelé lorsque le point de terminaison est appelé.

Vous pouvez coupler cette approche avec l’utilisation d’API de validation de prédéploiement comme mlflow.models.predict() pour vous assurer que l’assistant s’exécute de manière fiable lors de son déploiement pour la fourniture de services.

Pour voir un exemple de journalisation basée sur le code, consultez Exemple de création de ChatAgent sur Notebooks.

Déduire la signature du modèle pendant la journalisation

Remarque

Databricks recommande de créer un agent à l’aide de l’interface ChatAgent. Si vous utilisez ChatAgent, vous pouvez ignorer cette section ; MLflow déduit automatiquement une signature valide pour votre agent.

Si vous n’utilisez pas l’interface ChatAgent , vous devez utiliser l’une des méthodes suivantes pour spécifier la signature de modèle MLflow de votre agent au moment de la journalisation :

Définir manuellement la signature
Utilisez les fonctionnalités d’inférence de signature de modèle de MLflow pour générer automatiquement la signature de l’agent en fonction d’un exemple d’entrée que vous fournissez. Cette approche est plus pratique que de définir manuellement la signature.

La signature du modèle MLflow valide les entrées et les sorties pour vous assurer que l’agent interagit correctement avec les outils en aval tels que AI Playground et l’application de révision. Il guide également d’autres applications sur l’utilisation efficace de l’agent.

Les exemples LangChain et PyFunc ci-dessous utilisent l’inférence de signature de modèle.

Si vous préférez définir explicitement une signature de modèle vous-même au moment de la journalisation, consultez la documentation MLflow - Comment journaliser des modèles avec des signatures.

Journalisation basée sur le code avec LangChain

Les instructions et l’exemple de code suivants vous montrent comment consigner un agent avec LangChain.

Créez un notebook ou un fichier Python avec votre code. Pour cet exemple, le bloc-notes ou le fichier est nommé agent.py. Le notebook ou le fichier doit contenir un agent LangChain, appelé ici lc_agent.
Incluez mlflow.models.set_model(lc_agent) dans le bloc-notes ou le fichier.
Créez un notebook pour servir de notebook pilote (appelé driver.py dans cet exemple).
Dans le bloc-notes du pilote, utilisez le code suivant pour exécuter agent.py et journaliser les résultats sur un modèle MLflow :
```
mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
Le resources paramètre déclare les ressources gérées par Databricks nécessaires pour servir l’agent, comme un index de recherche vectorielle ou un point de terminaison servant un modèle de base. Pour plus d’informations, consultez Authentification pour les ressources Databricks.
Déployer le modèle. Consultez Déployer un agent pour les applications IA génératives.
Lorsque l’environnement de service est chargé, agent.py est exécuté.
Lorsqu’une requête de service arrive, lc_agent.invoke(...) est appelé.


import mlflow

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
  "messages": [
    {
      "role": "user",
      "content": "What is Retrieval-augmented Generation?",
    }
  ]
}

# example using langchain
with mlflow.start_run():
  logged_agent_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this configuration is used by agent code. The development_config is overwritten.
    artifact_path="agent", # This string is used as the path inside the MLflow model where artifacts are stored
    input_example=input_example, # Must be a valid input to the agent
    example_no_conversion=True, # Required
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.langchain.load_model(logged_agent_info.model_uri)
model.invoke(example)

Journalisation basée sur le code avec PyFunc

Les instructions et l’exemple de code suivants vous montrent comment journaliser un agent avec PyFunc.

Créez un notebook ou un fichier Python avec votre code. Pour cet exemple, le bloc-notes ou le fichier est nommé agent.py. Le notebook ou le fichier doit contenir une classe PyFunc, nommée PyFuncClass.
Incluez mlflow.models.set_model(PyFuncClass) dans le carnet ou le fichier.
Créez un notebook pour servir de notebook pilote (appelé driver.py dans cet exemple).
Dans le bloc-notes du pilote, utilisez le code suivant pour exécuter agent.py et utiliser log_model() pour consigner les résultats dans un modèle MLflow :
```
mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
Le resources paramètre déclare les ressources gérées par Databricks nécessaires pour servir l’agent, comme un index de recherche vectorielle ou un point de terminaison servant un modèle de base. Pour plus d’informations, consultez Authentification pour les ressources Databricks.
Déployer le modèle. Consultez Déployer un agent pour les applications IA génératives.
Lorsque l’environnement de service est chargé, agent.py est exécuté.
Lorsqu’une requête de service arrive, PyFuncClass.predict(...) est appelé.

import mlflow
from mlflow.models.resources import (
    DatabricksServingEndpoint,
    DatabricksVectorSearchIndex,
)

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
  "messages": [
    {
      "role": "user",
      "content": "What is Retrieval-augmented Generation?",
    }
  ]
}

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    artifact_path="agent",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
    resources=[
      DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct"),
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
    ]
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.pyfunc.load_model(logged_agent_info.model_uri)
model.invoke(example)

Authentification pour les ressources Databricks

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

Comme décrit dans l’authentification pour les ressources dépendantes, Model Serving prend en charge l’authentification auprès des ressources gérées par Databricks et externes lorsque vous déployez l’agent.

Model Serving prend en charge deux types d’authentification différents pour les ressources gérées par Databricks :

Authentification système : permet au principal du service de l'agent d'accéder aux ressources dépendantes spécifiées lors de la connexion de l'agent. Cela est utile pour accéder à des ressources partagées ou non sensibles, par exemple un index de recherche vectorielle contenant la documentation publique
[Bêta] Authentification au nom de l’utilisateur : permet à l’agent d’utiliser les informations d’identification de l’utilisateur final pour accéder aux ressources Databricks. Cela est utile pour les scénarios où votre agent doit accéder aux données sensibles ou interroger des API distantes pour effectuer des actions par utilisateur

Spécifier des ressources pour le passage automatique de l’authentification (authentification système)

Pour des types de ressources Databricks les plus courantes, Databricks prend en charge et recommande de déclarer des dépendances de ressources pour l’assistant au cours de la journalisation. Cela permet l’authentification directe automatique lorsque vous déployez l’agent. Databricks provisionne, fait pivoter et gère automatiquement les informations d’identification de courte durée pour accéder en toute sécurité à ces dépendances de ressources à partir du point de terminaison de l’agent.

Pour activer la passe directe d’authentification automatique, spécifiez les ressources dépendantes à l’aide du resources paramètre de l’API log_model() , comme indiqué dans le code suivant.

import mlflow
from mlflow.models.resources import (
  DatabricksVectorSearchIndex,
  DatabricksServingEndpoint,
  DatabricksSQLWarehouse,
  DatabricksFunction,
  DatabricksGenieSpace,
  DatabricksTable,
  DatabricksUCConnection,
  DatabricksApp
)

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    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"),
    ]
  )

Databricks vous recommande de spécifier manuellement resources pour toutes les versions d'agent.

Remarque

Si vous ne spécifiez pas de ressources lors de la journalisation des agents LangChain à l’aide de mlflow.langchain.log_model(...), MLflow effectue une inférence automatique optimale des ressources. Cependant, toutes les dépendances ne sont pas capturées, ce qui entraîne des erreurs d’autorisation lors du traitement ou de l’interrogation de l’assistant.

Le tableau suivant répertorie les ressources Databricks qui prennent en charge la passe d’authentification automatique et la version minimale mlflow requise pour enregistrer la ressource.

Type de ressource	Version minimale `mlflow` requise pour journaliser la ressource
Index de recherche vectorielle	Nécessite la version 2.13.1 ou ultérieure de `mlflow`
Point de terminaison de mise en service	Nécessite la version 2.13.1 ou ultérieure de `mlflow`
Entrepôt SQL	Nécessite la version 2.16.1 ou ultérieure de `mlflow`
Fonction Unity Catalog	Nécessite la version 2.16.1 ou ultérieure de `mlflow`
Espace Génie	Nécessite la version 2.17.1 ou ultérieure de `mlflow`
Tableau du catalogue Unity	Nécessite `mlflow` la version 2.18.0 ou ultérieure
Connexion du catalogue Unity	Nécessite la version 2.17.1 ou ultérieure de `mlflow`

Authentification au nom de l’utilisateur

Important

Cette fonctionnalité est en version bêta.

Lors de la journalisation d’un agent qui utilise l’authentification au nom de l’utilisateur, spécifiez le jeu minimal d’étendues d’API Databricks nécessaires pour effectuer des actions en tant qu’utilisateur final dans le code de l’agent. Cela garantit que l’agent dispose d’un accès avec des privilèges minimum pour effectuer des actions pour le compte de l’utilisateur final lors du déploiement, en améliorant la sécurité en empêchant les actions non autorisées et en réduisant le risque d’utilisation incorrecte des jetons.

Vous trouverez ci-dessous une liste d’étendues requises pour accéder à plusieurs types courants de ressources Databricks.

Ressource Databricks	Étendues d’API requises
Index de recherche vectorielle	`serving.serving-endpoints`,`vectorsearch.vector-search-endpoints`,`vectorsearch.vector-search-indexes`
Point de terminaison de mise en service de modèles	`serving.serving-endpoints`
SQL Warehouse	`sql.statement-execution`,`sql.warehouses`
Connexions UC	`catalog.connections`
L'Espace des Génies	`dashboards.genie`

Pour activer l’authentification au nom de l’utilisateur, transmettez un MLflow AuthPolicy à log_model(), comme illustré dans l’exemple ci-dessous. Un MLflow AuthPolicy a deux composants :

system_auth_policy: spécifiez les ressources pour l’authentification système. En règle générale, les agents utilisent l’authentification système pour les ressources partagées (par exemple, pour interroger des points de terminaison de service de modèle) en combinaison avec l’authentification au nom de l’utilisateur pour accéder aux ressources sensibles ou aux API
user_auth_policy: Spécifier les étendues d’API dont votre agent a besoin pour l’authentification au nom de l’utilisateur

from mlflow.models.resources import DatabricksServingEndpoint
from mlflow.models.auth_policy import SystemAuthPolicy, UserAuthPolicy, AuthPolicy

resources = [
    DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct")
]
# Specify resources here for system authentication
system_auth_policy = SystemAuthPolicy(resources=resources)

# Specify the minimal set of API scopes needed for on-behalf-of-user authentication
	# When deployed, the agent can access Databricks resources and APIs
	# on behalf of the end user, but only via REST APIs that are covered by the list of
	# scopes below

user_auth_policy = UserAuthPolicy(
    api_scopes=[
        "serving.serving-endpoints",
        "vectorsearch.vector-search-endpoints",
        "vectorsearch.vector-search-indexes",
    ]
)

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        ...
        # Instead of passing `resources` (which only supports system authentication),
	        # pass an auth_policy to log_model to enable both system authentication and
	        # on-behalf-of-user authentication
        auth_policy=AuthPolicy(
            system_auth_policy=system_auth_policy,
            user_auth_policy=user_auth_policy
        )
    )

Authentification automatique pour les clients OpenAI

Si votre agent utilise le client OpenAI, utilisez le Kit de développement logiciel (SDK) Databricks pour s’authentifier automatiquement pendant le déploiement. Databricks SDK fournit un wrapper pour construire le client OpenAI avec l’autorisation configurée automatiquement. get_open_ai_client() Exécutez ce qui suit dans votre bloc-notes :

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

Inscrire l’agent dans le catalogue Unity

Avant de déployer l’agent, vous devez inscrire l’agent dans le catalogue Unity. L’inscription de l’assistant génère un package en tant que modèle dans Unity Catalog. Par conséquent, vous pouvez utiliser les autorisations du catalogue Unity pour l’autorisation des ressources dans l’agent.

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "agent_name"

model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=model_name)

Voir mlflow.register_model().