Skip to main content

Tracing Quickstart (TS/JS)

info

The Python quickstart is available here.

This quickstart guide will walk you through setting up a simple GenAI application with MLflow Tracing. In less than 10 minutes, you'll enable tracing, run a basic application, and explore the generated traces in the MLflow UI.

Prerequisites​

Install the required packages by running the following command:

npm install mlflow-openai
info

The code example in this guide uses the OpenAI SDK; however, the contents of this guide are applicable to any other LLM providers, such as Anthropic, Google, Bedrock, and more.

For other LLM providers that are not natively integrated with auto-tracing, install the mlflow-tracing package and wrap the LLM call with the mlflow.trace wrapper function instead.

Example of wrapping Anthropic API calls to generate a trace
import * as mlflow from "mlflow-tracing";

// Instantiate the Anthropic client
const anthropic = new Anthropic({
  apiKey: process.env.ANTHROPIC_API_KEY,
});

// Wrap the client method with the `mlflow.trace` wrapper.
const tracedAnthropic = mlflow.trace(
  anthropic.messages.create,
  { spanType: mlflow.SpanType.LLM }
);

// Call the wrapped function just like the original function.
// MLflow will generate a trace that captures inputs, outputs, latency, and exceptions.
const response = await tracedAnthropic({
  model: "claude-sonnet-4-20250514",
  max_tokens: 1024,
  messages: [{ role: "user", content: "Hello, Claude" }],
});

Step 1: Set up your environment​

Connect to MLflow​

MLflow logs Traces in a tracking server. Connect your application to the tracking server by one of the following methods.

The Typescript package does not contain the tracking server. To start a self-hosted tracking server from command line using MLflow Python package:

# Install mlflow python package
pip install mlflow -U

# Start MLflow tracking server locally
mlflow ui --port 5000

# This will start the server at http://127.0.0.1:5000
tip

If you are unsure about how to set up an MLflow tracking server, you can start with the cloud-based MLflow powered by Databricks: Sign up for free →

Create a new MLflow Experiment​

Create a new MLflow experiment in the MLflow UI, or choose an existing experiment.

Initialize the Tracing SDK​

import { init } from "mlflow-tracing";

init({
    trackingUri: "<your-tracking-server-uri>",
    experimentId: "<your-experiment-id>",
});

Set OpenAI API Key (or other LLM providers)​

export OPENAI_API_KEY="your-api-key-here"  # Replace with your actual API key

Step 2: Trace a single LLM call​

Let's start with a simple example of tracing a single LLM call. We first wrap the OpenAI client with the tracedOpenAI function. After that, every call to OpenAI API will generate a trace span, capturing the input, output, latency, token counts, and other metadata.

import { OpenAI } from "openai";
import { tracedOpenAI } from "mlflow-openai";

// Wrap the OpenAI client with the tracedOpenAI function
const client = tracedOpenAI(new OpenAI());

// Invoke the client as usual
const response = await client.chat.completions.create({
    model: "o4-mini",
    messages: [
        {"role": "system", "content": "You are a helpful weather assistant."},
        {"role": "user", "content": "What's the weather like in Seattle?"},
    ],
})

After running the code above, go to the MLflow UI and select the "Traces" tab. It should show the newly created trace.

Single Trace

The table view shows the primary metadata of the trace, such as the trace ID, execution duration, token count, source system and status. You can add or remove displayed columns by selecting the columns in the drop down.

By clicking on the request row (the linked request text), you can view the detailed spans in the trace.

Single Trace Detail

The "Chat" view in the above screenshot shows the full chat messages exchanged between the user and the model. By clicking other tables such as "Inputs / Outputs" or "Attributes", you can see different aspects of the trace, including the raw input payload, token usage breakdown, and more.

Step 3: Trace a tool calling agent​

Next, let's add a bit more complexity to the application. To get the real-time weather information, we will use an external weather API as a tool. The application will include a tool calling flow, not only a simple LLM call. To instrument that custom Python flow, we will use the mlflow.trace wrapper function.

import * as mlflow from "mlflow-tracing";

// Wrap the tool function with the `mlflow.trace` wrapper. The wrapped function will be automatically traced and logged to MLflow.
const getWeather = mlflow.trace(
    // A tool function that fetches the weather information from a weather API
    async function getWeather(latitude: number, longitude: number) {
        const response = await fetch(`https://api.open-meteo.com/v1/forecast?latitude=${latitude}&longitude=${longitude}&current=temperature_2m,wind_speed_10m&hourly=temperature_2m,relative_humidity_2m,wind_speed_10m`);
        const data = await response.json();
        return data['current']['temperature_2m'];
    },
     // Set the span type to TOOL. You can also set other span configuration here.
    { spanType: mlflow.SpanType.TOOL }
);

