MLflow OpenAI 自动日志记录

MLflow 的 OpenAI 风格支持自动日志记录，以确保您的想法的实验、测试和验证可以动态捕获，而无需用日志记录样板代码来包装您的代码。

注意

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

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

• 通过 client.chat.completions.create() 进行聊天补全
• 通过 client.completions.create() 进行补全（旧版）
• 通过 client.embeddings.create() 进行嵌入

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

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

• 快速入门
• 使用 OpenAI 自动日志记录的示例
• OpenAI Swarm 的自动跟踪
• 常见问题

快速入门

要开始使用 MLflow 的 OpenAI 自动日志记录，您只需在脚本或笔记本的开头调用 mlflow.openai.autolog()。启用无参数覆盖的自动日志记录将默认启用跟踪。

提示

当自动日志记录被激活时，唯一默认启用的元素是跟踪信息的记录。您可以在此处阅读更多关于 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 自动日志记录的示例

import os

import mlflow
import openai

API_KEY = os.environ.get("OPENAI_API_KEY")
EXPERIMENT_NAME = "OpenAI Autologging Demonstration"

# set an active model for linking traces
mlflow.set_active_model(name="OpenAI_Model")

mlflow.openai.autolog()

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 客户端时使用的跟踪，如下图所示

OpenAI Swarm 的自动跟踪

MLflow 2.17.1 引入了对 OpenAI Swarm 的内置跟踪功能，这是一个来自 OpenAI 的多代理编排框架。该框架提供了一个清晰的接口，可以在 OpenAI 的函数调用能力和交接和例程模式概念之上构建多代理系统。

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 中查看，如下图所示

常见问题

自动日志记录是否支持异步 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)