Haystack RAG pipeline tutorial: query routing

Query routing is one of the highest-leverage improvements you can make to a haystack rag pipeline query classification system. Instead of throwing every question at a single retriever, you classify the intent first — then send code questions to a code index, documentation questions to a docs index, and general knowledge questions to a broad corpus. This tutorial builds a complete, runnable pipeline that does exactly that.

Prerequisites

Python 3.10+
An OpenAI API key (for embeddings and classification)
A working knowledge of Haystack 2.x concepts: components, pipelines, and document stores

Install the dependencies:

pip install haystack-ai==2.6.1 \
  openai==1.40.0 \
  python-dotenv==1.0.1 \
  datasets==2.18.0

Create a .env file with your key:

OPENAI_API_KEY=sk-...

Architecture overview

The pipeline has three stages:

Classify — An LLM-based router labels each query as code, docs, or general
Route — A conditional component sends the query to the matching retriever
Retrieve + Generate — The selected retriever fetches context; a generator produces the answer

Query → Classifier → Router → [Code Retriever | Docs Retriever | General Retriever] → Generator → Answer

Each retriever backs a separate document store with domain-tuned embeddings. The classifier uses a lightweight prompt with few-shot examples — fast, cheap, and deterministic enough for production.

Prepare sample data

We’ll use three tiny in-memory datasets so the tutorial runs without external infrastructure. In production you’d swap these for your real corpora.

# data_prep.py
from datasets import Dataset
from haystack import Document

code_docs = [
    Document(content="def fibonacci(n):\n    if n <= 1: return n\n    return fibonacci(n-1) + fibonacci(n-2)", meta={"source": "algorithms.py"}),
    Document(content="class LRUCache:\n    def __init__(self, capacity):\n        self.capacity = capacity\n        self.cache = {}\n        self.order = []", meta={"source": "cache.py"}),
    Document(content="async def fetch_user(session, user_id):\n    async with session.get(f'/users/{user_id}') as resp:\n        return await resp.json()", meta={"source": "api.py"}),
]

docs_docs = [
    Document(content="Haystack 2.x uses a component-based architecture. Pipelines connect components via typed inputs and outputs.", meta={"source": "architecture.md"}),
    Document(content="The DocumentStore protocol defines write_documents, delete_documents, and filter_documents methods.", meta={"source": "document_store.md"}),
    Document(content="Retrievers implement the Retriever protocol with a run method accepting query and filters.", meta={"source": "retriever.md"}),
]

general_docs = [
    Document(content="Python is a high-level, interpreted programming language created by Guido van Rossum in 1991.", meta={"source": "wiki_python.md"}),
    Document(content="The capital of France is Paris. Population: ~2.1 million.", meta={"source": "wiki_france.md"}),
    Document(content="Water boils at 100°C at standard atmospheric pressure.", meta={"source": "physics.md"}),
]

def to_dataset(docs, label):
    return Dataset.from_dict({
        "content": [d.content for d in docs],
        "meta": [d.meta for d in docs],
        "label": [label] * len(docs),
    })

code_ds = to_dataset(code_docs, "code")
docs_ds = to_dataset(docs_docs, "docs")
general_ds = to_dataset(general_docs, "general")

Run it to verify:

python data_prep.py

No output means success — the datasets are in memory.

Build the document stores and retrievers

We’ll use InMemoryDocumentStore with OpenAI embeddings. Each store gets its own index name so they stay isolated.

# stores.py
import os
from haystack import Document
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.embedders import OpenAIDocumentEmbedder
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from data_prep import code_ds, docs_ds, general_ds

os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

def build_store(dataset, index_name):
    store = InMemoryDocumentStore()
    embedder = OpenAIDocumentEmbedder(model="text-embedding-3-small")
    docs = [Document(content=row["content"], meta=row["meta"]) for row in dataset]
    docs_with_embeddings = embedder.run(docs)["documents"]
    store.write_documents(docs_with_embeddings, policy="overwrite")
    retriever = InMemoryEmbeddingRetriever(document_store=store, top_k=2)
    return store, retriever

code_store, code_retriever = build_store(code_ds, "code")
docs_store, docs_retriever = build_store(docs_ds, "docs")
general_store, general_retriever = build_store(general_ds, "general")

Verify the stores have vectors:

