追踪 LiteLLM🚄

LiteLLM 是一个开源的 LLM 网关，允许通过统一接口访问 100+ LLM。

MLflow Tracing 为 LiteLLM 提供了自动追踪能力。通过调用 mlflow.litellm.autolog() 函数启用 LiteLLM 的自动追踪后，MLflow 将捕获 LLM 调用的追踪并将其记录到当前活跃的 MLflow Experiment 中。

import mlflow

mlflow.litellm.autolog()

MLflow trace 自动捕获关于 LiteLLM 调用的以下信息：

• Prompts 和完成响应
• 延迟
• 关于 LLM 提供者的元数据，例如模型名称和端点 URL
• Token 用量和成本
• 缓存命中
• 发生的任何异常

基础示例

import mlflow
import litellm

# Enable auto-tracing for LiteLLM
mlflow.litellm.autolog()

# Optional: Set a tracking URI and an experiment
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("LiteLLM")

# Call Anthropic API via LiteLLM
response = litellm.completion(
    model="claude-3-5-sonnet-20240620",
    messages=[{"role": "user", "content": "Hey! how's it going?"}],
)

异步 API

MLflow 支持追踪 LiteLLM 的异步 API。

mlflow.litellm.autolog()

response = await litellm.acompletion(
    model="claude-3-5-sonnet-20240620",
    messages=[{"role": "user", "content": "Hey! how's it going?"}],
)

流式传输

MLflow 支持追踪 LiteLLM 的同步和异步流式传输 API。

mlflow.litellm.autolog()

response = litellm.completion(
    model="claude-3-5-sonnet-20240620",
    messages=[{"role": "user", "content": "Hey! how's it going?"}],
    stream=True,
)
for chunk in response:
    print(chunk.choices[0].delta.content, end="|")

MLflow 会将流块的连接输出记录为跨度（span）输出。

禁用自动追踪

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