embeddings.py

"""High-dimensional embedding generation with optional hyperbolic projection."""
import numpy as np
from typing import List, Dict, Any, Optional, Union
from dataclasses import dataclass
import torch
from tqdm import tqdm


@dataclass
class EmbeddedChunk:
    """A chunk with its embedding vector."""
    chunk_id: str
    content: str
    embedding: np.ndarray
    metadata: Dict[str, Any]
    
    # Optional hyperbolic embedding
    hyperbolic_embedding: Optional[np.ndarray] = None


class EmbeddingEngine:
    """
    High-dimensional embedding engine using Sentence Transformers.
    
    Supports:
    - Dense semantic embeddings (BGE, E5, etc.)
    - Optional hyperbolic projection for hierarchical data
    - Batch processing for 800+ page books
    """
    
    def __init__(
        self,
        model_name: str = "BAAI/bge-large-en-v1.5",
        device: str = "cuda",
        normalize: bool = True,
        use_hyperbolic: bool = False
    ):
        self.model_name = model_name
        self.device = device if torch.cuda.is_available() else "cpu"
        self.normalize = normalize
        self.use_hyperbolic = use_hyperbolic
        
        self._model = None
        self._hyperbolic_mapper = None
        
    @property
    def model(self):
        """Lazy load the embedding model."""
        if self._model is None:
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer(self.model_name, device=self.device)
            print(f"Loaded embedding model: {self.model_name}")
            print(f"Embedding dimension: {self._model.get_sentence_embedding_dimension()}")
        return self._model
    
    @property
    def dimension(self) -> int:
        """Get the embedding dimension."""
        return self.model.get_sentence_embedding_dimension()
    
    def embed_texts(
        self,
        texts: List[str],
        batch_size: int = 32,
        show_progress: bool = True
    ) -> np.ndarray:
        """
        Generate embeddings for a list of texts.
        
        Args:
            texts: List of text strings to embed
            batch_size: Batch size for processing
            show_progress: Whether to show progress bar
            
        Returns:
            numpy array of shape (len(texts), embedding_dim)
        """
        # Add instruction prefix for BGE models (improves retrieval)
        if "bge" in self.model_name.lower():
            texts = [f"Represent this cybersecurity document for retrieval: {t}" for t in texts]
        
        embeddings = self.model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=show_progress,
            normalize_embeddings=self.normalize,
            convert_to_numpy=True
        )
        
        return embeddings
    
    def embed_query(self, query: str) -> np.ndarray:
        """
        Generate embedding for a query (with query-specific prefix).
        
        Args:
            query: The search query
            
        Returns:
            numpy array of shape (embedding_dim,)
        """
        # BGE models use different prefixes for queries
        if "bge" in self.model_name.lower():
            query = f"Represent this question for searching relevant passages: {query}"
        
        embedding = self.model.encode(
            query,
            normalize_embeddings=self.normalize,
            convert_to_numpy=True
        )
        
        return embedding
    
    def project_to_hyperbolic(self, embeddings: np.ndarray) -> np.ndarray:
        """
        Project Euclidean embeddings to Poincaré ball (hyperbolic space).
        
        This is useful for hierarchical cybersecurity concepts where
        relationships form tree-like structures.
        """
        if not self.use_hyperbolic:
            return embeddings
        
        try:
            import geoopt
        except ImportError:
            print("Warning: geoopt not installed, skipping hyperbolic projection")
            return embeddings
        
        # Create Poincaré ball manifold
        ball = geoopt.PoincareBall()
        
        # Convert to tensor
        emb_tensor = torch.tensor(embeddings, dtype=torch.float32)
        
        # Project to hyperbolic space using exponential map
        # First normalize to unit ball
        norms = torch.norm(emb_tensor, dim=1, keepdim=True)
        normalized = emb_tensor / (norms + 1e-8)
        
        # Scale to fit in Poincaré ball (radius < 1)
        scaled = normalized * 0.9  # Leave margin from boundary
        
        # Apply exponential map from origin
        origin = torch.zeros_like(scaled[0])
        hyperbolic = ball.expmap(origin, scaled)
        
        return hyperbolic.numpy()


