LiteLLM 是一个开源 LLM 网关，允许访问统一接口中的 100 多个 LLM。
MLflow 跟踪 为 LiteLLM 提供自动跟踪功能。 通过调用 mlflow.litellm.autolog 函数为 LiteLLM 启用自动跟踪，MLflow 将捕获 LLM 调用的跟踪并将其记录到活动 MLflow 试验。
import mlflow
mlflow.litellm.autolog()
MLflow 跟踪自动捕获有关 LiteLLM 调用的以下信息：
- 提示和补全应答
- 潜伏期
- 有关 LLM 提供程序的元数据，例如模型名称和终结点 URL
- 令牌使用情况和成本
- 缓存命中
- 引发的任何异常
先决条件
在运行以下示例之前，请确保具备：
Databricks 凭据已配置：如果您在 Databricks 环境外运行，请设置环境变量：
export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com" export DATABRICKS_TOKEN="your-personal-access-token"小窍门
如果在 Databricks 笔记本中运行，系统会自动设置这些内容。
LLM 提供程序 API 密钥：为正在使用的 LLM 提供程序设置 API 密钥：
export ANTHROPIC_API_KEY="your-anthropic-api-key" # For Anthropic models export OPENAI_API_KEY="your-openai-api-key" # For OpenAI models # Add other provider keys as needed
基本示例
import mlflow
import litellm
# Enable auto-tracing for LiteLLM
mlflow.litellm.autolog()
# Set up MLflow tracking on Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/litellm-demo")
# Call Anthropic API via LiteLLM
response = litellm.completion(
model="claude-3-5-sonnet-20241022",
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-20241022",
messages=[{"role": "user", "content": "Hey! how's it going?"}],
)
流媒体
MLflow 支持跟踪 LiteLLM 的同步和异步流式处理 API：
mlflow.litellm.autolog()
response = litellm.completion(
model="claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hey! how's it going?"}],
stream=True,
)
for chunk in response:
print(chunk.choices[0].delta.content, end="|")
MLflow 会将流区块中的串联输出记录为范围输出。
禁用自动跟踪
可以通过调用 mlflow.litellm.autolog(disable=True) 或 mlflow.autolog(disable=True)调用来全局禁用 LiteLLM 的自动跟踪。