追踪 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 调用的以下信息

• 提示和完成响应
• 延迟
• 模型名称
• 如果指定，还包括额外的元数据，如 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("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 对 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 的自动追踪。