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)