Journaliser et inscrire des agents IA

Journaliser des agents IA à l’aide de Mosaic AI Agent Framework. La journalisation d’un agent 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 que vous puissiez évaluer la qualité de la configuration.

Spécifications

Créez un agent IA avant de le journaliser.

Journalisation basée sur le code et journalisation basée sur la sérialisation

Vous pouvez utiliser la journalisation MLflow basée sur le code ou la journalisation MLflow basée sur la sérialisation. Databricks vous recommande d’utiliser la journalisation basée sur le code.

Journalisation MLflow basée sur le code : le code de la chaîne est capturé sous la forme d’un fichier Python. L’environnement Python est capturé sous la forme d’une liste de packages. Lorsque la chaîne est déployée, l’environnement Python est restauré, et le code de la chaîne est exécuté pour charger la chaîne en mémoire afin qu’elle puisse être appelée lorsque le point de terminaison est appelé.

Journalisation MLflow basée sur la sérialisation : le code et l’état actuel de la chaîne dans l’environnement Python sont sérialisés sur disque, souvent à l’aide de bibliothèques telles que pickle ou joblib. Lorsque la chaîne est déployée, l’environnement Python est restauré, et l’objet sérialisé est chargé en mémoire afin de pouvoir être appelé lorsque le point de terminaison est appelé.

Le tableau montre les avantages et les inconvénients de chaque méthode.

Pour la journalisation basée sur le code, le code qui journalise votre agent ou chaîne doit se trouver dans un notebook distinct de votre code de chaîne. Ce notebook est appelé notebook pilote. Pour obtenir un exemple de notebook, consultez Exemples de notebooks.

Journalisation basée sur le code avec LangChain

1. Créez un notebook ou un fichier Python avec votre code. Pour les besoins de cet exemple, le notebook ou le fichier se nomme chain.py. Le notebook ou le fichier doit contenir une chaîne LangChain, appelée ici lc_chain.
2. Incluez mlflow.models.set_model(lc_chain) dans le notebook 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 mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”) pour exécuter chain.py et journaliser les résultats dans un modèle MLflow.
5. Déployer le modèle. Consultez Déployer un agent pour une application d’IA générative. Le déploiement de votre agent peut dépendre d’autres ressources Databricks, telles qu’un index de recherche vectorielle et des points de terminaison de service de modèle. Pour les agents LangChain :
   • MLflow log_model déduit les dépendances requises par la chaîne et les journalise dans le fichier MLmodel dans l’artefact de modèle journalisé. À compter de Mlflow version 2.17.0, vous pouvez remplacer ces dépendances déduites. Consultez Spécifier des ressources pour l’agent PyFunc ou LangChain.
   • Pendant le déploiement, databricks.agents.deploy crée automatiquement les jetons OAuth M2M requis pour accéder à ces dépendances de ressources déduites et communiquer avec ces dépendances de ressources déduites.
6. Lorsque l’environnement de service est chargé, chain.py est exécuté.
7. Lorsqu’une requête de service arrive, lc_chain.invoke(...) est appelé.

import mlflow

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

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using LangChain
with mlflow.start_run():
  logged_chain_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this is the configuration that is used for training the model. The development_config is overwritten.
    artifact_path="chain", # 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 your chain
    example_no_conversion=True, # Required
  )

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

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

Journalisation basée sur le code avec PyFunc

1. Créez un notebook ou un fichier Python avec votre code. Pour les besoins de cet exemple, le notebook ou le fichier se nomme chain.py. Le notebook ou le fichier doit contenir une classe PyFunc, appelée ici PyFuncClass.
2. Incluez mlflow.models.set_model(PyFuncClass) dans le notebook 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 mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”, resources=”/path/to/resources.yaml”) pour exécuter chain.py et journaliser les résultats dans un modèle MLflow. Le resources paramètre déclare toutes les ressources nécessaires pour servir le modèle, comme un index de recherche vectorielle ou un point de terminaison servant un modèle de base. Pour obtenir un exemple de fichier de ressources pour PyFunc, consultez Spécifier des ressources pour l’agent PyFunc ou LangChain.
5. Déployer le modèle. Consultez Déployer un agent pour une application d’IA générative.
6. Lorsque l’environnement de service est chargé, chain.py est exécuté.
7. Lorsqu’une requête de service arrive, PyFuncClass.predict(...) est appelé.

import mlflow

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

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using PyFunc model

resources_path = "/Workspace/Users/first.last/resources.yml"

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
  )

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

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

Spécifier des ressources pour l’agent PyFunc ou LangChain

Vous pouvez spécifier des ressources, telles qu’un index de recherche vectorielle et un point de terminaison de service, qui sont nécessaires pour servir le modèle.

Pour LangChain, les ressources sont automatiquement détectées et enregistrées avec le modèle à l’aide d’une approche optimale. À compter de MLflow version 2.17.0, vous pouvez remplacer ces ressources déduites automatiquement à l’aide du code similaire à celui indiqué ci-dessous. Cela est recommandé pour les cas d’utilisation de production, car il vous permet de vous assurer que les agents sont enregistrés avec les dépendances nécessaires.

Lors du déploiement d’un agent pyfunc aromatisé, vous devez ajouter manuellement toutes les dépendances de ressources de l’agent déployé. Un jeton OAuth M2M avec accès à toutes les ressources spécifiées dans le paramètre resources est créé et fourni à l’agent déployé.

Le code suivant spécifie les dépendances à l’aide du resources paramètre.

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

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    example_no_conversion=True,
    resources=[
      DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
      DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
      DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
      DatabricksFunction(function_name="ml.tools.python_exec"),
    ]
  )

Vous pouvez également ajouter des ressources en les spécifiant dans un fichier resources.yaml. Vous pouvez référencer ce chemin d’accès de fichier dans le paramètre resources. Un jeton OAuth M2M avec accès à toutes les ressources spécifiées dans resources.yaml est créé et fourni à l’agent déployé.

Voici un exemple de fichier resources.yaml qui définit des points de terminaison de service de modèle et un index de recherche vectorielle.

api_version: "1"
databricks:
  vector_search_index:
    - name: "catalog.schema.my_vs_index"
  serving_endpoint:
    - name: databricks-dbrx-instruct
    - name: databricks-bge-large-en

Inscrire la chaîne auprès de Unity Catalog

Avant de déployer la chaîne, vous devez l’inscrire auprès de Unity Catalog. Lorsque vous inscrivez la chaîne, elle est empaquetée en tant que modèle dans Unity Catalog, et vous pouvez utiliser des autorisations Unity Catalog pour l’autorisation pour les ressources contenues dans la chaîne.

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "chain_name"

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

Étapes suivantes

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