rag/vector_store.py

"""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