MLflow OpenAI 自动日志记录

MLflow 的 OpenAI 特性支持自动日志记录，以确保你的想法的实验、测试和验证能够被动态捕获，而无需用日志记录的样板代码来包装你的代码。

注意

自动日志记录仅支持 OpenAI SDK 的 1.17 及更高版本。

用于 OpenAI SDK 的 MLflow 自动日志记录支持以下接口：

• 聊天补全 (Chat Completions) 通过 client.chat.completions.create()
• 补全 (Completions) (旧版) 通过 client.completions.create()
• 嵌入 (Embeddings) 通过 client.embeddings.create()

其中 client 是 openai.OpenAI() 的一个实例。

在本指南中，我们将讨论自动日志记录功能中的一些关键特性。

• 快速入门
• OpenAI 自动日志记录配置
• 使用 OpenAI 自动日志记录的示例
• OpenAI Swarm 的自动追踪
• 常见问题 (FAQ)

快速入门

要开始使用 MLflow 的 OpenAI 自动日志记录，你只需要在脚本或笔记本的开头调用 mlflow.openai.autolog()。启用自动日志记录时，如果不覆盖任何参数，将按照下一节表格中的 default 配置行为。覆盖这些设置中的任何一项都可以让你记录更多元素。

提示

当自动日志记录激活时，唯一默认启用的元素是追踪信息的记录。你可以在此处阅读更多关于 MLflow 追踪的信息。

import os
import openai
import mlflow

# Enables trace logging by default
mlflow.openai.autolog()

openai_client = openai.OpenAI()

messages = [
    {
        "role": "user",
        "content": "What does turning something up to 11 refer to?",
    }
]

# The input messages and the response will be logged as a trace to the active experiment
answer = openai_client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    temperature=0.99,
)

注意

使用 OpenAI SDK 时，确保你的访问令牌已分配给环境变量 OPENAI_API_KEY。

OpenAI 自动日志记录配置

MLflow OpenAI 自动日志记录可以记录关于模型及其推理的各种信息。默认情况下，只启用追踪日志记录，但你可以通过在调用 mlflow.openai.autolog() 时设置相应的参数来启用其他信息的自动日志记录。

可用选项及其默认值如下所示。要了解更多附加参数，请参阅 API 文档。

目标	默认	参数	描述
追踪	true	log_traces	是否为模型生成并记录追踪。有关追踪功能的更多详细信息，请参阅 MLflow 追踪。
模型工件	false	log_models	如果设置为 True，则在调用 OpenAI 模型时会记录它。
模型签名	false	log_model_signatures	如果设置为 True，描述模型输入和输出的 ModelSignature 会在推理期间与 OpenAI 模型工件一起被收集和记录。此选项仅在启用 log_models 时可用。
输入示例	false	log_input_examples	如果设置为 True，来自推理数据的输入示例会在推理期间与 OpenAI 模型工件一起被收集和记录。此选项仅在启用 log_models 时可用。

例如，要禁用追踪日志记录，转而启用模型日志记录，运行以下代码：

import mlflow

mlflow.openai.autolog(
    log_traces=False,
    log_models=True,
)

使用 OpenAI 自动日志记录的示例

import os

import mlflow
import openai

API_KEY = os.environ.get("OPENAI_API_KEY")
EXPERIMENT_NAME = "OpenAI Autologging Demonstration"
REGISTERED_MODEL_NAME = "openai-auto"
MODEL_VERSION = 1

mlflow.openai.autolog(
    log_input_examples=True,
    log_model_signatures=True,
    log_models=True,
    log_traces=True,
    registered_model_name=REGISTERED_MODEL_NAME,
)

mlflow.set_experiment(EXPERIMENT_NAME)

openai_client = openai.OpenAI(api_key=API_KEY)

messages = [
    {
        "role": "user",
        "content": "State that you are responding to a test and that you are alive.",
    }
]

openai_client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    temperature=0.95,
)

在 UI 中查看记录的模型以及调用 OpenAI 客户端时使用的追踪，可以在下面的图像中看到：

可以通过记录和注册的模型使用 models uri 加载模型，并按照如下所示通过 pyfunc API 进行交互：

loaded_autologged_model = mlflow.pyfunc.load_model(
    f"models:/{REGISTERED_MODEL_NAME}/{MODEL_VERSION}"
)

loaded_autologged_model.predict(
    "How much relative time difference would occur between an astronaut travelling at 0.98c for 14 years "
    "as measured by an on-board clock on the spacecraft and humans on Earth, assuming constant speed?"
)

OpenAI Swarm 的自动追踪

MLflow 2.17.1 引入了针对 OpenAI Swarm（一个来自 OpenAI 的多智能体编排框架）的内置追踪功能。该框架提供了一个清晰的接口，用于在 OpenAI 的 Function Calling 能力之上构建多智能体系统，以及 handoff & routines 模式 的概念。

MLflow 的自动追踪能力提供对智能体之间交互、工具调用及其总体输出的无缝跟踪。你只需调用 mlflow.openai.autolog() 函数即可为 OpenAI Swarm 启用自动追踪。

import mlflow
from swarm import Swarm, Agent

# Calling the autolog API will enable trace logging by default.
mlflow.openai.autolog()

mlflow.set_experiment("OpenAI Swarm")

client = Swarm()


def transfer_to_agent_b():
    return agent_b


agent_a = Agent(
    name="Agent A",
    instructions="You are a helpful agent.",
    functions=[transfer_to_agent_b],
)

agent_b = Agent(
    name="Agent B",
    instructions="Only speak in Haikus.",
)

response = client.run(
    agent=agent_a,
    messages=[{"role": "user", "content": "I want to talk to agent B."}],
)
print(response)

与 OpenAI Swarm 实验相关联的记录追踪，可以在 MLflow UI 中看到，如下所示：

常见问题 (FAQ)

自动日志记录是否支持异步 API？

MLflow OpenAI 自动日志记录功能不支持异步 API 用于记录模型或追踪。

保存你的异步实现最好通过使用模型即代码功能来完成。

如果你想记录异步 OpenAI API 的追踪事件，下面是记录流式异步请求追踪的简化示例：

import openai
import mlflow
import asyncio

# Activate an experiment for logging traces to
mlflow.set_experiment("OpenAI")


async def fetch_openai_response(messages, model="gpt-4o", temperature=0.99):
    """
    Asynchronously gets a response from the OpenAI API using the provided messages and streams the response.

    Args:
        messages (list): List of message dictionaries for the OpenAI API.
        model (str): The model to use for the OpenAI API. Default is "gpt-4o".
        temperature (float): The temperature to use for the OpenAI API. Default is 0.99.

    Returns:
        None
    """
    client = openai.AsyncOpenAI()

    # Create the response stream
    response_stream = await client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature,
        stream=True,
    )

    # Manually log traces using the tracing fluent API
    with mlflow.start_span() as trace:
        trace.set_inputs(messages)
        full_response = []

        async for chunk in response_stream:
            content = chunk.choices[0].delta.content
            if content is not None:
                print(content, end="")
                full_response.append(content)

        trace.set_outputs("".join(full_response))


messages = [
    {
        "role": "user",
        "content": "How much additional hydrogen mass would Jupiter require to ignite a sustainable fusion cycle?",
    }
]

await fetch_openai_response(messages)