跟踪 AG2

MLflow 跟踪 与 AG2 （前 AutoGen 0.2）集成，以捕获多代理对话和工作流的统一跟踪。 集成自动检测代理循环和工具执行 - 只需调用 mlflow.ag2.autolog。

import mlflow

mlflow.ag2.autolog()

集成提供对以下功能的全面可见性：

• 在不同轮次调用哪个代理
• 代理之间传递的消息
• 每个代理的 LLM 和工具调用记录，按照代理和轮次进行整理
• 潜伏期
• 引发的任何异常

先决条件

若要将 MLflow 跟踪与 AG2 配合使用，需要安装 MLflow 和相关 AG2 （AutoGen） 包。

开发

对于开发环境，请使用 Databricks Extras 和 AutoGen 安装完整的 MLflow 包：

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

完整 mlflow[databricks] 包包括用于 Databricks 的本地开发和试验的所有功能。

生产

对于生产环境部署，请安装 mlflow-tracing 和 AutoGen。

pip install --upgrade mlflow-tracing pyautogen

包 mlflow-tracing 已针对生产用途进行优化。

在运行示例之前，需要配置环境：

对于不使用 Databricks 笔记本的用户：设置 Databricks 环境变量：

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

对于 Databricks 笔记本中的用户：这些凭据会自动为您设置。

基本示例

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

跟踪令牌使用情况

MLflow 3.2.0+ 支持 AG2 的令牌使用情况跟踪。 每次调用的使用情况记录在 mlflow.chat.tokenUsage span 属性中，总使用量显示在跟踪信息中。

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)

禁用自动跟踪

禁用 AG2 自动跟踪 mlflow.ag2.autolog(disable=True) 或禁用所有自动日志记录 mlflow.autolog(disable=True)。