AI Engineer Portal

Your personal operating system for career transition.

Private mode

Exercise

Build a multi-stage RAG pipeline with query routing

Build a Multi-Stage RAG Pipeline with Query Routing

Not all queries are the same. "What is our refund policy?" is a simple factual lookup. "Compare our Q3 2023 and Q4 2023 revenue performance" requires retrieving from multiple documents and synthesizing across them. Using the same retrieval pipeline for both wastes resources and produces worse answers for synthesis queries.

Query routing classifies incoming queries and directs them to specialized retrieval strategies optimized for each query type.

What you are building

Create a RoutedRAGPipeline with:

A query classifier — injectable function that returns a QueryType enum: FACTUAL, SYNTHESIS, or OUT_OF_SCOPE.
Factual retrieval — retrieves top-3 chunks for precision-optimized answering.
Synthesis retrieval — retrieves top-8 chunks, groups by source metadata field, deduplicates within each source, returns a multi-source context map for comparative answering.
Out-of-scope handling — returns a canned response without any retrieval or LLM calls.
Context formatting — factual uses a numbered list with source labels; synthesis uses per-source block headers (### Source: {name}).
Returns a RoutedResponse with: answer, query_type, routing_reason, chunks_used, and latency breakdown.

Constraints

Injectable classifier and retrieve functions. QueryType as a Python Enum. Type-hint everything.

Retrieval / hard / Step 15 of 15

Practice stage

Retrieval quality and ranking

Hint

Do not trust a single score blindly. Combine ranking logic with metadata and think about why a candidate deserves to rise.

Next drill

This is the end of the current mini-sequence.

Success criteria

- Ranking rule is explainable
- Supports future iteration without rewriting everything
- Feels tied to product trust rather than pure math trivia

Review checklist

- Did I account for both relevance and metadata?
- Would I be able to explain the ranking rule in an interview?
- Does the code make future tuning easy?

from __future__ import annotations

import time
from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable

class QueryType(Enum):
    FACTUAL = "factual"
    SYNTHESIS = "synthesis"
    OUT_OF_SCOPE = "out_of_scope"

@dataclass
class Chunk:
    id: str
    text: str
    score: float
    metadata: dict = field(default_factory=dict)

@dataclass
class RoutedResponse:
    answer: str
    query_type: QueryType
    routing_reason: str
    chunks_used: list[Chunk]
    latency_ms: dict[str, float]

def format_factual_context(chunks: list[Chunk]) -> str:
    # TODO: numbered list: "[1] (Source: {source})\n{text}"
    raise NotImplementedError

def format_synthesis_context(source_map: dict[str, list[Chunk]]) -> str:
    # TODO: per-source blocks: "### Source: {name}\n- {text}\n..."
    raise NotImplementedError

class RoutedRAGPipeline:
    def __init__(
        self,
        classify_fn: Callable[[str], QueryType],
        retrieve_fn: Callable[[str, int], list[Chunk]],
        llm_fn: Callable[[str, str], str],  # (question, context) -> answer
        scope_description: str = "our product documentation",
    ):
        self.classify_fn = classify_fn
        self.retrieve_fn = retrieve_fn
        self.llm_fn = llm_fn
        self.scope_description = scope_description

def query(self, question: str) -> RoutedResponse:
        """Classify, route to appropriate strategy, generate answer."""
        # TODO: implement
        raise NotImplementedError

Practice

Generate a variation

Generate a new exercise variation to deepen understanding or practice a related concept.

Attempt history

Recent submissions

Before you submit, decide what a strong answer should make obvious to the reviewer.

No attempts yet.