Tracing OpenAI Agent Responses utilizing MLFlow

MLflow is an open-source platform for managing and monitoring machine studying experiments. When used with the OpenAI Brokers SDK, MLflow routinely:

• Logs all agent interactions and API calls
• Captures software utilization, enter/output messages, and intermediate choices
• Tracks runs for debugging, efficiency evaluation, and reproducibility

That is particularly helpful if you’re constructing multi-agent programs the place completely different brokers collaborate or name capabilities dynamically

On this tutorial, we’ll stroll by way of two key examples: a easy handoff between brokers, and the usage of agent guardrails — all whereas tracing their habits utilizing MLflow.

Establishing the dependencies

Putting in the libraries

pip set up openai-agents mlflow pydantic pydotenv

OpenAI API Key

To get an OpenAI API key, go to https://platform.openai.com/settings/group/api-keys and generate a brand new key. When you’re a brand new consumer, it’s possible you’ll want so as to add billing particulars and make a minimal cost of $5 to activate API entry.

As soon as the bottom line is generated, create a .env file and enter the next:

Substitute with the important thing you generated.

Multi-Agent System (multi_agent_demo.py) 

On this script (multi_agent_demo.py), we construct a easy multi-agent assistant utilizing the OpenAI Brokers SDK, designed to route consumer queries to both a coding skilled or a cooking skilled. We allow mlflow.openai.autolog(), which routinely traces and logs all agent interactions with the OpenAI API — together with inputs, outputs, and agent handoffs — making it simple to observe and debug the system. MLflow is configured to make use of an area file-based monitoring URI (./mlruns) and logs all exercise beneath the experiment title “Agent‑Coding‑Cooking“.

import mlflow, asyncio
from brokers import Agent, Runner
import os
from dotenv import load_dotenv
load_dotenv()

mlflow.openai.autolog()                           # Auto‑hint each OpenAI name
mlflow.set_tracking_uri("./mlruns")
mlflow.set_experiment("Agent‑Coding‑Cooking")

coding_agent = Agent(title="Coding agent",
                     directions="You solely reply coding questions.")

cooking_agent = Agent(title="Cooking agent",
                      directions="You solely reply cooking questions.")

triage_agent = Agent(
    title="Triage agent",
    directions="If the request is about code, handoff to coding_agent; "
                 "if about cooking, handoff to cooking_agent.",
    handoffs=[coding_agent, cooking_agent],
)

async def major():
    res = await Runner.run(triage_agent,
                           enter="How do I boil pasta al dente?")
    print(res.final_output)

if __name__ == "__main__":
    asyncio.run(major())

MLFlow UI

To open the MLflow UI and think about all of the logged agent interactions, run the next command in a brand new terminal:

This may begin the MLflow monitoring server and show a immediate indicating the URL and port the place the UI is accessible — normally http://localhost:5000 by default.

We will view the complete interplay movement within the Tracing part — from the consumer’s preliminary enter to how the assistant routed the request to the suitable agent, and at last, the response generated by that agent. This end-to-end hint supplies worthwhile perception into decision-making, handoffs, and outputs, serving to you debug and optimize your agent workflows.

Tracing Guardrails (guardrails.py) 

On this instance, we implement a guardrail-protected buyer help agent utilizing the OpenAI Brokers SDK with MLflow tracing. The agent is designed to assist customers with basic queries however is restricted from answering medical-related questions. A devoted guardrail agent checks for such inputs, and if detected, blocks the request. MLflow captures the complete movement — together with guardrail activation, reasoning, and agent response — offering full traceability and perception into security mechanisms.

import mlflow, asyncio
from pydantic import BaseModel
from brokers import (
    Agent, Runner,
    GuardrailFunctionOutput, InputGuardrailTripwireTriggered,
    input_guardrail, RunContextWrapper)

from dotenv import load_dotenv
load_dotenv()

mlflow.openai.autolog()
mlflow.set_tracking_uri("./mlruns")
mlflow.set_experiment("Agent‑Guardrails")

class MedicalSymptons(BaseModel):
    medical_symptoms: bool
    reasoning: str


guardrail_agent = Agent(
    title="Guardrail examine",
    directions="Examine if the consumer is asking you for medical symptons.",
    output_type=MedicalSymptons,
)


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

    return GuardrailFunctionOutput(
        output_info=end result.final_output,
        tripwire_triggered=end result.final_output.medical_symptoms,
    )


agent = Agent(
    title="Buyer help agent",
    directions="You're a buyer help agent. You assist prospects with their questions.",
    input_guardrails=[medical_guardrail],
)


async def major():
    strive:
        await Runner.run(agent, "Ought to I take aspirin if I am having a headache?")
        print("Guardrail did not journey - that is sudden")

    besides InputGuardrailTripwireTriggered:
        print("Medical guardrail tripped")


if __name__ == "__main__":
    asyncio.run(major())

This script defines a buyer help agent with an enter guardrail that detects medical-related questions. It makes use of a separate guardrail_agent to judge whether or not the consumer’s enter comprises a request for medical recommendation. If such enter is detected, the guardrail triggers and prevents the primary agent from responding. Your complete course of, together with guardrail checks and outcomes, is routinely logged and traced utilizing MLflow.

MLFlow UI

To open the MLflow UI and think about all of the logged agent interactions, run the next command in a brand new terminal:

On this instance, we requested the agent, “Ought to I take aspirin if I’m having a headache?”, which triggered the guardrail. Within the MLflow UI, we will clearly see that the enter was flagged, together with the reasoning supplied by the guardrail agent for why the request was blocked.

I’m a Civil Engineering Graduate (2022) from Jamia Millia Islamia, New Delhi, and I’ve a eager curiosity in Knowledge Science, particularly Neural Networks and their software in numerous areas.