Design a RAG System

Problem Statement & Requirements

Think of RAG like an open-book exam. Instead of relying solely on memorized knowledge (the LLM's training data), the system first looks up relevant information from a knowledge base, then formulates an answer using both its general knowledge and the retrieved context.

Functional Requirements

• Document ingestion — Ingest PDFs, HTML, Markdown, databases, and APIs
• Chunking & embedding — Split documents into retrievable chunks with vector representations
• Semantic search — Find relevant chunks for a user query
• LLM generation — Generate answers grounded in retrieved context
• Citation & provenance — Link every claim to its source document
• Feedback loop — Users can rate answers; improve retrieval over time

Non-Functional Requirements

• End-to-end latency — <3 seconds (retrieval + generation)
• Relevance — Top-5 retrieved chunks contain the answer 90%+ of the time
• Freshness — New documents searchable within 15 minutes of ingestion
• Scale — 10M+ document chunks, 1000+ concurrent queries

Back-of-Envelope Estimation

System API Design

# Query the RAG system
POST /api/v1/query
{
  "question": "What is our refund policy for enterprise customers?",
  "conversation_id": "conv_123",
  "filters": { "source": "policy_docs" },
  "top_k": 5,
  "stream": true
}
# Response
{
  "answer": "Enterprise customers can request a full refund...",
  "citations": [
    { "chunk_id": "c_456", "doc": "refund-policy-v3.pdf", "page": 4, "score": 0.92 }
  ],
  "confidence": 0.88
}

# Ingest documents
POST /api/v1/documents
{
  "source": "s3://docs/policy/",
  "collection": "policy_docs",
  "chunking": { "strategy": "semantic", "max_tokens": 512 }
}

# Submit feedback
POST /api/v1/feedback
{ "query_id": "q_789", "rating": "helpful", "comment": "Accurate and well-cited" }

Data Model

CREATE TABLE documents (
    doc_id        VARCHAR PRIMARY KEY,
    source_url    TEXT,
    title         TEXT,
    collection    VARCHAR,
    ingested_at   TIMESTAMP,
    metadata      JSONB
);
CREATE TABLE chunks (
    chunk_id      VARCHAR PRIMARY KEY,
    doc_id        VARCHAR REFERENCES documents,
    content       TEXT,
    embedding     vector(1536),  -- pgvector type
    chunk_index   INT,
    token_count   INT,
    metadata      JSONB
);
CREATE INDEX ON chunks
    USING ivfflat (embedding vector_cosine_ops)
    WITH (lists = 1000);
CREATE TABLE conversations (
    conversation_id VARCHAR PRIMARY KEY,
    user_id         VARCHAR,
    messages        JSONB[],
    created_at      TIMESTAMP
);

High-Level Architecture

RAG has two pipelines: ingestion (offline) and query (online).

Deep Dive: Core Components

Chunking Strategies

from sentence_transformers import SentenceTransformer
import numpy as np

def semantic_chunk(text, model, threshold=0.75):
    """Split text where semantic similarity drops."""
    sentences = text.split(". ")
    embeddings = model.encode(sentences)

    chunks, current = [], [sentences[0]]
    for i in range(1, len(sentences)):
        sim = np.dot(embeddings[i], embeddings[i-1])
        if sim < threshold:
            chunks.append(". ".join(current))
            current = [sentences[i]]
        else:
            current.append(sentences[i])
    chunks.append(". ".join(current))
    return chunks

Vector Databases

Retrieval Strategies

• Dense retrieval: Embed query, find nearest vectors (semantic match)
• Sparse retrieval (BM25): Keyword-based (exact term match)
• Hybrid: Combine dense + sparse with reciprocal rank fusion (RRF). Best of both worlds — catches semantic similarity AND exact keywords

Guardrails & Hallucination Mitigation

Scaling & Optimization

Embedding Caching

Cache query embeddings (many queries are repeated or similar). Use a two-level cache: L1 in-memory (last 10K queries) and L2 in Redis (last 1M queries). Cache hit rates of 30-50% are common.

Hierarchical Retrieval

For large corpora, use two-stage retrieval: (1) Retrieve relevant documents first using document-level summaries, then (2) Retrieve specific chunks within those documents. Reduces search space 10-100x.

Evaluation Framework

Practice Problems

Quick Reference