Skip to main content

Tracing a GenAI App (Notebook)

This quickstart helps you integrate your GenAI app with MLflow Tracing if you use a Jupyter Notebook as your development environment. If you use a local IDE, please use the IDE quickstart instead.

What you'll achieve​

By the end of this tutorial, you will have:

  • A Jupyter Notebook with a linked MLflow Experiment for your GenAI app
  • Used MLflow Tracing to instrument your app.

Prerequisites​

  • Python Environment: Python 3.8+ with pip installed
  • Jupyter Notebook: Access to Jupyter Notebook or JupyterLab
  • OpenAI API Key: Access to OpenAI API

Step 1: Create a Jupyter Notebook​

note

Creating a Jupyter Notebook and setting up an MLflow Experiment will be the container for your GenAI application. Learn more about the Experiment and what it contains in the data model section.

  1. Open Jupyter Notebook or JupyterLab
  2. Create a new notebook
  3. Create your first cell to set up the environment

Step 2: Install and configure MLflow​

Install the latest version of MLflow for the best experience with GenAI capabilities, including the most comprehensive tracing features and robust support.

You can install MLflow in your notebook by running:

!pip install --upgrade "mlflow>=3.1" openai python-dotenv
MLflow Version Recommendation

While tracing features are available in MLflow 2.15.0+, it is strongly recommended to install MLflow 3 (specifically 3.1 or newer) for the latest GenAI capabilities, including expanded tracing features and robust support.

Step 3: Instrument your application​

Select the appropriate integration for your application:

OpenAI Integration​

  1. Create a cell in your notebook to configure your environment and MLflow:

    import mlflow
    import openai
    from openai import OpenAI
    import os
    from dotenv import load_dotenv

    # Load environment variables
    load_dotenv()

    # Set up MLflow experiment
    mlflow.set_experiment("genai-tracing-quickstart")

  2. Create a cell in your notebook with the below code

    Here, we use the @mlflow.trace decorator that makes it easy to trace any Python application combined with the OpenAI automatic instrumentation to captures the details of the call to the OpenAI SDK.

    # Enable MLflow's autologging to instrument your application with Tracing
    mlflow.openai.autolog()

    # Create an OpenAI client
    client = OpenAI(
        api_key=os.getenv("OPENAI_API_KEY")  # Set your OpenAI API key in environment
    )


    # Use the trace decorator to capture the application's entry point
    @mlflow.trace
    def my_app(input: str):
        # This call is automatically instrumented by `mlflow.openai.autolog()`
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",  # You can replace with gpt-4, gpt-4-turbo, etc.
            messages=[
                {
                    "role": "system",
                    "content": "You are a helpful assistant.",
                },
                {
                    "role": "user",
                    "content": input,
                },
            ],
        )
        return response.choices[0].message.content


    my_app(input="What is MLflow?")

  3. Run the Notebook cell to execute the code

Alternative: Other LLM Providers​

# For other providers, you can use manual tracing
import mlflow


@mlflow.trace
def my_custom_app(input: str):
    # Your custom LLM integration here
    # This could be Anthropic, Hugging Face, local models, etc.

    with mlflow.start_span(name="llm_call") as span:
        span.set_inputs({"input": input})

        # Your LLM call here
        response = "Your LLM response"

        span.set_outputs({"response": response})
        return response


my_custom_app(input="What is MLflow?")

Step 4: View the Trace in MLflow​

  1. Start the MLflow UI by running in your terminal:

    mlflow ui

  2. Open your browser to http://localhost:5000

  3. Navigate to your experiment: "genai-tracing-quickstart"

  4. You will now see the generated trace in the Traces tab

  5. Click on the trace to view its details

Understanding the Trace​

The trace you've just created shows:

  • Root Span: Represents the inputs to the my_app(...) function
    • Child Span: Represents the OpenAI completion request
  • Attributes: Contains metadata like model name, token counts, and timing information
  • Inputs: The messages sent to the model
  • Outputs: The response received from the model

This simple trace already provides valuable insights into your application's behavior, such as:

  • What was asked
  • What response was generated
  • How long the request took
  • How many tokens were used (affecting cost)

Next Steps​

Congratulations! You've successfully built your first GenAI application with MLflow Tracing and OpenAI.

tip

For more complex applications like RAG systems or multi-step agents, MLflow Tracing provides even more value by revealing the inner workings of each component and step.