Journaliser et enregistrer des agents d’IA

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.

Requirements

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 les exemples de notebooks de création de ResponsesAgent.

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

Si vous n’utilisez pas l’interface ResponsesAgent , 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 :

1. Définir manuellement la signature
2. 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.

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

2. Incluez mlflow.models.set_model(lc_agent) dans le bloc-notes ou le fichier.

3. Créez un notebook pour servir de notebook pilote (appelé driver.py dans cet exemple).

4. 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 Implémenter la passe d’authentification automatique.

5. Déployer le modèle. Consultez Déployer un agent pour les applications IA génératives.

6. Lorsque l’environnement de service est chargé, agent.py il est exécuté.

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

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

2. Incluez mlflow.models.set_model(PyFuncClass) dans le carnet ou le fichier.

3. Créez un notebook pour servir de notebook pilote (appelé driver.py dans cet exemple).

4. 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 Implémenter la passe d’authentification automatique.

5. Déployer le modèle. Consultez Déployer un agent pour les applications IA génératives.

6. Lorsque l’environnement de service est chargé, agent.py il est exécuté.

7. 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 déployé peut avoir besoin d’accéder à un index Recherche vectorielle pour interroger des données non structurées ou d’accéder au Registre d’invites pour charger des invites dynamiques.

L’authentification automatique et l’authentification pour le compte de l’utilisateur nécessitent une configuration pendant la connexion de l’agent.

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

Étapes suivantes

• Ajouter des suivis à un agent IA.
• Déployer un agent IA.