Tracing Strands Agents SDK

Strands Agents SDK 是 AWS 开发的一款开源、模型驱动的 SDK，开发者只需在几行代码中定义模型、一组工具和提示，即可轻松创建自主 AI 代理。

MLflow Tracing 为 Strands Agents SDK 提供了自动跟踪功能。通过调用 mlflow.strands.autolog() 函数为 Strands Agents SDK 启用自动跟踪，MLflow 将捕获 Agent 调用跟踪并将其记录到活动 MLflow 实验中。

import mlflow

mlflow.strands.autolog()

MLflow 跟踪会自动捕获 Agentic 调用相关的以下信息：

• 提示和完成响应
• 延迟
• 有关不同 Agent 的元数据，例如函数名称
• Token 使用量和成本
• 缓存命中
• 如果抛出任何异常

基本示例

import mlflow

mlflow.strands.autolog()
mlflow.set_experiment("Strand Agent")

from strands import Agent
from strands.models.openai import OpenAIModel
from strands_tools import calculator

model = OpenAIModel(
    client_args={"api_key": "<api-key>"},
    # **model_config
    model_id="gpt-4o",
    params={
        "max_tokens": 2000,
        "temperature": 0.7,
    },
)

agent = Agent(model=model, tools=[calculator])
response = agent("What is 2+2")
print(response)

Token 用量

MLflow >= 3.4.0 支持对 Strand Agent SDK 进行 Token 使用量跟踪。每个 Agent 调用的 Token 使用量将记录在 mlflow.chat.tokenUsage 属性中。整个跟踪过程的总 Token 使用量可在跟踪信息对象的 token_usage 字段中找到。

response = agent("What is 2+2")
print(response)

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: 2629
  Output tokens: 31
  Total tokens: 2660

== Detailed usage for each LLM call: ==
chat_1:
  Input tokens: 1301
  Output tokens: 16
  Total tokens: 1317
chat_2:
  Input tokens: 1328
  Output tokens: 15
  Total tokens: 1343

禁用自动跟踪

可以通过调用 mlflow.strands.autolog(disable=True) 或 mlflow.autolog(disable=True) 来全局禁用 Strands Agent SDK 的自动跟踪。