""" Vector store module for Qdrant integration. Handles embedding generation, storage, and retrieval from Qdrant. """ import logging import os from typing import List, Optional from qdrant_client import QdrantClient from qdrant_client.models import Distance, VectorParams, PointStruct, HasIdCondition from sentence_transformers import SentenceTransformer from metadata_schema import Chunk from config import QDRANT_CONFIG, LLM_CONFIG logger = logging.getLogger(__name__) class VectorStore: """ Manages embeddings and vector storage in Qdrant. Handles both in-memory and network-based Qdrant instances. """ def __init__(self): """Initialize vector store client.""" self.client = self._init_qdrant_client() # Using sentence-transformers for embeddings (all-MiniLM-L6-v2) self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2") self.vector_size = 384 # all-MiniLM-L6-v2 produces 384-dimensional vectors def _init_qdrant_client(self) -> QdrantClient: """ Initialize Qdrant client based on configuration. Returns: QdrantClient instance """ mode = QDRANT_CONFIG.get("mode", "memory") try: if mode == "memory": # In-memory Qdrant for development logger.info("Initializing Qdrant in-memory mode") return QdrantClient(":memory:") elif mode == "local": # Local persistent storage path = QDRANT_CONFIG.get("path", "qdrant_storage") logger.info(f"Initializing Qdrant in local persistent mode at: {path}") # Ensure directory exists os.makedirs(path, exist_ok=True) return QdrantClient(path=path) elif mode == "url": # Network Qdrant url = QDRANT_CONFIG.get("url", "localhost:6333") api_key = QDRANT_CONFIG.get("api_key") logger.info(f"Initializing Qdrant with URL: {url}") return QdrantClient( url=url, api_key=api_key, timeout=30, ) else: logger.warning(f"Unknown Qdrant mode: {mode}, defaulting to memory") return QdrantClient(":memory:") except Exception as e: logger.error(f"Failed to initialize Qdrant: {str(e)}") # Fallback to memory mode return QdrantClient(":memory:") def create_collection(self, collection_name: str, vector_size: int = None) -> bool: """ Create a collection in Qdrant. Args: collection_name: Name of the collection vector_size: Size of vectors (default from config) Returns: True if successful, False otherwise """ if vector_size is None: vector_size = self.vector_size try: # Check if collection exists collections = self.client.get_collections() if any(c.name == collection_name for c in collections.collections): logger.info(f"Collection '{collection_name}' already exists") return True # Create new collection self.client.create_collection( collection_name=collection_name, vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE), ) logger.info(f"Created collection: {collection_name}") return True except Exception as e: logger.error(f"Error creating collection {collection_name}: {str(e)}") return False def embed_text(self, text: str) -> Optional[List[float]]: """ Generate embedding for text using SentenceTransformer. Args: text: Text to embed Returns: Embedding vector or None if error """ try: # Truncate if too long (max ~512 tokens for sentence-transformers) if len(text) > 30000: text = text[:30000] embedding = self.embedding_model.encode(text, convert_to_tensor=False) return embedding.tolist() except Exception as e: logger.error(f"Error generating embedding with SentenceTransformer: {str(e)}") return None def store_chunks( self, chunks: List[Chunk], collection_name: str, ) -> bool: """ Store chunks with embeddings in Qdrant. Args: chunks: List of Chunk objects collection_name: Target collection name Returns: True if successful """ try: # Ensure collection exists if not self.create_collection(collection_name): logger.error(f"Failed to create collection {collection_name}") return False # Generate embeddings and prepare points points = [] for chunk in chunks: # Generate embedding embedding = self.embed_text(chunk.text) if not embedding: logger.warning(f"Failed to embed chunk {chunk.id}") continue # Create point with metadata payload point = PointStruct( id=self._hash_id(chunk.id), vector=embedding, payload=chunk.to_qdrant_payload(), ) points.append(point) if not points: logger.warning(f"No points to store in {collection_name}") return True # Upload points to Qdrant self.client.upsert( collection_name=collection_name, points=points, ) logger.info(f"Stored {len(points)} chunks in collection {collection_name}") return True except Exception as e: logger.error(f"Error storing chunks in {collection_name}: {str(e)}") return False def search_with_filter( self, collection_name: str, query_embedding: List[float], access_roles: List[str], top_k: int = 5, score_threshold: float = 0.5, ) -> List[dict]: """ Search collection with RBAC filter. CRITICAL: This ensures only chunks accessible to the user are returned. Args: collection_name: Collection to search query_embedding: Query embedding vector access_roles: Roles the user has (determines what they can access) top_k: Number of results to return score_threshold: Minimum similarity score Returns: List of matching chunks with metadata """ try: from qdrant_client.models import Filter, FieldCondition, MatchAny # Build native Qdrant RBAC filter # Checks if chunk's access_roles field contains any of the user's roles rbac_filter = Filter( must=[ FieldCondition( key="access_roles", match=MatchAny(any=access_roles) ) ] ) # qdrant-client >= 1.14 uses query_points; legacy .search() was removed. query_response = self.client.query_points( collection_name=collection_name, query=query_embedding, query_filter=rbac_filter, limit=top_k, score_threshold=score_threshold, with_payload=True, ) results = getattr(query_response, "points", None) or [] filtered_results = [] for scored_point in results: payload = scored_point.payload or {} filtered_results.append({ "id": scored_point.id, "score": scored_point.score, "source_document": payload.get("source_document", "unknown"), "collection": payload.get("collection", "unknown"), "access_roles": payload.get("access_roles", []), "section_title": payload.get("section_title", ""), "subsection_title": payload.get("subsection_title", ""), "page_number": payload.get("page_number", 0), "chunk_type": payload.get("chunk_type", "text"), "text": payload.get("text", ""), "parent_chunk_id": payload.get("parent_chunk_id", ""), "parent_summary": payload.get("parent_summary", ""), }) logger.info( f"Retrieved {len(filtered_results)} chunks from {collection_name} " f"after RBAC filtering (user roles: {access_roles})" ) return filtered_results[:top_k] except Exception as e: logger.error(f"Error searching collection {collection_name}: {str(e)}") return [] def search_by_text( self, collection_name: str, query_text: str, access_roles: List[str], top_k: int = 5, score_threshold: float = 0.5, ) -> List[dict]: """ Search by text query (convenience wrapper). Args: collection_name: Collection to search query_text: Query text access_roles: User's accessible roles top_k: Number of results score_threshold: Minimum score Returns: List of matching chunks """ # Embed query query_embedding = self.embed_text(query_text) if not query_embedding: logger.error("Failed to embed query") return [] # Search with RBAC filter return self.search_with_filter( collection_name=collection_name, query_embedding=query_embedding, access_roles=access_roles, top_k=top_k, score_threshold=score_threshold, ) def list_collections(self) -> List[str]: """ Get list of all collections in vector store. Returns: List of collection names """ try: collections = self.client.get_collections() return [c.name for c in collections.collections] except Exception as e: logger.error(f"Error listing collections: {str(e)}") return [] def delete_collection(self, collection_name: str) -> bool: """ Delete a collection. Args: collection_name: Collection to delete Returns: True if successful """ try: self.client.delete_collection(collection_name=collection_name) logger.info(f"Deleted collection: {collection_name}") return True except Exception as e: logger.error(f"Error deleting collection {collection_name}: {str(e)}") return False def get_collection_stats(self, collection_name: str) -> Optional[dict]: """ Get statistics about a collection. Args: collection_name: Collection name Returns: Dictionary with collection stats, or zeros if the collection does not exist in Qdrant yet (e.g. not ingested). None only on unexpected errors. """ try: if not self.client.collection_exists(collection_name=collection_name): return { "name": collection_name, "points_count": 0, "vectors_count": 0, } info = self.client.get_collection(collection_name=collection_name) # Qdrant REST CollectionInfo has no `name` (we already have it) or top-level # `vectors_count`; use points_count and indexed_vectors_count. # Use points_count as the definitive total document count points = info.points_count if info.points_count is not None else 0 # indexed_vectors_count shows how many have been HNSW-indexed (can be 0 initially) indexed = info.indexed_vectors_count # For the summary 'vectors_count', we prefer the total points if indexing is still 0 vectors_count = indexed if indexed is not None and indexed > 0 else points return { "name": collection_name, "points_count": points, "vectors_count": vectors_count, } except Exception as e: logger.error(f"Error getting collection stats: {str(e)}") return None @staticmethod def _hash_id(text_id: str) -> int: """ Convert string ID to integer hash for Qdrant. Args: text_id: Text ID Returns: Integer hash """ return abs(hash(text_id)) % (2**63) # Global vector store instance _vector_store = None def get_vector_store() -> VectorStore: """ Get singleton vector store instance. Returns: VectorStore instance """ global _vector_store if _vector_store is None: _vector_store = VectorStore() return _vector_store