Traçage des Agents OpenAI

MLflow Tracing offre une capacité de traçage automatique pour le SDK Agents OpenAI, un framework multi-agents développé par OpenAI. En activant le traçage automatique pour OpenAI en appelant la fonction mlflow.openai.autolog, MLflow capturera les traces et les enregistrera dans l’Experimentation MLflow active.

Python

import mlflow

mlflow.openai.autolog()

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.

Prérequis

Pour utiliser MLflow Tracing avec le SDK OpenAI Agents, vous devez installer MLflow, le SDK OpenAI et la bibliothèque openai-agents.

Development

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

Bash

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

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 de production, installez mlflow-tracing, openai et openai-agents:

Bash

pip install --upgrade mlflow-tracing openai openai-agents

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

remarque

MLflow 3 est fortement recommandé pour la meilleure expérience de traçage avec les agents 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 API.

Bash

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

Exemple de base

L'exemple suivant montre comment utiliser le SDK OpenAI Agents avec le suivi MLflow pour des agents de chat multilingues simples. Les trois agents collaborent pour déterminer la langue de l’entrée et la transfèrent au sous-agent approprié qui parle cette langue. MLflow capture la manière dont les agents interagissent entre eux et effectuent des appels à l'API OpenAI.

Python

import mlflow
import asyncio
from agents import Agent, Runner
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 Agents SDK
mlflow.openai.autolog() # This covers agents if using openai module for LLM calls
# If agents have their own autolog, e.g., mlflow.agents.autolog(), prefer that.

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

# Define a simple multi-agent workflow
spanish_agent = Agent(
    name="Spanish agent",
    instructions="You only speak Spanish.",
)

english_agent = Agent(
    name="English agent",
    instructions="You only speak English",
)

triage_agent = Agent(
    name="Triage agent",
    instructions="Handoff to the appropriate agent based on the language of the request.",
    handoffs=[spanish_agent, english_agent],
)


async def main():
    result = await Runner.run(triage_agent, input="Hola, ¿cómo estás?")
    print(result.final_output)


# If you are running this code in a Jupyter notebook, replace this with `await main()`.
if __name__ == "__main__":
    asyncio.run(main())

Appel de fonction

L'SDK d'agents OpenAI prend en charge la définition de fonctions qui peuvent être appelées par l'agent. MLflow capture les appels de fonction et affiche les fonctions disponibles pour l'agent, lesquelles sont appelées, ainsi que les entrées et les sorties des appels de fonction.

Python

import asyncio
from agents import Agent, Runner, function_tool
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 Agents SDK
mlflow.openai.autolog() # Assuming underlying LLM calls are via OpenAI

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

@function_tool
def get_weather(city: str) -> str:
    return f"The weather in {city} is sunny."


agent = Agent(
    name="Hello world",
    instructions="You are a helpful agent.",
    tools=[get_weather],
)


async def main():
    result = await Runner.run(agent, input="What's the weather in Tokyo?")
    print(result.final_output)
    # The weather in Tokyo is sunny.


# If you are running this code in a Jupyter notebook, replace this with `await main()`.
if __name__ == "__main__":
    asyncio.run(main())

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.

Garde-fous

Le SDK OpenAI Agents prend en charge la définition de garde-fous qui peuvent être utilisés pour vérifier l'entrée et la sortie de l'agent. MLflow capture les vérifications des garde-fous et affiche la raison de la vérification ainsi que si le garde-fou a été déclenché.

Python

from pydantic import BaseModel
from agents import (
    Agent,
    GuardrailFunctionOutput,
    InputGuardrailTripwireTriggered,
    RunContextWrapper,
    Runner,
    TResponseInputItem,
    input_guardrail,
)
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 Agents SDK
mlflow.openai.autolog() # Assuming underlying LLM calls are via OpenAI

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

class MathHomeworkOutput(BaseModel):
    is_math_homework: bool
    reasoning: str


guardrail_agent = Agent(
    name="Guardrail check",
    instructions="Check if the user is asking you to do their math homework.",
    output_type=MathHomeworkOutput,
)


@input_guardrail
async def math_guardrail(
    ctx: RunContextWrapper[None], agent: Agent, input
) -> GuardrailFunctionOutput:
    result = await Runner.run(guardrail_agent, input, context=ctx.context)

    return GuardrailFunctionOutput(
        output_info=result.final_output,
        tripwire_triggered=result.final_output.is_math_homework,
    )


agent = Agent(
    name="Customer support agent",
    instructions="You are a customer support agent. You help customers with their questions.",
    input_guardrails=[math_guardrail],
)


async def main():
    # This should trip the guardrail
    try:
        await Runner.run(agent, "Hello, can you help me solve for x: 2x + 3 = 11?")
        print("Guardrail didn't trip - this is unexpected")

    except InputGuardrailTripwireTriggered:
        print("Math homework guardrail tripped")


# If you are running this code in a Jupyter notebook, replace this with `await main()`.
if __name__ == "__main__":
    asyncio.run(main())

Désactiver le suivi automatique

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