visual-rag-toolkit 0.1.1__py3-none-any.whl

visual_rag/qdrant_admin.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+from urllib.parse import urlparse
+
+
+@dataclass(frozen=True)
+class QdrantConnection:
+    url: str
+    api_key: Optional[str]
+
+
+def _maybe_load_dotenv() -> None:
+    try:
+        from dotenv import load_dotenv
+    except Exception:
+        return
+    try:
+        from pathlib import Path
+
+        if Path(".env").exists():
+            load_dotenv(".env")
+    except Exception:
+        return
+
+
+def _resolve_qdrant_connection(
+    *,
+    url: Optional[str] = None,
+    api_key: Optional[str] = None,
+) -> QdrantConnection:
+    import os
+
+    _maybe_load_dotenv()
+    resolved_url = (
+        url
+        or os.getenv("SIGIR_QDRANT_URL")
+        or os.getenv("DEST_QDRANT_URL")
+        or os.getenv("QDRANT_URL")
+    )
+    if not resolved_url:
+        raise ValueError(
+            "Qdrant URL not set (pass url= or set SIGIR_QDRANT_URL/DEST_QDRANT_URL/QDRANT_URL)."
+        )
+    resolved_key = (
+        api_key
+        or os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
+    return QdrantConnection(url=str(resolved_url), api_key=resolved_key)
+
+
+def _infer_grpc_port(url: str) -> Optional[int]:
+    try:
+        if urlparse(url).port == 6333:
+            return 6334
+    except Exception:
+        return None
+    return None
+
+
+class QdrantAdmin:
+    def __init__(
+        self,
+        *,
+        url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        prefer_grpc: bool = False,
+        timeout: int = 60,
+    ):
+        from qdrant_client import QdrantClient
+
+        conn = _resolve_qdrant_connection(url=url, api_key=api_key)
+        grpc_port = _infer_grpc_port(conn.url) if prefer_grpc else None
+        self.client = QdrantClient(
+            url=conn.url,
+            api_key=conn.api_key,
+            prefer_grpc=bool(prefer_grpc),
+            grpc_port=grpc_port,
+            timeout=int(timeout),
+            check_compatibility=False,
+        )
+
+    def get_collection_info(self, *, collection_name: str) -> Dict[str, Any]:
+        info = self.client.get_collection(collection_name)
+        try:
+            return info.model_dump()
+        except Exception:
+            try:
+                return info.dict()
+            except Exception:
+                return {"collection": str(collection_name), "raw": str(info)}
+
+    def modify_collection_config(
+        self,
+        *,
+        collection_name: str,
+        hnsw_config: Optional[Dict[str, Any]] = None,
+        collection_params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+    ) -> bool:
+        """
+        Patch collection-level config via Qdrant update_collection.
+
+        Supported keys:
+        - hnsw_config: dict for HnswConfigDiff (e.g. on_disk, m, ef_construct, full_scan_threshold)
+        - collection_params: dict for CollectionParamsDiff (e.g. on_disk_payload)
+        """
+        from qdrant_client.http import models as m
+
+        hnsw_diff = m.HnswConfigDiff(**hnsw_config) if isinstance(hnsw_config, dict) else None
+        params_diff = (
+            m.CollectionParamsDiff(**collection_params)
+            if isinstance(collection_params, dict)
+            else None
+        )
+        if hnsw_diff is None and params_diff is None:
+            raise ValueError("No changes provided (pass hnsw_config and/or collection_params).")
+        return bool(
+            self.client.update_collection(
+                collection_name=str(collection_name),
+                hnsw_config=hnsw_diff,
+                collection_params=params_diff,
+                timeout=int(timeout) if timeout is not None else None,
+            )
+        )
+
+    def modify_collection_vector_config(
+        self,
+        *,
+        collection_name: str,
+        vectors: Dict[str, Dict[str, Any]],
+        timeout: Optional[int] = None,
+    ) -> bool:
+        """
+        Patch vector params under params.vectors[vector_name] using Qdrant update_collection.
+
+        Supported keys per vector:
+        - on_disk: bool
+        - hnsw_config: dict with optional keys: m, ef_construct, full_scan_threshold, on_disk
+        """
+        from qdrant_client.http import models as m
+
+        collection_name = str(collection_name)
+        info = self.client.get_collection(collection_name)
+        existing = set()
+        try:
+            existing = set((info.config.params.vectors or {}).keys())
+        except Exception:
+            existing = set()
+
+        missing = [str(k) for k in (vectors or {}).keys() if existing and str(k) not in existing]
+        if missing:
+            raise ValueError(
+                f"Vectors do not exist in collection '{collection_name}': {missing}. Existing: {sorted(existing)}"
+            )
+
+        ok = True
+        for name, cfg in (vectors or {}).items():
+            if not isinstance(cfg, dict):
+                raise ValueError(f"vectors['{name}'] must be a dict, got {type(cfg)}")
+            hnsw_cfg = cfg.get("hnsw_config")
+            hnsw_diff = m.HnswConfigDiff(**hnsw_cfg) if isinstance(hnsw_cfg, dict) else None
+            vectors_diff = {
+                str(name): m.VectorParamsDiff(
+                    on_disk=cfg.get("on_disk", None),
+                    hnsw_config=hnsw_diff,
+                )
+            }
+
+            ok = (
+                bool(
+                    self.client.update_collection(
+                        collection_name=collection_name,
+                        vectors_config=vectors_diff,
+                        timeout=int(timeout) if timeout is not None else None,
+                    )
+                )
+                and ok
+            )
+
+        return ok
+
+    def ensure_collection_all_on_disk(
+        self,
+        *,
+        collection_name: str,
+        timeout: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        """
+        Ensure:
+        - All existing named vectors have on_disk=True and hnsw_config.on_disk=True
+        - Collection hnsw_config.on_disk=True
+        - Collection params.on_disk_payload=True
+        Returns the post-update collection info (dict).
+        """
+        collection_name = str(collection_name)
+        info = self.client.get_collection(collection_name)
+        vectors = {}
+        try:
+            existing = list((info.config.params.vectors or {}).keys())
+        except Exception:
+            existing = []
+        for vname in existing:
+            vectors[str(vname)] = {"on_disk": True, "hnsw_config": {"on_disk": True}}
+
+        if vectors:
+            self.modify_collection_vector_config(
+                collection_name=collection_name, vectors=vectors, timeout=timeout
+            )
+
+        self.modify_collection_config(
+            collection_name=collection_name,
+            hnsw_config={"on_disk": True},
+            collection_params={"on_disk_payload": True},
+            timeout=timeout,
+        )
+
+        return self.get_collection_info(collection_name=collection_name)

visual_rag/retrieval/__init__.py

@@ -0,0 +1,19 @@
+"""
+Retrieval module - Search and retrieval strategies.
+
+Components:
+- TwoStageRetriever: Pooled prefetch → MaxSim reranking (our novel contribution)
+- SingleStageRetriever: Direct multi-vector or pooled search
+"""
+
+from visual_rag.retrieval.multi_vector import MultiVectorRetriever
+from visual_rag.retrieval.single_stage import SingleStageRetriever
+from visual_rag.retrieval.three_stage import ThreeStageRetriever
+from visual_rag.retrieval.two_stage import TwoStageRetriever
+
+__all__ = [
+    "TwoStageRetriever",
+    "SingleStageRetriever",
+    "MultiVectorRetriever",
+    "ThreeStageRetriever",
+]

visual_rag/retrieval/multi_vector.py

@@ -0,0 +1,222 @@
+import os
+from typing import Any, Dict, List, Optional
+from urllib.parse import urlparse
+
+from visual_rag.embedding.visual_embedder import VisualEmbedder
+from visual_rag.retrieval.single_stage import SingleStageRetriever
+from visual_rag.retrieval.three_stage import ThreeStageRetriever
+from visual_rag.retrieval.two_stage import TwoStageRetriever
+
+
+class MultiVectorRetriever:
+    @staticmethod
+    def _maybe_load_dotenv() -> None:
+        try:
+            from dotenv import load_dotenv
+        except ImportError:
+            return
+        if os.path.exists(".env"):
+            load_dotenv(".env")
+
+    def __init__(
+        self,
+        collection_name: str,
+        model_name: str = "vidore/colSmol-500M",
+        qdrant_url: Optional[str] = None,
+        qdrant_api_key: Optional[str] = None,
+        prefer_grpc: bool = False,
+        request_timeout: int = 120,
+        max_retries: int = 3,
+        retry_sleep: float = 0.5,
+        qdrant_client=None,
+        embedder: Optional[VisualEmbedder] = None,
+    ):
+        if qdrant_client is None:
+            self._maybe_load_dotenv()
+            try:
+                from qdrant_client import QdrantClient
+            except ImportError as e:
+                raise ImportError(
+                    "Qdrant client not installed. Install with: pip install visual-rag-toolkit[qdrant]"
+                ) from e
+
+            qdrant_url = (
+                qdrant_url
+                or os.getenv("SIGIR_QDRANT_URL")
+                or os.getenv("DEST_QDRANT_URL")
+                or os.getenv("QDRANT_URL")
+            )
+            if not qdrant_url:
+                raise ValueError(
+                    "QDRANT_URL is required (pass qdrant_url or set env var). "
+                    "You can also set DEST_QDRANT_URL to override."
+                )
+
+            qdrant_api_key = (
+                qdrant_api_key
+                or os.getenv("SIGIR_QDRANT_KEY")
+                or os.getenv("SIGIR_QDRANT_API_KEY")
+                or os.getenv("DEST_QDRANT_API_KEY")
+                or os.getenv("QDRANT_API_KEY")
+            )
+
+            grpc_port = None
+            if prefer_grpc:
+                try:
+                    if urlparse(qdrant_url).port == 6333:
+                        grpc_port = 6334
+                except Exception:
+                    grpc_port = None
+
+            def _make_client(use_grpc: bool):
+                return QdrantClient(
+                    url=qdrant_url,
+                    api_key=qdrant_api_key,
+                    prefer_grpc=bool(use_grpc),
+                    grpc_port=grpc_port,
+                    timeout=int(request_timeout),
+                    check_compatibility=False,
+                )
+
+            qdrant_client = _make_client(prefer_grpc)
+            if prefer_grpc:
+                try:
+                    _ = qdrant_client.get_collections()
+                except Exception as e:
+                    msg = str(e)
+                    if (
+                        "StatusCode.PERMISSION_DENIED" in msg
+                        or "http2 header with status: 403" in msg
+                    ):
+                        qdrant_client = _make_client(False)
+                    else:
+                        raise
+
+        self.client = qdrant_client
+        self.collection_name = collection_name
+        self.embedder = embedder or VisualEmbedder(model_name=model_name)
+
+        self._two_stage = TwoStageRetriever(
+            self.client,
+            collection_name=self.collection_name,
+            request_timeout=int(request_timeout),
+            max_retries=int(max_retries),
+            retry_sleep=float(retry_sleep),
+        )
+        self._three_stage = ThreeStageRetriever(
+            self.client,
+            collection_name=self.collection_name,
+            request_timeout=int(request_timeout),
+            max_retries=int(max_retries),
+            retry_sleep=float(retry_sleep),
+        )
+        self._single_stage = SingleStageRetriever(
+            self.client,
+            collection_name=self.collection_name,
+            request_timeout=int(request_timeout),
+        )
+
+    def build_filter(
+        self,
+        year: Optional[Any] = None,
+        source: Optional[str] = None,
+        district: Optional[str] = None,
+        filename: Optional[str] = None,
+        has_text: Optional[bool] = None,
+    ):
+        return self._two_stage.build_filter(
+            year=year,
+            source=source,
+            district=district,
+            filename=filename,
+            has_text=has_text,
+        )
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 10,
+        mode: str = "single_full",
+        prefetch_k: Optional[int] = None,
+        stage1_mode: str = "pooled_query_vs_tiles",
+        filter_obj=None,
+        return_embeddings: bool = False,
+    ) -> List[Dict[str, Any]]:
+        q = self.embedder.embed_query(query)
+        try:
+            import torch
+        except ImportError:
+            torch = None
+        if torch is not None and isinstance(q, torch.Tensor):
+            query_embedding = q.detach().cpu().numpy()
+        else:
+            query_embedding = q.numpy()
+
+        return self.search_embedded(
+            query_embedding=query_embedding,
+            top_k=top_k,
+            mode=mode,
+            prefetch_k=prefetch_k,
+            stage1_mode=stage1_mode,
+            filter_obj=filter_obj,
+            return_embeddings=return_embeddings,
+        )
+
+    def search_embedded(
+        self,
+        *,
+        query_embedding,
+        top_k: int = 10,
+        mode: str = "single_full",
+        prefetch_k: Optional[int] = None,
+        stage1_mode: str = "pooled_query_vs_tiles",
+        stage1_k: Optional[int] = None,
+        stage2_k: Optional[int] = None,
+        filter_obj=None,
+        return_embeddings: bool = False,
+    ) -> List[Dict[str, Any]]:
+        if mode == "single_full":
+            return self._single_stage.search(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                strategy="multi_vector",
+                filter_obj=filter_obj,
+            )
+
+        if mode == "single_tiles":
+            return self._single_stage.search(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                strategy="tiles_maxsim",
+                filter_obj=filter_obj,
+            )
+
+        if mode == "single_global":
+            return self._single_stage.search(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                strategy="pooled_global",
+                filter_obj=filter_obj,
+            )
+
+        if mode == "two_stage":
+            return self._two_stage.search_server_side(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                prefetch_k=prefetch_k,
+                filter_obj=filter_obj,
+                stage1_mode=stage1_mode,
+            )
+
+        if mode == "three_stage":
+            s1 = int(stage1_k) if stage1_k is not None else 1000
+            s2 = int(stage2_k) if stage2_k is not None else 300
+            return self._three_stage.search_server_side(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                stage1_k=s1,
+                stage2_k=s2,
+                filter_obj=filter_obj,
+            )
+
+        raise ValueError(f"Unknown mode: {mode}")

visual_rag/retrieval/single_stage.py

@@ -0,0 +1,126 @@
+"""
+Single-Stage Retrieval for Visual Document Search.
+
+Provides direct search without the two-stage complexity.
+Use when:
+- Collection is small (<10K documents)
+- Latency is not critical
+- Maximum accuracy is required
+"""
+
+import logging
+from typing import Any, Dict, List, Union
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class SingleStageRetriever:
+    """
+    Single-stage visual document retrieval using native Qdrant search.
+
+    Supports strategies:
+    - multi_vector: Native MaxSim on full embeddings (using="initial")
+    - tiles_maxsim: Native MaxSim between query tokens and tile vectors (using="mean_pooling")
+    - pooled_tile: Pooled query vs tile vectors (using="mean_pooling")
+    - pooled_global: Pooled query vs global pooled doc vector (using="global_pooling")
+
+    Args:
+        qdrant_client: Connected Qdrant client
+        collection_name: Name of the Qdrant collection
+
+    Example:
+        >>> retriever = SingleStageRetriever(client, "my_collection")
+        >>> results = retriever.search(query, top_k=10)
+    """
+
+    def __init__(
+        self,
+        qdrant_client,
+        collection_name: str,
+        request_timeout: int = 120,
+    ):
+        self.client = qdrant_client
+        self.collection_name = collection_name
+        self.request_timeout = int(request_timeout)
+
+    def search(
+        self,
+        query_embedding: Union[torch.Tensor, np.ndarray],
+        top_k: int = 10,
+        strategy: str = "multi_vector",
+        filter_obj=None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Single-stage search with configurable strategy.
+
+        Args:
+            query_embedding: Query embeddings [num_tokens, dim]
+            top_k: Number of results
+            strategy: "multi_vector", "tiles_maxsim", "pooled_tile", or "pooled_global"
+            filter_obj: Qdrant filter
+
+        Returns:
+            List of results with scores and metadata
+        """
+        query_np = self._to_numpy(query_embedding)
+
+        if strategy == "multi_vector":
+            # Native multi-vector MaxSim
+            vector_name = "initial"
+            query_vector = query_np.tolist()
+            logger.debug(f"🎯 Multi-vector search on '{vector_name}'")
+
+        elif strategy == "tiles_maxsim":
+            # Native multi-vector MaxSim against tile vectors
+            vector_name = "mean_pooling"
+            query_vector = query_np.tolist()
+            logger.debug(f"🎯 Tile MaxSim search on '{vector_name}'")
+
+        elif strategy == "pooled_tile":
+            # Tile-level pooled
+            vector_name = "mean_pooling"
+            query_pooled = query_np.mean(axis=0)
+            query_vector = query_pooled.tolist()
+            logger.debug(f"🔍 Tile-pooled search on '{vector_name}'")
+
+        elif strategy == "pooled_global":
+            # Global pooled vector (single vector)
+            vector_name = "global_pooling"
+            query_pooled = query_np.mean(axis=0)
+            query_vector = query_pooled.tolist()
+            logger.debug(f"🔍 Global-pooled search on '{vector_name}'")
+
+        else:
+            raise ValueError(f"Unknown strategy: {strategy}")
+
+        results = self.client.query_points(
+            collection_name=self.collection_name,
+            query=query_vector,
+            using=vector_name,
+            query_filter=filter_obj,
+            limit=top_k,
+            with_payload=True,
+            with_vectors=False,
+            timeout=self.request_timeout,
+        ).points
+
+        return [
+            {
+                "id": r.id,
+                "score": r.score,
+                "score_final": r.score,
+                "payload": r.payload,
+            }
+            for r in results
+        ]
+
+    def _to_numpy(self, embedding: Union[torch.Tensor, np.ndarray]) -> np.ndarray:
+        """Convert embedding to numpy array."""
+        if isinstance(embedding, torch.Tensor):
+            if embedding.dtype == torch.bfloat16:
+                return embedding.cpu().float().numpy()
+            return embedding.cpu().numpy()
+        return np.array(embedding, dtype=np.float32)