app/services/qdrant_vector_store.py

import uuid

from qdrant_client import QdrantClient, models

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


class QdrantVectorStore:
    HEALTH_TENANT_ID = "qdrant-health"
    HEALTH_ENTRY_ID = "qdrant-health-probe"
    HEALTH_TEXT = "qdrant health probe"

    def __init__(
        self,
        url: str,
        collection_name: str,
        vector_size: int,
        timeout: int = 10,
        client: QdrantClient | None = None,
    ):
        self.client = client or QdrantClient(url=url, timeout=timeout)
        self.collection_name = collection_name
        self.vector_size = vector_size

    @staticmethod
    def _point_id(tenant_id: str, vector_id: str) -> str:
        return str(uuid.uuid5(uuid.NAMESPACE_URL, f"{tenant_id}:{vector_id}"))

    @staticmethod
    def _tenant_filter(tenant_id: str) -> models.Filter:
        return models.Filter(
            must=[
                models.FieldCondition(
                    key="tenant_id",
                    match=models.MatchValue(value=tenant_id),
                )
            ]
        )

    @staticmethod
    def _docs_filter(tenant_id: str, entry_ids: list[str]) -> models.Filter:
        return models.Filter(
            must=[
                models.FieldCondition(
                    key="tenant_id",
                    match=models.MatchValue(value=tenant_id),
                ),
                models.FieldCondition(
                    key="entry_id",
                    match=models.MatchAny(any=entry_ids),
                ),
            ]
        )

    @staticmethod
    def _normalize_cosine_score(score: float) -> float:
        return max(0.0, min(1.0, (score + 1.0) / 2.0))

    def _ensure_collection(self) -> None:
        if self.client.collection_exists(collection_name=self.collection_name):
            return

        try:
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=models.VectorParams(
                    size=self.vector_size,
                    distance=models.Distance.COSINE,
                ),
            )
        except Exception:
            if self.client.collection_exists(collection_name=self.collection_name):
                return
            raise

        self.client.create_payload_index(
            collection_name=self.collection_name,
            field_name="tenant_id",
            field_schema=models.PayloadSchemaType.KEYWORD,
            wait=True,
        )
        self.client.create_payload_index(
            collection_name=self.collection_name,
            field_name="entry_id",
            field_schema=models.PayloadSchemaType.KEYWORD,
            wait=True,
        )

    def upsert(self, document: VectorDocument) -> None:
        self._ensure_collection()
        payload = {
            key: value for key, value in document.metadata.items() if value is not None
        }
        payload.update(
            {
                "vector_id": document.vector_id or document.entry_id,
                "tenant_id": document.tenant_id,
                "entry_id": document.entry_id,
                "text": document.text,
                "embedding_model": document.embedding_metadata.embedding_model,
                "embedding_dim": document.embedding_metadata.embedding_dim,
            }
        )
        self.client.upsert(
            collection_name=self.collection_name,
            points=[
                models.PointStruct(
                    id=self._point_id(
                        document.tenant_id, document.vector_id or document.entry_id
                    ),
                    vector=document.embedding,
                    payload=payload,
                )
            ],
            wait=True,
        )

    def query(
        self, tenant_id: str, query_embedding: list[float], n_results: int
    ) -> list[SearchResult]:
        self._ensure_collection()
        response = self.client.query_points(
            collection_name=self.collection_name,
            query=query_embedding,
            query_filter=self._tenant_filter(tenant_id),
            limit=n_results,
            with_payload=True,
            with_vectors=False,
        )
        return [
            SearchResult(
                entry_id=str(point.payload.get("entry_id", point.id)),
                text=str(point.payload.get("text", "")),
                similarity=self._normalize_cosine_score(point.score),
                metadata={
                    key: value
                    for key, value in point.payload.items()
                    if key not in {"text", "vector_id"} and value is not None
                },
            )
            for point in response.points
        ]

    def delete_docs(self, tenant_id: str, entry_ids: list[str]) -> None:
        if not entry_ids:
            return
        self._ensure_collection()
        self.client.delete(
            collection_name=self.collection_name,
            points_selector=models.FilterSelector(
                filter=self._docs_filter(tenant_id, entry_ids)
            ),
            wait=True,
        )

    def delete_tenant(self, tenant_id: str) -> None:
        self._ensure_collection()
        self.client.delete(
            collection_name=self.collection_name,
            points_selector=models.FilterSelector(filter=self._tenant_filter(tenant_id)),
            wait=True,
        )

    def get_collection_dimension(self) -> int | None:
        """Return the actual vector size of the collection, or None if it does not exist."""
        if not self.client.collection_exists(collection_name=self.collection_name):
            return None
        info = self.client.get_collection(collection_name=self.collection_name)
        return info.config.params.vectors.size

    def health_check(self) -> None:
        probe_vector = [1.0] + [0.0] * (self.vector_size - 1)
        self.upsert(
            VectorDocument(
                tenant_id=self.HEALTH_TENANT_ID,
                entry_id=self.HEALTH_ENTRY_ID,
                text=self.HEALTH_TEXT,
                embedding=probe_vector,
                embedding_metadata=EmbeddingMetadata(
                    embedding_model="qdrant-health-check",
                    embedding_dim=self.vector_size,
                ),
            )
        )
        results = self.query(self.HEALTH_TENANT_ID, probe_vector, 1)
        if not results or results[0].entry_id != self.HEALTH_ENTRY_ID:
            raise RuntimeError("Qdrant health probe query did not return probe document")
        self.delete_docs(self.HEALTH_TENANT_ID, [self.HEALTH_ENTRY_ID])