RAG in Plain Python: Chunking, TF-IDF, Cosine Similarity and Prompting

A typical RAG tutorial has you install five libraries and call .from_documents(). Ten minutes later a demo works and nothing about it is understood. When retrieval returns garbage or the model hallucinates, you're stuck, because the library did the thinking for you.

Here the pipeline is built the opposite way, in plain Python, using only the math module and one plain-text LLM call at the very end. No embedding model, no FAISS, no Pinecone. Knowing the mechanics means that when you later reach for LangChain or a real vector database, you'll understand what happens underneath. By the end there's a working pipeline and, more usefully, a feel for every number it produces.

The problem RAG solves

A language model's knowledge is frozen when training ends and lives entirely in its weights. To have it answer questions about internal company documents, a codebase or this morning's news, the two obvious options both disappoint:

1. Retrain the model. That's slow and costly, and tomorrow's changes need another round.
2. Paste everything into the prompt. This doesn't scale. Even a million-token window can't hold 50,000 documents, and research on long-context models repeatedly reports a "lost in the middle" effect, where information buried deep in a huge prompt is used worse even though it fits. Extra tokens also mean more latency and cost per call.

RAG's real trick is narrower than it sounds. Documents stay in external storage, and at question time only the few chunks relevant to that question are retrieved and handed to the model. The LLM never contains your data. It receives a fresh, relevant slice each time.

So any RAG system is two systems joined together: retrieval, which finds the right text in a far larger collection for a given query, and generation, which passes that text plus the query to an LLM and gets an answer. The rest of this piece builds both halves from basic parts.

Step 1: chunking

Turning a whole 50-page PDF into one vector wipes out nearly all useful signal, because the detail that would answer a question gets averaged into nothing. The first job is cutting documents into small, focused chunks.

The simplest method is a fixed-size sliding window counted in words, with overlap between neighboring chunks:

def chunk_text(text: str, chunk_size: int = 200, overlap: int = 50) -> list[str]:
    if overlap >= chunk_size:
        raise ValueError("overlap must be smaller than chunk_size")

    words = text.split()
    chunks = []
    start = 0
    while start < len(words):
        end = start + chunk_size
        chunks.append(" ".join(words[start:end]))
        if end >= len(words):
            break
        start = end - overlap  # step back so chunks share context
    return chunks

That guard clause isn't decoration. If overlap reaches or exceeds chunk_size, the start position stops advancing and the loop never ends. Tuning overlap upward "for more context" makes this an easy slip, so the function refuses to run instead of hanging quietly.

Overlap matters more than it seems. Without it, a sentence that answers a question may straddle a chunk boundary and get split, so neither chunk holds a complete idea that can be retrieved. With overlap, the boundary region survives whole in at least one chunk.

One edge case to know: fixed word windows ignore sentence and paragraph structure, so they can cut a sentence or a table row in half. Sentence-aware or paragraph-aware splitting avoids that but creates uneven chunk sizes, since a 10-word sentence and a 150-word sentence each become one chunk. That distorts the token budget and how strongly any single chunk dominates the later retrieval math. No chunking approach is free of trade-offs, and the choice is really about which failure you'd rather debug.

Step 2: turning text into numbers with TF-IDF

A computer can't compare the similarity of two strings directly. Every later step needs vectors, lists of numbers on which arithmetic works. The classic dependency-free way to build them is TF-IDF, term frequency multiplied by inverse document frequency.

• TF counts how often a word appears in this one chunk. A high count for "kidney" hints that the chunk is about kidneys.
• IDF measures how rare the word is across the whole collection. Words like "the" and "is" turn up everywhere and say nothing about which chunk you're reading, so they score near zero. A rare word in a chunk is a much stronger signal.

The code builds a vocabulary from all chunks, computes IDF for each term, and then converts every chunk into a TF-IDF vector:

import math
from collections import Counter

def build_vocabulary(chunks: list[str]) -> list[str]:
    vocab = set()
    for chunk in chunks:
        vocab.update(chunk.lower().split())
    return sorted(vocab)

def compute_idf(chunks: list[str], vocab: list[str]) -> dict[str, float]:
    num_chunks = len(chunks)
    idf = {}
    for word in vocab:
        containing = sum(1 for c in chunks if word in c.lower().split())
        idf[word] = math.log(num_chunks / (1 + containing))
    return idf

def compute_tf(chunk: str, vocab: list[str]) -> dict[str, float]:
    words = chunk.lower().split()
    counts = Counter(words)
    total = len(words)
    return {word: counts.get(word, 0) / total for word in vocab}

def vectorize(chunk: str, vocab: list[str], idf: dict[str, float]) -> list[float]:
    tf = compute_tf(chunk, vocab)
    return [tf[word] * idf[word] for word in vocab]

Try it on three chunks and something becomes clear. A word present in nearly every chunk, such as "the", gets an IDF near zero, so repeating it adds nothing to the vector. A rare, topic-specific word receives a clearly positive IDF. TF-IDF is literally a formula for asking whether a word tells you anything distinctive about a chunk.

The downside is scale. With 100,000 documents, the vocabulary may hold more than 500,000 unique words, so every chunk vector, even for a three-word chunk, is 500,000 numbers long and mostly zeros. These sparse vectors are expensive to store and compare. That is the main reason production systems switch to dense neural embeddings of 384 to 1536 dimensions with hardly any zeros.

Step 3: measuring relevance with cosine similarity

With chunks as vectors, the next question is how relevant a chunk is to a query. Raw distance isn't the right tool, because for text the angle between vectors matters more than their length. A chunk that mentions "cat" once and one that repeats it twenty times are still about the same topic, and cosine similarity treats them alike without caring which is louder.

The implementation needs a dot product, a magnitude and the ratio between them:

