追踪 Smolagents

MLflow Tracing 在使用 Smolagents 时提供自动追踪功能。当通过调用 mlflow.smolagents.autolog() 函数启用 Smolagents 自动追踪时，Smolagents SDK 的使用将自动记录交互式开发期间生成的追踪。

请注意，仅支持同步调用，异步 API 和流式方法不受追踪。

示例用法

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?",
)

Token 用量

MLflow >= 3.2.0 支持 Smolagents 的 token 用量追踪。每次 LLM 调用的 token 用量将记录在 mlflow.chat.tokenUsage 属性中。整个追踪的 token 总用量将在追踪信息对象的 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']}")

== Total token usage: ==
  Input tokens: 4360
  Output tokens: 185
  Total tokens: 4545

== Detailed usage for each LLM call: ==
LiteLLMModel.__call___1:
  Input tokens: 2047
  Output tokens: 124
  Total tokens: 2171
LiteLLMModel.__call___2:
  Input tokens: 2313
  Output tokens: 61
  Total tokens: 2374

禁用自动跟踪

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