ai_ml / ml system design / 02_design_a_rag_assistant.md

Design: a RAG assistant over company documents

6 interview angles 5 min read source

Design: a RAG assistant over company documents

The most likely AI system-design question for a backend engineer in 2026. Worked end to end.

Clarify first

  • Who and what? Employees asking questions of internal docs — policies, runbooks, product specs.
  • Corpus size? Say 100k documents, growing, with edits.
  • Scale? 5k employees, ~10 queries/day each = 50k/day, ~1 QPS average, maybe 10 peak. Small. Don’t over-engineer.
  • Latency? Conversational — first token under 2s.
  • Accuracy bar? Wrong answers about policy are costly. Citations are mandatory, and “I don’t know” must be an acceptable answer.
  • Access control? Yes — documents have permissions. This is the requirement that shapes the design most.

That last point is the one to surface early. Retrieval that ignores permissions is a data breach, not a relevance bug.

Architecture

                    ┌─── ingestion (async) ──────────────────┐
  source systems -> │ parse -> chunk -> contextualise -> embed│ -> vector + keyword index
  (SharePoint,      └────────────────────────────────────────┘
   Confluence, S3)

  user query -> rewrite -> hybrid retrieve -> filter by ACL -> rerank
             -> assemble context -> generate w/ citations -> validate -> stream

Ingestion

Parsing is where quality is won or lost. PDFs with tables and multi-column layouts are the hard case; a document-layout model beats naive text extraction substantially. Preserve structure — headings become metadata you can filter and cite on.

Chunking: structure-aware, splitting on headings and sections rather than fixed character counts, ~500-1000 tokens with overlap.

Contextual retrieval: prepend a generated description of where the chunk sits before embedding. A chunk reading “the limit is 30 days” is useless without “from the 2026 Expense Policy, section 4”. This is the single highest-return ingestion improvement. See ../09_rag_embeddings/08_hybrid_search_and_reranking.md.

Metadata per chunk: doc_id, source, section, updated_at, and critically acl_groups.

Incremental updates: content-hash each chunk, re-embed only what changed. Re-embedding 100k documents nightly is waste.

Storage

At 100k documents (~1M chunks), this is not big data. pgvector on Postgres is likely the right answer: one datastore, transactional metadata, hybrid search via tsvector alongside vector, and ACL filtering as an ordinary SQL predicate.

Reach for a dedicated vector database at tens of millions of chunks or when you need features Postgres lacks. Proposing Pinecone for 1M chunks when the team already runs Postgres is over-engineering, and saying so is a good signal.

Retrieval

async def retrieve(query: str, user: User) -> list[Chunk]:
    q = await rewrite_query(query, history)           # resolve "what about that one?"

    dense, sparse = await asyncio.gather(
        vector_search(q, user.acl_groups, top_k=50),
        bm25_search(q, user.acl_groups, top_k=50),
    )
    fused = rrf(dense, sparse)                         # rank fusion, no score calibration
    return (await rerank(q, fused))[:5]                # cross-encoder

Four deliberate choices:

  • Query rewriting — multi-turn queries are meaningless standalone.
  • Hybrid — BM25 catches error codes and product names that embeddings miss.
  • ACL as a filter inside the query, not applied afterwards. Post-filtering leaks the existence of documents and can return an empty page when the user had accessible results deeper in the ranking.
  • Rerank to 5 — more chunks dilute attention. See ../06_transformers_llm/08_context_windows.md.

Generation

Context order matters: system prompt and tools first (stable, cacheable), then retrieved chunks with the strongest at the extremes, then history, then the question last.

[system + citation instructions]   <- cached across all requests
[retrieved chunks, numbered]
[conversation history]
[user question]                    <- last, adjacent to generation

Require citations by chunk number and verify them mechanically — a citation referencing a chunk that wasn’t retrieved is a fabrication and trivially detectable.

Instruct explicitly: answer only from the provided context, and say so when it’s insufficient.

Validation before display

  1. Schema/format check — free.
  2. Citation check — every cited ID was actually retrieved.
  3. Groundedness on a sample or for high-stakes topics — claim decomposition plus entailment.

Stream to the user while validating in parallel; block only actions, not text. See ../14_guardrails_safety/02_guardrails_and_output_validation.md.

Evaluation

  • Golden set of ~150 questions with known-correct chunks, including unanswerable ones.
  • Retrieval: recall@50 and MRR, measured separately from answer quality. If recall@50 is poor, no prompt work will fix the answers.
  • Generation: faithfulness, citation validity, answer relevance.
  • CI gate on prompt, model, chunking or index changes.

Cost and scale

At 1 QPS this is cheap. The levers if it grows:

  • Prefix caching — the system prompt is identical across every request, so structure the prompt to exploit it.
  • Cache identical queries — internal corpora produce heavy repetition (“what’s the holiday policy”).
  • Route by difficulty — a small model handles lookup-style questions.
  • Cap max_tokens.

Failure modes to raise

  • Stale index after a document edit — track updated_at, monitor index lag.
  • ACL drift — permissions change in the source system; re-sync, don’t trust cached ACLs.
  • Confabulation on unanswerable questions — the trust-destroying failure. Test for it explicitly.
  • Prompt injection via document content — a document containing instructions. The assistant is read-only, which bounds the damage; that’s a deliberate design choice worth stating.
  • Lost in the middle — mitigated by retrieving few chunks and ordering by relevance.

Interview angle

  • “Walk me through a RAG system for internal docs.” — ingestion (parse, structure-aware chunk, contextualise, embed, index) and query time (rewrite, hybrid retrieve with ACL filtering, rerank, assemble, generate with citations, validate). Surface access control early; it shapes retrieval more than anything else.
  • “Which vector database?” — at ~1M chunks, pgvector on existing Postgres: one datastore, transactional metadata, hybrid search and ACL filtering as a SQL predicate. Dedicated vector DBs earn their place at tens of millions of chunks.
  • “How do you enforce permissions?” — as a filter inside the retrieval query, never a post-filter. Post-filtering leaks document existence and can return an empty result when accessible matches existed further down.
  • “How do you stop it making things up?” — mandatory citations verified mechanically against retrieved chunk IDs, explicit instruction to answer only from context, groundedness checking by claim entailment, and unanswerable questions in the eval set so you actually measure refusal.
  • “Answers are wrong — where do you look first?” — recall@50. If the right chunk isn’t retrieved, the generation stage is irrelevant and you’ve been tuning the wrong thing.
  • “How would you cut the cost at 10x scale?” — prefix caching with stable content first, caching repeated queries (internal corpora repeat heavily), routing simple lookups to a small model, and capping output tokens.
Previous
ML system design framework