追踪 Strands 代理 SDK

Strands Agents SDK 是 AWS 提供的开源 SDK，用于创建可与外部工具和 API 交互的自治 AI 代理。

MLflow 跟踪 为 Strands 代理 SDK 提供自动跟踪功能。 通过调用 mlflow.strands.autolog 函数为 Strands 启用自动跟踪，MLflow 将捕获跟踪，并在调用代理时将其记录到活动的 MLflow 试验。

import mlflow

mlflow.strands.autolog()

MLflow 跟踪自动捕获有关 Strands 代理调用的以下信息：

• 提示和响应
• 潜伏期
• 代理元数据
• 令牌使用情况和成本
• 缓存命中信息
• 引发的任何异常

先决条件

若要将 MLflow 跟踪与 Strands 代理 SDK 配合使用，需要安装 MLflow、Strands SDK 和所需的依赖项。

开发

对于开发环境，请安装包含 Databricks 附加组件和 Strands 包的完整 MLflow 包：

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

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

生产

对于生产部署，请安装 mlflow-tracing 和 Strands 包。

pip install --upgrade mlflow-tracing strands strands-tools

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

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

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

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

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

API 密钥：确保配置 LLM 提供程序 API 密钥。 对于生产环境，请使用 马赛克 AI 网关或 Databricks 机密，而不要使用硬编码的值，以便更安全地管理 API 密钥。

export OPENAI_API_KEY="your-openai-api-key"
# Add other provider keys as needed

示例用法

以下示例演示如何将 Strands Agents SDK 与 MLflow 跟踪配合使用。 代理使用 OpenAI 模型，并且有权访问计算器工具来执行算术运算。

import mlflow
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured

# Enable auto tracing for Strands Agents SDK
mlflow.strands.autolog()

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/strands-agent-demo")

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

# Configure the OpenAI model
model = OpenAIModel(
    client_args={"api_key": os.environ.get("OPENAI_API_KEY")},
    model_id="gpt-4o",
    params={
        "max_tokens": 2000,
        "temperature": 0.7,
    },
)

# Create an agent with the calculator tool
agent = Agent(model=model, tools=[calculator])

# Run the agent
response = agent("What is 2+2?")
print(response)

令牌使用情况跟踪

使用 MLflow 3.4.0 或更高版本时，MLflow 会自动跟踪 Strands 代理的令牌使用情况。 令牌使用信息包括代理执行期间使用的输入令牌、输出令牌和总令牌。

import mlflow

mlflow.strands.autolog()

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

model = OpenAIModel(
    client_args={"api_key": os.environ.get("OPENAI_API_KEY")},
    model_id="gpt-4o",
    params={
        "max_tokens": 2000,
        "temperature": 0.7,
    },
)

agent = Agent(model=model, tools=[calculator])

# Run the agent and retrieve trace information
with mlflow.start_span(name="strands_agent_run") as span:
    response = agent("Calculate the sum of 15 and 27")
    print(response)

# Token usage is automatically logged and visible in the MLflow UI
trace_info = mlflow.get_last_active_trace()
print(f"Trace ID: {trace_info.request_id}")

令牌使用情况详细信息显示在 MLflow 跟踪 UI 中，允许监视和优化代理的性能和成本。

禁用自动跟踪

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