RAG Pipeline

Implement a retrieval-augmented generation pipeline: ingest documents, chunk them with overlap, embed each chunk, then retrieve the top-K most similar chunks for a query using dot-product similarity. The key design decision is separating ingest() from update() - making intentional overwrite explicit rather than silent.

“ingest() is create - fail loudly on duplicate. update() is an intentional overwrite. Never silently overwrite on ingest: data corruption is worse than a raised exception. retrieve() returns empty on an empty index - never raise, never hallucinate.”

from dataclasses import dataclass, field
from typing import List

@dataclass
class Document:
    doc_id: str
    content: str
    chunks: List[str] = field(default_factory=list)

class RAGPipeline:
    def __init__(self, chunk_size: int = 512, overlap: int = 50):
        self.chunk_size = chunk_size
        self.overlap = overlap
        self.index = {}  # doc_id -> list of (chunk, embedding)

    def ingest(self, doc: Document):
        if doc.doc_id in self.index:
            raise ValueError(f"Use update() to replace existing document '{doc.doc_id}'")
        self._index_document(doc)

    def update(self, doc: Document):
        self._index_document(doc)

    def _index_document(self, doc: Document):
        if not doc.content or not doc.doc_id:
            raise ValueError("Document must have non-empty id and content")
        chunks = self._chunk(doc.content)
        doc.chunks = chunks
        self.index[doc.doc_id] = [(chunk, self._embed(chunk)) for chunk in chunks]

    def _chunk(self, text: str) -> List[str]:
        chunks, start = [], 0
        while start < len(text):
            end = min(start + self.chunk_size, len(text))
            chunks.append(text[start:end])
            start += self.chunk_size - self.overlap
        return chunks

    def _embed(self, text: str) -> List[float]:
        # Stub: replace with real embedding API call
        return [float(ord(c)) for c in text[:10]]

    def retrieve(self, query: str, top_k: int = 3) -> List[str]:
        if not self.index:
            return []
        query_emb = self._embed(query)
        scored = []
        for chunks_embs in self.index.values():
            for chunk, emb in chunks_embs:
                score = sum(a * b for a, b in zip(query_emb, emb))
                scored.append((score, chunk))
        scored.sort(reverse=True)
        return [chunk for _, chunk in scored[:top_k]]


# Usage
pipeline = RAGPipeline(chunk_size=512, overlap=50)
pipeline.ingest(Document(doc_id="doc_1", content="RAG stands for Retrieval Augmented Generation..."))
pipeline.ingest(Document(doc_id="doc_2", content="Vector databases store embeddings efficiently..."))
results = pipeline.retrieve("What is RAG?", top_k=2)
print(results)

• 1ingest() vs update() - Explicit Intent - ingest() raises ValueError on duplicate doc_id. update() overwrites unconditionally. This makes the caller declare intent - accidental overwrites on ingest are a data correctness bug, so fail loudly. Common pattern in database write paths.
• 2Chunking with Overlap - _chunk() slides a window of chunk_size characters, advancing by (chunk_size - overlap) each step. Overlap ensures a concept split across a boundary appears fully in at least one chunk. Without overlap, relevant context at chunk edges is silently lost.
• 3Dot-Product Similarity (Stub) - retrieve() computes sum(a*b for a,b in zip(query_emb, emb)). For real unit-normalized embeddings this equals cosine similarity. The stub embedding (first 10 ASCII values) is declared upfront - interviewers appreciate knowing exactly where the real API call goes.
• 4_index_document() - Shared Private Method - Both ingest() and update() delegate to _index_document() after their respective pre-conditions. Shared logic lives in one place - if chunking or embedding changes, you change it once. Classic DRY applied correctly.