跟踪 DeepSeek

MLflow Tracing 通过 OpenAI SDK 集成为 Deepseek 模型提供了自动跟踪功能。由于 DeepSeek 使用 OpenAI 兼容的 API 格式，您可以使用 mlflow.openai.autolog() 来跟踪与 DeepSeek 模型的交互。

import mlflow

# Enable
mlflow.openai.autolog()

MLflow trace 自动捕获有关 DeepSeek 调用的以下信息

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

支持的 API

MLflow 通过 OpenAI 集成支持对以下 DeepSeek API 的自动跟踪

聊天完成	函数调用	流式传输	异步
✅	✅	✅ (*1)	✅ (*2)

(*1) 流式传输支持需要 MLflow 2.15.0 或更高版本。(*2) 异步支持需要 MLflow 2.21.0 或更高版本。

要请求支持更多 API，请在 GitHub 上提出功能请求。

基本示例

import openai
import mlflow

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

# Optional: Set a tracking URI and an experiment
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("DeepSeek")

# Initialize the OpenAI client with DeepSeek API endpoint
client = openai.OpenAI(
    base_url="https://api.deepseek.com", api_key="<your_deepseek_api_key>"
)

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "What is the capital of France?"},
]

response = client.chat.completions.create(
    model="deepseek-chat",
    messages=messages,
    temperature=0.1,
    max_tokens=100,
)

上述示例应在 MLflow UI 的 DeepSeek 实验中生成一个 trace

流式传输和异步支持

MLflow 支持对流式传输和异步 DeepSeek API 进行跟踪。请访问OpenAI 跟踪文档，获取通过 OpenAI SDK 跟踪流式传输和异步调用的示例代码片段。

高级示例：函数调用 Agent

MLflow Tracing 通过 OpenAI SDK 自动捕获来自 DeepSeek 模型的函数调用响应。响应中的函数指令将在 trace UI 中高亮显示。此外，您可以使用 @mlflow.trace 装饰器注解工具函数，为工具执行创建一个 span。

以下示例使用 DeepSeek 函数调用和 MLflow Tracing 实现了一个简单的函数调用 agent。

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

# Initialize the OpenAI client with DeepSeek API endpoint
client = OpenAI(base_url="https://api.deepseek.com", api_key="<your_deepseek_api_key>")


# 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="deepseek-chat",
        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="deepseek-chat", 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) 来全局禁用 DeepSeek 的自动跟踪（通过 OpenAI SDK）。