跟踪 Gemini

MLflow 跟踪为 Google Gemini 提供自动跟踪功能。通过调用 mlflow.gemini.autolog() 函数启用 Gemini 的自动跟踪，MLflow 将捕获嵌套跟踪并在调用 Gemini Python SDK 时将它们记录到活动的 MLflow 实验中。在 Typescript 中，您可以使用 tracedGemini 函数来包装 Gemini 客户端。

MLflow 跟踪自动捕获有关 Gemini 调用的以下信息

• 提示和完成响应
• 延迟
• 模型名称
• 如果指定，还包括额外的元数据，如 temperature、max_tokens。
• Token 用量（输入、输出和总 Token 数）
• 如果响应中返回函数调用
• 如果抛出任何异常

开始使用

1

安装依赖项

Python

bash

pip install 'mlflow[genai]' google-generativeai

JS / TS

bash

npm install mlflow-gemini @google/generative-ai

2

启动 MLflow 服务器

本地 (pip)

如果您有本地 Python 环境 >= 3.10，您可以使用 mlflow CLI 命令在本地启动 MLflow 服务器。

bash

mlflow server

本地 (docker)

MLflow 还提供了一个 Docker Compose 文件，用于启动带有 postgres 数据库和 minio 服务器的本地 MLflow 服务器。

bash

git clone --depth 1 --filter=blob:none --sparse https://github.com/mlflow/mlflow.git
cd mlflow
git sparse-checkout set docker-compose
cd docker-compose
cp .env.dev.example .env
docker compose up -d

有关更多详细信息，请参阅 说明，例如覆盖默认环境变量。

3

启用跟踪并进行 API 调用

Python

使用 mlflow.gemini.autolog() 启用跟踪，并像往常一样进行 API 调用。

python

import mlflow
import google.generativeai as genai
import os

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

# Set a tracking URI and an experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("Gemini")

# Configure your API key
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Use Gemini as usual - traces will be automatically captured
model = genai.GenerativeModel("gemini-2.5-flash")
response = model.generate_content("What is the capital of France?")
print(response.text)

JS / TS

使用 tracedGemini 函数包装 Gemini 客户端，并像往常一样进行 API 调用。

typescript

import { GoogleGenerativeAI } from "@google/generative-ai";
import { tracedGemini } from "mlflow-gemini";

// Wrap the Gemini client with the tracedGemini function
const genAI = tracedGemini(new GoogleGenerativeAI(process.env.GOOGLE_API_KEY));
const model = genAI.getGenerativeModel({ model: "gemini-2.5-flash" });

// Invoke the client as usual
const result = await model.generateContent("What is the capital of France?");
console.log(result.response.text());

4

在 MLflow UI 中查看跟踪

浏览到 MLflow UI，网址为 https://:5000（或您的 MLflow 服务器 URL），您应该会看到 Gemini API 调用的跟踪。

→ 查看 后续步骤，了解有关更多 MLflow 功能（如用户反馈跟踪、提示管理和评估）的信息。

注意

当前的 MLflow 跟踪集成支持新的 Google GenAI SDK 和旧的 Google AI Python SDK。但是，MLflow 可能会在不通知的情况下停止支持旧包，强烈建议您将用例迁移到新的 Google GenAI SDK。

支持的 API

MLflow 支持对以下 Gemini API 进行自动跟踪

Python

文本生成	聊天	函数调用	流式传输	异步	图像	视频
✅	✅	✅	-	✅ (*1)	-	-

(*1) 异步支持已在 MLflow 3.2.0 中添加。

TypeScript / JavaScript

内容生成	聊天	函数调用	流式传输	异步
✅	-	✅ (*2)	-	✅

(*2) 仅支持 models.generateContent()。捕获响应中的函数调用，并可以在 MLflow UI 中渲染。TypeScript SDK 本质上是异步的。

如需支持更多 API，请在 GitHub 上提交 功能请求。

示例

基本文本生成

Python

python

import mlflow
import google.genai as genai
import os

# Turn on auto tracing for Gemini
mlflow.gemini.autolog()

# Optional: Set a tracking URI and an experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("Gemini")


