跟踪 Smolagents

Smolagents 是一个轻型代理框架，强调极简主义和可组合性。

MLflow 跟踪 与 Smolagents 集成，以捕获轻型代理工作流的简化跟踪。 使用 mlflow.smolagents.autolog 启用它。

先决条件

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

开发

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

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

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

生产

对于生产部署，请安装和 mlflow-tracing Smolagents：

pip install --upgrade mlflow-tracing smolagents openai

包 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

示例用法

import mlflow

mlflow.smolagents.autolog()

from smolagents import CodeAgent, LiteLLMModel
import mlflow

# Turn on auto tracing for Smolagents by calling mlflow.smolagents.autolog()
mlflow.smolagents.autolog()

model = LiteLLMModel(model_id="openai/gpt-4o-mini", api_key=API_KEY)
agent = CodeAgent(tools=[], model=model, add_base_tools=True)

result = agent.run(
    "Could you give me the 118th number in the Fibonacci sequence?",
)

照常运行 Smolagents 工作流。 跟踪记录将显示在试验 UI 中。

令牌跟踪使用情况

MLflow 记录每次代理调用至 mlflow.chat.tokenUsage 的令牌使用情况。 整个跟踪的总令牌使用情况在跟踪信息对象的 token_usage 字段中提供。

import json
import mlflow

mlflow.smolagents.autolog()

model = LiteLLMModel(model_id="openai/gpt-4o-mini", api_key=API_KEY)
agent = CodeAgent(tools=[], model=model, add_base_tools=True)

result = agent.run(
    "Could you give me the 118th number in the Fibonacci sequence?",
)

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

禁用自动跟踪

使用 mlflow.smolagents.autolog(disable=True) 禁用，或使用 mlflow.autolog(disable=True) 全局禁用。