Traçage OpenAI

MLflow Tracing fournit une capacité de traçage automatique pour OpenAI. En activant le traçage automatique pour OpenAI en appelant la fonction mlflow.openai.autolog, MLflow capturera les traces pour l'invocation de LLM et les enregistrera dans l'Experimentation MLflow active.

La trace MLflow capture automatiquement les informations suivantes concernant les appels OpenAI :

• Prompts et réponses de complétion
• Latences
• Nom du modèle
• Métadonnées supplémentaires telles que temperature, max_tokens, si spécifié.
• Appel de fonction si renvoyé dans la réponse
• Toute exception si levée

Prérequis

Pour utiliser MLflow Tracing avec OpenAI, vous devez installer MLflow et le SDK OpenAI.

Development

Pour les environnements de développement, installez le package MLflow complet avec les extras Databricks et openai:

Bash

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

Le package mlflow[databricks] complet inclut toutes les fonctionnalités pour le développement local et l’expérimentation sur Databricks.

Production

Pour les déploiements en production, installez mlflow-tracing et openai:

Bash

pip install --upgrade mlflow-tracing openai

Le mlflow-tracing package est optimisé pour une utilisation en production.

remarque

MLflow 3 est fortement recommandé pour une meilleure expérience de traçage avec OpenAI.

Avant d'exécuter les exemples, vous devrez configurer votre environnement :

Pour les utilisateurs en dehors des notebooks Databricks : Définissez vos variables d'environnement Databricks :

Bash

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

Pour les utilisateurs dans les Notebooks Databricks : ces identifiants sont définis automatiquement pour vous.

Clés API : Assurez-vous que votre clé API OpenAI est configurée. Pour les environnements de production, utilisez AI Gateway ou les secrets Databricks au lieu de valeurs codées en dur pour une gestion sécurisée des clés d'API :

Bash

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

APIs prises en charge

MLflow prend en charge le traçage automatique pour les APIs OpenAI suivantes :

Achèvement du chat	Intégrations	Appel de fonction	Sorties structurées	Streaming	Asynchrone	Image	Audio
✅	✅	✅	✅ (*1)	✅ (*2)	✅ (*1)

Achèvement du chat	Intégrations	Appel de fonction	Sorties structurées	Streaming	Asynchrone	Image	Audio
✅	✅	✅	✅ (*1)	✅ (*2)	✅ (*1)

(*1) La prise en charge du streaming a été ajoutée dans MLflow 2.15.0.

(*2) La prise en charge de la sortie asynchrone et structurée a été ajoutée dans MLflow 2.21.0.

Pour demander la prise en charge d'APIs supplémentaires, veuillez ouvrir une demande de fonctionnalité sur GitHub.

Exemple de base

remarque

Sur les clusters de compute Serverless, la journalisation automatique pour les cadres de traçage genAI n'est pas activée automatiquement. Vous devez activer explicitement l'autologging en appelant la fonction mlflow.<library>.autolog() appropriée pour les intégrations spécifiques que vous souhaitez suivre.

Python

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

attention

Pour les environnements de production, utilisez AI Gateway ou les secrets Databricks au lieu de valeurs codées en dur pour la gestion sécurisée des clés API.

Exemple OpenAI avec secrets

Une fois votre secret stocké, utilisez-le dans votre code de suivi MLflow :

Python

import openai
import mlflow
import os

# Configure your secret scope and key names
secret_scope_name = "openai-secrets"
secret_key_name = "api-key"

# Retrieve the API key from Databricks secrets
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get(
    scope=secret_scope_name,
    key=secret_key_name
)

# Verify the API key was loaded successfully
assert os.environ["OPENAI_API_KEY"] is not None, "API key not loaded from secrets"

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

# Now you can use OpenAI as usual
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,
)

streaming

MLflow Tracing prend en charge l'API de streaming de l'OpenAI SDK. Avec la même configuration de traçage automatique, MLflow trace automatiquement la réponse de streaming et affiche la sortie concaténée dans l'interface utilisateur de l'étendue. Les fragments réels dans le Stream de réponse peuvent également être trouvés dans l'tab Event.

Python

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="")

Async

MLflow Tracing prend en charge l'API asynchrone du SDK OpenAI depuis MLflow 2.21.0. L'utilisation est la même que celle de l'API synchrone.

Python

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
)

Appel de fonction

MLflow Tracing capture automatiquement la réponse d'appel de fonction des modèles OpenAI. L’instruction de fonction dans la réponse sera mise en surbrillance dans l’interface utilisateur de la trace. De plus, vous pouvez annoter la fonction d’outil avec le décorateur @mlflow.trace pour créer un intervalle pour l’exécution de l’outil.

L'exemple suivant met en œuvre un agent d'appel de fonction simple utilisant OpenAI Function Calling et MLflow Tracing pour OpenAI.

Python

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)

Désactiver le suivi automatique

Le traçage automatique pour OpenAI peut être désactivé globalement en appelant mlflow.openai.autolog(disable=True) ou mlflow.autolog(disable=True).