Chroma for Local RAG

This is the hands-on post: a complete, runnable local RAG pipeline rather than fragments. The cover sets expectations — we use the modern Chroma API and build the canonical arc you'll repeat for every project: create a persistent collection, chunk and add documents (Chroma embeds them), retrieve the top-k relevant chunks for a question, and stuff those chunks into an LLM prompt as grounding context.

We'll also show how to attach metadata for filtering and citations, and how to swap the default local embedding model for an OpenAI one. Following it in order shows how the pieces compose, so afterward you can confidently point the same skeleton at your own documents.

Step one creates a persistent collection. PersistentClient(path='./rag_db') writes everything to that directory, so your index survives restarts — the right choice for anything beyond a throwaway test. get_or_create_collection is used deliberately instead of create_collection: it returns the existing collection if it's already there, so re-running the script doesn't raise a 'collection exists' error.

The metadata={'hnsw:space': 'cosine'} sets the distance space at creation time, matching what most text embedding models expect. count() prints 0 on the first run, confirming an empty collection. This single call establishes both where data lives and how similarity is measured — two decisions worth making explicitly rather than relying on defaults.

Step two chunks the document and adds it. The chunk() helper splits text into overlapping windows: size controls how much text per chunk, and overlap carries a few characters across each boundary so a sentence split between chunks still appears intact in at least one of them. This is a simple character-based chunker; real systems often chunk on tokens or sentence boundaries, but the principle is identical.

The add() call passes documents, matching ids, and a metadata dict per chunk tagging its source. Because the collection has an embedding function (the default MiniLM here), Chroma embeds every chunk automatically — you never handle vectors directly. Adding all chunks in one call rather than looping is what makes the embedding step efficient.

Step three is retrieval, wrapped in a small function so the rest of the pipeline stays clean. retrieve() embeds the question with the collection's embedding function, searches the HNSW index, and returns the top-k document texts. Here k=4, a reasonable starting point — enough context to answer without drowning the prompt in noise.

The result structure is worth noting: query() returns parallel lists nested one level per query, so res['documents'][0] is the list of chunk texts for the first (and only) query. The loop prints the first 60 characters of each retrieved chunk, which is the single most valuable debugging habit in RAG: always look at what was actually retrieved before blaming the LLM for a bad answer.

Step four is where retrieval becomes RAG: the retrieved chunks are injected into an LLM prompt as context. answer() calls retrieve(), joins the chunks into a context block, and builds a prompt that instructs the model to use ONLY that context — a simple but important guardrail that reduces hallucination by discouraging the model from answering from its own parametric memory.

The prompt then asks the question and the model responds grounded in the supplied passages. This is the essence of RAG: the LLM provides language and reasoning, while Chroma provides the facts. The quality ceiling of the answer is set by retrieval — if the right chunk isn't in the context, no prompting trick recovers it, which is why the earlier tuning matters so much.

This trace shows the expected behavior end to end. The question goes in, the comment notes that four chunks were retrieved (all from the handbook source), and the output is an answer grounded in those chunks — 'Full-time staff receive 20 days...'. The final comment makes the key property explicit: the answer is grounded in the handbook text, not invented.

Reading a trace like this is how you verify a RAG pipeline is actually doing retrieval-augmented generation rather than just letting the LLM freewheel. If the retrieved chunks didn't contain the answer, you'd see the model either refuse or hallucinate — and that signal tells you to fix retrieval (chunking, k, embeddings) rather than the prompt.

Step five shows how to swap the embedding model, which is the main lever for improving retrieval quality. chromadb.utils.embedding_functions provides ready-made functions; OpenAIEmbeddingFunction wraps OpenAI's embedding API. You pass it to a new collection via embedding_function, and from then on every add() and query() on that collection embeds through OpenAI instead of the local MiniLM model.

The critical discipline is that the embedding function is fixed per collection. You cannot embed some documents with MiniLM and others with OpenAI in the same collection and expect meaningful results — their vector spaces are unrelated. That's why the example creates a separate collection (handbook_oai) rather than changing the function on the existing one. Choose the embedding model when you create the collection.

This slide explains the chunking decision the code embodies, because it's the highest-leverage tuning knob in RAG. Embeddings and LLMs work best on focused passages: a chunk should be small enough that its vector represents one coherent idea, so similarity search can pinpoint it. That's why you split long documents rather than embedding them whole.

Overlap exists to handle boundary effects. Without it, a sentence or fact that happens to fall across a chunk boundary gets split and may not retrieve well from either side. A small overlap — a few dozen tokens — ensures such content appears intact in at least one chunk. Size and overlap interact with your embedding model and content, so treat them as parameters to tune empirically, not constants to set once and forget.

These practical notes cover the things that trip up first-time builders. get_or_create_collection prevents the 'already exists' error that breaks scripts on the second run — a constant annoyance with create_collection. Pinning one embedding function per collection is restated because violating it is the most common silent failure. The OpenAI embedder and the LLM both need OPENAI_API_KEY in the environment, and forgetting it surfaces as an auth error at query or generate time, not at setup.

The final note — store source and page metadata — pays off twice: it enables where-filtering to scope retrieval, and it lets you cite where an answer came from. Adding metadata at add() time is cheap; retrofitting it means re-importing the whole corpus, so do it from the start.

This flow diagram summarizes the whole post as a four-step loop: chunk, add, query, prompt. It mirrors the code exactly and is worth memorizing because it's the skeleton of essentially every local RAG program you'll write. Swap the chunker, swap the embedding function, swap the LLM, and the shape stays identical.

Seeing it as a flow also reinforces the ordering dependencies: you can't add before chunking, you can't query before adding, and you can't prompt usefully before retrieving good chunks. The diagram is the mental checklist you run before debugging — if an answer is wrong, walk the flow backward to find which stage produced the bad input to the next.

The teaser points to Day 90's Common Mistakes post. With a working pipeline in hand, the natural next step is learning the failure modes — the in-memory client that loses your data, mismatched embedding functions, bad chunk sizes, missing metadata, and unevaluated retrieval — so you can run this same pipeline reliably instead of being surprised by silent degradation.