"""ChromaDB persistence and LangChain ``Chroma`` vector store helpers. Collections are named per ingest target; documents are stored with UUID chunk ids. Telemetry is disabled at the client level for quieter logs in production. """ from datetime import datetime, timezone from pathlib import Path from uuid import uuid4 import chromadb from chromadb.config import Settings from langchain_chroma import Chroma from langchain_core.documents import Document from langchain_core.embeddings import Embeddings _CHROMA_CLIENT_SETTINGS = Settings(anonymized_telemetry=False) def _utc_now_iso() -> str: return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") def _chroma_client(persist_directory: str) -> chromadb.PersistentClient: Path(persist_directory).mkdir(parents=True, exist_ok=True) return chromadb.PersistentClient(path=persist_directory, settings=_CHROMA_CLIENT_SETTINGS) def get_vector_store( persist_directory: str, collection_name: str, embedding_function: Embeddings, ) -> Chroma: """Open or create a persisted Chroma collection wired to the given embedder.""" client = _chroma_client(persist_directory) try: client.get_collection(name=collection_name) except Exception: client.get_or_create_collection( name=collection_name, metadata={"created_at": _utc_now_iso()}, ) return Chroma( collection_name=collection_name, embedding_function=embedding_function, persist_directory=persist_directory, client_settings=_CHROMA_CLIENT_SETTINGS, ) def add_documents(vector_store: Chroma, chunks: list[Document]) -> list[str]: """Embed and insert chunks; return the generated vector ids.""" document_ids = [str(uuid4()) for _ in chunks] vector_store.add_documents(documents=chunks, ids=document_ids) return document_ids def list_collection_names(persist_directory: str) -> list[str]: """Sorted list of collection names in the persist directory.""" client = _chroma_client(persist_directory) return sorted(c.name for c in client.list_collections()) def delete_collection(persist_directory: str, collection_name: str) -> int: """Delete a collection and return the number of documents that were removed (best effort).""" client = _chroma_client(persist_directory) removed = 0 try: col = client.get_collection(name=collection_name) removed = int(col.count()) except Exception: removed = 0 client.delete_collection(name=collection_name) return removed def collection_document_count(persist_directory: str, collection_name: str) -> int: """Number of vectors in a collection, or 0 if the collection does not exist.""" client = _chroma_client(persist_directory) try: col = client.get_collection(name=collection_name) return int(col.count()) except Exception: return 0 def collection_created_at(persist_directory: str, collection_name: str) -> str | None: """Return collection metadata ``created_at`` if present (Chroma-specific).""" client = _chroma_client(persist_directory) try: col = client.get_collection(name=collection_name) meta = getattr(col, "metadata", None) or {} if isinstance(meta, dict): raw = meta.get("created_at") or meta.get("created") if raw is not None: return str(raw) except Exception: pass return None def ensure_collection_created_at( persist_directory: str, collection_name: str, *, fallback: str | None = None, ) -> str | None: """Persist ``created_at`` on the Chroma collection when missing; never overwrites an existing value.""" client = _chroma_client(persist_directory) try: col = client.get_collection(name=collection_name) except Exception: return None meta = getattr(col, "metadata", None) or {} if not isinstance(meta, dict): meta = {} raw = meta.get("created_at") or meta.get("created") if raw is not None: return str(raw) value = fallback or _utc_now_iso() updated = dict(meta) updated["created_at"] = value col.modify(metadata=updated) return value