ChromaDB persistence and idempotent ingest

An in-memory vector store is not a deploy. We want the index on disk so a pod restart does not wipe it. ChromaDB ships a PersistentClient that stores the collection in a directory you choose via CHROMA_PERSIST_DIR. In Kubernetes, that directory will be a mounted volume.

pipeline/index.py
python
import os
from dataclasses import dataclass
from typing import Dict, List, Optional
import chromadb


@dataclass
class RetrievedDoc:
    id: str
    text: str
    score: float
    metadata: Dict


class ChromaIndex:
    def __init__(
        self,
        collection_name: str = 'enterprise_rag',
        persist_directory: Optional[str] = None,
    ):
        self.collection_name = collection_name
        self.persist_directory = persist_directory or os.getenv(
            'CHROMA_PERSIST_DIR', './chroma_db'
        )
        self.client = chromadb.PersistentClient(path=self.persist_directory)
        self.collection = self.client.get_or_create_collection(name=collection_name)

    def add(
        self,
        ids: List[str],
        texts: List[str],
        embeddings: List[List[float]],
        metadatas: Optional[List[Dict]] = None,
    ) -> int:
        if not ids:
            return 0
        metas = metadatas or [{'id': i} for i in ids]
        self.collection.upsert(
            ids=ids,
            embeddings=embeddings,
            documents=texts,
            metadatas=metas,
        )
        return len(ids)

    def query(self, embedding: List[float], top_k: int = 4) -> List[RetrievedDoc]:
        res = self.collection.query(query_embeddings=[embedding], n_results=top_k)
        out: List[RetrievedDoc] = []
        if not res.get('ids') or not res['ids'][0]:
            return out
        ids = res['ids'][0]
        docs = res['documents'][0] if res.get('documents') else [''] * len(ids)
        metas = res['metadatas'][0] if res.get('metadatas') else [{}] * len(ids)
        dists = res['distances'][0] if res.get('distances') else [0.0] * len(ids)
        for i in range(len(ids)):
            # Chroma returns distance; invert to a similarity-style score
            score = 1.0 - float(dists[i]) if dists[i] is not None else 0.0
            out.append(RetrievedDoc(id=ids[i], text=docs[i] or '', score=score, metadata=metas[i] or {}))
        return out

    def count(self) -> int:
        return self.collection.count()

upsert is the key choice on the write side. If a document already exists with the same id, it is replaced, so re-running ingest on the same batch is safe. On the read side, query takes a pre-computed embedding and returns RetrievedDoc objects carrying id, text, and a similarity-style score (Chroma reports distance, so we invert it). The grounded answer path consumes exactly this shape.

Idempotent ingest flow

Why upsert is the right choice for a service that will be retried by Kubernetes and client code.

Quiz: Quiz

Loading practice…