chroma.py

"""Singleton ChromaDB client and collection.

All modules that need the vector store import get_collection() from here.
The PersistentClient and collection are created once on first call and
reused for the lifetime of the process.

COST: ZERO external API tokens.
Embeddings use sentence-transformers 'all-MiniLM-L6-v2' running locally —
the same model used by the semantic query cache.
"""
from pathlib import Path
from typing import Optional

import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction

VECTOR_STORE_DIR = Path(__file__).parent / "vector_store"
COLLECTION_NAME = "rag_documents"

_client: Optional[chromadb.Client] = None
_collection: Optional[chromadb.Collection] = None


def get_collection() -> chromadb.Collection:
    """Return the shared ChromaDB collection, creating it on first call."""
    global _client, _collection
    if _collection is None:
        _client = chromadb.PersistentClient(path=str(VECTOR_STORE_DIR))
        ef = SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
        _collection = _client.get_or_create_collection(
            name=COLLECTION_NAME,
            embedding_function=ef,
        )
    return _collection