View MLflow traces

To view MLflow traces after running a generative AI agent, use one of the following options:

• Inline visualization: See traces rendered directly in Databricks notebook cell outputs.

• Workspace experiments: In your Databricks workspace, go to Experiments, select an experiment, and then click Traces.

• MLflow run: Traces appear on the Run page of the MLflow UI when an agent runs under an active MLflow Run.

• Agent Evaluation UI: In Agent Evaluation, click See detailed trace view in the evaluation result.

Search for traces

Use the UI or API to search and filter traces on properties such as trace name, tags, and status. You can also extract specific fields from spans using the extract_field argument.

UI search

The MLflow UI supports a SQL-like query language to filter traces by properties.

Find traces for a specific function:

name = 'get_weather'

Get traces with a specific status:

status = 'ERROR'

Combine filters:

name = 'get_weather' AND status = 'ERROR'

API search

MLflow provides two APIs to search traces programmatically, mlflow.search_traces() and MlflowClient.search_traces().

Use mlflow.search_traces() to get a list of traces in the pandas DataFrame format. The DataFrame format is helpful for analysis and evaluation.

Python

# High-level API - returns a pandas DataFrame
traces_df = mlflow.search_traces(
    filter_string="trace.name = 'agent_function' AND trace.status = 'OK'",
    order_by=["timestamp DESC"],
    # You can also extract specific span fields as dataframe columns
    extract_fields=[
        "<span_name>.inputs",
        "<span_name>.outputs",
    ],
    max_results=100
)

Use MlflowClient.search_traces() to get a list of traces in a Python list format.

Python

# Client API - returns Trace objects
from mlflow import MlflowClient

client = MlflowClient()
traces = client.search_traces(
    experiment_ids=[experiment.experiment_id],
    filter_string="tag.user_id = '12345'",
)

For more information about the trace filter syntax, see the MLflow Tracing documentation - filtering traces.

Annotate traces with tags

MLflow trace tags help you filter and search traces in the UI. Tags are key-value pairs that allow you to add custom metadata to traces, such as a conversation ID, a user ID, Git commit hash, etc.

You can set tags on ongoing or completed traces using MLflow APIs or the UI. The following example adds a tag to an ongoing trace using the mlflow.update_current_trace() API:

Python

@mlflow.trace
def my_func(x):
    mlflow.update_current_trace(tags={"fruit": "apple"})
    return x + 1

To learn more about tagging traces and how to use them to filter and search traces, see MLflow documentation - Setting Trace Tags.

Evaluate agents using traces

Trace data serves as a valuable resource for evaluating your agents. By capturing detailed information about the execution of your models, MLflow Tracing is instrumental in offline evaluation. You can use the trace data to evaluate your agent's performance against a golden dataset, identify issues, and improve your agent's performance.

Python

%pip install -U mlflow databricks-agents
%restart_python

Python

import mlflow

# Get the recent 50 successful traces from the experiment
traces = mlflow.search_traces(
    max_results=50,
    filter_string="status = 'OK'",
)

traces.drop_duplicates("request", inplace=True) # Drop duplicate requests.
traces["trace"] = traces["trace"].apply(lambda x: x.to_json()) # Convert the trace to JSON format.

# Evaluate the agent with the trace data
mlflow.evaluate(data=traces, model_type="databricks-agent")

To learn more about agent evaluation see Run an evaluation and view the results.

Monitor deployed agents with inference tables

After an agent is deployed to Mosaic AI Model Serving, you can use inference tables to monitor the agent.

The inference tables contain detailed logs of requests, responses, agent traces, and agent feedback from the review app. This information lets you debug issues, monitor performance, and create a golden dataset for offline evaluation.

To enable inference tables for agent deployments, see Enable inference tables for AI agents.

Use a notebook to query the inference table and analyze the results.

To visualize traces, run display(<the request logs table>) and select rows to inspect:

Python

# Query the inference table
df = spark.sql("SELECT * FROM <catalog.schema.my-inference-table-name>")
display(df)

Next steps

See Monitor apps deployed using Agent Framework.