Tracing a GenAI App (IDE)

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

What you'll achieve

By the end of this tutorial, you will have:

• A MLflow Experiment for your GenAI app
• Your local development environment connected to MLflow
• A simple GenAI application instrumented with MLflow Tracing
• A trace from that app in your MLflow Experiment

Prerequisites

• Databricks Workspace: Access to a Databricks workspace.

Step 1: Install MLflow

When working in your local IDE, you need to install MLflow with Databricks connectivity.

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

Step 2: Create a new MLflow Experiment

An MLflow Experiment is the container for your GenAI application. Learn more about the Experiment and what it contains in the data model section.

1. Open your Databricks workspace
2. Go to Experiments in the left sidebar under Machine Learning
3. At the top of the Experiments page, click on New GenAI Experiment

Step 3: Connect your environment to MLflow

Use environment variables

1. Click Generate API Key

2. Copy and run the generated code in your terminal.

   export DATABRICKS_TOKEN=<databricks-personal-access-token>
   export DATABRICKS_HOST=https://<workspace-name>.cloud.databricks.com
   export MLFLOW_TRACKING_URI=databricks
   export MLFLOW_EXPERIMENT_ID=<experiment-id>
   

Use a .env file

1. Click Generate API Key

2. Copy the generated code to a .env file in your project root

   DATABRICKS_TOKEN=<databricks-personal-access-token>
   DATABRICKS_HOST=https://<workspace-name>.cloud.databricks.com
   MLFLOW_TRACKING_URI=databricks
   MLFLOW_EXPERIMENT_ID=<experiment-id>
   

3. Install the python-dotenv package

   pip install python-dotenv
   

4. Load environment variables in your code

   # At the beginning of your Python script
   from dotenv import load_dotenv
   
   # Load environment variables from .env file
   load_dotenv()
   

Step 4: Create and instrument your application

The below example quickstart uses the OpenAI SDK to connect to a Databricks hosted LLM. If you want to use your own OpenAI key, update the client = OpenAI(...) line.

If you prefer to use one of the other 20+ LLM SDKs (Anthropic, Bedrock, etc) or GenAI authoring frameworks (LangGraph, etc) that MLflow supports, follow the instructions in the MLflow Experiment UI in the previous step.

1. Create a Python file named app.py in your project directory.

   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.

   The code snippet below uses Anthropic's Claude Sonnet LLM. You can select another LLM from the list of supported foundational models.

   import mlflow
   from openai import OpenAI
   
   # Enable MLflow's autologging to instrument your application with Tracing
   mlflow.openai.autolog()
   
   # Connect to a Databricks LLM via OpenAI using the same credentials as MLflow
   # Alternatively, you can use your own OpenAI credentials here
   mlflow_creds = mlflow.utils.databricks_utils.get_databricks_host_creds()
   client = OpenAI(
       api_key=mlflow_creds.token,
       base_url=f"{mlflow_creds.host}/serving-endpoints"
   )
   
   # 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="databricks-claude-sonnet-4",  # This example uses a Databricks hosted LLM - you can replace this with any AI Gateway or Model Serving endpoint. If you provide your own OpenAI credentials, replace with a valid OpenAI model e.g., gpt-4o, etc.
       messages=[
         {
           "role": "system",
           "content": "You are a helpful assistant.",
         },
         {
           "role": "user",
           "content": input,
         },
       ],
     )
     return response.choices[0].message.content
   
   result = my_app(input="What is MLflow?")
   print(result)
   

2. Run the application

   python app.py
   

Step 5: View the Trace in MLflow

1. Navigate back to the MLflow Experiment UI
2. You will now see the generated trace in the Traces tab
3. 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

Continue your journey with these recommended actions and tutorials.

• Evaluate your app's quality - Measure and improve your GenAI app's quality with MLflow's evaluation capabilities
• Collect human feedback - Learn how to collect feedback from users and ___domain experts
• Track users & sessions - Add user and conversation context to your traces

Reference guides

Explore detailed documentation for concepts and features mentioned in this guide.

• Tracing concepts - Understand the fundamentals of MLflow Tracing
• Tracing data model - Learn about traces, spans, and how MLflow structures observability data
• Manual tracing APIs - Explore advanced tracing techniques for custom instrumentation