# verify_stores.py
from stores import code_store, docs_store, general_store

for name, store in [("code", code_store), ("docs", docs_store), ("general", general_store)]:
    count = store.count_documents()
    sample = store.filter_documents()[:1]
    has_embedding = sample[0].embedding is not None if sample else False
    print(f"{name}: {count} docs, embeddings={has_embedding}")

Expected output:

code: 3 docs, embeddings=True
docs: 3 docs, embeddings=True
general: 3 docs, embeddings=True

Create the query classifier

The classifier is a prompt-driven component that returns a single label. We’ll wrap it in a custom component so it plugs cleanly into the pipeline.

# classifier.py
from typing import Literal
from haystack import component
from haystack.components.generators import OpenAIGenerator
from haystack.dataclasses import ChatMessage

@component
class QueryClassifier:
    def __init__(self, model: str = "gpt-4o-mini"):
        self.generator = OpenAIGenerator(model=model)
        self.system_prompt = """Classify the user query into exactly ONE category:
- code: programming questions, syntax, algorithms, debugging, libraries
- docs: questions about Haystack, framework usage, APIs, configuration
- general: general knowledge, facts, definitions, non-technical topics

Return ONLY the category name. No explanation."""

        self.few_shots = [
            ChatMessage.from_user("How do I implement LRU cache in Python?"),
            ChatMessage.from_assistant("code"),
            ChatMessage.from_user("What is the DocumentStore protocol in Haystack?"),
            ChatMessage.from_assistant("docs"),
            ChatMessage.from_user("What is the capital of France?"),
            ChatMessage.from_assistant("general"),
            ChatMessage.from_user("Explain async/await in Python"),
            ChatMessage.from_assistant("code"),
            ChatMessage.from_user("How do I connect a custom retriever in Haystack?"),
            ChatMessage.from_assistant("docs"),
        ]

    @component.output_types(label=str)
    def run(self, query: str):
        messages = [
            ChatMessage.from_system(self.system_prompt),
            *self.few_shots,
            ChatMessage.from_user(query),
        ]
        response = self.generator.run(messages)
        label = response["replies"][0].strip().lower()
        # Guard against hallucinated labels
        if label not in ("code", "docs", "general"):
            label = "general"
        return {"label": label}

Test it in isolation:

# test_classifier.py
from classifier import QueryClassifier

clf = QueryClassifier()
tests = [
    "How do I write a binary search in Python?",
    "What does the Haystack Retriever protocol look like?",
    "Who won the World Cup in 2022?",
]

for q in tests:
    result = clf.run(query=q)
    print(f"'{q}' → {result['label']}")

Expected output:

'How do I write a binary search in Python?' → code
'What does the Haystack Retriever protocol look like?' → docs
'Who won the World Cup in 2022?' → general

Build the routing pipeline

Haystack’s ConditionalRouter component handles the branching. We define routes as a list of (condition, output_name) tuples where the condition is a Jinja2 expression against the classifier’s output.

# pipeline.py
from haystack import Pipeline
from haystack.components.routers import ConditionalRouter
from haystack.components.joiners import DocumentJoiner
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator
from haystack.dataclasses import ChatMessage
from classifier import QueryClassifier
from stores import code_retriever, docs_retriever, general_retriever

# Classifier
classifier = QueryClassifier()

# Router: sends query to exactly one retriever based on label
routes = [
    {"condition": "{{label == 'code'}}", "output": "code_query", "output_type": str},
    {"condition": "{{label == 'docs'}}", "output": "docs_query", "output_type": str},
    {"condition": "{{label == 'general'}}", "output": "general_query", "output_type": str},
]
router = ConditionalRouter(routes)

# Retrievers (each receives the query string on its named input)
# We'll connect router outputs to retriever inputs via pipeline connections

# Joiner merges documents from whichever retriever fired
joiner = DocumentJoiner()

# Prompt + Generator
prompt_template = """Answer the question using only the provided context.
If the context doesn't contain the answer, say you don't know.

Context:
{% for doc in documents %}
{{doc.content}}
{% endfor %}

Question: {{query}}
Answer:"""

prompt_builder = PromptBuilder(template=prompt_template)
generator = OpenAIGenerator(model="gpt-4o-mini")

