"""Qdrant database manager for hybrid (dense + sparse) vector storage.""" from __future__ import annotations from typing import Any from qdrant_client import QdrantClient, models from qdrant_client.models import ( Distance, PointStruct, SparseVector, SparseVectorParams, VectorParams, ) from src.config import ( EMBEDDING_DIMENSION, QDRANT_API_KEY, QDRANT_COLLECTION_NAME, QDRANT_TIMEOUT, QDRANT_URL, ) # Batch size for upsert operations _UPSERT_BATCH_SIZE = 64 def _build_client() -> QdrantClient: """Create a Qdrant client, preferring cloud when credentials are available.""" if QDRANT_URL and QDRANT_API_KEY and "your_" not in QDRANT_URL: print(f"[QDRANT] Connecting to cloud: {QDRANT_URL} (timeout={QDRANT_TIMEOUT}s)") return QdrantClient( url=QDRANT_URL, api_key=QDRANT_API_KEY, timeout=QDRANT_TIMEOUT, ) # Fallback: local persistent storage local_path = "data/qdrant_local" print(f"[QDRANT] Using local storage: {local_path}") return QdrantClient(path=local_path, timeout=QDRANT_TIMEOUT) class QdrantManager: """Manage a Qdrant collection that stores both dense and sparse vectors.""" def __init__( self, collection_name: str = QDRANT_COLLECTION_NAME, client: QdrantClient | None = None, ) -> None: self.collection_name = collection_name self.client = client or _build_client() # ------------------------------------------------------------------ # Collection lifecycle # ------------------------------------------------------------------ def init_collection(self, *, recreate: bool = False) -> None: """Create the collection if it doesn't exist. If *recreate* is ``True``, delete and recreate it. """ exists = self.client.collection_exists(self.collection_name) if exists and not recreate: print(f"[QDRANT] Collection '{self.collection_name}' already exists — skipping creation.") return if exists and recreate: self.client.delete_collection(self.collection_name) print(f"[QDRANT] Deleted existing collection '{self.collection_name}'.") self.client.create_collection( collection_name=self.collection_name, vectors_config={ "dense": VectorParams( size=EMBEDDING_DIMENSION, distance=Distance.COSINE, ), }, sparse_vectors_config={ "sparse": SparseVectorParams(), }, ) print(f"[QDRANT] Created collection '{self.collection_name}' (dense={EMBEDDING_DIMENSION}d + sparse).") # ------------------------------------------------------------------ # Upsert # ------------------------------------------------------------------ def upsert_chunks( self, *, chunk_ids: list[str], contents: list[str], metadatas: list[dict[str, Any]], dense_vectors: list[list[float]], sparse_vectors: list[dict[str, Any]], ) -> int: """Upload chunks with both dense and sparse vectors to Qdrant. Returns the number of points upserted. """ n = len(chunk_ids) if not (n == len(contents) == len(metadatas) == len(dense_vectors) == len(sparse_vectors)): raise ValueError("All input lists must have the same length.") total_upserted = 0 for start in range(0, n, _UPSERT_BATCH_SIZE): end = min(start + _UPSERT_BATCH_SIZE, n) points: list[PointStruct] = [] for i in range(start, end): # Qdrant requires integer or UUID point IDs. # We use a deterministic hash of the chunk_id string. point_id = _stable_id(chunk_ids[i]) payload = { "chunk_id": chunk_ids[i], "content": contents[i], **metadatas[i], } sv = sparse_vectors[i] point = PointStruct( id=point_id, vector={ "dense": dense_vectors[i], "sparse": SparseVector( indices=sv["indices"], values=sv["values"], ), }, payload=payload, ) points.append(point) self.client.upsert(collection_name=self.collection_name, points=points) total_upserted += len(points) print(f"[QDRANT] Upserted batch {start}–{end} ({len(points)} points)") return total_upserted # ------------------------------------------------------------------ # Search # ------------------------------------------------------------------ def dense_search( self, *, dense_vector: list[float], limit: int = 5, ) -> list[models.ScoredPoint]: """Search by dense vector only (cosine similarity).""" results = self.client.query_points( collection_name=self.collection_name, query=dense_vector, using="dense", limit=limit, ) return results.points def sparse_search( self, *, sparse_vector: dict[str, Any], limit: int = 5, ) -> list[models.ScoredPoint]: """Search by sparse (BM25) vector only.""" sv = SparseVector( indices=sparse_vector["indices"], values=sparse_vector["values"], ) results = self.client.query_points( collection_name=self.collection_name, query=sv, using="sparse", limit=limit, ) return results.points def hybrid_search( self, *, dense_vector: list[float], sparse_vector: dict[str, Any], limit: int = 5, dense_limit: int | None = None, sparse_limit: int | None = None, fusion_limit: int | None = None, ) -> list[models.ScoredPoint]: """Run a hybrid query using Qdrant's built-in RRF fusion.""" sv = SparseVector( indices=sparse_vector["indices"], values=sparse_vector["values"], ) dense_prefetch_limit = dense_limit or limit sparse_prefetch_limit = sparse_limit or limit final_limit = fusion_limit or limit results = self.client.query_points( collection_name=self.collection_name, prefetch=[ models.Prefetch( query=dense_vector, using="dense", limit=dense_prefetch_limit, ), models.Prefetch( query=sv, using="sparse", limit=sparse_prefetch_limit, ), ], query=models.FusionQuery(fusion=models.Fusion.RRF), limit=final_limit, ) return results.points # ------------------------------------------------------------------ # Info # ------------------------------------------------------------------ def delete_by_doc_id(self, doc_id: str) -> None: """Delete all points whose payload doc_id matches.""" self.client.delete( collection_name=self.collection_name, points_selector=models.FilterSelector( filter=models.Filter( must=[ models.FieldCondition( key="doc_id", match=models.MatchValue(value=doc_id), ) ] ) ), ) def count(self) -> int: """Return the number of points in the collection.""" info = self.client.get_collection(self.collection_name) return int(info.points_count or 0) def _stable_id(chunk_id: str) -> int: """Produce a positive 64-bit integer from a chunk_id string.""" import hashlib digest = hashlib.sha256(chunk_id.encode()).hexdigest() return int(digest[:16], 16)