MLflow for LLMs: Tracking Prompts and Evaluations

Introduction to MLflow for LLMs

As Large Language Models (LLMs) become central to modern applications, managing the complexity of prompts, model versions, and evaluation metrics has emerged as a critical challenge. MLflow, originally designed for traditional machine learning lifecycle management, has extended its capabilities to support LLM workflows. This tutorial covers how to use MLflow to track prompts, log LLM interactions, and evaluate model outputs systematically.

What is MLflow for LLMs?

MLflow for LLMs is a set of features within the MLflow ecosystem specifically designed to handle the unique requirements of generative AI workflows. It provides tools for tracking prompts as first-class artifacts, logging model parameters and responses, and running evaluations against custom or built-in metrics. Unlike traditional ML where inputs and outputs are numeric, LLM workflows deal with free-form text, making tracking and evaluation significantly more nuanced.

Why It Matters

LLM development is inherently experimental. Small changes to a prompt can dramatically alter output quality, and without systematic tracking, teams risk losing valuable insights. MLflow addresses several pain points:

• Reproducibility: Every prompt, parameter, and model version is logged, enabling exact reproduction of results.
• Comparison: Multiple prompt variants or models can be compared side by side using consistent evaluation metrics.
• Collaboration: Teams can share tracked experiments, making it easier to review and iterate on prompt engineering.
• Production Readiness: Tracked prompts and models can be deployed with confidence, backed by evaluation data.

Setting Up MLflow for LLM Tracking

Before diving into tracking, you need to install MLflow and set up a tracking environment. MLflow supports local tracking out of the box, but for team collaboration, you may want to set up a tracking server.

Installation

Install MLflow along with the LLM-related dependencies. The following command installs MLflow with support for popular LLM providers:

pip install mlflow openai langchain

For local development, MLflow stores data in a local directory by default. To start a tracking server for team use, run:

mlflow server --host 0.0.0.0 --port 5000

Then configure your Python environment to point to the server:

import mlflow

mlflow.set_tracking_uri("http://localhost:5000")

Tracking Prompts and Model Interactions

The core of MLflow for LLMs is the ability to log prompts, model parameters, and responses as part of an experiment run. Let's walk through a practical example using OpenAI's API.

Basic Prompt Tracking

Start by creating an experiment and logging a simple LLM interaction. The key is to treat the prompt as a parameter and the response as an artifact or metric.

import mlflow
from openai import OpenAI

# Set up the experiment
mlflow.set_experiment("llm_prompt_tracking")

# Initialize the OpenAI client
client = OpenAI()

# Define the prompt
system_prompt = "You are a helpful assistant that summarizes text."
user_prompt = "Summarize the following text in one sentence: Machine learning is a subset of artificial intelligence that enables systems to learn and improve from experience without being explicitly programmed."

# Start an MLflow run
with mlflow.start_run(run_name="summarization_v1"):
    # Log prompts as parameters
    mlflow.log_param("system_prompt", system_prompt)
    mlflow.log_param("user_prompt", user_prompt)
    mlflow.log_param("model", "gpt-4")
    mlflow.log_param("temperature", 0.7)
    mlflow.log_param("max_tokens", 150)

    # Make the API call
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=0.7,
        max_tokens=150,
    )

    # Log the response
    output_text = response.choices[0].message.content
    mlflow.log_text(output_text, "response.txt")
    
    # Log token usage as metrics
    mlflow.log_metric("prompt_tokens", response.usage.prompt_tokens)
    mlflow.log_metric("completion_tokens", response.usage.completion_tokens)
    mlflow.log_metric("total_tokens", response.usage.total_tokens)

    print(f"Response: {output_text}")

This example logs the system and user prompts as parameters, the model configuration, the response as a text artifact, and token usage as metrics. Everything is captured in a single run, making it easy to review later.

Using Prompt Templates

For more complex applications, you often work with prompt templates that include variables. MLflow allows you to log these templates and the variables used in each run.