# Assemble pipeline
pipe = Pipeline()
pipe.add_component("classifier", classifier)
pipe.add_component("router", router)
pipe.add_component("code_retriever", code_retriever)
pipe.add_component("docs_retriever", docs_retriever)
pipe.add_component("general_retriever", general_retriever)
pipe.add_component("joiner", joiner)
pipe.add_component("prompt", prompt_builder)
pipe.add_component("generator", generator)

# Connections
pipe.connect("classifier.label", "router.label")
pipe.connect("router.code_query", "code_retriever.query")
pipe.connect("router.docs_query", "docs_retriever.query")
pipe.connect("router.general_query", "general_retriever.query")

# Each retriever outputs documents → joiner
pipe.connect("code_retriever.documents", "joiner.documents")
pipe.connect("docs_retriever.documents", "joiner.documents")
pipe.connect("general_retriever.documents", "joiner.documents")

# Joiner → prompt (documents) + classifier (query for prompt variable)
pipe.connect("joiner.documents", "prompt.documents")
pipe.connect("classifier.label", "prompt.query")  # reuse label? no, need original query
# We'll pass query separately at runtime via pipeline.run()

There’s a small issue above: the prompt needs the original query, not the label. Fix by adding a QueryPassThrough component or just passing the query at runtime. Let’s do the clean thing — add a tiny component that forwards the query.

# query_passthrough.py
from haystack import component

@component
class QueryPassThrough:
    @component.output_types(query=str)
    def run(self, query: str):
        return {"query": query}

Update pipeline.py to include it:

# pipeline.py (updated)
from haystack import Pipeline
from haystack.components.routers import ConditionalRouter
from haystack.components.joiners import DocumentJoiner
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator
from classifier import QueryClassifier
from query_passthrough import QueryPassThrough
from stores import code_retriever, docs_retriever, general_retriever

classifier = QueryClassifier()

routes = [
    {"condition": "{{label == 'code'}}", "output": "code_query", "output_type": str},
    {"condition": "{{label == 'docs'}}", "output": "docs_query", "output_type": str},
    {"condition": "{{label == 'general'}}", "output": "general_query", "output_type": str},
]
router = ConditionalRouter(routes)

joiner = DocumentJoiner()

prompt_template = """Answer the question using only the provided context.
If the context doesn't contain the answer, say you don't know.

Context:
{% for doc in documents %}
{{doc.content}}
{% endfor %}

Question: {{query}}
Answer:"""

prompt_builder = PromptBuilder(template=prompt_template)
generator = OpenAIGenerator(model="gpt-4o-mini")
query_pass = QueryPassThrough()

pipe = Pipeline()
pipe.add_component("classifier", classifier)
pipe.add_component("router", router)
pipe.add_component("code_retriever", code_retriever)
pipe.add_component("docs_retriever", docs_retriever)
pipe.add_component("general_retriever", general_retriever)
pipe.add_component("joiner", joiner)
pipe.add_component("prompt", prompt_builder)
pipe.add_component("generator", generator)
pipe.add_component("query_pass", query_pass)

pipe.connect("classifier.label", "router.label")
pipe.connect("router.code_query", "code_retriever.query")
pipe.connect("router.docs_query", "docs_retriever.query")
pipe.connect("router.general_query", "general_retriever.query")

pipe.connect("code_retriever.documents", "joiner.documents")
pipe.connect("docs_retriever.documents", "joiner.documents")
pipe.connect("general_retriever.documents", "joiner.documents")

pipe.connect("joiner.documents", "prompt.documents")
pipe.connect("query_pass.query", "prompt.query")
pipe.connect("prompt.prompt", "generator.prompt")

Run end-to-end tests

# run_pipeline.py
from pipeline import pipe

questions = [
    "How do I implement LRU cache in Python?",
    "What is the DocumentStore protocol in Haystack?",
    "What is the capital of France?",
    "Explain async/await in Python",
    "How do I connect a custom retriever in Haystack?",
]

for q in questions:
    print(f"\n=== Question: {q} ===")
    result = pipe.run({
        "classifier": {"query": q},
        "query_pass": {"query": q},
    })
    answer = result["generator"]["replies"][0]
    print(f"Answer: {answer}")

Expected output (abbreviated for space):

