From Fragmented Evaluation to Unified Agent Evaluation with MLflow for AI Observability

Building the initial version of an agent is easier than ever thanks to modern agent development toolkits and coding assistants. However, real value comes from successfully deploying these systems in the real world, which continues to be a challenge. Evaluation is a key lever in closing this gap with reports showing that companies that adopt it bring 6× more AI projects into production.

Subject matter expert (SME) evaluation is considered the gold standard for assessing agent quality. However, human evaluation is expensive, slow, and inconsistent due to subjective preferences, limiting its scalability. As a result, we increasingly rely on LLM-as-a-Judge approaches, where language models evaluate agents outputs across multiple quality dimensions. This approach enables scalable, inexpensive, and reproducible evaluation.

Several evaluation frameworks, such as DeepEval, RAGAS, Arize Phoenix, and MLflow, have emerged to provide LLM-as-a-judge capabilities with metrics specialised for different aspects of agent evaluation, including retrieval quality, hallucination detection, and response quality. However, using multiple frameworks often creates evaluation silos, making it difficult to combine metrics, compare results side by side, or track performance trends across iterations without custom integrations.

MLflow simplifies this by integrating DeepEval, RAGAS, and Arize Phoenix through the MLflow Scorer API, providing a unified interface for evaluation via mlflow.genai.evaluate. This enables teams to use more than 50 metrics from these frameworks through a single API and UI, while seamlessly integrating evaluations with existing MLflow experiments and traces.

In this story, I’ll walk you through the process of building a complete workflow for building and evaluating AI agents leveraging **MLflow 3.10.0’s unified evaluation framework **as the foundation for evaluation-driven development.

The Evaluation Driven Development loop: users interact with the agent, traces are captured, ground-truth expectations are added, systematic evaluation is run, and the results feed back into agent iteration and improvement.

Building and Evaluating a School Policy RAG Agent with MLflow

To understand how these new MLflow 3.10.0 capabilities come together, let’s walk through a complete workflow to build, evaluate, and iteratively improve a school policy RAG agent following the evaluation-driven development cycle. The agent answers students’ questions about school rules and procedures, such as “What happens if I miss too many classes?” or “Can I use my phone at school?”, by retrieving relevant sections from the school’s policy document and generating clear, child-friendly responses using an LLM.

The school policy RAG agent is built on three core components:

• **LangChain**: An open-source framework that provides AI developers with composable building blocks for connecting LLMs to external data sources, tools, and memory.
• **Amazon Bedrock**: A fully managed service that provides access to foundation models from leading AI providers through a unified API.
• FAISS Vector Store: An in-memory vector store that enables fast similarity search during retrieval. To build the FAISS vector store, the school policy document is loaded, split into overlapping chunks, embedded using Amazon Bedrock Titan Embed Text v2, and stored in a FAISS index for semantic retrieval.

To build a comprehensive evaluation pipeline for the school policy RAG agent’s answers, we leverage MLflow to orchestrate multiple judges from different evaluation frameworks (RAGAS and Arize Phoenix) through a single interface, alongside with use-case-specific custom scorers and judges.

📦 All source code is available in the rag-agent-mlflow-evaluation repository, organized in a modular structure with ready-to-use scripts for each step. In this story, we focus on the MLflow-related code and skip detailed explanations of utilities.

Set up your MLflow environment

Each step in the evaluation pipeline begins by configuring the MLflow tracking server and setting the RAG agent experiment. This ensures all agent assets are logged to a single experiment, making it easy to track and manage everything from the MLflow UI.

MLflow UI displaying the centralized experiment for the RAG agent, where evaluation dashboards, traces for observability, prompt registry, and evaluation runs can all be accessed in one place.

Step 1: Register the Prompt

We start by registering the RAG agent chat prompt into MLflow’s Prompt Registry, which enables full prompt lifecycle management with Git-inspired commit-based versioning and aliasing.

The agent’s chat prompt includes a system message that instructs the LLM to act as a school assistant providing correct, child-appropriate explanations, and a user message with placeholders for the retrieved context and the student’s question.

# scripts/register_prompt.py

# ── Register the chat prompt ────────────────────────────────────────────────
# Creates a new version of the prompt in the Prompt Registry.
registered_prompt = mlflow.genai.register_prompt(
    name=PromptConfig.PROMPT_NAME,
    template=[
        {
            "role": "system",
            "content": (
                "You are a school assistant that answers questions from students.\n"
                "Use ONLY the information provided in the context to answer the question.\n"
                "Your goal is to give a correct and child-appropriate explanation "
                "of the school rules or procedures."
            ),
        },
        {
            "role": "user",
            "content": "Context:\n{{ context }}\n\nStudent Question:\n{{ question }}",
        },
    ],
    tags={
        "author": "joana.mesquita@test.com",
        "task": "RAG chat prompt",
        "language": "en",
    },
)

