# semantic_search/vector_store.py import logging import numpy as np from typing import List, Dict, Any, Tuple, Optional from dataclasses import dataclass import json logger = logging.getLogger(__name__) @dataclass class Document: """Represents a document with its embedding and metadata.""" id: str text: str embedding: np.ndarray metadata: Dict[str, Any] def to_dict(self) -> Dict[str, Any]: """Convert to dictionary (excluding embedding for serialization).""" return { 'id': self.id, 'text': self.text, 'metadata': self.metadata } class VectorStore: """ In-memory vector storage with similarity search capabilities. Future versions can use Faiss, ChromaDB, or other vector databases. """ def __init__(self, dimension: int = 768): """ Initialize vector store. Args: dimension: Embedding dimension size """ self.dimension = dimension self.documents: List[Document] = [] self.embeddings: Optional[np.ndarray] = None self.index_built = False logger.info(f"Initialized VectorStore with dimension: {dimension}") def add_documents(self, ids: List[str], texts: List[str], embeddings: np.ndarray, metadata: Optional[List[Dict[str, Any]]] = None) -> int: """ Add documents to the vector store. Args: ids: Document IDs texts: Document texts embeddings: Document embeddings (N x dimension) metadata: Optional metadata for each document Returns: Number of documents added """ if len(ids) != len(texts) or len(ids) != embeddings.shape[0]: raise ValueError("Mismatched lengths for ids, texts, and embeddings") if metadata and len(metadata) != len(ids): raise ValueError("Metadata length doesn't match document count") # Add documents for i in range(len(ids)): doc = Document( id=ids[i], text=texts[i], embedding=embeddings[i], metadata=metadata[i] if metadata else {} ) self.documents.append(doc) # Rebuild index self._build_index() logger.info(f"Added {len(ids)} documents to vector store. Total: {len(self.documents)}") return len(ids) def _build_index(self): """Build or rebuild the embedding index.""" if not self.documents: self.embeddings = None self.index_built = False return # Stack all embeddings into a single array self.embeddings = np.vstack([doc.embedding for doc in self.documents]) self.index_built = True logger.debug(f"Built index with shape: {self.embeddings.shape}") def search(self, query_embedding: np.ndarray, top_k: int = 10, threshold: Optional[float] = None) -> List[Tuple[Document, float]]: """ Search for similar documents using cosine similarity. Args: query_embedding: Query embedding vector top_k: Number of results to return threshold: Optional similarity threshold (0-1) Returns: List of (Document, similarity_score) tuples """ if not self.index_built or self.embeddings is None: logger.warning("No documents in vector store") return [] # Ensure query is 2D if len(query_embedding.shape) == 1: query_embedding = query_embedding.reshape(1, -1) # Compute cosine similarities (assuming normalized embeddings) similarities = np.dot(self.embeddings, query_embedding.T).squeeze() # Apply threshold if specified if threshold is not None: valid_indices = np.where(similarities >= threshold)[0] if len(valid_indices) == 0: logger.info(f"No documents above threshold {threshold}") return [] similarities = similarities[valid_indices] valid_docs = [self.documents[i] for i in valid_indices] else: valid_docs = self.documents # Get top-k indices top_k = min(top_k, len(valid_docs)) if top_k == 0: return [] # Use argpartition for efficiency with large arrays if len(similarities) > top_k: top_indices = np.argpartition(similarities, -top_k)[-top_k:] top_indices = top_indices[np.argsort(similarities[top_indices])[::-1]] else: top_indices = np.argsort(similarities)[::-1] # Create results results = [] for idx in top_indices: doc = valid_docs[idx] if threshold else self.documents[idx] score = float(similarities[idx]) results.append((doc, score)) logger.info(f"Search returned {len(results)} results (top_k={top_k})") return results def hybrid_search(self, query_embedding: np.ndarray, keyword_scores: Dict[str, float], top_k: int = 10, alpha: float = 0.5) -> List[Tuple[Document, float]]: """ Hybrid search combining vector similarity and keyword scores. Args: query_embedding: Query embedding vector keyword_scores: Document ID to keyword relevance score mapping top_k: Number of results to return alpha: Weight for vector similarity (1-alpha for keyword score) Returns: List of (Document, combined_score) tuples """ if not self.index_built: logger.warning("No documents in vector store") return [] # Get vector similarities vector_results = self.search(query_embedding, top_k=len(self.documents)) # Combine scores combined_scores = [] for doc, vector_score in vector_results: keyword_score = keyword_scores.get(doc.id, 0.0) # Normalize keyword score to 0-1 range if needed if keyword_score > 1.0: keyword_score = keyword_score / max(keyword_scores.values()) combined_score = alpha * vector_score + (1 - alpha) * keyword_score combined_scores.append((doc, combined_score)) # Sort by combined score and return top-k combined_scores.sort(key=lambda x: x[1], reverse=True) results = combined_scores[:top_k] logger.info(f"Hybrid search returned {len(results)} results") return results def clear(self): """Clear all documents from the store.""" self.documents = [] self.embeddings = None self.index_built = False logger.info("Cleared vector store") def size(self) -> int: """Get number of documents in store.""" return len(self.documents) def get_by_id(self, doc_id: str) -> Optional[Document]: """Get document by ID.""" for doc in self.documents: if doc.id == doc_id: return doc return None def get_stats(self) -> Dict[str, Any]: """Get statistics about the vector store.""" stats = { 'num_documents': len(self.documents), 'dimension': self.dimension, 'index_built': self.index_built, 'memory_usage_mb': 0 } if self.embeddings is not None: # Estimate memory usage memory_bytes = self.embeddings.nbytes for doc in self.documents: memory_bytes += len(doc.text.encode('utf-8')) memory_bytes += len(json.dumps(doc.metadata).encode('utf-8')) stats['memory_usage_mb'] = memory_bytes / (1024 * 1024) return stats