Trace MLflow AI Gateway

MLflow AI Gateway 是一个统一的、集中的接口，用于访问多个 LLM 提供商。它简化了 API 密钥管理，在不同提供商之间提供了统一的 API，并实现了 OpenAI、Anthropic、Google 等提供商的模型之间的无缝切换。

由于 MLflow AI Gateway 暴露了一个与 OpenAI 兼容的 API，您可以使用 MLflow 的自动跟踪集成来捕获 LLM 交互的详细跟踪记录。

集成选项

有两种方法可以跟踪通过 MLflow AI Gateway 的 LLM 调用

方法	描述	最适合
服务器端跟踪 (即将推出)	Gateway 自动记录所有请求	对通过 Gateway 的所有请求进行集中跟踪
客户端跟踪	使用 OpenAI SDK 配合 MLflow 自动日志记录	将 LLM 跟踪与您的代理或应用程序跟踪相结合

即将推出

MLflow AI Gateway 的服务器端跟踪尚不可用。请关注更新！

先决条件

使用 AI Gateway 启动 MLflow 服务器

要使用 AI Gateway 启动 MLflow 服务器，您需要安装 mlflow[genai] 包。

bash
pip install mlflow[genai]

然后像往常一样启动 MLflow 服务器，无需额外配置。

bash
mlflow server

创建端点

在 MLflow AI Gateway 中创建一个端点，将请求路由到您的 LLM 提供商。有关详细的设置说明，请参阅 AI Gateway 快速入门。

查询 Gateway

您可以使用以下任何方法来跟踪通过 MLflow AI Gateway 的 LLM 调用

OpenAI SDK
Anthropic SDK
Gemini SDK
LangChain

由于 MLflow AI Gateway 暴露了一个与 OpenAI 兼容的 API，您可以使用 MLflow 的 OpenAI 自动跟踪集成来跟踪调用。

python
import mlflow
from openai import OpenAI

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Set MLflow tracking URI and experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("MLflow AI Gateway")

# Point OpenAI client to MLflow AI Gateway
client = OpenAI(
    base_url="https://:5000/gateway/openai/v1",
    api_key="dummy",  # API key not needed, configured server-side
)

response = client.chat.completions.create(
    model="my-endpoint", messages=[{"role": "user", "content": "Hello!"}]
)
print(response.choices[0].message.content)

您可以使用 MLflow 的 Anthropic 自动跟踪集成来跟踪通过 MLflow AI Gateway 的调用。

python
import mlflow
import anthropic

# Enable auto-tracing for Anthropic
mlflow.anthropic.autolog()

# Set MLflow tracking URI and experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("MLflow AI Gateway")

# Point Anthropic client to MLflow AI Gateway
client = anthropic.Anthropic(
    base_url="https://:5000/gateway/anthropic",
    api_key="dummy",  # API key not needed, configured server-side
)

# Make API calls - traces will be captured automatically
response = client.messages.create(
    max_tokens=1004,
    model="<your-endpoint-name>",  # Use your endpoint name
    messages=[{"role": "user", "content": "Hello world"}],
)
print(response.content[0].text)

您可以使用 MLflow 的 Gemini 自动跟踪集成来跟踪通过 MLflow AI Gateway 的调用。

python
import mlflow
from google import genai

# Enable auto-tracing for Gemini
mlflow.gemini.autolog()

# Set MLflow tracking URI and experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("MLflow AI Gateway")

# Point Gemini client to MLflow AI Gateway
client = genai.Client(
    http_options={"base_url": "https://:5000/gateway/gemini"},
    api_key="dummy",  # API key not needed, configured server-side
)

# Make API calls - traces will be captured automatically
response = client.models.generate_content(
    model="<your-endpoint-name>",  # Use your endpoint name
    contents={"text": "Hello!"},
)
print(response.text)

您可以使用 MLflow 的 LangChain 自动跟踪集成来跟踪通过 MLflow AI Gateway 的调用。

python
import mlflow
from langchain.agents import create_agent
from langchain_openai import ChatOpenAI

# Enable auto-tracing for LangChain
mlflow.langchain.autolog()

# Set MLflow tracking URI and experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("MLflow AI Gateway")

# Point LangChain to MLflow AI Gateway
llm = ChatOpenAI(
    base_url="https://:5000/gateway/mlflow/v1",
    model="<your-endpoint-name>",  # Use your endpoint name
    api_key="dummy",  # API key not needed, configured server-side
)


def get_weather(city: str) -> str:
    """Get weather for a given city."""
    return f"It's always sunny in {city}!"


# Run the agent as usual
agent = create_agent(
    llm,
    tools=[get_weather],
    system_prompt="You are a helpful assistant",
)
agent.invoke({"input": "What's the weather in San Francisco?"})