ChromaVectorDatabase

On this page

Table of contents

Overview

A wrapper around ChromaDB supporting in-memory or persistent storage, automatic batching for large upserts, optional multi-tenant isolation via a tenant_id field injected into every document's metadata, and both a modern (DocumentChunk-list) and legacy (explicit ids/documents/metadatas) calling convention for add().

normalize_metadata (module-level helper)

Python
def normalize_metadata(metadata: Any) -> Dict

Converts None → {}, passes a dict through unchanged, converts a dataclass instance via asdict(), and raises TypeError for any other type. Not listed in __all__ / vector_database/__init__.py, so it's an internal helper of this module rather than a package-level public export — documented here because add() uses it directly on chunk metadata before merging in doc_id.

Class attributes

Python
SUPPORTED_METRICS = {"cosine": "cosine", "l2": "l2", "ip": "ip"}

(A dict rather than a set — keys are the metric names this wrapper accepts; values are what's passed through to ChromaDB's hnsw:space collection setting.)

Constructor

Python · 9 lines
ChromaVectorDatabase(
    embedder: Optional[Any] = None,
    collection_name: str = "default_collection",
    persist_directory: Optional[str] = None,
    distance_metric: str = "cosine",
    tenant_id: Optional[str] = None,
    batch_size: int = 500,
    strict_mode: bool = True,
)
ParameterTypeDefaultNotes
embedderOptional[Any]None.encode(texts) -> np.ndarray interface.
collection_namestr"default_collection"ChromaDB collection to create or reuse.
persist_directoryOptional[str]NoneIf given, uses chromadb.PersistentClient(path=...); otherwise chromadb.Client() (in-memory).
distance_metricstr"cosine"Must be a key of SUPPORTED_METRICS.
tenant_idOptional[str]NoneIf set, injected into every document's metadata and into every query's where filter.
batch_sizeint500Clamped to at least 1 via max(1, batch_size).
strict_modeboolTrueWhen True, add() validates that every ID is a non-empty string.

Raises: ImportError if chromadb isn't installed; ValueError for an unsupported distance_metric; RuntimeError if client or collection initialization fails.

Initialization workflow

  1. super().__init__(embedder).
  2. Validate distance_metric, map it to self.metric_space.
  3. _initialize_client(persist_directory) — persistent or in-memory client.
  4. _initialize_collection() — tries get_collection(name=...); on any exception, falls back to create_collection(name=..., metadata={"hnsw:space": self.metric_space}).

[!NOTE] Existing collections keep their original distance metric Because _initialize_collection only sets hnsw:space when creating a new collection, reopening an existing collection with a different distance_metric argument does not change the collection's actual metric — Chroma keeps whatever metric it was created with. self.distance_metric on the Python object reflects the constructor argument, which can silently diverge from the collection's real configuration in that scenario.

Thread safety

Not addressed in source.

Properties

Python · 4 lines
@property
def client(self) -> Any          # raises RuntimeError if not initialized
@property
def collection(self) -> Collection  # raises RuntimeError if not initialized

Both simply guard against None and return the underlying ChromaDB objects, giving escape-hatch access to the native client/collection API for anything this wrapper doesn't expose.

Add / upsert operations

add

Python
def add(self, chunks: List[Any], embeddings: Optional[np.ndarray] = None) -> None

Supports two calling styles, detected by inspecting the first element of chunks:

Unified style (chunks is a list of DocumentChunk-like objects with a chunk_id attribute):

Python
db.add(chunks)  # list[DocumentChunk]
  • ids = [c.chunk_id for c in chunks]
  • documents = [c.text for c in chunks]
  • metadatas = [sanitize_metadata({**normalize_metadata(c.metadata), "doc_id": c.doc_id}) for c in chunks]
  • if embeddings is None and every chunk already has .embedding set, those are reused directly instead of calling the embedder again.

Legacy style (backward-compatible, keyword-argument based):

Python
db.add(ids=[...], documents=[...], metadatas=[...])

[!DANGER] Legacy calling style is broken in this version of the source In the legacy branch, the function body reads:

Python · 3 lines
ids = chunks              # first positional arg treated as ids
documents = documents     # references a name that is NOT a parameter of add()
metadatas = [sanitize_metadata(m) for m in (metadatas or [])]   # same issue

documents and metadatas are not parameters of add() — there is no documents=/metadatas= parameter in the signature (def add(self, chunks, embeddings=None)), so referencing those names raises NameError unless a variable with that exact name happens to already exist in an enclosing scope. In practice, calling add() in the documented legacy style as shown in the docstring will raise NameError, not perform an upsert. This is a defect in the source as provided, not a documentation gap — see Troubleshooting for a full explanation and workaround.

