追踪 AutoGen

AutoGen AgentChat 是一个用于构建对话式单智能体和多智能体应用程序的开源框架。

MLflow 追踪为 AutoGen AgentChat 提供了自动追踪功能。通过调用 mlflow.autogen.autolog() 函数为 AutoGen 启用自动追踪后，MLflow 将在智能体执行时捕获嵌套的追踪信息，并将其记录到活动的 MLflow 实验中。

import mlflow

mlflow.autogen.autolog()

注意

请注意，mlflow.autogen.autolog() 应在导入被追踪的 AutoGen 类之后调用。ChatCompletionClient 的子类（如 OpenAIChatCompletionClient 或 AnthropicChatCompletionClient）和 BaseChatAgent 的子类（如 AssistantAgent 或 CodeExecutorAgent）应在调用 mlflow.autogen.autolog() 之前导入。另请注意，此集成适用于 AutoGen 0.4.9 或更高版本。如果您使用的是 AutoGen 0.2，请改用 AG2 集成。

MLflow 捕获有关 AutoGen 智能体的以下信息

• 传递给智能体的消息，包括图像
• 来自智能体的响应
• 每个智能体进行的 LLM 和工具调用
• 延迟
• 如果抛出任何异常

支持的 API

MLflow 支持对以下 AutoGen API 进行自动追踪。它不支持对异步生成器的追踪。异步流式 API（如 run_stream 或 on_messages_stream）不会被追踪。

• ChatCompletionClient.create
• BaseChatAgent.run
• BaseChatAgent.on_messages

基本示例

import os

# Imports of autogen classes should happen before calling autolog.
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient

import mlflow

# Turn on auto tracing for AutoGen
mlflow.autogen.autolog()

# Optional: Set a tracking URI and an experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("AutoGen")

model_client = OpenAIChatCompletionClient(
    model="gpt-4.1-nano",
    # api_key="YOUR_API_KEY",
)

agent = AssistantAgent(
    name="assistant",
    model_client=model_client,
    system_message="You are a helpful assistant.",
)

result = await agent.run(task="Say 'Hello World!'")
print(result)

工具智能体

import os

# Imports of autogen classes should happen before calling autolog.
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient

import mlflow

# Turn on auto tracing for AutoGen
mlflow.autogen.autolog()

# Optional: Set a tracking URI and an experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("AutoGen")

model_client = OpenAIChatCompletionClient(
    model="gpt-4.1-nano",
    # api_key="YOUR_API_KEY",
)


def add(a: int, b: int) -> int:
    """add two numbers"""
    return a + b


agent = AssistantAgent(
    name="assistant",
    model_client=model_client,
    system_message="You are a helpful assistant.",
    tools=[add],
)

await agent.run(task="1+1")

令牌用量

MLflow >= 3.2.0 支持对 AutoGen 的令牌用量进行追踪。每次 LLM 调用的令牌用量将记录在 mlflow.chat.tokenUsage 属性中。整个追踪过程中的总令牌用量将在追踪信息对象的 token_usage 字段中提供。

import json
import mlflow

mlflow.autogen.autolog()

# Run the tool calling agent defined in the previous section
await agent.run(task="1+1")

# Get the trace object just created
last_trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(trace_id=last_trace_id)

# Print the token usage
total_usage = trace.info.token_usage
print("== Total token usage: ==")
print(f"  Input tokens: {total_usage['input_tokens']}")
print(f"  Output tokens: {total_usage['output_tokens']}")
print(f"  Total tokens: {total_usage['total_tokens']}")

# Print the token usage for each LLM call
print("\n== Detailed usage for each LLM call: ==")
for span in trace.data.spans:
    if usage := span.get_attribute("mlflow.chat.tokenUsage"):
        print(f"{span.name}:")
        print(f"  Input tokens: {usage['input_tokens']}")
        print(f"  Output tokens: {usage['output_tokens']}")
        print(f"  Total tokens: {usage['total_tokens']}")

== Total token usage: ==
  Input tokens: 65
  Output tokens: 17
  Total tokens: 82

== Detailed usage for each LLM call: ==
create:
  Input tokens: 65
  Output tokens: 17
  Total tokens: 82

支持的 API

MLflow 支持对以下 AutoGen API 进行自动令牌用量追踪。它不支持对异步生成器的令牌用量追踪。异步流式 API（如 run_stream 或 on_messages_stream）的用量不会被追踪。

• ChatCompletionClient.create
• BaseChatAgent.run
• BaseChatAgent.on_messages

禁用自动跟踪

可以通过调用 mlflow.autogen.autolog(disable=True) 或 mlflow.autolog(disable=True) 来全局禁用对 AutoGen 的自动追踪。