Rastreando DeepSeek

O MLflow Tracing fornece capacidade de rastreamento automático para modelos Deepseek por meio da integração do OpenAI SDK. Como o DeepSeek usa um formato de API compatível com OpenAI, você pode usar mlflow.openai.autolog() para rastrear interações com modelos DeepSeek.

import mlflow

mlflow.openai.autolog()

O rastreamento MLflow captura automaticamente as seguintes informações sobre chamadas DeepSeek:

• Prompts e respostas de conclusão
• Latências
• Nome do modelo
• Metadados adicionais, como temperature, max_tokens, se especificado.
• Chamada de função se for retornada na resposta
• Qualquer exceção, se for levantada

Pré-requisitos

Para usar o MLflow Tracing com o DeepSeek (por meio de sua API compatível com OpenAI), você precisa instalar o MLflow e o OpenAI SDK.

Desenvolvimento

Para ambientes de desenvolvimento, instale o pacote MLflow completo com os extras do Databricks e openai:

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

O pacote completo mlflow[databricks] inclui todos os recursos para desenvolvimento local e experimentação no Databricks.

Produção

Para implantações de produção, instale mlflow-tracing e openai:

pip install --upgrade mlflow-tracing openai

O mlflow-tracing pacote é otimizado para uso em produção.

Observação

O MLflow 3 é altamente recomendado para a melhor experiência de rastreamento.

Antes de executar os exemplos, você precisará configurar seu ambiente:

Para utilizadores fora dos notebooks Databricks: defina as suas variáveis de ambiente Databricks:

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

Para utilizadores dentro de blocos de notas Databricks: estas credenciais são definidas automaticamente para si.

Chaves de API: certifique-se de que sua chave de API do DeepSeek esteja configurada. Para uso em produção, use segredos do Mosaic AI Gateway ou Databricks em vez de valores codificados:

export DEEPSEEK_API_KEY="your-deepseek-api-key"

APIs suportadas

O MLflow suporta rastreamento automático para as seguintes APIs do DeepSeek por meio da integração OpenAI:

Conclusão do bate-papo	Chamada de função	Serviço de streaming	Assíncrono
✅	✅	✅ (*1)	✅ (*2)

(*1) O suporte de streaming requer MLflow 2.15.0 ou posterior. (*2) O suporte assíncrono requer MLflow 2.21.0 ou posterior.

Para solicitar suporte para APIs adicionais, abra uma solicitação de recurso no GitHub.

Exemplo básico

import openai
import mlflow
import os

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

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

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

# Initialize the OpenAI client with DeepSeek API endpoint and your key
client = openai.OpenAI(
    base_url="https://api.deepseek.com",
    api_key=os.environ.get("DEEPSEEK_API_KEY") # Or directly pass your key string
)

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,
)

O exemplo acima deve gerar um registo no experimento na interface do utilizador do MLflow.

Advertência

Para ambientes de produção, use segredos do Mosaic AI Gateway ou Databricks em vez de valores codificados para gerenciamento seguro de chaves de API.

Suporte a streaming e assíncrono

O MLflow suporta rastreamento para streaming e APIs DeepSeek assíncronas. Visite a documentação do OpenAI Tracing para obter trechos de código para rastrear streaming e chamadas assíncronas por meio do OpenAI SDK.

Exemplo avançado: Agente de chamada de função

O MLflow Tracing captura automaticamente respostas de chamada de função de modelos DeepSeek através do OpenAI SDK. A instrução de função na resposta será realçada na interface do usuário de rastreamento. Além disso, pode-se adicionar uma anotação à função da ferramenta com o decorador @mlflow.trace para criar um intervalo para a execução da ferramenta.

O exemplo a seguir implementa um agente de chamada de função simples usando DeepSeek Function Calling e MLflow Tracing.

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

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

# Initialize the OpenAI client with DeepSeek API endpoint and your key
client = OpenAI(
    base_url="https://api.deepseek.com",
    api_key=os.environ.get("DEEPSEEK_API_KEY") # Or directly pass your key string
)

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

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

# 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)

Desativar o rastreamento automático

O rastreamento automático para DeepSeek (através do OpenAI SDK) pode ser desativado globalmente chamando mlflow.openai.autolog(disable=True) ou mlflow.autolog(disable=True).