RAG Architecture: Components, Timing & Design Patterns › Building a Complete RAG Pipeline

from dataclasses import dataclass

import numpy as np
from rank_bm25 import BM25Okapi


@dataclass
class RetrievedDocument:
    """A document with its retrieval score."""

    content: str
    score: float
    doc_id: int


class BM25Retriever:
    """BM25-based sparse retriever."""

    def __init__(self, documents: list[str]):
        self.documents = documents
        self.tokenized_docs = [doc.lower().split() for doc in documents]
        self.bm25 = BM25Okapi(self.tokenized_docs)

    def retrieve(self, query: str, k: int = 3) -> list[RetrievedDocument]:
        """Retrieve top-k documents for a query."""
        tokenized_query = query.lower().split()
        scores = self.bm25.get_scores(tokenized_query)

        top_indices = np.argsort(scores)[::-1][:k]

        return [
            RetrievedDocument(
                content=self.documents[idx], score=scores[idx], doc_id=idx
            )
            for idx in top_indices
        ]

class SimpleGenerator:
    """Placeholder generator (in practice, use an LLM API or local model)."""

    def __init__(self):
        pass

    def generate(self, prompt: str, max_tokens: int = 256) -> str:
        """Generate a response given a prompt.

        This is a placeholder. In practice, you would:
        - Call an API (OpenAI, Anthropic, etc.)
        - Run a local model (llama.cpp, vLLM, etc.)
        """
        return "[Generated response based on retrieved context]"

class RAGPipeline:
    """Complete RAG pipeline combining retriever and generator."""

    def __init__(
        self,
        retriever: BM25Retriever,
        generator: SimpleGenerator,
        k: int = 3,
        include_scores: bool = False,
    ):
        self.retriever = retriever
        self.generator = generator
        self.k = k
        self.include_scores = include_scores

    def format_context(self, docs: list[RetrievedDocument]) -> str:
        """Format retrieved documents into context string."""
        parts = []
        for i, doc in enumerate(docs, 1):
            if self.include_scores:
                parts.append(f"[{i}] (score: {doc.score:.2f}) {doc.content}")
            else:
                parts.append(f"[{i}] {doc.content}")
        return "\n\n".join(parts)

    def build_prompt(self, query: str, context: str) -> str:
        """Construct the full prompt for the generator."""
        return f"""Answer the question based only on the following context. If the context doesn't contain the answer, say "I don't have enough information to answer this question."

Context:
{context}

Question: {query}

Answer:"""

    def answer(self, query: str) -> dict:
        """Run the full RAG pipeline."""
        retrieved_docs = self.retriever.retrieve(query, k=self.k)
        context = self.format_context(retrieved_docs)
        prompt = self.build_prompt(query, context)
        response = self.generator.generate(prompt)

        return {
            "query": query,
            "retrieved_documents": retrieved_docs,
            "prompt": prompt,
            "response": response,
        }

# Build a knowledge base about ML concepts
knowledge_base = [
    "Transformers use self-attention mechanisms that allow each token to attend to all other tokens in the sequence. This enables capturing long-range dependencies without the sequential processing limitations of RNNs.",
    "The attention mechanism computes compatibility scores between queries and keys, then uses these scores to create weighted sums of values. This can be expressed as Attention(Q,K,V) = softmax(QK^T / sqrt(d_k))V.",
    "BERT (Bidirectional Encoder Representations from Transformers) uses masked language modeling for pre-training, randomly masking 15% of input tokens and training the model to predict them.",
    "GPT models are autoregressive, generating one token at a time while attending only to previously generated tokens. This causal attention prevents information leakage from future tokens.",
    "Retrieval-augmented generation grounds language models in external knowledge by retrieving relevant documents before generation. This reduces hallucination and enables access to up-to-date information.",
    "The key innovation of transformers is replacing recurrence with self-attention. While RNNs process sequences step-by-step, transformers process all positions in parallel during training.",
    "Multi-head attention runs several attention operations in parallel, letting the model to capture different types of relationships at different positions.",
    "Position encodings are necessary in transformers because self-attention is permutation-invariant. Sinusoidal encodings and learned position embeddings are common approaches.",
]

retriever = BM25Retriever(knowledge_base)
generator = SimpleGenerator()
rag = RAGPipeline(retriever, generator, k=3, include_scores=True)

result = rag.answer("What is the key innovation of transformers?")

def check_retrieval_quality(
    docs: list[RetrievedDocument], score_threshold: float = 1.0
) -> dict:
    """Assess whether retrieved documents are likely relevant."""
    if not docs:
        return {"status": "no_documents", "action": "fallback_to_parametric"}

    max_score = max(doc.score for doc in docs)
    avg_score = sum(doc.score for doc in docs) / len(docs)

    if max_score < score_threshold:
        return {
            "status": "low_confidence",
            "max_score": max_score,
            "action": "warn_user_or_expand_search",
        }

    return {
        "status": "good",
        "max_score": max_score,
        "avg_score": avg_score,
        "action": "proceed_with_generation",
    }


# Test with a query that has good matches
good_docs = retriever.retrieve("How does attention work in transformers?", k=3)
quality_good = check_retrieval_quality(good_docs)

# Test with a query outside the knowledge base
poor_docs = retriever.retrieve("What is the capital of France?", k=3)
quality_poor = check_retrieval_quality(poor_docs)