Skip to content

ADK 智能体的 MLflow AI 网关

Supported in ADKPython

MLflow AI Gateway 是一个基于数据库的 LLM 代理，内置于 MLflow 跟踪服务器（MLflow ≥ 3.0）。它为数十家提供商提供统一的兼容 OpenAI 的 API，包括 Gemini、Anthropic、Mistral、Bedrock、Ollama 等，具有内置的密钥管理、回退/重试、流量拆分和预算跟踪功能，全部通过 MLflow UI 配置。

由于 MLflow AI Gateway 公开了一个兼容 OpenAI 的端点，你可以使用 LiteLLM 模型连接器将 ADK 智能体连接到它。

使用场景

  • 多提供商路由：切换 LLM 提供商而无需更改智能体代码
  • 密钥管理：提供商 API 密钥在服务器上加密存储；你的应用程序不发送提供商密钥
  • 回退和重试：失败时自动故障转移到备份模型
  • 预算跟踪：每个端点或每个用户的令牌预算
  • 流量拆分：将请求百分比路由到不同模型以进行 A/B 测试
  • 使用跟踪：每次调用自动记录为 MLflow 跟踪

前置条件

  • MLflow 3.0 或更高版本
  • 在你的环境中安装了 Google ADK 和 LiteLLM

设置

安装依赖项：

pip install mlflow[genai] google-adk litellm

启动 MLflow 服务器：

mlflow server --host 127.0.0.1 --port 5000

MLflow UI 将在 http://localhost:5000 上可用。

通过导航到 http://localhost:5000 的 MLflow UI 创建网关端点，然后转到 AI 网关 → 创建端点。选择提供商（例如 Google Gemini）和模型（例如 gemini-flash-latest），然后输入你的提供商 API 密钥，该密钥将在服务器上加密存储。

MLflow AI Gateway - Create Endpoint

See the MLflow AI Gateway documentation for more details on endpoint configuration.

与智能体配合使用

使用 LiteLlm 包装器，并将 api_base 指向 MLflow 网关的端点。model 参数应使用 openai/ 前缀后跟你的网关端点名称。

from google.adk.agents import LlmAgent
from google.adk.models.lite_llm import LiteLlm

# 指向 MLflow AI Gateway 端点。
# "my-chat-endpoint" 是你在 MLflow UI 中创建的端点名称。
agent = LlmAgent(
    model=LiteLlm(
        model="openai/my-chat-endpoint",
        api_base="http://localhost:5000/gateway/openai/v1",
        api_key="unused",  # 提供商密钥由 MLflow 服务器管理
    ),
    name="gateway_agent",
    instruction="你是一个由 MLflow AI Gateway 驱动的得力助手。",
)

你可以在 MLflow UI 中重新配置网关端点来随时切换底层 LLM 提供商，无需更改 ADK 智能体代码。

提示

  • api_key 参数是 LiteLLM 必需的，但网关不会验证它。将其设置为任何非空字符串。
  • 在代理后面或远程主机上时，将 localhost:5000 替换为你的服务器地址。
  • 结合 MLflow Tracing 为你的 ADK 智能体提供端到端的可观测性。

资源