=== Question: How do I implement LRU cache in Python? ===
Answer: Based on the context, here's an LRU cache implementation in Python:
class LRUCache:
    def __init__(self, capacity):
        self.capacity = capacity
        self.cache = {}
        self.order = []

=== Question: What is the DocumentStore protocol in Haystack? ===
Answer: The DocumentStore protocol defines write_documents, delete_documents, and filter_documents methods.

=== Question: What is the capital of France? ===
Answer: The capital of France is Paris. Population: ~2.1 million.

=== Question: Explain async/await in Python ===
Answer: Based on the context:
async def fetch_user(session, user_id):
    async with session.get(f'/users/{user_id}') as resp:
        return await resp.json()

=== Question: How do I connect a custom retriever in Haystack? ===
Answer: Retrievers implement the Retriever protocol with a run method accepting query and filters.

Each question hits the correct retriever. The code question retrieved cache.py, the docs question retrieved document_store.md, and the general question retrieved wiki_france.md.

Inspect routing decisions

Add logging to see which branch fired:

# debug_routing.py
from pipeline import pipe
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("haystack.components.routers.conditional_router")
logger.setLevel(logging.DEBUG)

result = pipe.run({
    "classifier": {"query": "How do I write a binary search?"},
    "query_pass": {"query": "How do I write a binary search?"},
})

You’ll see log lines like:

DEBUG - haystack.components.routers.conditional_router - Condition '{{label == 'code'}}' evaluated to True, routing to output 'code_query'

Production considerations

Classifier latency

The LLM classifier adds ~200-400ms per request. For high-throughput systems, consider:

Caching classifier results for repeated queries
Using a smaller model (gpt-4o-mini is already fast) or a fine-tuned BERT classifier
Running classification asynchronously and speculatively executing the most likely retriever

Embedding costs

Three separate stores means three embedding passes at index time. At query time, only one retriever runs, so query embedding cost stays constant. If you have millions of documents, use a vector database (Qdrant, Weaviate, Pinecone) per domain instead of in-memory stores.

Fallback behavior

The classifier defaults to general on unexpected labels. You might want a stricter fallback — e.g., run a hybrid search across all stores when confidence is low. Haystack’s ConditionalRouter supports an else route for this:

routes = [
    {"condition": "{{label == 'code'}}", "output": "code_query", "output_type": str},
    {"condition": "{{label == 'docs'}}", "output": "docs_query", "output_type": str},
    {"condition": "{{label == 'general'}}", "output": "general_query", "output_type": str},
    {"condition": "True", "output": "fallback_query", "output_type": str},  # catches anything else
]

Then connect fallback_query to a hybrid retriever or a multi-store search.

Observability

Log the classifier label, retriever used, and document scores for every request. This lets you measure routing accuracy and detect drift. A simple structlog line:

import structlog
log = structlog.get_logger()
log.info("rag_query_routed", query=query, label=label, retriever=retriever_name, doc_scores=[d.score for d in docs])

Using n4n.ai as the model gateway

If you’re routing across multiple model providers for the generator or classifier, n4n.ai gives you one OpenAI-compatible endpoint that addresses 240+ models with automatic fallback when a provider is rate-limited or degraded. You’d point OpenAIGenerator and OpenAIDocumentEmbedder at the n4n.ai base URL and pass routing directives in headers to control which upstream model handles each call.

Extending the pattern

This architecture scales to dozens of domains. Common extensions:

Domain	Retriever type	Example use case
Code	Embedding + BM25 hybrid	Internal libraries, Stack Overflow dumps
Docs	Embedding + section-aware chunking	API references, tutorials
Tickets	Embedding + metadata filters	Jira/Linear history, resolution patterns
Logs	Sparse/dense hybrid	Incident debugging, error signatures
General	Dense vector	Wikipedia, knowledge bases

Each domain gets its own chunking strategy, embedding model, and retrieval hyperparameters. The classifier prompt grows linearly with domains — still manageable at 20-30 categories with good few-shot examples.

Summary

You built a haystack rag pipeline query classification system that:

Classifies queries with a few-shot LLM prompt
Routes to domain-specific retrievers via ConditionalRouter
Joins results and generates grounded answers

The pattern is production-ready. Swap the in-memory stores for your vector databases, add caching and observability, and you have a routing layer that cuts latency and improves accuracy by keeping each retriever focused on what it knows best.