import mlflow
from openai import OpenAI

client = OpenAI()

# Define a prompt template
template = """
You are an expert in {domain}.
Answer the following question concisely:
Question: {question}
"""

mlflow.set_experiment("prompt_template_tracking")

with mlflow.start_run(run_name="domain_qa_v1"):
    # Log the template as a text artifact
    mlflow.log_text(template, "prompt_template.txt")
    
    # Log the variables
    variables = {
        "domain": "data engineering",
        "question": "What is the difference between ETL and ELT?"
    }
    mlflow.log_param("template_variables", str(variables))
    
    # Fill in the template
    filled_prompt = template.format(**variables)
    
    # Call the model
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": filled_prompt}],
        temperature=0.3,
    )
    
    output = response.choices[0].message.content
    mlflow.log_text(output, "response.txt")
    mlflow.log_metric("total_tokens", response.usage.total_tokens)
    
    print(output)

Logging LLM Models with MLflow Flavors

MLflow provides specific flavors for logging LLM models, which package the model along with its configuration so it can be loaded and used later. The mlflow.openai flavor is one of the most commonly used.

Logging an OpenAI Model

import mlflow
from mlflow.models import infer_signature

mlflow.set_experiment("llm_model_logging")

# Define the model configuration
model_config = {
    "model": "gpt-4",
    "temperature": 0.5,
    "max_tokens": 200,
}

# Define input and output examples
input_example = {
    "messages": [
        {"role": "user", "content": "What is the capital of France?"}
    ]
}

output_example = {
    "content": "The capital of France is Paris."
}

signature = infer_signature(input_example, output_example)

with mlflow.start_run(run_name="openai_gpt4_logged"):
    # Log the model
    mlflow.openai.log_model(
        model="gpt-4",
        task="chat.completions",
        artifact_path="model",
        config=model_config,
        signature=signature,
        input_example=input_example,
    )
    
    # Log additional metadata
    mlflow.set_tag("use_case", "general_qa")
    mlflow.set_tag("team", "ml_platform")
    mlflow.log_param("prompt_strategy", "zero_shot")

Loading and Using a Logged Model

import mlflow

# Load the model from a specific run
run_id = "your_run_id_here"
model_uri = f"runs:/{run_id}/model"

loaded_model = mlflow.pyfunc.load_model(model_uri)

# Use the model
response = loaded_model.predict({
    "messages": [
        {"role": "user", "content": "Explain recursion in one sentence."}
    ]
})

print(response)

Evaluating LLM Outputs

Evaluation is where MLflow truly shines for LLM workflows. Traditional metrics like accuracy do not directly apply to generative tasks. MLflow provides an evaluation framework that supports both built-in and custom metrics for assessing LLM outputs.

Built-in LLM Evaluation Metrics

MLflow includes several pre-defined metrics for LLM evaluation, such as relevance, toxicity, and factual correctness. These metrics use an LLM as a judge to score outputs.

import mlflow
import pandas as pd

mlflow.set_experiment("llm_evaluation")

# Prepare evaluation data
eval_data = pd.DataFrame({
    "inputs": [
        "What is machine learning?",
        "Explain the concept of recursion.",
        "What are the benefits of cloud computing?",
    ],
    "outputs": [
        "Machine learning is a branch of AI that allows systems to learn from data.",
        "Recursion is when a function calls itself to solve smaller instances of a problem.",
        "Cloud computing offers scalability, cost savings, and accessibility.",
    ],
    "targets": [
        "Machine learning is a subset of artificial intelligence that enables systems to learn patterns from data and make predictions.",
        "Recursion is a programming technique where a function calls itself to break down a problem into smaller subproblems.",
        "Cloud computing provides on-demand resources, scalability, reduced costs, and global accessibility.",
    ],
})

