app/services/vector_store.py

import chromadb

from app.services.vector_store_contract import EmbeddingMetadata, SearchResult, VectorDocument


class ChromaReadinessError(RuntimeError):
    """Raised when the ChromaDB persistent store cannot pass a write/read/delete probe."""


class VectorStore:
    READINESS_TENANT_ID = "chroma-readiness"
    READINESS_ENTRY_ID = "__chroma_readiness_probe__"
    READINESS_TEXT = "__chroma_readiness_probe__"
    READINESS_EMBEDDING_DIMENSION = 3072

    def __init__(self, persist_path: str):
        self.client = chromadb.PersistentClient(path=persist_path)

    def _get_collection(self, tenant_id: str) -> chromadb.Collection:
        collection = self.client.get_or_create_collection(name=f"kb_{tenant_id}")

        return collection

    @staticmethod
    def _metadata_for_document(
        document: VectorDocument,
    ) -> dict[str, str | int | float | bool | None]:
        custom_metadata = {
            key: value for key, value in document.metadata.items() if value is not None
        }
        return {
            **custom_metadata,
            "vector_id": document.vector_id or document.entry_id,
            "tenant_id": document.tenant_id,
            "entry_id": document.entry_id,
            "embedding_model": document.embedding_metadata.embedding_model,
            "embedding_dim": document.embedding_metadata.embedding_dim,
        }

    @staticmethod
    def _legacy_document(
        tenant_id: str, entry_id: str, text: str, embedding: list[float]
    ) -> VectorDocument:
        return VectorDocument(
            tenant_id=tenant_id,
            entry_id=entry_id,
            text=text,
            embedding=embedding,
            embedding_metadata=EmbeddingMetadata(
                embedding_model="unknown",
                embedding_dim=len(embedding),
            ),
        )

    @staticmethod
    def _similarity_from_distance(distance: float | None) -> float:
        if distance is None:
            return 0.0
        return max(0.0, min(1.0, 1.0 / (1.0 + distance)))

    @staticmethod
    def _search_results_from_chroma(results: dict) -> list[SearchResult]:
        docs = results.get("documents") or []
        ids = results.get("ids") or []
        distances = results.get("distances") or []
        metadatas = results.get("metadatas") or []

        first_docs = docs[0] if docs else []
        first_ids = ids[0] if ids else []
        first_distances = distances[0] if distances else []
        first_metadatas = metadatas[0] if metadatas else []

        search_results = []
        for index, entry_id in enumerate(first_ids):
            text = first_docs[index] if index < len(first_docs) else ""
            distance = first_distances[index] if index < len(first_distances) else None
            metadata = first_metadatas[index] if index < len(first_metadatas) else {}
            search_results.append(
                SearchResult(
                    entry_id=str(metadata.get("entry_id", entry_id)),
                    text=text,
                    similarity=VectorStore._similarity_from_distance(distance),
                    metadata={
                        key: value
                        for key, value in (metadata or {}).items()
                        if key != "vector_id"
                    },
                )
            )

        return search_results

    def add(
        self, tenant_id: str, entry_id: str, text: str, embedding: list[float]
    ) -> None:
        self.upsert(self._legacy_document(tenant_id, entry_id, text, embedding))

    def query(
        self, tenant_id: str, query_embedding: list[float], n_results: int
    ) -> list[SearchResult]:

        collection = self._get_collection(tenant_id=tenant_id)

        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=n_results,
            include=["documents", "metadatas", "distances"],
        )

        return self._search_results_from_chroma(results)

    def upsert(
        self,
        document: VectorDocument | None = None,
        *,
        tenant_id: str | None = None,
        entry_id: str | None = None,
        text: str | None = None,
        embedding: list[float] | None = None,
    ) -> None:
        if document is None:
            if (
                tenant_id is None
                or entry_id is None
                or text is None
                or embedding is None
            ):
                raise TypeError(
                    "upsert requires a VectorDocument or legacy keyword arguments"
                )
            document = self._legacy_document(tenant_id, entry_id, text, embedding)

        collection = self._get_collection(tenant_id=document.tenant_id)

        collection.upsert(
            ids=[document.vector_id or document.entry_id],
            embeddings=[document.embedding],
            documents=[document.text],
            metadatas=[self._metadata_for_document(document)],
        )

    def delete(self, tenant_id: str, entry_id: str) -> None:
        self.delete_docs(tenant_id=tenant_id, entry_ids=[entry_id])

    def delete_docs(self, tenant_id: str, entry_ids: list[str]) -> None:
        if not entry_ids:
            return
        collection = self._get_collection(tenant_id=tenant_id)

        collection.delete(where={"entry_id": {"$in": entry_ids}})

    def delete_tenant(self, tenant_id: str) -> None:
        try:
            self.client.delete_collection(name=f"kb_{tenant_id}")
        except Exception:
            collection_names = {
                collection.name for collection in self.client.list_collections()
            }
            if f"kb_{tenant_id}" in collection_names:
                raise

    def health_check(self) -> None:
        self.verify_readiness()

    def verify_readiness(self) -> None:
        """Verify ChromaDB can write, read, and delete from the persistent store."""
        collection = self._get_collection(tenant_id=self.READINESS_TENANT_ID)
        probe_embedding = [0.0] * self.READINESS_EMBEDDING_DIMENSION

        try:
            collection.upsert(
                ids=[self.READINESS_ENTRY_ID],
                embeddings=[probe_embedding],
                documents=[self.READINESS_TEXT],
            )
            results = collection.query(
                query_embeddings=[probe_embedding],
                n_results=1,
            )

            docs = results["documents"] or []
            ids = results["ids"] or []
            first_docs = docs[0] if docs else []
            first_ids = ids[0] if ids else []

            if (
                self.READINESS_ENTRY_ID not in first_ids
                or self.READINESS_TEXT not in first_docs
            ):
                raise ChromaReadinessError(
                    "ChromaDB readiness probe query did not return the probe document"
                )

            collection.delete(ids=[self.READINESS_ENTRY_ID])
        except ChromaReadinessError:
            raise
        except Exception as exc:
            raise ChromaReadinessError("ChromaDB readiness probe failed") from exc