retrieval/qdrant_client.py

"""Qdrant vector database manager with RBAC-aware operations."""

from __future__ import annotations

import contextlib
import uuid
from typing import Any

from qdrant_client import QdrantClient, models
from qdrant_client.http.models import (
    Distance,
    PointStruct,
    SparseVector,
    SparseVectorParams,
    VectorParams,
)

from config.settings import settings
from ingestion.metadata import SensitivityLevel, UserContext, sensitivity_to_int
from utils.logging import get_logger

logger = get_logger(__name__)


class QdrantManager:
    """Manages Qdrant vector database operations including collection lifecycle and document upsert.



    Provides methods for collection management and RBAC-aware document storage.



    Args:

        url: Qdrant server URL. Defaults to settings.qdrant_url.

        collection_name: Target collection name. Defaults to settings.qdrant_collection.

        api_key: Optional API key for Qdrant Cloud authentication.

    """

    def __init__(

        self,

        url: str | None = None,

        collection_name: str | None = None,

        api_key: str | None = None,

    ) -> None:
        """Initialize the Qdrant manager.



        Args:

            url: Qdrant server URL. Falls back to settings.qdrant_url.

            collection_name: Collection name. Falls back to settings.qdrant_collection.

            api_key: API key for authentication. Falls back to settings.qdrant_api_key.

        """
        self._url = url if url is not None else settings.qdrant_url
        self._collection_name = (
            collection_name if collection_name is not None else settings.qdrant_collection
        )
        self._api_key = api_key if api_key is not None else settings.qdrant_api_key

        self._client = QdrantClient(
            url=self._url,
            api_key=self._api_key,
            timeout=30,
        )
        # Per-tenant manager cache. In multi-tenant mode each `for_org(org_id)`
        # call previously created a fresh QdrantManager (new HTTP client +
        # extra `get_collections` round-trip via `ensure_collection`). Caching
        # by collection name turns repeat calls into pure dict lookups so the
        # per-request overhead disappears. Stays bound to *this* root manager
        # — distinct roots (different URLs) keep distinct caches.
        self._tenant_cache: dict[str, QdrantManager] = {}

        logger.info(
            "qdrant_manager_initialized",
            url=self._url,
            collection=self._collection_name,
        )

    @property
    def collection_name(self) -> str:
        """Return the current collection name."""
        return self._collection_name

    @property
    def client(self) -> QdrantClient:
        """Return the underlying QdrantClient instance."""
        return self._client

    def for_org(self, org_id: str) -> QdrantManager:
        """Return a QdrantManager scoped to an organization-specific collection.



        When ``settings.multi_tenant_collections`` is True, this returns a

        per-org manager bound to ``documents_{org_id}``. Each tenant collection

        is created the first time it is requested (with the same dense + sparse

        vector configuration as the global collection — sparse isolation is

        therefore structural: org A's sparse vectors live in

        ``documents_acme_corp.sparse``, org B's in ``documents_partner_inc.sparse``,

        and Qdrant cannot cross collections in a single query) and the manager

        is cached on the root instance so repeat requests are O(1) dict lookups

        rather than fresh HTTP-client + ``get_collections`` round-trips.



        When ``multi_tenant_collections`` is False, returns ``self``.



        Args:

            org_id: Organization identifier.



        Returns:

            A QdrantManager instance (new, cached, or self).

        """
        if not settings.multi_tenant_collections:
            return self
        from retrieval.multitenancy import get_collection_name

        org_collection = get_collection_name(org_id)
        if org_collection == self._collection_name:
            return self
        cached = self._tenant_cache.get(org_collection)
        if cached is not None:
            return cached
        mgr = QdrantManager(
            url=self._url,
            collection_name=org_collection,
            api_key=self._api_key,
        )
        mgr.ensure_collection()
        self._tenant_cache[org_collection] = mgr
        logger.info(
            "tenant_collection_cached",
            collection=org_collection,
            cache_size=len(self._tenant_cache),
        )
        return mgr

    def for_session(self, session_id: str) -> QdrantManager:
        """Return a QdrantManager bound to a BYOK visitor's session collection.



        Mirrors ``for_org`` but uses the session-scoped naming convention

        ``documents_sess_<sanitized_session>``. The collection is created

        on first request with the same dense + sparse vector configuration

        as the base collection, then cached on this root manager so repeat

        requests are O(1) dict lookups.



        BYOK uploads live in the visitor's session collection only; the

        24-hour purge cron drops abandoned collections so 1 GB Qdrant

        Cloud quota stays bounded.



        Args:

            session_id: Per-visitor session UUID (BYOK mode).



        Returns:

            A QdrantManager scoped to the session collection.

        """
        if not session_id:
            return self
        from retrieval.multitenancy import get_collection_name

        # Force BYOK-style naming even when settings.byok_mode is False so
        # tests can exercise the path without flipping the global flag.
        base = self._collection_name
        sanitized = "".join(c if c.isalnum() else "_" for c in session_id)
        sess_collection = f"{base}_sess_{sanitized}"
        # Honour the canonical helper when both flags align so a future rename
        # of the prefix has a single source of truth.
        if settings.byok_mode:
            with contextlib.suppress(Exception):
                sess_collection = get_collection_name(session_id=session_id)

        if sess_collection == self._collection_name:
            return self
        cached = self._tenant_cache.get(sess_collection)
        if cached is not None:
            return cached
        mgr = QdrantManager(
            url=self._url,
            collection_name=sess_collection,
            api_key=self._api_key,
        )
        mgr.ensure_collection()
        # Qdrant Cloud requires explicit payload indexes on filterable fields.
        # Mirror the indexes created on the base demo collection.
        for field, schema in (
            ("org_id", "keyword"),
            ("sensitivity_level_int", "integer"),
            ("roles", "keyword"),
            ("user_id", "keyword"),
            ("source_file", "keyword"),
            # Per-upload group key used by the BYOK delete endpoint to drop
            # all chunks of one upload. Qdrant Cloud refuses Filter()
            # predicates on un-indexed payload keys, so this is mandatory.
            ("source_file_id", "keyword"),
        ):
            # Index may already exist; safe to ignore.
            with contextlib.suppress(Exception):
                mgr._client.create_payload_index(
                    collection_name=sess_collection,
                    field_name=field,
                    field_schema=schema,
                )
        # Stamp the collection with a creation timestamp so the 24h purge cron
        # can actually find and drop it later (Qdrant has no writable collection
        # metadata slot — see retrieval/session_purge.write_session_sentinel).
        with contextlib.suppress(Exception):
            from retrieval.session_purge import write_session_sentinel

            write_session_sentinel(mgr._client, sess_collection, settings.embedding_dim)
        self._tenant_cache[sess_collection] = mgr
        logger.info(
            "byok_session_collection_cached",
            collection=sess_collection,
            cache_size=len(self._tenant_cache),
        )
        return mgr

    def ensure_collection(self, vector_size: int | None = None) -> None:
        """Create the collection if it does not already exist.



        Creates both dense and sparse vector configurations so that hybrid

        search (dense + sparse) works out of the box.



        Args:

            vector_size: Dimension of the embedding vectors.

                Defaults to settings.embedding_dim.

        """
        size = vector_size if vector_size is not None else settings.embedding_dim

        try:
            collections = self._client.get_collections().collections
            existing_names = {c.name for c in collections}

            if self._collection_name in existing_names:
                logger.info(
                    "collection_already_exists",
                    collection=self._collection_name,
                )
                return

            sparse_name = getattr(settings, "sparse_vector_name", "sparse")
            self._client.create_collection(
                collection_name=self._collection_name,
                vectors_config=VectorParams(
                    size=size,
                    distance=Distance.COSINE,
                ),
                sparse_vectors_config={sparse_name: SparseVectorParams()},
            )
            logger.info(
                "collection_created",
                collection=self._collection_name,
                vector_size=size,
                distance="Cosine",
                sparse_vector=sparse_name,
            )

        except Exception as exc:
            logger.error(
                "collection_ensure_failed",
                collection=self._collection_name,
                error=str(exc),
            )
            raise

    async def upsert_documents(

        self,

        chunks: list[str],

        embeddings: list[list[float]],

        metadatas: list[dict],

        sparse_vectors: list[SparseVector] | None = None,

    ) -> list[str]:
        """Upsert document chunks with embeddings and metadata into Qdrant.



        Generates UUID for each point and stores the chunk text in the payload

        alongside the provided metadata. When *sparse_vectors* are supplied

        they are written to the named sparse vector field configured by

        ``settings.sparse_vector_name``.



        Args:

            chunks: List of text chunks.

            embeddings: Corresponding dense embedding vectors.

            metadatas: Corresponding metadata dictionaries.

            sparse_vectors: Optional sparse vectors for hybrid search.



        Returns:

            List of point ID strings (UUIDs).



        Raises:

            ValueError: If input lists have mismatched lengths.

            Exception: On Qdrant upsert failure.

        """
        if not (len(chunks) == len(embeddings) == len(metadatas)):
            raise ValueError(
                f"Input length mismatch: chunks={len(chunks)}, "
                f"embeddings={len(embeddings)}, metadatas={len(metadatas)}"
            )
        if sparse_vectors is not None and len(sparse_vectors) != len(chunks):
            raise ValueError(
                f"Sparse vector length mismatch: sparse={len(sparse_vectors)}, chunks={len(chunks)}"
            )

        if not chunks:
            return []

        point_ids: list[str] = []
        points: list[PointStruct] = []
        sparse_name = getattr(settings, "sparse_vector_name", "sparse")
        has_sparse = sparse_vectors is not None

        for idx, (chunk_text, embedding, metadata) in enumerate(
            zip(chunks, embeddings, metadatas, strict=False)
        ):
            point_id = str(uuid.uuid4())
            point_ids.append(point_id)

            payload = {
                "text": chunk_text,
                **metadata,
            }

            # Defensive: ensure sensitivity_level_int present even if caller
            # passed metadata not produced by DocumentMetadata.to_qdrant_payload.
            if "sensitivity_level_int" not in payload:
                sl = payload.get("sensitivity_level")
                if sl is not None:
                    try:
                        payload["sensitivity_level_int"] = sensitivity_to_int(SensitivityLevel(sl))
                    except (ValueError, KeyError):
                        payload["sensitivity_level_int"] = 1

            vector: dict[str, Any] | list[float] = embedding
            if has_sparse:
                vector = {
                    "": embedding,
                    sparse_name: sparse_vectors[idx],
                }

            points.append(
                PointStruct(
                    id=point_id,
                    vector=vector,
                    payload=payload,
                )
            )

        try:
            self._client.upsert(
                collection_name=self._collection_name,
                points=points,
            )
            logger.info(
                "documents_upserted",
                collection=self._collection_name,
                count=len(points),
                has_sparse=has_sparse,
            )
        except Exception as exc:
            logger.error(
                "upsert_failed",
                collection=self._collection_name,
                count=len(points),
                error=str(exc),
            )
            raise

        return point_ids

    def get_collection_info(self) -> dict | None:
        """Retrieve information about the current collection.



        Returns:

            Dictionary with collection info, or None if collection doesn't exist.

        """
        try:
            info = self._client.get_collection(self._collection_name)
            # vectors_count was removed from CollectionInfo in qdrant-client >= 1.10;
            # use getattr so this stays forward-compatible.
            return {
                "name": self._collection_name,
                "points_count": info.points_count,
                "vectors_count": getattr(info, "vectors_count", info.points_count),
                "status": info.status.value if info.status else None,
            }
        except Exception as exc:
            logger.warning(
                "collection_info_failed",
                collection=self._collection_name,
                error=str(exc),
            )
            return None

    def delete_collection(self) -> None:
        """Delete the current collection from Qdrant.



        Logs a warning if the collection doesn't exist.

        """
        try:
            self._client.delete_collection(self._collection_name)
            logger.info("collection_deleted", collection=self._collection_name)
        except Exception as exc:
            logger.warning(
                "collection_delete_failed",
                collection=self._collection_name,
                error=str(exc),
            )

    def build_rbac_filter(self, user_context: UserContext) -> models.Filter:
        """Build a Qdrant filter that enforces role-based access control.



        The filter ensures:

        - User belongs to the same organization as the document.

        - Document sensitivity level is within the user's clearance.

        - At least one of the user's roles matches the document's roles.



        Args:

            user_context: Authenticated user context with org, roles, and clearance.



        Returns:

            A Qdrant Filter object ready for use in search queries.

        """
        must_conditions = [
            models.FieldCondition(
                key="org_id",
                match=models.MatchValue(value=user_context.org_id),
            ),
            models.FieldCondition(
                key="sensitivity_level_int",
                range=models.Range(lte=user_context.clearance_level),
            ),
            models.FieldCondition(
                key="roles",
                match=models.MatchAny(any=user_context.roles),
            ),
        ]
        return models.Filter(must=must_conditions)

    def build_combined_filter(

        self,

        user_context: UserContext,

        extra_conditions: list[dict[str, Any]] | None = None,

    ) -> models.Filter:
        """Build a Qdrant filter combining RBAC with self-query conditions.



        Args:

            user_context: Authenticated user context for RBAC.

            extra_conditions: List of condition dicts from

                ``self_query.build_qdrant_filter_conditions``.



        Returns:

            A Qdrant Filter with RBAC must-conditions plus any extra conditions.

        """
        rbac = self.build_rbac_filter(user_context)
        if not extra_conditions:
            return rbac

        combined_must = list(rbac.must or [])
        for cond in extra_conditions:
            if "match" in cond:
                combined_must.append(
                    models.FieldCondition(
                        key=cond["key"],
                        match=cond["match"],
                    )
                )
            elif "range" in cond:
                combined_must.append(
                    models.FieldCondition(
                        key=cond["key"],
                        range=cond["range"],
                    )
                )
        return models.Filter(must=combined_must)

    def search_with_rbac(

        self,

        query_embedding: list[float],

        user_context: UserContext,

        top_k: int | None = None,

        score_threshold: float | None = None,

        extra_filter: models.Filter | None = None,

    ) -> list[models.ScoredPoint]:
        """Search the collection with RBAC filter applied.



        Args:

            query_embedding: Query vector for similarity search.

            user_context: Authenticated user context for RBAC filtering.

            top_k: Maximum number of results. Defaults to settings.top_k.

            score_threshold: Minimum score threshold. Defaults to None.



        Returns:

            List of scored points matching the query with RBAC constraints.

        """
        k = top_k if top_k is not None else settings.top_k
        rbac_filter = extra_filter or self.build_rbac_filter(user_context)

        try:
            # qdrant-client >= 1.13 replaced .search() with .query_points()
            # which returns a QueryResponse wrapping a list of ScoredPoint.
            response = self._client.query_points(
                collection_name=self._collection_name,
                query=query_embedding,
                query_filter=rbac_filter,
                limit=k,
                score_threshold=score_threshold,
            )
            results = response.points
            logger.info(
                "search_with_rbac_completed",
                collection=self._collection_name,
                results_count=len(results),
                user_id=user_context.user_id,
                org_id=user_context.org_id,
            )
            return results
        except Exception as exc:
            logger.error(
                "search_with_rbac_failed",
                collection=self._collection_name,
                error=str(exc),
            )
            return []

    def search_sparse_with_rbac(

        self,

        sparse_vector: models.SparseVector,

        user_context: UserContext,

        top_k: int | None = None,

        score_threshold: float | None = None,

        extra_filter: models.Filter | None = None,

    ) -> list[models.ScoredPoint]:
        """Search the sparse vector field with RBAC filter applied.



        Args:

            sparse_vector: Query sparse vector (indices + values).

            user_context: Authenticated user context for RBAC filtering.

            top_k: Maximum number of results. Defaults to settings.top_k.

            score_threshold: Minimum score threshold. Defaults to None.

            extra_filter: Optional additional Qdrant filter.



        Returns:

            List of scored points from the sparse vector index.

        """
        k = top_k if top_k is not None else settings.top_k
        rbac_filter = extra_filter or self.build_rbac_filter(user_context)
        sparse_name = getattr(settings, "sparse_vector_name", "sparse")

        try:
            response = self._client.query_points(
                collection_name=self._collection_name,
                query=sparse_vector,
                using=sparse_name,
                query_filter=rbac_filter,
                limit=k,
                score_threshold=score_threshold,
            )
            results = response.points
            logger.info(
                "search_sparse_with_rbac_completed",
                collection=self._collection_name,
                results_count=len(results),
                user_id=user_context.user_id,
                org_id=user_context.org_id,
            )
            return results
        except Exception as exc:
            logger.error(
                "search_sparse_with_rbac_failed",
                collection=self._collection_name,
                error=str(exc),
            )
            return []

    def search_without_rbac(

        self,

        query_embedding: list[float],

        top_k: int | None = None,

        score_threshold: float | None = None,

        admin_context: UserContext | None = None,

    ) -> list[models.ScoredPoint]:
        """Search the collection without RBAC filtering (admin/debug use).



        Requires admin role for security. Logs a warning when invoked.



        Args:

            query_embedding: Query vector for similarity search.

            top_k: Maximum number of results. Defaults to settings.top_k.

            score_threshold: Minimum score threshold. Defaults to None.

            admin_context: UserContext that must contain 'admin' role.



        Returns:

            List of scored points matching the query.



        Raises:

            PermissionError: If admin_context is missing or lacks admin role.

        """
        if admin_context is None or "admin" not in admin_context.roles:
            logger.warning(
                "search_without_rbac_called_without_admin",
                admin_context_provided=admin_context is not None,
            )
            raise PermissionError("Admin role required for unfiltered search")

        logger.warning(
            "search_without_rbac_invoked",
            user_id=admin_context.user_id,
            org_id=admin_context.org_id,
        )

        k = top_k if top_k is not None else settings.top_k

        try:
            response = self._client.query_points(
                collection_name=self._collection_name,
                query=query_embedding,
                limit=k,
                score_threshold=score_threshold,
            )
            results = response.points
            logger.info(
                "search_without_rbac_completed",
                collection=self._collection_name,
                results_count=len(results),
            )
            return results
        except Exception as exc:
            logger.error(
                "search_without_rbac_failed",
                collection=self._collection_name,
                error=str(exc),
            )
            return []

    def get_document_count(self) -> int:
        """Return total number of points in the collection.



        Returns:

            Integer count of documents, or 0 if collection info unavailable.

        """
        try:
            info = self._client.get_collection(self._collection_name)
            return info.points_count or 0
        except Exception as exc:
            logger.warning(
                "get_document_count_failed",
                collection=self._collection_name,
                error=str(exc),
            )
            return 0

    def scroll_documents(

        self,

        filter_: models.Filter | None = None,

        limit: int = 100,

    ) -> list[models.Record]:
        """Scroll/list documents from the collection with optional filtering.



        Args:

            filter_: Optional Qdrant filter to apply.

            limit: Maximum number of documents to return.



        Returns:

            List of point records from the collection.

        """
        try:
            results, _ = self._client.scroll(
                collection_name=self._collection_name,
                scroll_filter=filter_,
                limit=limit,
            )
            return results
        except Exception as exc:
            logger.error(
                "scroll_documents_failed",
                collection=self._collection_name,
                error=str(exc),
            )
            return []

    def delete_documents_by_filter(

        self,

        filter_: models.Filter | None = None,

    ) -> int:
        """Delete documents matching the given filter.



        If no filter is provided, deletes ALL documents in the collection.

        Use with caution.



        Args:

            filter_: Qdrant filter to match documents for deletion.



        Returns:

            Number of documents deleted.

        """
        try:
            result = self._client.delete(
                collection_name=self._collection_name,
                points_selector=models.FilterSelector(filter=filter_)
                if filter_
                else models.PointIdsList(points=[]),
            )
            deleted = getattr(result, "operation_id", 0)
            logger.info(
                "documents_deleted",
                collection=self._collection_name,
                deleted=deleted,
                filter_applied=filter_ is not None,
            )
            return deleted
        except Exception as exc:
            logger.error(
                "delete_documents_failed",
                collection=self._collection_name,
                error=str(exc),
            )
            return 0

    def delete_document_by_id(self, point_id: str) -> bool:
        """Delete a single document by its point ID.



        Args:

            point_id: The UUID of the point to delete.



        Returns:

            True if deletion was successful, False otherwise.

        """
        try:
            self._client.delete(
                collection_name=self._collection_name,
                points_selector=models.PointIdsList(points=[point_id]),
            )
            logger.info("document_deleted", point_id=point_id)
            return True
        except Exception as exc:
            logger.error("delete_document_failed", point_id=point_id, error=str(exc))
            return False

    def update_document_metadata(

        self,

        point_id: str,

        metadata: dict,

    ) -> bool:
        """Update metadata for a specific document.



        Args:

            point_id: The UUID of the point to update.

            metadata: Dict of metadata fields to update.



        Returns:

            True if update was successful, False otherwise.

        """
        try:
            # Ensure sensitivity_level_int is updated if sensitivity_level changed
            if "sensitivity_level" in metadata and "sensitivity_level_int" not in metadata:
                try:
                    metadata["sensitivity_level_int"] = sensitivity_to_int(
                        SensitivityLevel(metadata["sensitivity_level"])
                    )
                except (ValueError, KeyError):
                    metadata["sensitivity_level_int"] = 1

            self._client.set_payload(
                collection_name=self._collection_name,
                payload=metadata,
                points=[point_id],
            )
            logger.info("document_metadata_updated", point_id=point_id)
            return True
        except Exception as exc:
            logger.error(
                "update_document_metadata_failed",
                point_id=point_id,
                error=str(exc),
            )
            return False

    def get_documents_by_source(

        self,

        source_file: str,

        org_id: str | None = None,

    ) -> list[models.Record]:
        """Get all documents originating from a specific source file.



        Args:

            source_file: The source filename to search for.

            org_id: Optional org_id filter.



        Returns:

            List of matching point records.

        """
        conditions = [
            models.FieldCondition(
                key="source_file",
                match=models.MatchValue(value=source_file),
            ),
        ]
        if org_id:
            conditions.append(
                models.FieldCondition(
                    key="org_id",
                    match=models.MatchValue(value=org_id),
                )
            )
        filter_ = models.Filter(must=conditions)
        return self.scroll_documents(filter_=filter_, limit=1000)