vector_store.py

"""
Vector store interface and implementations for RAG.

This module provides pluggable vector store backends for storing and retrieving
chunk embeddings. All implementations are local-first (no external APIs).
"""

import json
import pickle
from abc import ABC, abstractmethod
from pathlib import Path
from typing import List, Dict, Any, Tuple, Optional

import numpy as np
from tqdm import tqdm

from logger import get_logger

logger = get_logger()


class VectorStore(ABC):
    """
    Abstract base class for vector stores.

    All implementations must be local-first (no external APIs required).
    """

    @abstractmethod
    def add(self, embeddings: np.ndarray, chunks: List[Dict[str, Any]], ids: List[str]):
        """
        Add embeddings and associated chunks to the store.

        Args:
            embeddings: (N, D) array of embedding vectors
            chunks: List of chunk dictionaries with metadata
            ids: List of unique chunk IDs
        """
        pass

    @abstractmethod
    def search(
        self,
        query_embedding: np.ndarray,
        top_k: int = 5,
        filters: Optional[Dict[str, Any]] = None
    ) -> List[Tuple[Dict[str, Any], float]]:
        """
        Search for similar chunks.

        Args:
            query_embedding: Query vector
            top_k: Number of results to return
            filters: Optional metadata filters (e.g., {"chapter": "Chapter 1"})

        Returns:
            List of (chunk, similarity_score) tuples, sorted by similarity
        """
        pass

    @abstractmethod
    def save(self, path: Path):
        """Save the vector store to disk."""
        pass

    @abstractmethod
    def load(self, path: Path):
        """Load the vector store from disk."""
        pass

    @abstractmethod
    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about the vector store."""
        pass


class InMemoryVectorStore(VectorStore):
    """
    Simple in-memory vector store using numpy for similarity search.

    Features:
    - Fast cosine similarity search
    - Metadata-based filtering
    - Persistent storage via pickle
    - No external dependencies beyond numpy

    Best for:
    - Small to medium datasets (< 100k chunks)
    - Development and testing
    - Fully local deployment
    """

    def __init__(self):
        """Initialize empty vector store."""
        self.embeddings = None  # (N, D) array
        self.chunks = []  # List of chunk dicts
        self.ids = []  # List of chunk IDs
        self.embedding_dim = None

    def add(self, embeddings: np.ndarray, chunks: List[Dict[str, Any]], ids: List[str]):
        """
        Add embeddings and chunks to the store.

        Args:
            embeddings: (N, D) array of embedding vectors
            chunks: List of chunk dictionaries
            ids: List of unique chunk IDs
        """
        if len(embeddings) != len(chunks) or len(embeddings) != len(ids):
            raise ValueError("Embeddings, chunks, and ids must have same length")

        # Initialize or validate embedding dimension
        if self.embedding_dim is None:
            self.embedding_dim = embeddings.shape[1]
        elif embeddings.shape[1] != self.embedding_dim:
            raise ValueError(f"Embedding dimension mismatch: {embeddings.shape[1]} != {self.embedding_dim}")

        # Append to existing data
        if self.embeddings is None:
            self.embeddings = embeddings
        else:
            self.embeddings = np.vstack([self.embeddings, embeddings])

        self.chunks.extend(chunks)
        self.ids.extend(ids)

    def search(
        self,
        query_embedding: np.ndarray,
        top_k: int = 5,
        filters: Optional[Dict[str, Any]] = None
    ) -> List[Tuple[Dict[str, Any], float]]:
        """
        Search for similar chunks using cosine similarity.

        Args:
            query_embedding: Query vector (D,)
            top_k: Number of results to return
            filters: Optional metadata filters

        Returns:
            List of (chunk, similarity_score) tuples
        """
        if self.embeddings is None or len(self.chunks) == 0:
            return []

        # Validate query embedding dimension
        if len(query_embedding) != self.embedding_dim:
            raise ValueError(f"Query embedding dimension mismatch: {len(query_embedding)} != {self.embedding_dim}")

        # Compute cosine similarities
        similarities = self._cosine_similarity_batch(query_embedding, self.embeddings)

        # Apply metadata filters
        if filters:
            valid_indices = self._apply_filters(filters)
            # Set invalid similarities to -inf
            mask = np.ones(len(similarities), dtype=bool)
            mask[valid_indices] = False
            similarities[mask] = -np.inf

        # Get top-k indices
        top_indices = np.argsort(similarities)[::-1][:top_k]

        # Filter out -inf scores (filtered out chunks)
        results = []
        for idx in top_indices:
            score = similarities[idx]
            if score == -np.inf:
                continue
            results.append((self.chunks[idx], float(score)))

        return results

    def save(self, path: Path):
        """
        Save vector store to disk using pickle.

        Args:
            path: Output file path (will create .pkl file)
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        data = {
            "embeddings": self.embeddings,
            "chunks": self.chunks,
            "ids": self.ids,
            "embedding_dim": self.embedding_dim,
        }

        with open(path, "wb") as f:
            pickle.dump(data, f)

    def load(self, path: Path):
        """
        Load vector store from disk.

        Args:
            path: Input file path (.pkl file)
        """
        path = Path(path)

        if not path.exists():
            raise FileNotFoundError(f"Vector store file not found: {path}")

        with open(path, "rb") as f:
            data = pickle.load(f)

        self.embeddings = data["embeddings"]
        self.chunks = data["chunks"]
        self.ids = data["ids"]
        self.embedding_dim = data["embedding_dim"]

    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about the vector store."""
        stats = {
            "num_chunks": len(self.chunks),
            "embedding_dim": self.embedding_dim,
        }

        # Count chunks by role
        if self.chunks:
            role_counts = {}
            for chunk in self.chunks:
                role = chunk.get("metadata", {}).get("pedagogical_role", "unknown")
                role_counts[role] = role_counts.get(role, 0) + 1
            stats["chunks_by_role"] = role_counts

            # Count chunks by section
            section_counts = {}
            for chunk in self.chunks:
                section = chunk.get("metadata", {}).get("section", "unknown")
                section_counts[section] = section_counts.get(section, 0) + 1
            stats["chunks_by_section"] = section_counts

        return stats

    def _cosine_similarity_batch(self, query: np.ndarray, embeddings: np.ndarray) -> np.ndarray:
        """
        Compute cosine similarity between query and batch of embeddings.

        Args:
            query: Query vector (D,)
            embeddings: Batch of vectors (N, D)

        Returns:
            Similarity scores (N,) in range [0, 1]
        """
        # Normalize query
        query_norm = query / (np.linalg.norm(query) + 1e-8)

        # Normalize embeddings
        embeddings_norm = embeddings / (np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-8)

        # Compute dot products
        similarities = np.dot(embeddings_norm, query_norm)

        # Normalize from [-1, 1] to [0, 1]
        similarities = (similarities + 1) / 2

        return similarities

    def _apply_filters(self, filters: Dict[str, Any]) -> List[int]:
        """
        Find indices of chunks matching metadata filters.

        Args:
            filters: Dictionary of metadata filters

        Returns:
            List of valid chunk indices
        """
        valid_indices = []

        for idx, chunk in enumerate(self.chunks):
            metadata = chunk.get("metadata", {})

            # Check if all filter conditions are met
            match = True
            for key, value in filters.items():
                if metadata.get(key) != value:
                    match = False
                    break

            if match:
                valid_indices.append(idx)

        return valid_indices


def build_vector_store_from_chunks(
    chunks_path: Path,
    embedding_engine,
    output_path: Path,
    show_progress: bool = True
) -> InMemoryVectorStore:
    """
    Build a vector store from a chunks JSON file.

    Args:
        chunks_path: Path to chunks.json or chunks_semantic.json
        embedding_engine: EmbeddingEngine instance for generating embeddings
        output_path: Path to save the vector store
        show_progress: Whether to show progress bar

    Returns:
        Populated InMemoryVectorStore
    """
    # Load chunks
    with open(chunks_path, "r", encoding="utf-8") as f:
        chunks = json.load(f)

    if not chunks:
        raise ValueError("No chunks found in file")

    logger.info(f"Building vector store from {len(chunks)} chunks...")

    # Extract chunk IDs and texts
    ids = [chunk["id"] for chunk in chunks]
    texts = [chunk["text"] for chunk in chunks]

    # Generate embeddings (uses caching automatically)
    embeddings = embedding_engine.embed_batch(
        texts,
        batch_size=32,
        show_progress=show_progress
    )

    # Create and populate vector store
    store = InMemoryVectorStore()
    store.add(embeddings, chunks, ids)

    # Save to disk
    store.save(output_path)

    logger.info(f"Vector store saved to {output_path}")
    logger.info(f"Total chunks indexed: {len(chunks)}")

    return store