追踪 Gemini

MLflow Tracing 为 Google Gemini 提供了自动追踪功能。通过调用 mlflow.gemini.autolog() 函数为 Gemini 启用自动追踪后，MLflow 将在调用 Gemini Python SDK 时捕获嵌套的追踪信息，并将其记录到活动的 MLflow 实验中。

import mlflow

mlflow.gemini.autolog()

注意

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

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

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

支持的 API

MLflow 支持对以下 Anthropic API 进行自动追踪

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

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

如需请求对其他 API 的支持，请在 GitHub 上提交功能请求。

基本示例

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"
)

多轮聊天交互

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

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 3.2.0 起，MLflow Tracing 支持 Gemini SDK 的异步 API。其用法与同步 API 相同。

# 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"
)

嵌入

MLflow Tracing for Gemini SDK 支持嵌入 API

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

禁用自动跟踪

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