with mlflow.start_run(run_name="eval_v1"):
    # Evaluate using built-in metrics
    results = mlflow.evaluate(
        data=eval_data,
        targets="targets",
        predictions="outputs",
        extra_metrics=[
            mlflow.metrics.genai.answer_similarity("targets"),
            mlflow.metrics.genai.answer_correctness("targets"),
            mlflow.metrics.latency(),
        ],
        evaluator_config={
            "col_mapping": {
                "inputs": "inputs",
                "predictions": "outputs",
                "targets": "targets",
            }
        },
    )
    
    # Print the evaluation results
    print(f"Metrics: {results.metrics}")
    
    # Access the per-row results table
    results_table = results.tables["eval_results_table"]
    print(results_table)

Custom Evaluation Metrics

For use cases where built-in metrics are insufficient, you can define custom evaluation functions. This is particularly useful for domain-specific quality criteria.

import mlflow
from mlflow.metrics import MetricValue

def response_length_metric(eval_df, built_in_metrics=None):
    """Custom metric that checks if response length is within a target range."""
    predictions = eval_df["prediction"]
    lengths = predictions.apply(lambda x: len(str(x).split()))
    
    # Score: 1.0 if between 10 and 50 words, 0.5 otherwise
    scores = lengths.apply(lambda l: 1.0 if 10 <= l <= 50 else 0.5)
    
    return MetricValue(
        scores=scores,
        aggregate_results={
            "mean_score": scores.mean(),
            "within_range_pct": (scores == 1.0).mean(),
        },
    )

def keyword_presence_metric(eval_df, built_in_metrics=None):
    """Custom metric that checks for presence of required keywords."""
    predictions = eval_df["prediction"]
    required_keywords = ["important", "because", "therefore"]
    
    scores = predictions.apply(
        lambda p: sum(1 for kw in required_keywords if kw in str(p).lower()) / len(required_keywords)
    )
    
    return MetricValue(
        scores=scores,
        aggregate_results={
            "mean_keyword_coverage": scores.mean(),
        },
    )

# Use the custom metrics in evaluation
eval_data = pd.DataFrame({
    "inputs": ["Why is data validation important?"],
    "outputs": ["Data validation is important because it ensures data quality and therefore prevents errors downstream."],
})

with mlflow.start_run(run_name="custom_eval_v1"):
    results = mlflow.evaluate(
        data=eval_data,
        predictions="outputs",
        extra_metrics=[
            response_length_metric,
            keyword_presence_metric,
        ],
    )
    
    print(f"Custom Metrics: {results.metrics}")

LLM-as-a-Judge Evaluation

One powerful pattern is using an LLM to evaluate the outputs of another LLM. MLflow supports this through the genai metrics module, which allows you to define custom grading prompts.

import mlflow

# Define a custom grading prompt for the LLM judge
grading_prompt = """
You are an expert evaluator. Score the following response on a scale of 1-5.

Question: {input}
Response: {prediction}
Reference Answer: {target}

Scoring Criteria:
- 5: Excellent, comprehensive and accurate
- 4: Good, mostly accurate with minor gaps
- 3: Acceptable, covers the basics
- 2: Poor, significant inaccuracies
- 1: Unacceptable, completely wrong

Provide only the numeric score.
"""

mlflow.set_experiment("llm_judge_eval")

eval_data = pd.DataFrame({
    "inputs": ["What is a primary key in a database?"],
    "outputs": ["A primary key is a unique identifier for each record in a table."],
    "targets": ["A primary key is a column or set of columns that uniquely identifies each row in a database table, ensuring no duplicate records exist."],
})

with mlflow.start_run(run_name="llm_judge_v1"):
    custom_metric = mlflow.metrics.genai.make_genai_metric(
        name="expert_score",
        definition="Scores response quality from 1-5 based on accuracy and completeness.",
        grading_prompt=grading_prompt,
        model="gpt-4",
        examples=[
            {
                "input": "What is an API?",
                "prediction": "An API is a set of rules for building software.",
                "target": "An API is an interface that allows different software applications to communicate with each other.",
                "score": 3,
            }
        ],
        version="v1",
        aggregations=["mean", "median"],
        greater_is_better=True,
    )
    
    results = mlflow.evaluate(
        data=eval_data,
        targets="targets",
        predictions="outputs",
        extra_metrics=[custom_metric],
    )
    
    print(f"Judge Scores: {results.metrics}")

