vector_store.py

"""Vector database with hybrid search support."""
import json
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
import numpy as np


@dataclass
class SearchResult:
    """A search result with score and metadata."""
    chunk_id: str
    content: str
    score: float
    metadata: Dict[str, Any]
    parent_content: Optional[str] = None  # For parent-document retrieval


class VectorStore:
    """
    ChromaDB-based vector store with hybrid search capabilities.
    
    Features:
    - Dense vector similarity search
    - Metadata filtering (by chapter, severity, CVE, etc.)
    - Parent-document retrieval
    - Reciprocal Rank Fusion for hybrid results
    """
    
    def __init__(
        self,
        collection_name: str = "cybersecurity_book",
        persist_directory: str = "./vector_db",
        distance_metric: str = "cosine"
    ):
        self.collection_name = collection_name
        self.persist_directory = persist_directory
        self.distance_metric = distance_metric
        
        self._client = None
        self._collection = None
        self._parent_store: Dict[str, str] = {}  # parent_id -> content
        
    @property
    def client(self):
        """Lazy load ChromaDB client."""
        if self._client is None:
            import chromadb
            from chromadb.config import Settings
            
            self._client = chromadb.PersistentClient(
                path=self.persist_directory,
                settings=Settings(anonymized_telemetry=False)
            )
        return self._client
    
    @property
    def collection(self):
        """Get or create the collection."""
        if self._collection is None:
            self._collection = self.client.get_or_create_collection(
                name=self.collection_name,
                metadata={"hnsw:space": self.distance_metric}
            )
        return self._collection
    
    def add_chunks(
        self,
        chunk_ids: List[str],
        embeddings: np.ndarray,
        contents: List[str],
        metadatas: List[Dict[str, Any]],
        batch_size: int = 100
    ):
        """
        Add chunks to the vector store.
        
        Args:
            chunk_ids: Unique IDs for each chunk
            embeddings: numpy array of embeddings
            contents: Text content of each chunk
            metadatas: Metadata for each chunk
            batch_size: Batch size for insertion
        """
        # Store parent content separately
        for i, meta in enumerate(metadatas):
            if meta.get('is_parent', False):
                self._parent_store[chunk_ids[i]] = contents[i]
        
        # Filter out parent chunks from vector store (we only search children)
        child_indices = [i for i, m in enumerate(metadatas) if not m.get('is_parent', False)]
        
        if not child_indices:
            return
        
        child_ids = [chunk_ids[i] for i in child_indices]
        child_embeddings = embeddings[child_indices]
        child_contents = [contents[i] for i in child_indices]
        child_metadatas = [self._clean_metadata(metadatas[i]) for i in child_indices]
        
        # Add in batches
        for i in range(0, len(child_ids), batch_size):
            end_idx = min(i + batch_size, len(child_ids))
            
            self.collection.add(
                ids=child_ids[i:end_idx],
                embeddings=child_embeddings[i:end_idx].tolist(),
                documents=child_contents[i:end_idx],
                metadatas=child_metadatas[i:end_idx]
            )
        
        # Save parent store
        self._save_parent_store()
        
        print(f"Added {len(child_ids)} child chunks to vector store")
        print(f"Stored {len(self._parent_store)} parent chunks for retrieval")
    
    def _clean_metadata(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
        """Clean metadata for ChromaDB (only supports str, int, float, bool)."""
        cleaned = {}
        for key, value in metadata.items():
            if isinstance(value, (str, int, float, bool)):
                cleaned[key] = value
            elif isinstance(value, list):
                # Convert lists to comma-separated strings
                cleaned[key] = ",".join(str(v) for v in value)
            elif value is None:
                cleaned[key] = ""
            else:
                cleaned[key] = str(value)
        return cleaned
    
    def _save_parent_store(self):
        """Save parent store to disk."""
        parent_file = Path(self.persist_directory) / "parent_store.json"
        parent_file.parent.mkdir(parents=True, exist_ok=True)
        with open(parent_file, 'w') as f:
            json.dump(self._parent_store, f)
    
    def _load_parent_store(self):
        """Load parent store from disk."""
        parent_file = Path(self.persist_directory) / "parent_store.json"
        if parent_file.exists():
            with open(parent_file, 'r') as f:
                self._parent_store = json.load(f)
    
    def search(
        self,
        query_embedding: np.ndarray,
        n_results: int = 10,
        where: Optional[Dict[str, Any]] = None,
        include_parent: bool = True
    ) -> List[SearchResult]:
        """
        Search for similar chunks.
        
        Args:
            query_embedding: Query vector
            n_results: Number of results to return
            where: Metadata filter (e.g., {"chapter": "Network Security"})
            include_parent: Whether to include parent content in results
            
        Returns:
            List of SearchResult objects
        """
        # Load parent store if needed
        if include_parent and not self._parent_store:
            self._load_parent_store()
        
        results = self.collection.query(
            query_embeddings=[query_embedding.tolist()],
            n_results=n_results,
            where=where,
            include=["documents", "metadatas", "distances"]
        )
        
        search_results = []
        
        for i in range(len(results['ids'][0])):
            chunk_id = results['ids'][0][i]
            content = results['documents'][0][i]
            metadata = results['metadatas'][0][i]
            distance = results['distances'][0][i]
            
            # Convert distance to similarity score
            score = 1 - distance if self.distance_metric == "cosine" else 1 / (1 + distance)
            
            # Get parent content if available
            parent_content = None
            if include_parent:
                parent_id = metadata.get('parent_id', '')
                if parent_id and parent_id in self._parent_store:
                    parent_content = self._parent_store[parent_id]
            
            search_results.append(SearchResult(
                chunk_id=chunk_id,
                content=content,
                score=score,
                metadata=metadata,
                parent_content=parent_content
            ))
        
        return search_results
    
    def hybrid_search(
        self,
        query_embedding: np.ndarray,
        query_text: str,
        n_results: int = 10,
        dense_weight: float = 0.7,
        where: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """
        Hybrid search combining dense and sparse (keyword) retrieval.
        
        Uses Reciprocal Rank Fusion (RRF) to merge results.
        """
        # Dense search
        dense_results = self.search(
            query_embedding=query_embedding,
            n_results=n_results * 2,  # Get more for fusion
            where=where,
            include_parent=False
        )
        
        # Sparse search (keyword matching via ChromaDB's where_document)
        try:
            sparse_results = self.collection.query(
                query_texts=[query_text],
                n_results=n_results * 2,
                where=where,
                include=["documents", "metadatas", "distances"]
            )
            
            sparse_search_results = []
            for i in range(len(sparse_results['ids'][0])):
                sparse_search_results.append(SearchResult(
                    chunk_id=sparse_results['ids'][0][i],
                    content=sparse_results['documents'][0][i],
                    score=1.0,  # Keyword matches are binary
                    metadata=sparse_results['metadatas'][0][i]
                ))
        except:
            sparse_search_results = []
        
        # Reciprocal Rank Fusion
        fused = self._reciprocal_rank_fusion(
            dense_results=dense_results,
            sparse_results=sparse_search_results,
            dense_weight=dense_weight,
            k=60  # RRF constant
        )
        
        # Load parent content for final results
        self._load_parent_store()
        for result in fused[:n_results]:
            parent_id = result.metadata.get('parent_id', '')
            if parent_id in self._parent_store:
                result.parent_content = self._parent_store[parent_id]
        
        return fused[:n_results]
    
    def _reciprocal_rank_fusion(
        self,
        dense_results: List[SearchResult],
        sparse_results: List[SearchResult],
        dense_weight: float,
        k: int = 60
    ) -> List[SearchResult]:
        """
        Merge results using Reciprocal Rank Fusion.
        
        RRF score = sum(1 / (k + rank))
        """
        scores: Dict[str, float] = {}
        result_map: Dict[str, SearchResult] = {}
        
        # Score dense results
        for rank, result in enumerate(dense_results, start=1):
            rrf_score = dense_weight * (1 / (k + rank))
            scores[result.chunk_id] = scores.get(result.chunk_id, 0) + rrf_score
            result_map[result.chunk_id] = result
        
        # Score sparse results
        for rank, result in enumerate(sparse_results, start=1):
            rrf_score = (1 - dense_weight) * (1 / (k + rank))
            scores[result.chunk_id] = scores.get(result.chunk_id, 0) + rrf_score
            if result.chunk_id not in result_map:
                result_map[result.chunk_id] = result
        
        # Sort by fused score
        sorted_ids = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)
        
        # Update scores in results
        fused = []
        for chunk_id in sorted_ids:
            result = result_map[chunk_id]
            result.score = scores[chunk_id]
            fused.append(result)
        
        return fused
    
    def filter_by_category(
        self,
        query_embedding: np.ndarray,
        category: str,
        n_results: int = 10
    ) -> List[SearchResult]:
        """
        Search within a specific cybersecurity category.
        
        Categories: network_security, web_security, cryptography, malware,
                   authentication, incident_response, compliance, cloud_security
        """
        where = {"categories": {"$contains": category}}
        return self.search(query_embedding, n_results, where)
    
    def filter_by_severity(
        self,
        query_embedding: np.ndarray,
        severity: str,
        n_results: int = 10
    ) -> List[SearchResult]:
        """
        Search for chunks with specific severity level.
        
        Severities: critical, high, medium, low
        """
        where = {"severity": severity}
        return self.search(query_embedding, n_results, where)
    
    def get_chunk_count(self) -> int:
        """Get the number of chunks in the store."""
        return self.collection.count()
    
    def clear(self):
        """Clear all data from the store."""
        self.client.delete_collection(self.collection_name)
        self._collection = None
        self._parent_store = {}
        print(f"Cleared collection: {self.collection_name}")