Common workflow after ids/documents/metadatas are resolved:

  1. _validate_inputs(ids, documents, metadatas, embeddings) — raises ValueError if any of ids/documents/metadatas is empty, if their lengths mismatch, if embeddings length mismatches, or (when strict_mode=True) if any ID is not a non-empty string.
  2. If embeddings is None, requires self.embedder (else ValueError), then calls self.embedder.encode(documents), wrapping failures in RuntimeError.
  3. Cast embeddings to float32, run _validate_embeddings (checks 2-D shape and consistent dimension across calls — the first call's dimension becomes self.embedding_dim).
  4. Iterate in batches of self.batch_size, apply _enforce_tenant_isolation (injects tenant_id into each batch's metadata copies), then self.collection.upsert(ids=..., documents=..., metadatas=..., embeddings=batch_emb.tolist()).

Raises: ValueError for invalid input shapes; RuntimeError wrapping upsert failures (per-batch, naming the failing batch's start index).

Complexity

Upsert cost is whatever ChromaDB's own HNSW-backed collection charges per batch; this wrapper adds only O(n) bookkeeping (validation, batching, tenant-tag injection) on top.

Search operations

Python · 8 lines
def search(
    self,
    query: Union[str, np.ndarray],
    top_k: int = 5,
    score_threshold: Optional[float] = None,
    filters: Optional[Dict] = None,
    **kwargs: Any,
) -> List[Tuple[DocumentChunk, float]]
ParameterTypeDefaultNotes
querystr | np.ndarray—Text requires self.embedder.
top_kint5Must be > 0.
score_thresholdOptional[float]NoneApplied after converting distance → similarity.
filtersOptional[Dict]NoneChroma where filter. Combined with tenant isolation via $and if tenant_id is set.
**kwargs——Accepted but unused in the current implementation.

Workflow:

  1. Validate top_k > 0.
  2. Resolve query_embedding (encode text, or cast an array).
  3. If self.embedding_dim is already known (from a prior add()), validate the query's dimension against it.
  4. If tenant_id is set, AND the caller's filters (if any) with {"tenant_id": tenant_id}.
  5. self.collection.query(query_embeddings=[...], n_results=top_k, where=filters).
  6. For each hit: convert Chroma's distance to a similarity score via _distance_to_score; drop it if below score_threshold; reconstruct a DocumentChunk (doc_id taken from metadata if present, else falls back to the Chroma document ID itself; the returned chunk's metadata excludes the doc_id key since it's promoted to the DocumentChunk.doc_id field).

Returns: List[Tuple[DocumentChunk, float]].

Raises: ValueError for bad top_k, missing embedder on a text query, or dimension mismatch; RuntimeError wrapping query failures.

_distance_to_score (internal)

Python
def _distance_to_score(self, distance: float) -> float
distance_metricFormula
"cosine"1.0 - distance
"l2"-distance
"ip"distance (used as-is)

This means the numeric range and meaning of the returned "score" is metric-dependent — a score_threshold tuned for cosine similarity ([0, 1]-ish) is not comparable to one tuned for negated L2 distance.

Delete operations

delete_by_ids

Python
def delete_by_ids(self, ids: List[str]) -> int

Empty list → returns 0 immediately. Otherwise self.collection.delete(ids=ids) and returns len(ids) — note this is the requested count, not a confirmed-deleted count from Chroma.

delete_by_filter

Python
def delete_by_filter(self, filters: Dict) -> int

Empty filters → warns and returns 0. Applies tenant isolation the same way search does, fetches matching IDs via self.collection.get(where=filters), deletes them, and returns the actual number of IDs that matched (this one is a confirmed count, unlike delete_by_ids).

clear_collection

Python
def clear_collection(self) -> int

Reads the current count, fetches every ID via collection.get(), deletes them all, and returns the pre-deletion count. This cannot be undone (per the source's own docstring warning).

Statistics & utilities

stats

Python
def stats(self) -> Dict[str, Any]
Python · 9 lines
{
    "collection_name": str,
    "total_vectors": int,          # self.collection.count()
    "distance_metric": str,
    "embedding_dimension": int | None,
    "tenant_id": str | None,
    "batch_size": int,
    "strict_mode": bool,
}

get_by_ids

Python
def get_by_ids(self, ids: List[str]) -> List[Tuple[str, Dict]]

Returns (document_text, metadata) pairs (note: not DocumentChunk objects, unlike every other cross-backend method — this one keeps Chroma's native shape).

Dunder methods

Python · 2 lines
def __len__(self) -> int      # self.collection.count()
def __repr__(self) -> str     # collection, metric, vectors, tenant

Unified interface aliases

Python · 8 lines
def remove_by_doc_id(self, doc_id: str) -> int:
    return self.delete_by_filter({"doc_id": doc_id})

def get_stats(self) -> Dict[str, Any]:
    base = self.stats()
    base.setdefault("total_vectors", base.get("total_vectors", len(self)))
    base.setdefault("has_embedder", self.embedder is not None)
    return base

remove_by_doc_id only works correctly if doc_id was actually stored as a metadata key — which it is, whenever chunks were added via the unified DocumentChunk-list style (add() injects "doc_id": c.doc_id into metadata). If chunks were added via the legacy style (setting aside that it currently raises NameError — see above) with metadata that never included doc_id, remove_by_doc_id would match nothing.

Context manager support

batch_operation

Python · 2 lines
@contextmanager
def batch_operation(self)

A thin try/except/finally wrapper around a yield self that logs start/success/failure/ finalization — it does not buffer or defer writes; calls made inside the with block still hit ChromaDB immediately. Its value is centralized logging/error propagation around a group of calls, not batching in the performance sense.

Async API

Python · 3 lines
async def asearch(self, query, top_k=10, **kwargs)
async def aadd(self, chunks, **kwargs)
async def __aenter__(self) / async def __aexit__(self, *_)

All are asyncio.to_thread(...) wrappers, same pattern as the other backends.

Common mistakes

  • Calling add() in the legacy keyword style as shown in its own docstring — currently raises NameError. Use the DocumentChunk-list style instead. See Troubleshooting.
  • Comparing score_threshold values across different distance_metric settings — the scale and direction of "score" changes with the metric (see _distance_to_score).
  • Assuming get_by_ids returns DocumentChunk objects like other backends' equivalent methods — it returns (text, metadata) tuples instead.
  • Expecting distance_metric passed to the constructor to retroactively change an already-existing collection's actual similarity metric — it doesn't.

Related pages