Tracking Prompt Iterations and A/B Testing

A common workflow in LLM development is iterating on prompts and comparing results. MLflow makes this straightforward by allowing you to log multiple runs under the same experiment and compare them.

Comparing Prompt Variants

import mlflow
from openai import OpenAI

client = OpenAI()
mlflow.set_experiment("prompt_ab_testing")

# Define multiple prompt variants
prompt_variants = {
    "v1_direct": "List three benefits of exercise.",
    "v2_contextual": "You are a health expert. List three benefits of regular exercise and explain each briefly.",
    "v3_structured": "List three benefits of exercise in the following format:\n1. [Benefit]: [Explanation]\n2. [Benefit]: [Explanation]\n3. [Benefit]: [Explanation]",
}

for variant_name, prompt in prompt_variants.items():
    with mlflow.start_run(run_name=f"prompt_{variant_name}"):
        mlflow.log_param("prompt_variant", variant_name)
        mlflow.log_param("prompt_text", prompt)
        mlflow.log_param("model", "gpt-4")
        mlflow.log_param("temperature", 0.5)
        
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.5,
        )
        
        output = response.choices[0].message.content
        mlflow.log_text(output, f"response_{variant_name}.txt")
        mlflow.log_metric("response_length", len(output))
        mlflow.log_metric("total_tokens", response.usage.total_tokens)
        
        print(f"\n--- {variant_name} ---")
        print(output)

After running these experiments, you can compare them in the MLflow UI or programmatically:

from mlflow.tracking import MlflowClient

client = MlflowClient()
experiment = client.get_experiment_by_name("prompt_ab_testing")

runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    order_by=["metrics.response_length DESC"],
)

for run in runs:
    print(f"Run: {run.info.run_name}")
    print(f"  Variant: {run.data.params.get('prompt_variant')}")
    print(f"  Response Length: {run.data.metrics.get('response_length')}")
    print(f"  Total Tokens: {run.data.metrics.get('total_tokens')}")
    print()

Best Practices

To get the most out of MLflow for LLM tracking and evaluation, consider the following best practices:

• Version your prompts: Treat prompts like code. Log every version and use descriptive run names so you can trace which prompt produced which results.
• Log full context: Include not just the final prompt but also the system prompt, few-shot examples, and any preprocessing steps. This ensures full reproducibility.
• Use tags for organization: Tags like use_case, team, and environment make it easier to filter and find relevant runs in large organizations.
• Establish baseline metrics: Before iterating on prompts, establish baseline evaluation metrics. This gives you a reference point to measure improvements against.
• Combine automated and human evaluation: LLM-as-a-judge metrics are useful for quick iteration, but human review remains important for nuanced quality assessment. Log human ratings as additional metrics when available.
• Track cost alongside quality: Log token counts and estimated costs as metrics. This helps balance quality improvements against budget constraints.
• Use the MLflow UI for comparison: The web-based UI provides visual comparison tools that are especially helpful when evaluating multiple prompt variants side by side.
• Store evaluation datasets as artifacts: Log the evaluation dataset used for each run so that future evaluations can use the same data for consistent comparisons.

Conclusion

MLflow provides a robust framework for managing the experimental nature of LLM development. By systematically tracking prompts, model configurations, and evaluation metrics, teams can move from ad-hoc prompt tweaking to a disciplined, reproducible workflow. The ability to log every interaction, compare variants, and evaluate outputs with both built-in and custom metrics makes MLflow an essential tool for any team building LLM-powered applications. As the generative AI landscape continues to evolve, having a solid tracking and evaluation foundation will be key to shipping reliable, high-quality LLM features to production.