Tracing Strands Agents SDK

Strands Agents SDK 是 AWS 开发的一款开源、模型驱动的 SDK，它使开发人员只需用几行代码定义模型、一组工具和提示，即可创建自主 AI 智能体。

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

python

import mlflow

mlflow.strands.autolog()

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

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

基本示例

python

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 支持 Strands Agent SDK 的 Token 用量跟踪。每次 Agent 调用的 Token 用量将记录在 mlflow.chat.tokenUsage 属性中。整个跟踪的总 Token 用量可在跟踪信息对象的 token_usage 字段中查看。

python

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']}")

bash

== 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 的自动跟踪。