def dot_product(a, b):
    return sum(x * y for x, y in zip(a, b))

def magnitude(v):
    return math.sqrt(sum(x * x for x in v))

def cosine_similarity(a, b):
    denom = magnitude(a) * magnitude(b)
    return 0.0 if denom == 0 else dot_product(a, b) / denom

Cosine similarity in principle ranges from -1 to 1, and with this TF-IDF setup it stays near the 0 to 1 band nearly always:

• 1.0 means the vectors point in exactly the same direction, maximum similarity.
• 0.0 means they're orthogonal and share no meaningful dimension, which in plain terms says the query and chunk have no non-trivial words in common.

A score of 0 isn't a quirk. It's the system correctly reporting that the chunk has nothing to do with the query.

One precision point: the smoothed IDF formula can give a slightly negative value to a word that appears in every chunk, so cosine similarity can occasionally dip a little below zero. For real text this barely matters, but claiming it's never negative would overstate things.

The blind spot of TF-IDF with cosine similarity is synonyms. A query that says "automobile" against a document that says "car" scores near zero, because only literal vocabulary overlap counts and there's no notion of meaning. Dense embedding models, which place synonyms close together, exist to close exactly that gap.

Step 4: the retriever

With vectorization and similarity ready, retrieval means vectorizing each chunk once, vectorizing the incoming query, scoring every chunk against it, sorting and keeping the top k:

class VectorStore:
    def __init__(self):
        self.chunks, self.vocab, self.idf, self.vectors = [], [], {}, []

    def index(self, documents, chunk_size=50, overlap=10):
        all_chunks = []
        for doc in documents:
            all_chunks.extend(chunk_text(doc, chunk_size, overlap))
        self.chunks = all_chunks
        self.vocab = build_vocabulary(all_chunks)
        self.idf = compute_idf(all_chunks, self.vocab)
        self.vectors = [vectorize(c, self.vocab, self.idf) for c in all_chunks]

    def search(self, query, top_k=3):
        query_vector = vectorize(query, self.vocab, self.idf)
        scored = [(c, cosine_similarity(query_vector, v))
                  for c, v in zip(self.chunks, self.vectors)]
        scored.sort(key=lambda pair: pair[1], reverse=True)
        return scored[:top_k]

One detail can silently ruin everything. The query has to be converted using the vocabulary and IDF values computed from the documents, and a vocabulary derived from the query by itself is a mistake. Cosine similarity presumes equal-length vectors where position i stands for the same word in each. Give the query its own vocabulary and position 5 might be "heart" in one vector and "and" in the other, which leaves the similarity score structurally meaningless, not merely slightly off.

Step 5: augmenting the prompt

This step involves no math, only discipline. You combine the retrieved chunks with the user's question into a prompt that tells the model exactly how to behave:

def build_augmented_prompt(query: str, retrieved_chunks: list[str]) -> str:
    context = "\n\n".join(
        f"[Source {i+1}]: {chunk}" for i, chunk in enumerate(retrieved_chunks)
    )
    return f"""Answer the question using ONLY the context below.
If the answer is not contained in the context, say "I don't know based on the provided context."
Cite which source number(s) you used.

Context:
{context}

Question: {query}

Answer:"""

Two details here aren't cosmetic.

"Using ONLY the context below." Without that explicit rule, the model blends the retrieved chunks with whatever it remembers from training, and you can no longer tell which part of an answer came from trusted documents and which is guessing. Delivered confidently, such a blend does more damage than skipping RAG entirely, since it appears to have a source when it doesn't.

Numbered source labels and a citation instruction. These make answers traceable. When an answer is wrong, which will happen, labeled sources show which chunk misled the model, whether the chunk itself was stale or incorrect, or whether the model misread a correct one. With no labels, tracking down the cause of a bad answer is largely guesswork.

Step 6: wiring it together

The complete flow runs retrieval, then augmentation, then one LLM call. The key design decision is to keep indexing and querying as two separate lifecycles instead of one function doing both:

class RagPipeline:
    def __init__(self, client, model="openai/gpt-oss-20b"):
        self.store = VectorStore()
        self.client = client
        self.model = model

    def build_index(self, documents: list[str]):
        """Run once, whenever documents are loaded or updated."""
        self.store.index(documents, chunk_size=50, overlap=10)

    def answer(self, query: str, top_k: int = 3) -> str:
        """Run on every user query — never re-indexes."""
        results = self.store.search(query, top_k=top_k)
        chunks = [c for c, _ in results]
        prompt = build_augmented_prompt(query, chunks)

        response = self.client.chat.completions.create(
            model=self.model,
            max_tokens=500,
            messages=[{"role": "user", "content": prompt}],
        )
        return response.choices[0].message.content

Were indexing placed inside the per-query function, every query would split the text again, rebuild the vocabulary and recompute all the TF-IDF vectors for the whole collection, although none of it changes between queries. Across 10,000 documents and 1,000 queries that wastes 999 identical rebuilds. Building the index once and then querying many times is the layout FAISS, Chroma and Pinecone also nudge you toward. It's the correct architecture from the start, not a premature optimization.

What's wrong with this system

The limitations deserve candor, because they aren't bugs. They follow directly from the design choices above.

• No semantic understanding. "Car" and "automobile" never match, since TF-IDF sees only literal vocabulary overlap.
• Sparse, very wide vectors. Vocabulary size tracks the corpus, so each vector widens along with it. Beyond a few thousand documents this gets unwieldy.
• Nothing in the code verifies that retrieval found something useful. Even when every top-k chunk scores close to zero, the pipeline still passes them to the model. The prompt's request to admit ignorance relies wholly on the LLM cooperating, so a weak retrieval can still flow into generation unchecked.