Spaces:
Running
Running
| import os | |
| import re | |
| import uuid | |
| from typing import List, Dict, Any, Optional | |
| from qdrant_client import QdrantClient | |
| from qdrant_client.models import ( | |
| Distance, | |
| VectorParams, | |
| PointStruct, | |
| Filter, | |
| FieldCondition, | |
| MatchValue, | |
| PayloadSchemaType | |
| ) | |
| from sentence_transformers import SentenceTransformer, CrossEncoder | |
| from rank_bm25 import BM25Okapi | |
| from core.config import settings | |
| from core.logger import logger | |
| class VectorStore: | |
| _instance = None | |
| def __new__(cls): | |
| if cls._instance is None: | |
| cls._instance = super(VectorStore, cls).__new__(cls) | |
| try: | |
| cls._instance.client = QdrantClient( | |
| url=settings.QDRANT_URL, | |
| api_key=settings.QDRANT_API_KEY, | |
| timeout=5.0 | |
| ) | |
| cls._instance.client.get_collections() | |
| logger.info(f"Connected to online Qdrant cloud database at {settings.QDRANT_URL}") | |
| except Exception as e: | |
| logger.warning( | |
| f"Remote Qdrant Cloud cluster unreachable ({e}). " | |
| f"Falling back to local in-memory Qdrant instance." | |
| ) | |
| cls._instance.client = QdrantClient(location=":memory:") | |
| cls._instance.collection_name = "ecommerce_chunks" | |
| cls._instance._ensure_collection() | |
| logger.info(f"Initializing SentenceTransformer embedding model on CPU ({settings.EMBEDDING_MODEL_NAME})") | |
| cls._instance.embedding_model = SentenceTransformer(settings.EMBEDDING_MODEL_NAME) | |
| cls._instance.embedding_model.to('cpu') | |
| logger.info(f"Initializing CrossEncoder reranker model on CPU ({settings.RERANKER_MODEL_NAME})") | |
| cls._instance.reranker = CrossEncoder(settings.RERANKER_MODEL_NAME) | |
| return cls._instance | |
| def _ensure_collection(self): | |
| try: | |
| collections = self.client.get_collections().collections | |
| exists = any(c.name == self.collection_name for c in collections) | |
| if not exists: | |
| logger.info(f"Creating Qdrant collection: {self.collection_name}") | |
| self.client.create_collection( | |
| collection_name=self.collection_name, | |
| vectors_config=VectorParams(size=384, distance=Distance.COSINE) | |
| ) | |
| # Ensure payload index on session_id for fast isolated filtering | |
| try: | |
| self.client.create_payload_index( | |
| collection_name=self.collection_name, | |
| field_name="session_id", | |
| field_schema=PayloadSchemaType.KEYWORD | |
| ) | |
| except Exception: | |
| pass # Index might already exist | |
| except Exception as e: | |
| logger.error(f"Error ensuring Qdrant collection: {str(e)}") | |
| raise | |
| def add_chunks(self, chunks: List[Dict[str, Any]], session_id: Optional[str] = None): | |
| """ | |
| Add chunks to Qdrant collection with session isolation and parent-document linkage. | |
| """ | |
| if not chunks: | |
| return | |
| active_session = session_id or "global" | |
| logger.info(f"Generating embeddings and indexing {len(chunks)} chunks in Qdrant (Session: {active_session})") | |
| texts = [chunk["document"] for chunk in chunks] | |
| # Generate dense embeddings | |
| embeddings = self.embedding_model.encode(texts, convert_to_numpy=True, show_progress_bar=False) | |
| points = [] | |
| for i, chunk in enumerate(chunks): | |
| point_id = chunk["id"] | |
| if len(point_id) < 36: | |
| point_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, point_id)) | |
| chunk_session = chunk.get("session_id") or active_session | |
| payload = { | |
| "document": chunk["document"], | |
| "raw_text": chunk.get("raw_text", chunk["document"]), | |
| "source": chunk.get("source", ""), | |
| "chunk_type": chunk.get("chunk_type", "child"), | |
| "created_at": chunk.get("created_at", ""), | |
| "scraped_timestamp": chunk.get("scraped_timestamp", ""), | |
| "sequence": chunk.get("sequence", None), | |
| "parent_id": chunk.get("parent_id", None), | |
| "parent_content": chunk.get("parent_content", chunk["document"]), | |
| "session_id": chunk_session, | |
| "structured_data": chunk.get("structured_data", {}) | |
| } | |
| points.append(PointStruct( | |
| id=point_id, | |
| vector=embeddings[i].tolist(), | |
| payload=payload | |
| )) | |
| self.client.upsert( | |
| collection_name=self.collection_name, | |
| points=points | |
| ) | |
| logger.info(f"Successfully indexed {len(points)} chunks in Qdrant for session '{active_session}'") | |
| def search( | |
| self, | |
| query_text: str, | |
| k: int = 5, | |
| session_id: Optional[str] = None | |
| ) -> List[Dict[str, Any]]: | |
| """ | |
| Hybrid retrieval with dense vectors, BM25, Cross-Encoder reranking, and parent document resolution. | |
| """ | |
| logger.info(f"Hybrid SOTA Search for: '{query_text[:50]}' (Session: {session_id or 'all'})") | |
| try: | |
| query_vector = self.embedding_model.encode( | |
| query_text, | |
| convert_to_numpy=True, | |
| show_progress_bar=False | |
| ).tolist() | |
| candidate_pool_size = max(k * 4, 15) | |
| query_filter = None | |
| if session_id and session_id != "global": | |
| query_filter = Filter(must=[ | |
| FieldCondition(key="session_id", match=MatchValue(value=session_id)) | |
| ]) | |
| response = self.client.query_points( | |
| collection_name=self.collection_name, | |
| query=query_vector, | |
| query_filter=query_filter, | |
| limit=candidate_pool_size | |
| ) | |
| hits = response.points | |
| if not hits and query_filter is not None: | |
| logger.info("No session-specific chunks found; falling back to broader corpus") | |
| response = self.client.query_points( | |
| collection_name=self.collection_name, | |
| query=query_vector, | |
| limit=candidate_pool_size | |
| ) | |
| hits = response.points | |
| if not hits: | |
| return [] | |
| candidates = [] | |
| for hit in hits: | |
| payload = hit.payload or {} | |
| candidates.append({ | |
| "id": hit.id, | |
| "document": payload.get("document", ""), | |
| "raw_text": payload.get("raw_text", payload.get("document", "")), | |
| "parent_content": payload.get("parent_content", payload.get("document", "")), | |
| "parent_id": payload.get("parent_id") or hit.id, | |
| "metadata": { | |
| "source": payload.get("source", ""), | |
| "chunk_type": payload.get("chunk_type", ""), | |
| "created_at": payload.get("created_at", ""), | |
| "scraped_timestamp": payload.get("scraped_timestamp", ""), | |
| "sequence": payload.get("sequence", None), | |
| "parent_id": payload.get("parent_id", None), | |
| "session_id": payload.get("session_id", ""), | |
| "structured_data": payload.get("structured_data", {}) | |
| }, | |
| "dense_score": float(hit.score) | |
| }) | |
| tokenized_corpus = [ | |
| re.findall(r'\w+', (c["document"] + " " + str(c["metadata"].get("structured_data", ""))).lower()) | |
| for c in candidates | |
| ] | |
| tokenized_query = re.findall(r'\w+', query_text.lower()) | |
| if tokenized_corpus and tokenized_query: | |
| bm25 = BM25Okapi(tokenized_corpus) | |
| bm25_scores = bm25.get_scores(tokenized_query) | |
| bm25_ranked_indices = sorted(range(len(candidates)), key=lambda i: bm25_scores[i], reverse=True) | |
| bm25_ranks = {candidates[idx]["id"]: rank for rank, idx in enumerate(bm25_ranked_indices)} | |
| dense_ranked_indices = sorted(range(len(candidates)), key=lambda i: candidates[i]["dense_score"], reverse=True) | |
| dense_ranks = {candidates[idx]["id"]: rank for rank, idx in enumerate(dense_ranked_indices)} | |
| for c in candidates: | |
| cid = c["id"] | |
| d_rank = dense_ranks.get(cid, len(candidates)) | |
| s_rank = bm25_ranks.get(cid, len(candidates)) | |
| c["rrf_score"] = (1.0 / (60.0 + d_rank)) + (1.0 / (60.0 + s_rank)) | |
| candidates.sort(key=lambda x: x.get("rrf_score", 0.0), reverse=True) | |
| logger.info(f"Applied BM25 + Dense RRF fusion over {len(candidates)} candidates") | |
| if self.reranker: | |
| try: | |
| pairs = [(query_text, c["document"][:800]) for c in candidates] | |
| rerank_scores = self.reranker.predict(pairs) | |
| for i, score in enumerate(rerank_scores): | |
| candidates[i]["rerank_score"] = float(score) | |
| candidates.sort(key=lambda x: x.get("rerank_score", -999.0), reverse=True) | |
| logger.info("Successfully re-ranked candidates using Cross-Encoder") | |
| except Exception as e: | |
| logger.error(f"Cross-encoder reranking error: {e}") | |
| seen_parent_ids = set() | |
| parent_results = [] | |
| for c in candidates: | |
| parent_id = c["parent_id"] | |
| if parent_id in seen_parent_ids: | |
| continue | |
| seen_parent_ids.add(parent_id) | |
| full_parent_doc = c["parent_content"] or c["document"] | |
| parent_results.append({ | |
| "id": c["id"], | |
| "parent_id": parent_id, | |
| "document": full_parent_doc, | |
| "child_match_snippet": c["document"][:300], | |
| "metadata": c["metadata"], | |
| "score": c.get("rerank_score", c.get("rrf_score", c["dense_score"])) | |
| }) | |
| if len(parent_results) >= k: | |
| break | |
| logger.info(f"Retrieved and resolved {len(parent_results)} parent documents after SOTA re-ranking") | |
| return parent_results | |
| except Exception as e: | |
| logger.error(f"Qdrant SOTA search error: {str(e)}") | |
| return [] | |