ChromaDB indexing and retrieval

Chunks alone are not searchable. We need vector embeddings so semantic queries can find the right chunk even when the user phrases a question differently from the source document. ChromaDB persists the embeddings on disk and runs a local similarity search. No external vector service, no credentials, perfect for a workshop.

rag_utils.py
python
import os
import uuid
import chromadb
from chromadb.config import Settings as ChromaSettings

CHROMA_PATH = os.getenv('CHROMA_PATH', 'chroma_db')
COLLECTION_NAME = os.getenv('CHROMA_COLLECTION', 'insurance_policies')
EMBED_MODEL_NAME = os.getenv('EMBEDDING_MODEL', 'all-MiniLM-L6-v2')

_client = None
_embed_model = None


def _get_collection():
    global _client
    if _client is None:
        _client = chromadb.PersistentClient(
            path=CHROMA_PATH,
            settings=ChromaSettings(anonymized_telemetry=False),
        )
    return _client.get_or_create_collection(name=COLLECTION_NAME)


def _embed(texts: list[str]) -> list[list[float]]:
    global _embed_model
    if _embed_model is None:
        from sentence_transformers import SentenceTransformer
        _embed_model = SentenceTransformer(EMBED_MODEL_NAME)
    return _embed_model.encode(texts, convert_to_tensor=False).tolist()


def index_policy_pdf(path: str, source_name: str | None = None) -> int:
    text = _load_pdf_markdown(path)
    chunks = _split_text(text)
    if not chunks:
        return 0
    source = source_name or os.path.basename(path)
    ids = [f'{source}-{uuid.uuid4().hex[:8]}-{i}' for i in range(len(chunks))]
    metas = [{'source': source, 'chunk_index': i} for i in range(len(chunks))]
    collection = _get_collection()
    collection.add(ids=ids, documents=chunks, metadatas=metas, embeddings=_embed(chunks))
    return len(chunks)

Lazy singletons for the Chroma client and embedding model mean importing this file does not spin up a model or open the DB. You pay the cost on first query, not on import. The source name in metadata is what lets the subagent cite the original PDF filename in its answer.

rag_utils.py
python
def search_policies(query: str, k: int = 4) -> list[dict]:
    """Semantic search over indexed policy chunks."""
    collection = _get_collection()
    if collection.count() == 0:
        return []
    embedded = _embed([query])
    results = collection.query(query_embeddings=embedded, n_results=k)
    docs = results.get('documents', [[]])[0]
    metas = results.get('metadatas', [[]])[0]
    distances = results.get('distances', [[]])[0] if results.get('distances') else [None] * len(docs)
    out = []
    for doc, meta, dist in zip(docs, metas, distances):
        out.append({
            'content': doc,
            'source': (meta or {}).get('source', 'policy'),
            'distance': dist,
        })
    return out

The return shape is deliberate: content, source, and distance. Every subagent downstream depends on this exact structure. If you return raw Chroma results, you leak the vector store into your agent code and cannot swap Chroma for Qdrant or Pinecone later.

Flashcards: Flashcards

Loading practice…

Quiz: Quiz

Loading practice…