Aller au contenu principal
Dernière mise à jour le

Traçage AG2

MLflow Tracing s'intègre à AG2 (anciennement AutoGen 0.2) pour capturer les traces unifiées des conversations multi-agents et des workflows. L'intégration instrumente automatiquement les boucles d'agent et l'exécution d'outils — il suffit d'appeler mlflow.ag2.autolog.

Python
import mlflow

mlflow.ag2.autolog()
remarque

Sur les clusters de compute serverless, l'autologging n'est pas activé automatiquement. Vous devez appeler explicitement mlflow.ag2.autolog() pour activer le traçage automatique de cette intégration.

L'intégration offre une visibilité complète sur :

  • Quel agent est appelé à différentes étapes
  • Messages échangés entre les agents
  • Appels de LLM et d'outils effectués par chaque agent, organisés par agent et par tour
  • Latences
  • Toute exception si levée

Prérequis​

Pour utiliser MLflow Tracing avec AG2, vous devez installer MLflow et les packages AG2 (AutoGen) pertinents.

Pour les environnements de développement, installez le package MLflow complet avec les extras Databricks et AutoGen :

Bash
pip install --upgrade "mlflow[databricks]>=3.1" pyautogen

Le package mlflow[databricks] complet inclut toutes les fonctionnalités pour le développement local et l’expérimentation sur Databricks.

remarque

MLflow 3 est recommandé pour une meilleure expérience de traçage avec AG2.

Avant d'exécuter les exemples, vous devrez configurer votre environnement :

Pour les utilisateurs en dehors des notebooks Databricks : Définissez vos variables d'environnement Databricks :

Bash
export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"

Pour les utilisateurs dans les Notebooks Databricks : ces identifiants sont définis automatiquement pour vous.

Exemple de base​

Python
import os
from typing import Annotated, Literal

from autogen import ConversableAgent
import mlflow

# Enable auto-tracing for AG2
mlflow.ag2.autolog()

# Track to Databricks (optional if already configured)
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/ag2-tracing-demo")

# Define a simple multi-agent workflow using AG2 (AutoGen 0.2)
config_list = [
    {
        "model": "gpt-4o-mini",
        # Requires OPENAI_API_KEY in env for this example
        "api_key": os.environ.get("OPENAI_API_KEY"),
    }
]

Operator = Literal["+", "-", "*", "/"]


def calculator(a: int, b: int, operator: Annotated[Operator, "operator"]) -> int:
    if operator == "+":
        return a + b
    elif operator == "-":
        return a - b
    elif operator == "*":
        return a * b
    elif operator == "/":
        return int(a / b)
    else:
        raise ValueError("Invalid operator")


assistant = ConversableAgent(
    name="Assistant",
    system_message=(
        "You are a helpful AI assistant. You can help with simple calculations. "
        "Return 'TERMINATE' when the task is done."
    ),
    llm_config={"config_list": config_list},
)

user_proxy = ConversableAgent(
    name="Tool Agent",
    llm_config=False,
    is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
)

assistant.register_for_llm(name="calculator", description="A simple calculator")(calculator)
user_proxy.register_for_execution(name="calculator")(calculator)

response = user_proxy.initiate_chat(
    assistant, message="What is (44231 + 13312 / (230 - 20)) * 4?"
)

Suivi de l'utilisation des jetons​

MLflow 3.2.0+ prend en charge le suivi de l'utilisation des jetons pour AG2. L'utilisation par appel est enregistrée dans l'attribut de portée mlflow.chat.tokenUsage ; l'utilisation totale apparaît dans les informations de trace.

Python
import mlflow

last_trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(trace_id=last_trace_id)

total = trace.info.token_usage
print("Input:", total["input_tokens"], "Output:", total["output_tokens"], "Total:", total["total_tokens"])

for span in trace.data.spans:
    usage = span.get_attribute("mlflow.chat.tokenUsage")
    if usage:
        print(span.name, usage)

Désactiver le suivi automatique​

Désactiver l'auto-tracing AG2 avec mlflow.ag2.autolog(disable=True) ou désactiver tout l'autologging avec mlflow.autolog(disable=True).

Sur cette page