# ── Assign a stable alias ──────────────────────────────────────────────────
# Points the alias (e.g. "production") to the newly registered version
mlflow.genai.set_prompt_alias(
    name=PromptConfig.PROMPT_NAME,
    alias=PromptConfig.PROMPT_ALIAS,
    version=registered_prompt.version,
)

The prompt is now versioned in MLflow and can be loaded at inference time with using an alias, enabling prompt iteration without redeploying the agent.

MLflow Prompt Registry UI displaying versioned RAG agent prompts with metadata tracking, and lifecycle management through alias.

Step 2: Define and Register the Agent

In this step, we define our agent by building a SchoolRAGAgent class that extends [mlflow.pyfunc.ResponsesAgent](https://mlflow.org/docs/latest/api_reference/python_api/mlflow.pyfunc.html#mlflow.pyfunc.ResponsesAgent). This base class wraps agent logic in a standardized serving interface and implements the [mlflow.pyfunc.PythonModel](https://mlflow.org/docs/latest/api_reference/python_api/mlflow.pyfunc.html#mlflow.pyfunc.PythonModel) interface, allowing the agent to integrate directly with MLflow’s logging, versioning, and deployment capabilities.

The SchoolRAGAgent class exposes four key methods:

• __init__: Initialises the Bedrock LLM client with the provided model ID and loads the versioned prompt from MLflow’s Prompt Registry using mlflow.genai.load_prompt with the prompt name and alias.
• load_context: Called automatically by MLflow when the model is loaded. It initialises the FAISS vector store from the pre-built index on disk (or creates it from the school policy document if the index does not yet exist).
• retrieve_relevant_documents: Performs a similarity search against the FAISS vector store using the input question and returns the top-k most relevant document chunks. This method is decorated with @mlflow.trace(span_type=SpanType.RETRIEVER), which automatically creates a retriever span in the MLflow trace and records the retrieved documents as span outputs.
• predict: Orchestrates the full RAG pipeline: it retrieves relevant documents using retrieve_relevant_documents, joins the retrieved chunks into a single context string, formats the prompt with the context and the user’s question, invokes the Bedrock LLM, and returns the generated response as a ResponsesAgentResponse.

# src/core/agent_model.py

class SchoolRAGAgent(ResponsesAgent):
    """Retrieval-Augmented Generation agent using Bedrock LLM and FAISS."""

    def __init__(self, model_id: str) -> None:
        """Initialise the agent's LLM client and load the versioned prompt."""
        self.model = ChatBedrockConverse(
            model_id=model_id,
            region_name=AWSConfig.REGION_NAME,
            temperature=ChatModelConfig.TEMPERATURE,
            max_tokens=ChatModelConfig.MAX_TOKENS,
            credentials_profile_name=AWSConfig.AWS_PROFILE,
        )
        self.vector_store = None
        self.prompt = mlflow.genai.load_prompt(
            f"prompts:/{PromptConfig.PROMPT_NAME}@{PromptConfig.PROMPT_ALIAS}"
        )

    def load_context(self, context) -> None:
        """Initialise the FAISS vector store (called by MLflow on model load)."""
        # Setup vector store
        self.vector_store = setup_vector_store(_logger, DataConfig.DOCUMENT_PATH)

    @mlflow.trace(span_type=SpanType.RETRIEVER)
    def retrieve_relevant_documents(self, question: str) -> List[Document]:
        """Retrieve the top-k most relevant chunks from the FAISS index."""
        # Get documents from the search store
        docs = self.vector_store.similarity_search(question, k=3)

        # Get the current active span (created by @mlflow.trace)
        span = mlflow.get_current_active_span()

        if span is not None:
            # Set the outputs of the span
            outputs = [
                Document(
                    page_content=doc.page_content,
                    metadata={"source": doc.metadata.get("source", "")},
                )
                for doc in docs
            ]
            span.set_outputs(outputs)

        return docs

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        """Orchestrate the full RAG pipeline and return a response."""
        _logger.info(f"Received question: {request.input[0].content}")

        # Retrieve relevant documents from the vector store
        _logger.info("Retrieving relevant documents from vector store...")
        docs = self.retrieve_relevant_documents(request.input[0].content)

        # Build the context from retrieved documents
        context = "\n\n".join(doc.page_content for doc in docs)

        # Format the prompt with context and question
        _logger.info("Formatting prompt with context and question...")
        messages = self.prompt.format(
            context=context, question=request.input[0].content
        )

        # Generate the answer using the model
        _logger.info("Generating answer using the model...")

        answer = self.model.invoke(messages).content
        return ResponsesAgentResponse(
            output=[
                self.create_text_output_item(
                    text=answer,
                    id="msg-1",
                ),
            ],
            custom_outputs={
                "trace_id": (
                    mlflow.get_active_trace_id()
                    if mlflow.get_active_trace_id()
                    else None
                )
            },
        )

After defining the SchoolRAGAgent class, we log and register the model using MLflow's Models from Code feature. The first step is calling mlflow.models.set_model to designate the SchoolRAGAgent class as the model of interest for MLflow.

# src/core/agent_model.py

# Enable MLflow autologging for LangChain so every LLM call is traced.
mlflow.langchain.autolog()

# Instantiate the agent and set model
# This is required so that mlflow.pyfunc.log_model can discover the
# agent when logging it to the Model Registry.
agent = SchoolRAGAgent(model_id=ChatModelConfig.MODEL_ID)
set_model(agent)

Finally, we log and register the model using the standard mlflow.pyfunc.log_model API, specifying the path to the file containing the SchoolRAGAgent definition and using the code_paths parameter to include the agent's dependencies.

# scripts/register_model.py

# ── Log Model ──────────────────────────────────────────────────────────────
# Creates a new version of the model in the Model Registry.
with mlflow.start_run():
    logged_model_info = mlflow.pyfunc.log_model(
        python_model="src/core/agent_model.py",
        name="rag-agent",
        registered_model_name=ModelRegistryConfig.REGISTERED_MODEL_NAME,
        code_paths=["src/", "config/"],
        input_example=ResponsesAgentRequest(
            input=[
                {
                    "role": "user",
                    "content": "What happens if I miss too many classes?",
                }
            ],
        ),
    )

# ── Set Model Alias ─────────────────────────────────────────────────────────
# Assigns an alias to the newly registered model version for easy retrieval.
client.set_registered_model_alias(
    ModelRegistryConfig.REGISTERED_MODEL_NAME,
    ModelRegistryConfig.MODEL_ALIAS,
    logged_model_info.registered_model_version,
)

Step 3: Create and Register a Custom LLM Judge

While the agent architecture itself is straightforward, evaluating it requires assessing multiple quality dimensions: Is the retrieved context relevant? Is the response factually accurate? And most importantly for this use case, is the answer appropriate for a school-aged child?

In this step, we define a custom LLM-as-a-judge using mlflow.genai.judges.make_judge to evaluate whether the agent’s responses are appropriate for a school-aged audience.

# src/evaluation/custom_judges.py

child_appropriateness_judge = make_judge(
    name=EvaluationConfig.CUSTOM_JUDGE_NAME,
    instructions=(
        "Evaluate whether the response in {{ outputs }} is appropriate for a school-aged child "
        "asking the question in {{ inputs }}.\n\n"
        "Rate as:\n"
        "- excellent: Perfectly suited for children, engaging and clear\n"
        "- good: Appropriate with minor improvements possible\n"
        "- ok: Acceptable but could be more child-friendly\n"
        "- bad: Inappropriate or too complex for children"
    ),
    model=EvaluationConfig.CUSTOM_JUDGE_MODEL_ID,
    feedback_value_type=Literal["excellent", "good", "ok", "bad"],
)

With the custom judge defined, the next step is to register it under the RAG agent experiment, so it can be versioned and retrieved during the evaluation step.

# scripts/register_judge.py

# ── Judge registration ──────────────────────────────────────────────────────────
child_appropriateness_judge.register(
    experiment_id=client.get_experiment_by_name(
        ModelRegistryConfig.EXPERIMENT_NAME
    ).experiment_id,
)

Step 4: Collect Traces with Ground-Truth Expectations

In this step, we run the school policy RAG agent on a sample of student questions to generate responses and capture execution traces. For each trace, we attach ground-truth expectations using mlflow.log_expectations. These expectations serve as references during evaluation, allowing scorers and judges to compare the agent’s outputs against the expected results.

Each trace is annotated with two expectations:

• expected_output: The ideal answer the agent should generate.
• expected_context: The specific document sections that should have been retrieved to support the answer.

# scripts/create_traces.py

# ── Load registered model ───────────────────────────────────────────────────────
model_uri: str = (
    f"models:/{ModelRegistryConfig.REGISTERED_MODEL_NAME}@{ModelRegistryConfig.MODEL_ALIAS}"
)
model: Any = mlflow.pyfunc.load_model(model_uri)

# ── Generate traces and log expectations ───────────────────────────────────────
evaluation_question: Dict[str, Any]
for evaluation_question in EVALUATION_DATASET:
    question: str = evaluation_question["question"]
    expected_output: str = evaluation_question["expected_output"]
    expected_context: List[str] = evaluation_question["expected_context"]

    request = {
        "input": [{"role": "user", "content": question}],
    }
    answer = model.predict(request)

    # Record the human-authored source for the logged expectations.
    source: AssessmentSource = AssessmentSource(
        source_type=AssessmentSourceType.HUMAN,
        source_id=EvaluationConfig.EVALUATION_SOURCE,
    )

    # Log the ideal answer expected for this trace.
    mlflow.log_expectation(
        trace_id=answer["custom_outputs"]["trace_id"],
        name="expected_output",
        value=expected_output,
        source=source,
    )

    # Log the context passages that should have been retrieved.
    mlflow.log_expectation(
        trace_id=answer["custom_outputs"]["trace_id"],
        name="expected_context",
        value=expected_context,
        source=source,
    )

At this point, MLflow contains a set of traces, each with the agent’s retrieval spans, generated outputs, and SME-provided expectations.

MLflow Traces UI displaying captured RAG agent execution traces and expectations.

Step 5: Evaluate with Unified Multi-Framework Scorers

We now run evaluation across the collected traces using MLflow’s unified evaluation interface. MLflow allows us to use scorers from multiple frameworks, including RAGAS and Arize Phoenix, together with custom scorers and the previously registered ChildAppropriateness judge through a single API. This removes the need to integrate multiple evaluation frameworks separately and keeps the entire evaluation workflow centralized in MLflow.

The following code defines the scorers used in the evaluation pipeline. We load the previously registered custom judge from the experiment and combine it with out-of-the-box scorers from RAGAS and Phoenix, as well as a custom scorer that validates retrieval quality.

from mlflow.genai.scorers.phoenix import QA
from mlflow.genai.scorers.ragas import FactualCorrectness
from src.evaluation.custom_scorers import ContextRetrieval

# ── Configure scorers and judge ────────────────────────────────────────────────
# Load the registered custom judge from the experiment.
child_appropriateness_judge: Any = get_scorer(
    name=EvaluationConfig.CUSTOM_JUDGE_NAME, experiment_id=experiment_id
)

scorers: List[Scorer] = [
    FactualCorrectness(model=EvaluationConfig.JUDGE_MODEL_ID),
    QA(model=EvaluationConfig.JUDGE_MODEL_ID),
    ContextRetrieval,
    child_appropriateness_judge,
]

Each scorer evaluates a different quality dimension of the agent’s behavior:

• Factual Correctness (RAGAS): Measures whether the generated response is factually consistent with the expected output.
• QA (Arize Phoenix): Evaluates the overall quality of the answer to the user’s question.
• Context Retrieval (Custom scorer): Validates that the expected policy sections appear in the retrieved documents.
• Child Appropriateness (Custom LLM judge): Assesses whether the answer is suitable for a school-aged audience.

Once the scorers are configured, we retrieve the traces generated earlier and run evaluation for all scorers through a single mlflow.genai.evaluate call:

# ── Retrieve traces ─────────────────────────────────────────────────────────────
# Search for recent traces to evaluate.
traces: List[Trace] = mlflow.search_traces(
    max_results=10,
)

# ── Run unified evaluation ──────────────────────────────────────────────────────
_logger.info("Running evaluation")
results: Any = mlflow.genai.evaluate(
    data=traces,
    scorers=scorers,
)

The results are automatically logged into an evaluation run and can be explored in the MLflow UI alongside the agent’s traces:

MLflow UI displaying evaluation run result.

They are also displayed in the Agents Dashboard Overview tab linked to the RAG agent experiment, providing instant, out-of-the-box visibility into the agent health and quality metrics derived from MLflow scorers:

MLflow Agents Dashboard Overview tab displaying RAG agent health and quality metrics derived from MLflow scorers.

Conclusion

In this story, we built a workflow for building, evaluating, and iterating on the school policy RAG agent using MLflow's unified evaluation framework. By bringing multiple evaluation frameworks under a single interface and integrating evaluation across the entire agent lifecycle, from registration and scoring to monitoring and observability, MLflow removes the complexity of comprehensive agent evaluation.

The results of evaluation feed into an iteration loop, allowing teams to refine prompts, tools, and overall system behaviour. This evaluation-driven development builds the confidence needed to bridge the gap between building agents and successfully deploying them in production.

However, evaluation is only as good as the judges running it. If the judges are unreliable, teams risk iterating in the wrong direction. While LLM-as-a-judge approaches correlate well with SME evaluations for general tasks, their reliability can decrease in domain-specific scenarios that require specialised knowledge. In our school policy assistant, this means the ChildAppropriateness judge may miss subtle issues, such as vocabulary that is technically correct but too formal or abstract for a child to understand.

In the next story, we will close the evaluation loop by aligning the ChildAppropriateness judge with SME feedback using **MemAlign**, a framework introduced in MLflow that enables LLM judges to learn domain-specific evaluation standards from a small number of expert examples, bringing us closer to evaluation systems that scale while remaining aligned with human expertise.