Harshavard21/FinRAG
2
1"""2src/vectorstore/qdrant_store.py3================================4Qdrant vector store operations — collection management, upsert, search.5 6Running in LOCAL mode (no Docker/server needed):7 QdrantClient(path="./data/qdrant_store") stores everything to disk.8 This is perfectly fine for our scale (~15K chunks).9 10HNSW Index:11 m=16 — controls graph connectivity (higher = better recall, more memory)12 ef_construct=200 — build-time search depth (higher = better index quality)13 These values are production defaults used by major RAG systems.14 15Payload Filtering:16 Every chunk is stored with its metadata as a Qdrant "payload".17 This allows queries like:18 "only search in TCS documents"19 "only search FY2025 data"20 "only search table chunks"21 Without payload filtering, multi-company RAG would mix up contexts.22"""23 24from __future__ import annotations25 26from dataclasses import dataclass27from typing import Optional28 29import numpy as np30 31from qdrant_client import QdrantClient32from qdrant_client.models import (33 Distance,34 VectorParams,35 HnswConfigDiff,36 PointStruct,37 Filter,38 FieldCondition,39 MatchValue,40 MatchAny,41 ScoredPoint,42)43 44from config.settings import settings45from src.utils.logger import logger46from src.chunking.hierarchical_chunker import Chunk47 48 49@dataclass50class SearchResult:51 """A single result from vector search."""52 chunk_id: str53 content: str54 score: float55 company: str56 ticker: str57 source_file: str58 fiscal_year: str59 page_number: int60 section: str61 subsection: str62 content_type: str63 chunk_level: str64 parent_chunk_id: Optional[str] = None65 66 67class QdrantStore:68 """69 Manages the Qdrant vector store for the financial RAG system.70 Handles collection creation, chunk upsert, and similarity search.71 """72 73 def __init__(self):74 # Connect to Cloud if URL provided, else fallback to Local disk75 if settings.qdrant_url and settings.qdrant_api_key:76 self.client = QdrantClient(77 url=settings.qdrant_url,78 port=443,79 api_key=settings.qdrant_api_key,80 timeout=60,81 )82 mode_msg = f"Cloud | url: {settings.qdrant_url}"83 else:84 self.client = QdrantClient(85 path=str(settings.qdrant_path),86 timeout=60,87 )88 mode_msg = f"Local | path: {settings.qdrant_path}"89 90 self.collection_name = settings.qdrant_collection91 self.dim = settings.embedding_dim92 logger.info(f"QdrantStore initialized | collection: {self.collection_name} | {mode_msg}")93 94 # ---------------------------------------------------------------- #95 # Collection Management96 # ---------------------------------------------------------------- #97 98 def create_collection(self, recreate: bool = False) -> None:99 """100 Create the Qdrant collection with HNSW index.101 102 Args:103 recreate: If True, drop and recreate the collection (fresh start).104 Use this if you change embedding dims or want to re-index.105 """106 existing = [c.name for c in self.client.get_collections().collections]107 108 if self.collection_name in existing:109 if recreate:110 logger.warning(f"Dropping existing collection: {self.collection_name}")111 self.client.delete_collection(self.collection_name)112 else:113 logger.info(f"Collection '{self.collection_name}' already exists — skipping creation")114 return115 116 logger.info(f"Creating collection '{self.collection_name}' | dim={self.dim} | HNSW")117 118 self.client.create_collection(119 collection_name=self.collection_name,120 vectors_config=VectorParams(121 size=self.dim,122 distance=Distance.COSINE, # cosine similarity (L2 normalized → same as dot product)123 ),124 hnsw_config=HnswConfigDiff(125 m=settings.hnsw_m,126 ef_construct=settings.hnsw_ef_construct,127 full_scan_threshold=10000, # use exact search below this count128 on_disk=False, # keep index in RAM for speed129 ),130 )131 132 logger.info(f"Collection created successfully | HNSW m={settings.hnsw_m}, ef_construct={settings.hnsw_ef_construct}")133 134 def collection_exists(self) -> bool:135 existing = [c.name for c in self.client.get_collections().collections]136 return self.collection_name in existing137 138 def get_collection_info(self) -> dict:139 """Get collection stats (useful for debugging)."""140 if not self.collection_exists():141 return {"status": "does not exist", "vectors_count": 0}142 info = self.client.get_collection(self.collection_name)143 # qdrant-client API varies by version — handle gracefully144 try:145 # Newer qdrant-client uses points_count146 points = info.points_count or 0147 except AttributeError:148 points = 0149 try:150 vectors = info.vectors_count or points151 except AttributeError:152 vectors = points153 return {154 "vectors_count": vectors,155 "points_count": points,156 "status": str(info.status),157 }158 159 # ---------------------------------------------------------------- #160 # Upsert161 # ---------------------------------------------------------------- #162 163 def upsert_chunks(self, chunks: list[Chunk], embeddings: np.ndarray) -> None:164 """165 Store chunks with their embeddings in Qdrant.166 167 Args:168 chunks: List of Chunk objects169 embeddings: np.ndarray of shape (len(chunks), 768)170 """171 if len(chunks) != len(embeddings):172 raise ValueError(f"Chunk count ({len(chunks)}) != embedding count ({len(embeddings)})")173 174 points = []175 for i, (chunk, vector) in enumerate(zip(chunks, embeddings)):176 # Qdrant requires a numeric or UUID point ID177 # We use a hash of chunk_id converted to int178 point_id = int(chunk.chunk_id, 16) % (2**63) # SHA256 hex → int179 180 payload = {181 "chunk_id": chunk.chunk_id,182 "content": chunk.content,183 "content_type": chunk.content_type,184 "chunk_level": chunk.chunk_level,185 "parent_chunk_id": chunk.parent_chunk_id or "",186 "company": chunk.company,187 "ticker": chunk.ticker,188 "source_file": chunk.source_file,189 "fiscal_year": chunk.fiscal_year,190 "page_number": chunk.page_number,191 "section": chunk.section,192 "subsection": chunk.subsection,193 "token_count": chunk.token_count,194 "char_count": chunk.char_count,195 }196 197 points.append(PointStruct(198 id=point_id,199 vector=vector.tolist(),200 payload=payload,201 ))202 203 # Batch upsert in groups of 100 (Qdrant recommended batch size)204 batch_size = 100205 for i in range(0, len(points), batch_size):206 batch = points[i: i + batch_size]207 self.client.upsert(208 collection_name=self.collection_name,209 points=batch,210 wait=True,211 )212 213 logger.info(f"Upserted {len(points)} chunks into Qdrant")214 215 # ---------------------------------------------------------------- #216 # Search217 # ---------------------------------------------------------------- #218 219 def search(220 self,221 query_vector: np.ndarray,222 top_k: int = 30,223 company_filter: Optional[str | list[str]] = None,224 fiscal_year_filter: Optional[str] = None,225 content_type_filter: Optional[str] = None,226 chunk_level_filter: Optional[str] = "child", # default: search child chunks227 ) -> list[SearchResult]:228 """229 Dense vector similarity search with optional payload filters.230 231 Args:232 query_vector: L2-normalized query embedding (768-dim)233 top_k: Number of results to return234 company_filter: Filter by company name (str) or list of companies235 fiscal_year_filter: Filter by fiscal year (e.g. "FY2025")236 content_type_filter: Filter by "text", "table", or "ocr_text"237 chunk_level_filter: "child" | "atomic" | "parent" | None (no filter)238 239 Returns:240 List of SearchResult objects sorted by score descending241 """242 must_conditions = []243 244 # Company filter245 if company_filter:246 if isinstance(company_filter, list):247 must_conditions.append(FieldCondition(248 key="company",249 match=MatchAny(any=company_filter),250 ))251 else:252 must_conditions.append(FieldCondition(253 key="company",254 match=MatchValue(value=company_filter),255 ))256 257 # Fiscal year filter258 if fiscal_year_filter:259 must_conditions.append(FieldCondition(260 key="fiscal_year",261 match=MatchValue(value=fiscal_year_filter),262 ))263 264 # Content type filter265 if content_type_filter:266 must_conditions.append(FieldCondition(267 key="content_type",268 match=MatchValue(value=content_type_filter),269 ))270 271 # Chunk level filter (default: child chunks only for retrieval)272 if chunk_level_filter:273 must_conditions.append(FieldCondition(274 key="chunk_level",275 match=MatchValue(value=chunk_level_filter),276 ))277 278 query_filter = Filter(must=must_conditions) if must_conditions else None279 280 # qdrant-client >= 1.7: use query_points() instead of deprecated search()281 response = self.client.query_points(282 collection_name=self.collection_name,283 query=query_vector.tolist(),284 limit=top_k,285 query_filter=query_filter,286 score_threshold=settings.score_threshold,287 with_payload=True,288 )289 hits = response.points # list[ScoredPoint]290 291 results = []292 for hit in hits:293 p = hit.payload294 results.append(SearchResult(295 chunk_id=p.get("chunk_id", ""),296 content=p.get("content", ""),297 score=hit.score,298 company=p.get("company", ""),299 ticker=p.get("ticker", ""),300 source_file=p.get("source_file", ""),301 fiscal_year=p.get("fiscal_year", ""),302 page_number=p.get("page_number", 0),303 section=p.get("section", ""),304 subsection=p.get("subsection", ""),305 content_type=p.get("content_type", ""),306 chunk_level=p.get("chunk_level", ""),307 parent_chunk_id=p.get("parent_chunk_id") or None,308 ))309 310 return results311 312 def search_tables_only(313 self,314 query_vector: np.ndarray,315 top_k: int = 10,316 company_filter: Optional[str | list[str]] = None,317 ) -> list[SearchResult]:318 """Convenience: search only table (atomic) chunks."""319 return self.search(320 query_vector=query_vector,321 top_k=top_k,322 company_filter=company_filter,323 chunk_level_filter="atomic",324 )325 326 def fetch_parent_chunk(self, parent_chunk_id: str) -> Optional[str]:327 """328 Fetch a parent chunk's content given a parent_chunk_id.329 Used in parent-child retrieval: find child → return parent content to LLM.330 """331 # qdrant-client >= 1.7: scroll() uses query_filter (not scroll_filter)332 try:333 results = self.client.scroll(334 collection_name=self.collection_name,335 scroll_filter=Filter(336 must=[FieldCondition(337 key="chunk_id",338 match=MatchValue(value=parent_chunk_id),339 )]340 ),341 limit=1,342 with_payload=True,343 )344 points, _ = results345 except TypeError:346 # Newer API uses query_filter instead of scroll_filter347 results = self.client.scroll(348 collection_name=self.collection_name,349 query_filter=Filter(350 must=[FieldCondition(351 key="chunk_id",352 match=MatchValue(value=parent_chunk_id),353 )]354 ),355 limit=1,356 with_payload=True,357 )358 points, _ = results359 360 if points:361 return points[0].payload.get("content", "")362 return None363