Building a citation-aware RAG pipeline in Haystack

Most RAG demos return text and hope the user trusts it. A citation-aware rag pipeline haystack implementation forces the model to ground every claim in retrieved documents and emit traceable references, which is what production systems actually require.

Step 1: Install dependencies and import components

Haystack 2.x provides composable components for embedding, retrieval, and generation. Install the core package plus the OpenAI generator extra:

pip install haystack-ai openai

Import the pieces you need for an in-memory store and a standard embedding retriever:

from haystack import Document, Pipeline
from haystack.document_stores import InMemoryDocumentStore
from haystack.components.embedders import (
    SentenceTransformersDocumentEmbedder,
    SentenceTransformersTextEmbedder,
)
from haystack.components.retrievers import InMemoryEmbeddingRetriever
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator

This baseline avoids external services until the generation step, keeping the citation logic fully testable offline.

Step 2: Load and chunk source documents

Citations are only useful if the retrieved unit maps cleanly to a source span. Split long texts into coherent chunks of 200–500 words before indexing.

raw_texts = [
    "Haystack is a framework for building NLP pipelines. It supports retrieval-augmented generation.",
    "Citations in RAG improve traceability. Models should reference document IDs when answering.",
    "OpenAI-compatible endpoints allow swapping LLM providers without changing client code.",
]

docs = [Document(content=t, meta={"src": f"doc_{i}"}) for i, t in enumerate(raw_texts)]

Attach metadata (src) now; you will surface it in the citation output later.

Step 3: Embed and index into the document store

Use a lightweight sentence transformer for embeddings. Warm up the embedder, transform the documents, and write them to the store.

store = InMemoryDocumentStore(embedding_similarity_function="cosine")
doc_embedder = SentenceTransformersDocumentEmbedder(
    model="sentence-transformers/all-MiniLM-L6-v2"
)
doc_embedder.warm_up()
embedded_docs = doc_embedder.run(documents=docs)["documents"]
store.write_documents(embedded_docs)

The InMemoryEmbeddingRetriever will query against these vectors. Set top_k=3 to keep citation lists short and verifiable.

retriever = InMemoryEmbeddingRetriever(store, top_k=3)
text_embedder = SentenceTransformersTextEmbedder(
    model="sentence-transformers/all-MiniLM-L6-v2"
)
text_embedder.warm_up()

Step 4: Build the citation prompt template

The prompt is where the citation contract is enforced. Number the retrieved documents and instruct the model to use bracketed indices.

template = """
You answer questions using ONLY the provided documents.
Cite every factual claim with the document number in square brackets, e.g. [1].
If the answer is not in the documents, say "No citation found".

Documents:
{% for doc in documents %}
[{{ loop.index }}] {{ doc.content }} (source: {{ doc.meta.src }})
{% endfor %}

Question: {{ question }}
Answer:
"""

prompt_builder = PromptBuilder(template=template)

This citation-aware rag pipeline haystack design makes the retriever output directly visible to the model with stable indices.

Step 5: Configure the generator with a fallback-friendly endpoint

Use OpenAIGenerator and point it at any OpenAI-compatible API. If you point it at an OpenAI-compatible endpoint such as n4n.ai, you keep a single code path while gaining automatic fallback when a provider is rate-limited or degraded.

generator = OpenAIGenerator(
    api_key="your-key",
    api_base="https://api.n4n.ai/v1",  # optional: swap for OpenAI base if preferred
    model="gpt-4o-mini",
    generation_kwargs={"temperature": 0.0},
)

Temperature zero reduces hallucinated citations. The gateway forwards provider cache-control hints, so repeated document prefixes cost less.

Step 6: Assemble the pipeline

Wire the text embedder to the retriever, then the retriever output into the prompt builder, and finally the prompt into the generator.

pipe = Pipeline()
pipe.add_component("text_embedder", text_embedder)
pipe.add_component("retriever", retriever)
pipe.add_component("prompt_builder", prompt_builder)
pipe.add_component("generator", generator)

pipe.connect("text_embedder.embedding", "retriever.query_embedding")
pipe.connect("retriever.documents", "prompt_builder.documents")
pipe.connect("prompt_builder.prompt", "generator.prompt")

The citation-aware rag pipeline haystack graph is now a single callable. No custom nodes required.

Step 7: Run a query and parse citations

Execute the pipeline and extract the reply plus the retrieved documents for validation.

question = "What does Haystack support?"
result = pipe.run(
    {
        "text_embedder": {"text": question},
        "prompt_builder": {"question": question},
    }
)

answer = result["generator"]["replies"][0]
retrieved = result["retriever"]["documents"]
print(answer)

A correct run prints something like:

Haystack supports retrieval-augmented generation [1]. It is a framework for building NLP pipelines [1].

Step 8: Verify the citation contract

Write a small assertion that every [n] in the answer maps to a retrieved document index. This is how you confirm the pipeline works end to end.

import re

def verify_citations(answer, docs):
    cited = {int(i) for i in re.findall(r"\[(\d+)\]", answer)}
    assert cited.issubset(set(range(1, len(docs) + 1))), "Citation out of range"
    return True

assert verify_citations(answer, retrieved)

If the assertion passes, the citation-aware rag pipeline haystack setup is enforcing grounding. Add this check to your integration tests so regressions in prompt formatting break the build.

Handling multi-document overlaps

When chunks from the same source overlap, duplicate [n] indices confuse users. Deduplicate by meta.src before prompt assembly:

seen = set()
unique_docs = []
for d in retrieved:
    if d.meta["src"] not in seen:
        seen.add(d.meta["src"])
        unique_docs.append(d)

Pass unique_docs to the prompt builder via a custom function component if needed. The core principle stays: one index, one verifiable span.

Why this pattern holds up in production

A citation-aware rag pipeline haystack approach shifts trust from the model to the data layer. The retriever is deterministic; the generator is constrained by the template. When you meter usage per token at the gateway, you can attribute cost to specific document sets and tune top_k accordingly.

Keep the prompt strict, keep chunk sizes sane, and treat the citation parser as a unit-tested boundary. That is the difference between a demo and a system you can ship.