# Configure the SDK with your API key.
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

# Use the generate_content method to generate responses to your prompts.
response = client.models.generate_content(
    model="gemini-1.5-flash", contents="The opposite of hot is"
)

JS / TS

typescript

import { GoogleGenAI } from "@google/genai";
import { tracedGemini } from "mlflow-gemini";

const client = tracedGemini(new GoogleGenAI({ apiKey: process.env.GEMINI_API_KEY }));

const response = await client.models.generateContent({
    model: "gemini-2.5-flash",
    contents: "What is the capital of France?"
});

多轮聊天交互

MLflow 支持跟踪与 Gemini 的多轮对话

python

import mlflow

mlflow.gemini.autolog()

chat = client.chats.create(model="gemini-1.5-flash")
response = chat.send_message(
    "In one sentence, explain how a computer works to a young child."
)
print(response.text)
response = chat.send_message(
    "Okay, how about a more detailed explanation to a high schooler?"
)
print(response.text)

异步

MLflow 跟踪自 MLflow 3.2.0 起支持 Gemini SDK 的异步 API。用法与同步 API 相同。

Python

python

# Configure the SDK with your API key.
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

# Async API is invoked through the `aio` namespace.
response = await client.aio.models.generate_content(
    model="gemini-1.5-flash", contents="The opposite of hot is"
)

JS / TS

Gemini Typescript / Javascript SDK 本质上是异步的。请参阅上面的基本示例。

Embeddings

MLflow 对 Gemini SDK 的跟踪支持嵌入式 API（仅限 Python）

python

result = client.models.embed_content(model="text-embedding-004", contents="Hello world")

Token 用量

MLflow >= 3.4.0 支持 Gemini 的 Token 使用量跟踪。每次 LLM 调用的 Token 使用量将记录在 mlflow.chat.tokenUsage 属性中。跟踪信息对象中的 token_usage 字段将显示整个跟踪过程中的总 Token 使用量。

Python

python

import json
import mlflow

mlflow.gemini.autolog()

client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

# Use the generate_content method to generate responses to your prompts.
response = client.models.generate_content(
    model="gemini-1.5-flash", contents="The opposite of hot is"
)

# Get the trace object just created
trace = mlflow.get_trace(mlflow.get_last_active_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']}")

JS / TS

typescript

import * as mlflow from "mlflow-tracing";

// After your Gemini call completes, flush and fetch the trace
await mlflow.flushTraces();
const lastTraceId = mlflow.getLastActiveTraceId();

if (lastTraceId) {
  const client = new mlflow.MlflowClient({ trackingUri: "https://:5000" });
  const trace = await client.getTrace(lastTraceId);

  // Total token usage on the trace
  console.log("== Total token usage: ==");
  console.log(trace.info.tokenUsage); // { input_tokens, output_tokens, total_tokens }

  // Per-span usage (if provided by the provider)
  console.log("\n== Detailed usage for each LLM call: ==");
  for (const span of trace.data.spans) {
    const usage = span.attributes?.["mlflow.chat.tokenUsage"];
    if (usage) {
      console.log(`${span.name}:`, usage);
    }
  }
}

bash

== Total token usage: ==
  Input tokens: 5
  Output tokens: 2
  Total tokens: 7

== Detailed usage for each LLM call: ==
Models.generate_content:
  Input tokens: 5
  Output tokens: 2
  Total tokens: 7
Models._generate_content:
  Input tokens: 5
  Output tokens: 2
  Total tokens: 7

Token 使用量跟踪同时支持 Python 和 TypeScript/JavaScript 实现。

禁用自动跟踪

可以通过调用 mlflow.gemini.autolog(disable=True) 或 mlflow.autolog(disable=True) 全局禁用 Gemini 的自动跟踪。

后续步骤

跟踪用户反馈

记录用户对跟踪的反馈，以跟踪用户满意度。

了解反馈 →

管理提示

了解如何使用 MLflow 的提示注册表管理提示。

管理提示 →

评估跟踪

使用 LLM 裁判评估跟踪，以了解和改进 AI 应用的行为。

评估跟踪 →