src/embedding_manager.py

"""
Embedding management system using sentence-transformers with caching and optimization.
"""

import os
import pickle
import hashlib
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any, Union
import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
import time

from .error_handler import EmbeddingError, ResourceError
from .cache_manager import CacheManager


class EmbeddingManager:
    """Manages document embeddings with caching and batch processing."""
    
    def __init__(self, config: Dict[str, Any], cache_manager: Optional[CacheManager] = None):
        self.config = config
        self.model_config = config.get("models", {}).get("embedding", {})
        self.cache_config = config.get("cache", {})
        
        self.model_name = self.model_config.get("name", "sentence-transformers/all-MiniLM-L6-v2")
        self.max_seq_length = self.model_config.get("max_seq_length", 256)
        self.batch_size = self.model_config.get("batch_size", 32)
        self.device = self._get_device()
        
        self.model: Optional[SentenceTransformer] = None
        self.cache_manager = cache_manager
        self._model_loaded = False
        
        # Performance tracking
        self.stats = {
            "embeddings_generated": 0,
            "cache_hits": 0,
            "total_time": 0,
            "batch_count": 0
        }
    
    def _get_device(self) -> str:
        """Determine the best device for computation."""
        device_config = self.model_config.get("device", "auto")
        
        if device_config == "auto":
            if torch.cuda.is_available():
                return "cuda"
            elif torch.backends.mps.is_available():  # Apple Silicon
                return "mps"
            else:
                return "cpu"
        else:
            return device_config
    
    def _load_model(self) -> None:
        """Lazy load the sentence transformer model."""
        if self._model_loaded and self.model is not None:
            return
        
        try:
            print(f"Loading embedding model: {self.model_name}")
            start_time = time.time()
            
            self.model = SentenceTransformer(self.model_name)
            
            # Set device
            if self.device != "cpu":
                self.model = self.model.to(self.device)
            
            # Set max sequence length
            if hasattr(self.model, 'max_seq_length'):
                self.model.max_seq_length = self.max_seq_length
            
            load_time = time.time() - start_time
            print(f"Model loaded in {load_time:.2f}s on device: {self.device}")
            
            self._model_loaded = True
            
        except Exception as e:
            raise EmbeddingError(f"Failed to load embedding model: {str(e)}") from e
    
    def generate_embeddings(
        self, 
        texts: List[str], 
        show_progress: bool = True
    ) -> np.ndarray:
        """
        Generate embeddings for a list of texts with caching.
        
        Args:
            texts: List of text strings to embed
            show_progress: Whether to show progress bar
            
        Returns:
            Array of embeddings with shape (len(texts), embedding_dim)
        """
        if not texts:
            return np.array([])
        
        start_time = time.time()
        
        # Check cache for existing embeddings
        cached_embeddings, missing_indices, missing_texts = self._check_cache(texts)
        
        # Generate embeddings for missing texts
        if missing_texts:
            self._load_model()
            new_embeddings = self._generate_batch_embeddings(missing_texts, show_progress)
            
            # Cache new embeddings
            self._cache_embeddings(missing_texts, new_embeddings)
        else:
            new_embeddings = np.array([])
        
        # Combine cached and new embeddings
        all_embeddings = self._combine_embeddings(texts, cached_embeddings, missing_indices, new_embeddings)
        
        # Update stats
        generation_time = time.time() - start_time
        self.stats["total_time"] += generation_time
        self.stats["embeddings_generated"] += len(texts)
        
        return all_embeddings
    
    def _check_cache(self, texts: List[str]) -> Tuple[Dict[str, np.ndarray], List[int], List[str]]:
        """Check cache for existing embeddings."""
        cached_embeddings = {}
        missing_indices = []
        missing_texts = []
        
        if not self.cache_manager:
            return cached_embeddings, list(range(len(texts))), texts
        
        for i, text in enumerate(texts):
            cache_key = self._get_cache_key(text)
            cached_embedding = self.cache_manager.get(f"embedding_{cache_key}")
            
            if cached_embedding is not None:
                cached_embeddings[text] = cached_embedding
                self.stats["cache_hits"] += 1
            else:
                missing_indices.append(i)
                missing_texts.append(text)
        
        return cached_embeddings, missing_indices, missing_texts
    
    def _generate_batch_embeddings(self, texts: List[str], show_progress: bool) -> np.ndarray:
        """Generate embeddings in batches."""
        try:
            embeddings = []
            
            # Process in batches
            batches = [texts[i:i + self.batch_size] for i in range(0, len(texts), self.batch_size)]
            
            if show_progress and len(batches) > 1:
                batches = tqdm(batches, desc="Generating embeddings")
            
            for batch in batches:
                try:
                    # Generate embeddings for batch
                    batch_embeddings = self.model.encode(
                        batch,
                        convert_to_numpy=True,
                        show_progress_bar=False,
                        batch_size=len(batch)
                    )
                    
                    embeddings.append(batch_embeddings)
                    self.stats["batch_count"] += 1
                    
                except Exception as e:
                    raise EmbeddingError(f"Failed to generate embeddings for batch: {str(e)}") from e
            
            if not embeddings:
                return np.array([])
            
            # Combine all batch embeddings
            all_embeddings = np.vstack(embeddings)
            return all_embeddings
            
        except torch.cuda.OutOfMemoryError as e:
            raise ResourceError(
                "GPU memory insufficient for embedding generation. "
                "Try reducing batch_size or using CPU."
            ) from e
        except Exception as e:
            raise EmbeddingError(f"Failed to generate embeddings: {str(e)}") from e
    
    def _cache_embeddings(self, texts: List[str], embeddings: np.ndarray) -> None:
        """Cache generated embeddings."""
        if not self.cache_manager or embeddings.size == 0:
            return
        
        for text, embedding in zip(texts, embeddings):
            cache_key = self._get_cache_key(text)
            self.cache_manager.set(f"embedding_{cache_key}", embedding)
    
    def _combine_embeddings(
        self, 
        original_texts: List[str],
        cached_embeddings: Dict[str, np.ndarray],
        missing_indices: List[int],
        new_embeddings: np.ndarray
    ) -> np.ndarray:
        """Combine cached and newly generated embeddings."""
        if not original_texts:
            return np.array([])
        
        # Get embedding dimension
        if new_embeddings.size > 0:
            embedding_dim = new_embeddings.shape[1]
        elif cached_embeddings:
            embedding_dim = next(iter(cached_embeddings.values())).shape[0]
        else:
            # Fallback - load model to get dimension
            self._load_model()
            sample_embedding = self.model.encode(["sample"], convert_to_numpy=True)
            embedding_dim = sample_embedding.shape[1]
        
        # Initialize result array
        result = np.zeros((len(original_texts), embedding_dim))
        
        # Fill in cached embeddings
        for i, text in enumerate(original_texts):
            if text in cached_embeddings:
                result[i] = cached_embeddings[text]
        
        # Fill in new embeddings
        if new_embeddings.size > 0:
            for i, original_idx in enumerate(missing_indices):
                result[original_idx] = new_embeddings[i]
        
        return result
    
    def _get_cache_key(self, text: str) -> str:
        """Generate cache key for text."""
        # Include model name and config in hash for cache invalidation
        cache_input = f"{self.model_name}_{self.max_seq_length}_{text}"
        return hashlib.md5(cache_input.encode()).hexdigest()
    
    def get_embedding_dimension(self) -> int:
        """Get the dimension of embeddings."""
        self._load_model()
        
        # Generate a sample embedding to get dimensions
        sample_embedding = self.model.encode(["sample"], convert_to_numpy=True)
        return sample_embedding.shape[1]
    
    def compute_similarity(self, query_embedding: np.ndarray, doc_embeddings: np.ndarray) -> np.ndarray:
        """Compute cosine similarity between query and document embeddings."""
        if query_embedding.ndim == 1:
            query_embedding = query_embedding.reshape(1, -1)
        
        # Normalize vectors
        query_norm = query_embedding / np.linalg.norm(query_embedding, axis=1, keepdims=True)
        doc_norm = doc_embeddings / np.linalg.norm(doc_embeddings, axis=1, keepdims=True)
        
        # Compute cosine similarity
        similarities = np.dot(query_norm, doc_norm.T).flatten()
        return similarities
    
    def clear_cache(self) -> None:
        """Clear embedding cache."""
        if self.cache_manager:
            # Clear all embedding entries
            keys_to_remove = []
            for key in self.cache_manager._memory_cache.keys():
                if key.startswith("embedding_"):
                    keys_to_remove.append(key)
            
            for key in keys_to_remove:
                self.cache_manager.delete(key)
    
    def get_stats(self) -> Dict[str, Any]:
        """Get performance statistics."""
        stats = self.stats.copy()
        
        if stats["embeddings_generated"] > 0:
            stats["avg_time_per_embedding"] = stats["total_time"] / stats["embeddings_generated"]
        else:
            stats["avg_time_per_embedding"] = 0
        
        if stats["batch_count"] > 0:
            stats["avg_batch_size"] = stats["embeddings_generated"] / stats["batch_count"]
        else:
            stats["avg_batch_size"] = 0
        
        stats["cache_hit_rate"] = (
            stats["cache_hits"] / (stats["cache_hits"] + stats["embeddings_generated"]) 
            if (stats["cache_hits"] + stats["embeddings_generated"]) > 0 else 0
        )
        
        stats["model_loaded"] = self._model_loaded
        stats["device"] = self.device
        
        return stats
    
    def warmup(self) -> None:
        """Warm up the model with a sample embedding."""
        self._load_model()
        
        # Generate a sample embedding to warm up the model
        sample_texts = ["This is a sample text for model warmup."]
        self.model.encode(sample_texts, convert_to_numpy=True, show_progress_bar=False)
        
        print("Embedding model warmed up successfully")
    
    def unload_model(self) -> None:
        """Unload the model to free memory."""
        if self.model is not None:
            del self.model
            self.model = None
            self._model_loaded = False
            
            # Clear GPU cache if using CUDA
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            print("Embedding model unloaded")