追踪 OpenAI

MLflow 跟踪 为 OpenAI 提供自动跟踪功能。 通过调用 mlflow.openai.autolog 函数为 OpenAI 启用自动跟踪，MLflow 将捕获 LLM 调用的跟踪并将其记录到活动 MLflow 试验。

MLflow 跟踪自动捕获有关 OpenAI 调用的以下信息：

• 提示和补全应答
• 潜伏期
• 模型名称
• 其他元数据（例如 temperature， max_tokens如果指定）。
• 在响应中返回时进行函数调用
• 引发的任何异常

先决条件

若要将 MLflow 跟踪与 OpenAI 配合使用，需要安装 MLflow 和 OpenAI SDK。

开发

对于开发环境，请使用 Databricks 附加组件安装完整的 MLflow 包，如下所示：

pip install --upgrade "mlflow[databricks]>=3.1" openai

完整 mlflow[databricks] 包包括用于 Databricks 的本地开发和试验的所有功能。

生产

对于生产部署，请安装mlflow-tracing和openai。

pip install --upgrade mlflow-tracing openai

包 mlflow-tracing 已针对生产用途进行优化。

在运行示例之前，需要配置环境：

对于不使用 Databricks 笔记本的用户：设置 Databricks 环境变量：

export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"

对于 Databricks 笔记本中的用户：这些凭据会自动为您设置。

API 密钥：确保已设置 OpenAI API 密钥：

export OPENAI_API_KEY="your-openai-api-key"

受支持的 API

MLflow 支持以下 OpenAI API 的自动跟踪：

(*1)MLflow 2.15.0 中添加了流式处理支持。

(*2)MLflow 2.21.0 中添加了支持的异步和结构化输出。

若要请求对其他 API 的支持，请在 GitHub 上打开 功能请求 。

基本示例

import openai
import mlflow
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/openai-tracing-demo")

openai_client = openai.OpenAI()

messages = [
    {
        "role": "user",
        "content": "What is the capital of France?",
    }
]

response = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    temperature=0.1,
    max_tokens=100,
)

流媒体

MLflow 跟踪功能支持 OpenAI SDK 的流式 API。 通过相同的自动跟踪配置，MLflow 会自动记录流式响应，并在 Span 用户界面中呈现合并后的输出。 响应流中的实际区块也可以在选项卡中找到 Event 。

import openai
import mlflow
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured

# Enable trace logging
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/openai-streaming-demo")

client = openai.OpenAI()

stream = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "user", "content": "How fast would a glass of water freeze on Titan?"}
    ],
    stream=True,  # Enable streaming response
)
for chunk in stream:
    print(chunk.choices[0].delta.content or "", end="")

异步

MLflow 自 2.21.0 版本起支持 OpenAI SDK 的异步 API 跟踪。 用法与同步 API 相同。

import openai
import mlflow
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured

# Enable trace logging
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/openai-async-demo")

client = openai.AsyncOpenAI()

response = await client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "user", "content": "How fast would a glass of water freeze on Titan?"}
    ],
    # Async streaming is also supported
    # stream=True
)

函数调用

MLflow 跟踪会自动捕获从 OpenAI 模型函数调用产生的响应。 响应中的函数指令将在跟踪 UI 中突出显示。 此外，可以使用@mlflow.trace修饰器为工具函数添加批注，以便为工具执行创建跨度。

以下示例使用 OpenAI 函数调用和 OpenAI 的 MLflow 跟踪实现简单的函数调用代理。

import json
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/openai-function-agent-demo")

# Assuming autolog is enabled globally or called earlier
# mlflow.openai.autolog()

client = OpenAI()


# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
def get_weather(city: str) -> str:
    if city == "Tokyo":
        return "sunny"
    elif city == "Paris":
        return "rainy"
    return "unknown"


tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "parameters": {
                "type": "object",
                "properties": {"city": {"type": "string"}},
            },
        },
    }
]

_tool_functions = {"get_weather": get_weather}


# Define a simple tool calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str):
    messages = [{"role": "user", "content": question}]

    # Invoke the model with the given question and available tools
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        tools=tools,
    )
    ai_msg = response.choices[0].message
    messages.append(ai_msg)

    # If the model request tool call(s), invoke the function with the specified arguments
    if tool_calls := ai_msg.tool_calls:
        for tool_call in tool_calls:
            function_name = tool_call.function.name
            if tool_func := _tool_functions.get(function_name):
                args = json.loads(tool_call.function.arguments)
                tool_result = tool_func(**args)
            else:
                raise RuntimeError("An invalid tool is returned from the assistant!")

            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": tool_result,
                }
            )

        # Sent the tool results to the model and get a new response
        response = client.chat.completions.create(
            model="gpt-4o-mini", messages=messages
        )

    return response.choices[0].message.content


# Run the tool calling agent
question = "What's the weather like in Paris today?"
answer = run_tool_agent(question)

禁用自动跟踪

可以通过调用 mlflow.openai.autolog(disable=True) 或 mlflow.autolog(disable=True)调用来全局禁用 OpenAI 的自动跟踪。