Rastreamento de Modelos Fundamentais do Databricks

O Rastreamento do MLflow oferece capacidade de rastreamento automático para os Modelos Fundamentais do Databricks. Como os Modelos do Databricks Foundation usam uma API compatível com OpenAI, você pode habilitar o rastreamento automático chamando a função mlflow.openai.autolog, e o MLflow capturará rastreamentos para invocação de LLM e os registrará no Experimento de MLflow ativo.

import mlflow

mlflow.openai.autolog()

O rastreamento do MLflow captura automaticamente as seguintes informações sobre chamadas do Databricks Foundation Model:

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

Observação

Em clusters de computação sem servidor, o registro automático não é habilitado automaticamente. Você deve chamar mlflow.openai.autolog() explicitamente para habilitar o rastreamento automático para essa integração.

Pré-requisitos

Para usar o Rastreamento do MLflow com modelos do Databricks Foundation, você precisa instalar o MLflow e o SDK do OpenAI (já que os Modelos do Databricks Foundation usam uma API compatível com OpenAI).

Desenvolvimento

Para ambientes de desenvolvimento, instale o pacote completo do MLflow com extras do Databricks e o SDK do 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 o SDK do 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 com o Databricks Foundation Models.

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

Para usuários fora dos notebooks do Databricks: defina suas variáveis de ambiente do Databricks:

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

Para usuários dentro de blocos de anotações do Databricks: essas credenciais são definidas automaticamente para você.

APIs com suporte

O MLflow dá suporte ao rastreamento automático para as seguintes APIs de Modelo do Databricks Foundation:

Conclusão do chat	Chamada de função	Transmissão ao vivo	Async
✅	✅	✅	✅

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

Exemplo básico

import mlflow
import os
from openai import OpenAI

# Databricks Foundation Model APIs use Databricks authentication.

# Enable auto-tracing for OpenAI (which will trace Databricks Foundation Model API calls)
mlflow.openai.autolog()

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

# Create OpenAI client configured for Databricks
client = OpenAI(
    api_key=os.environ.get("DATABRICKS_TOKEN"),
    base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)

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

response = client.chat.completions.create(
    model="databricks-llama-4-maverick",
    messages=messages,
    temperature=0.1,
    max_tokens=100,
)

Transmissão ao vivo

O MLflow Tracing dá suporte à API de streaming de Modelos do Databricks Foundation. Com a mesma configuração de rastreamento automático, o MLflow rastreia automaticamente a resposta de streaming e renderiza a saída concatenada na interface do usuário do componente span.

import mlflow
import os
from openai import OpenAI

# Enable auto-tracing for OpenAI (which will trace Databricks Foundation Model API calls)
mlflow.openai.autolog()

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

# Create OpenAI client configured for Databricks
client = OpenAI(
    api_key=os.environ.get("DATABRICKS_TOKEN"),
    base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)

stream = client.chat.completions.create(
    model="databricks-llama-4-maverick",
    messages=[
        {"role": "user", "content": "Explain the benefits of using Databricks Foundation Models"}
    ],
    stream=True,  # Enable streaming response
)

for chunk in stream:
    print(chunk.choices[0].delta.content or "", end="")

Chamada de função

O Rastreamento de MLflow captura automaticamente a resposta de chamada de função dos Modelos Fundamentais da Databricks. A instrução de função na resposta será destacada na UI de rastreamento. Além disso, você pode anotar a função de ferramenta com o @mlflow.trace decorador para criar um intervalo para a execução da ferramenta.

O exemplo a seguir implementa um agente de chamada de função simples usando os Modelos do Databricks Foundation e o Rastreamento de MLflow.

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

# Enable auto-tracing for OpenAI (which will trace Databricks Foundation Model API calls)
mlflow.openai.autolog()

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

# Create OpenAI client configured for Databricks
client = OpenAI(
    api_key=os.environ.get("DATABRICKS_TOKEN"),
    base_url=f"{os.environ.get('DATABRICKS_HOST')}/serving-endpoints"
)

# 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="databricks-llama-4-maverick",
        messages=messages,
        tools=tools,
    )
    ai_msg = response.choices[0].message

    # If the model requests 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,
                }
            )

        # Send the tool results to the model and get a new response
        response = client.chat.completions.create(
            model="databricks-llama-4-maverick", 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)

Modelos disponíveis

O Databricks Foundation Models fornece acesso a uma variedade de modelos de última geração, incluindo Llama, Antropic e outros modelos de base líderes.

Para obter a lista completa e mais up-tode datas dos modelos disponíveis e suas IDs de modelo, consulte a documentação do Databricks Foundation Models.

Desabilitar o rastreamento automático

O rastreamento automático de modelos do Databricks Foundation pode ser desabilitado globalmente chamando mlflow.openai.autolog(disable=True) ou mlflow.autolog(disable=True).