To pass the function as a tool to the LLM, we need to define the JSON schema for the function.

const tools = [{
    type: "function",
    function: {
        name: "get_weather" as const,
        description: "Get current temperature for provided coordinates in celsius.",
        parameters: {
            type: "object",
            properties: {
                latitude: { type: "number" },
                longitude: { type: "number" }
            },
            required: ["latitude", "longitude"],
            additionalProperties: false
        },
        strict: true
    }
}];

Lastly, define a simple flow that first asks the LLM to get instructions for calling the tool, then invokes the tool function, and lastly returns the result to the LLM.. This example uses the mlflow.withSpan API to create a parent span for the agent flow, but you can also achieve the same by using the mlflow.trace API like the previous example.

async function runToolAgent(question: string) {
    console.log(`\n🤖 Running tool agent with question: "${question}"`);

    return await mlflow.withSpan(
        async () => {
            const client = tracedOpenAI(new OpenAI());
            const messages: any[] = [{"role": "user", "content": question}];

            // First LLM call to get tool instructions
            const response = await client.chat.completions.create({
                model: "o4-mini",
                messages: messages,
                tools: tools,
            });

            const aiMsg = response.choices[0].message;
            messages.push(aiMsg);

            // If the model requests tool call(s), invoke the function
            if (aiMsg.tool_calls && aiMsg.tool_calls.length > 0) {
                for (const toolCall of aiMsg.tool_calls) {
                    const functionName = toolCall.function.name;

                    if (functionName === "get_weather") {
                        // Invoke the tool function with the provided arguments
                        const args = JSON.parse(toolCall.function.arguments);
                        const toolResult = await getWeather(args.latitude, args.longitude);
                        messages.push({
                            "role": "tool",
                            "tool_call_id": toolCall.id,
                            "content": String(toolResult),
                        });
                    } else {
                        throw new Error(`Invalid tool returned: ${functionName}`);
                    }
                }

                // Second LLM call with tool results
                const finalResponse = await client.chat.completions.create({
                    model: "o4-mini",
                    messages: messages
                });
                return finalResponse.choices[0].message.content;
            }
            return aiMsg.content;
        },
        {
            name: "run_tool_agent",
            spanType: mlflow.SpanType.AGENT,
        }
    );
}

Now we can run the application.

# Run the tool calling agent
const question = "What's the weather like in Seattle?";
await runToolAgent(question);

Step 4: Explore Traces in the UI​

After running the application, you can explore the traces in the MLflow UI.

Tool Calling Trace

The trace shows all LLM invocations and tool calls, organized in a tree structure. You can also inspect the timeline breakdown by clicking the timeline icon next to the tree view. This helps you understand where the time is spent in the application.

Tool Calling Trace Timeline

Step 5: Attach Feedbacks on Traces​

As a last step of this quickstart, let's attach feedback on the generated traces. In real world development, human feedback is critical to improve the quality of any LLM-powered application.

To add a feedback to a trace, you can open the trace detail page and click the "Add new Assessment" button on the top right. It will open an input form where you can provide various feedback values and metadata. For example, we can add feedback called "Quality" with an integer value (1~5), indicating how good the answer is. We can also put the detailed rationale behind the score for future reference.

Feedback Input Form

When you submit the form with "Create" button, the feedback will be attached to the trace.

Feedback List

The aggregated score in the experiment can be seen in the Trace list. You can do slice-and-dice by various criteria, such as timestamp, source, tags, and it will update the aggregated score in real-time.

Feedback Aggregated

Summary​

Congratulations! You've successfully:

  • ✅ Set up MLflow Tracing for a GenAI application
  • ✅ Enabled automatic tracing for OpenAI API calls
  • ✅ Generated and explored traces in the MLflow UI
  • ✅ Learned how to add custom tracing using decorators
  • ✅ Learned how to attach feedback on traces

MLflow Tracing provides powerful observability for your GenAI applications, helping you monitor performance, debug issues, and understand user interactions. Continue exploring the advanced features to get the most out of your tracing setup!

Next Steps​

Now that you have basic tracing working, explore these advanced features: