追踪 Gemini

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

import mlflow

mlflow.gemini.autolog()

注意

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

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

• Prompt 和完成响应
• 延迟
• 模型名称
• 附加元数据，例如 temperature、max_tokens（如果已指定）。
• 响应中返回的函数调用（如果适用）
• 抛出的任何异常（如果适用）

注意

目前，MLflow Gemini 集成仅支持对文本交互的同步调用进行追踪。异步 API 不会被追踪，对于多模态输入，可能无法记录完整的输入。

基本示例

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("http://localhost: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)

Embedding

MLflow 对 Gemini SDK 的追踪支持 embedding API

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

禁用自动追踪

通过调用 mlflow.gemini.autolog(disable=True) 或 mlflow.autolog(disable=True)，可以在全局范围内禁用 Gemini 的自动追踪。