class ContextualEmbedder:
    """
    Embedder that adds contextual information to improve retrieval.
    
    Implements the "Contextual Breadcrumbs" strategy from the plan.
    """
    
    def __init__(self, base_engine: EmbeddingEngine):
        self.engine = base_engine
    
    def create_contextual_text(
        self,
        content: str,
        chapter: Optional[str] = None,
        section: Optional[str] = None,
        summary: Optional[str] = None,
        categories: Optional[List[str]] = None
    ) -> str:
        """
        Create enriched text with contextual information prepended.
        
        This "anchors" the embedding in the semantic space by adding
        structural context that helps with retrieval.
        """
        context_parts = []
        
        # Add hierarchical breadcrumbs
        if chapter:
            context_parts.append(f"Chapter: {chapter}")
        if section:
            context_parts.append(f"Section: {section}")
        
        # Add categories as tags
        if categories:
            context_parts.append(f"Topics: {', '.join(categories)}")
        
        # Add summary if available
        if summary:
            context_parts.append(f"Summary: {summary}")
        
        # Combine context with content
        if context_parts:
            context = " | ".join(context_parts)
            return f"[{context}]\n\n{content}"
        
        return content
    
    def embed_chunks(
        self,
        chunks: List[Any],  # List of Chunk objects
        batch_size: int = 32
    ) -> List[EmbeddedChunk]:
        """
        Embed chunks with contextual enrichment.
        """
        # Create contextual texts
        texts = []
        for chunk in chunks:
            contextual_text = self.create_contextual_text(
                content=chunk.content,
                chapter=chunk.chapter,
                section=chunk.section,
                summary=chunk.summary,
                categories=chunk.metadata.get('categories')
            )
            texts.append(contextual_text)
        
        # Generate embeddings
        embeddings = self.engine.embed_texts(texts, batch_size=batch_size)
        
        # Optionally project to hyperbolic space
        hyperbolic_embeddings = None
        if self.engine.use_hyperbolic:
            hyperbolic_embeddings = self.engine.project_to_hyperbolic(embeddings)
        
        # Create EmbeddedChunk objects
        embedded_chunks = []
        for i, chunk in enumerate(chunks):
            embedded = EmbeddedChunk(
                chunk_id=chunk.id,
                content=chunk.content,
                embedding=embeddings[i],
                metadata=chunk.to_dict(),
                hyperbolic_embedding=hyperbolic_embeddings[i] if hyperbolic_embeddings is not None else None
            )
            embedded_chunks.append(embedded)
        
        return embedded_chunks


class VectorArithmetic:
    """
    Enables semantic vector arithmetic like: King - Man + Woman = Queen
    
    For cybersecurity: Zero Trust + Cloud = Cloud-Native Identity
    """
    
    def __init__(self, engine: EmbeddingEngine):
        self.engine = engine
        self._concept_cache: Dict[str, np.ndarray] = {}
    
    def get_concept_vector(self, concept: str) -> np.ndarray:
        """Get or compute vector for a concept."""
        if concept not in self._concept_cache:
            self._concept_cache[concept] = self.engine.embed_query(concept)
        return self._concept_cache[concept]
    
    def compute(self, expression: str) -> np.ndarray:
        """
        Compute vector arithmetic from expression.
        
        Example: "Zero Trust + Cloud Architecture - On-Premise"
        """
        # Parse the expression
        expression = expression.strip()
        
        # Split by + and -
        parts = []
        current = ""
        sign = 1
        
        for char in expression + "+":  # Add + to flush last term
            if char in "+-":
                if current.strip():
                    parts.append((sign, current.strip()))
                sign = 1 if char == "+" else -1
                current = ""
            else:
                current += char
        
        # Compute result vector
        result = np.zeros(self.engine.dimension)
        for sign, concept in parts:
            vec = self.get_concept_vector(concept)
            result += sign * vec
        
        # Normalize result
        result = result / (np.linalg.norm(result) + 1e-8)
        
        return result
    
    def analogy(self, a: str, b: str, c: str) -> np.ndarray:
        """
        Compute analogy: A is to B as C is to ?
        
        Example: "SQL Injection" is to "Web Application" as "?" is to "Database"
        Formula: B - A + C
        """
        vec_a = self.get_concept_vector(a)
        vec_b = self.get_concept_vector(b)
        vec_c = self.get_concept_vector(c)
        
        result = vec_b - vec_a + vec_c
        result = result / (np.linalg.norm(result) + 1e-8)
        
        return result