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 instrument utilization, enter/output messages, and intermediate selections
• Tracks runs for debugging, efficiency evaluation, and reproducibility

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

On this tutorial, we’ll stroll by means of two key examples: a easy handoff between brokers, and the usage of agent guardrails — all whereas tracing their conduct 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. In case you’re a brand new consumer, chances are you’ll want so as to add billing particulars and make a minimal fee of $5 to activate API entry.

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

Exchange 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 watch and debug the system. MLflow is configured to make use of a neighborhood file-based monitoring URI (./mlruns) and logs all exercise below the experiment identify “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(identify="Coding agent",
                     directions="You solely reply coding questions.")

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

triage_agent = Agent(
    identify="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 essential():
    res = await Runner.run(triage_agent,
                           enter="How do I boil pasta al dente?")
    print(res.final_output)

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

MLFlow UI

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

This can 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 whole interplay move within the Tracing part — from the consumer’s preliminary enter to how the assistant routed the request to the suitable agent, and eventually, the response generated by that agent. This end-to-end hint gives useful 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 assist agent utilizing the OpenAI Brokers SDK with MLflow tracing. The agent is designed to assist customers with normal 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 whole move — 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(
    identify="Guardrail verify",
    directions="Test 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:
    outcome = await Runner.run(guardrail_agent, enter, context=ctx.context)

    return GuardrailFunctionOutput(
        output_info=outcome.final_output,
        tripwire_triggered=outcome.final_output.medical_symptoms,
    )


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


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

    besides InputGuardrailTripwireTriggered:
        print("Medical guardrail tripped")


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

This script defines a buyer assist 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 principle agent from responding. All the course of, together with guardrail checks and outcomes, is routinely logged and traced utilizing MLflow.

MLFlow UI

To open the MLflow UI and examine 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 are able to 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 Information Science, particularly Neural Networks and their utility in numerous areas.