AIPass/src/aipass/memory/apps/handlers/vector/embedder.py

# =================== AIPass ====================
# Name: embedder.py
# Description: Vector Embedding Handler
# Version: 0.2.0
# Created: 2025-11-16
# Modified: 2026-03-06
# =============================================

"""
Vector Embedding Handler

Generates semantic embeddings using sentence-transformers/all-MiniLM-L6-v2.
Implements production best practices from research.

Purpose:
    Convert text memories into 384-dimensional vectors for semantic search.
    Optimized for batch processing (100 lines during rollover).

Best Practices Applied:
    - Pre-sort by length (30% padding reduction)
    - Built-in normalization (L2 distance requirement)
    - GPU memory cleanup (prevent VRAM leaks)
    - Batch size optimization (64 GPU, 16 CPU)
    - Singleton pattern (model loaded once)

Dependencies (optional):
    - sentence-transformers
    - torch
"""

import logging
from typing import List, Dict, Any
from pathlib import Path

from aipass.prax.apps.modules.logger import get_system_logger

logger = get_system_logger()

# No service imports - handlers are pure workers (3-tier architecture)
# No module imports (handler independence)


# =============================================================================
# EMBEDDING SERVICE (Singleton)
# =============================================================================

class EmbeddingService:
    """
    Production-ready embedding service

    Implements best practices:
    - Batch size optimization (64 GPU, 16 CPU)
    - Pre-sorting by length (reduces padding waste 30%)
    - Built-in normalization (critical for L2 distance)
    - GPU memory cleanup (prevents VRAM leaks)
    """

    def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
        """
        Initialize embedding service

        Args:
            model_name: HuggingFace model identifier

        Raises:
            ImportError: If sentence-transformers or torch are not installed
        """
        # Late imports (heavy optional dependencies)
        try:
            import torch
            from sentence_transformers import SentenceTransformer
        except ImportError as e:
            raise ImportError(
                f"Embedding requires sentence-transformers and torch. "
                f"Install with: pip install sentence-transformers torch. "
                f"Original error: {e}"
            )

        self.model_name = model_name
        self.model = SentenceTransformer(model_name)

        # GPU optimization if available
        self.use_gpu = torch.cuda.is_available()
        if self.use_gpu:
            self.model = self.model.to('cuda')
            self.batch_size = 64
        else:
            self.batch_size = 16

        self.dimension = 384  # all-MiniLM-L6-v2 output dimension


    def encode_batch(self, texts: List[str]) -> Dict[str, Any]:
        """
        Encode batch of texts with all optimizations

        Best practices applied:
        1. Pre-sort by length (reduces padding waste)
        2. Batch processing (optimal batch size)
        3. Built-in normalization (L2 distance requirement)
        4. GPU cleanup (prevent VRAM leaks)

        Args:
            texts: List of text strings to encode

        Returns:
            Dict with embeddings and metadata
        """
        import torch

        if not texts:
            return {
                "embeddings": [],
                "count": 0,
                "dimension": self.dimension
            }

        # Pre-sort by length (reduces padding waste by 30%)
        sorted_pairs = sorted(enumerate(texts), key=lambda x: len(x[1]))
        sorted_indices, sorted_texts = zip(*sorted_pairs)

        # Encode with optimal settings
        embeddings = self.model.encode(
            sorted_texts,
            batch_size=self.batch_size,
            convert_to_tensor=False,  # Return numpy for Chroma
            normalize_embeddings=True,  # Critical for L2 distance
            show_progress_bar=False
        )

        # Restore original order
        ordered_embeddings = [None] * len(texts)
        for original_idx, sorted_idx in enumerate(sorted_indices):
            ordered_embeddings[sorted_idx] = embeddings[original_idx]

        # Cleanup GPU memory if used
        if self.use_gpu:
            torch.cuda.empty_cache()

        return {
            "embeddings": ordered_embeddings,
            "count": len(ordered_embeddings),
            "dimension": self.dimension
        }


# Global service instance (singleton pattern)
_embedding_service = None


def _get_service() -> EmbeddingService:
    """
    Get or create embedding service singleton

    Lazy initialization - model loaded on first use
    """
    global _embedding_service
    if _embedding_service is None:
        _embedding_service = EmbeddingService()
    return _embedding_service


# =============================================================================
# PUBLIC API
# =============================================================================

def encode_batch(texts: List[str]) -> Dict[str, Any]:
    """
    Encode batch of texts to embeddings

    This is the main public API. Delegates to singleton service
    to avoid reloading the model.

    Args:
        texts: List of text strings to encode

    Returns:
        Dict with embeddings and metadata

    Example:
        result = encode_batch(["memory 1", "memory 2"])
        if result['success']:
            embeddings = result['embeddings']
            # Each embedding is 384-dim numpy array
    """
    if not texts:
        return {
            'success': True,
            'embeddings': [],
            'count': 0,
            'message': 'No texts provided'
        }

    try:
        service = _get_service()
        result = service.encode_batch(texts)

        return {
            'success': True,
            **result
        }

    except Exception as e:
        return {
            'success': False,
            'error': f"Encoding failed: {e}"
        }


def encode_memories(memories: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Encode memory entries to embeddings

    Extracts text from memory entries and encodes them.
    Preserves original memory structure for metadata.

    Args:
        memories: List of memory entry dicts (from extraction)

    Returns:
        Dict with embeddings and original memories

    Example:
        memories = [{"content": "...", "timestamp": "..."}]
        result = encode_memories(memories)
        embeddings = result['embeddings']
        original = result['memories']
    """
    if not memories:
        return {
            'success': True,
            'embeddings': [],
            'memories': [],
            'count': 0,
            'message': 'No memories provided'
        }

    # Extract text content from memories
    texts = []
    for memory in memories:
        # Try common fields for text content
        text = (
            memory.get('content') or
            memory.get('text') or
            memory.get('message') or
            str(memory)  # Fallback to string representation
        )
        texts.append(text)

    # Encode texts
    encode_result = encode_batch(texts)

    if not encode_result['success']:
        return encode_result

    # Combine embeddings with original memories
    return {
        'success': True,
        'embeddings': encode_result['embeddings'],
        'memories': memories,
        'count': len(memories),
        'dimension': encode_result['dimension']
    }


def get_model_info() -> Dict[str, Any]:
    """
    Get embedding model information

    Returns:
        Dict with model metadata
    """
    try:
        service = _get_service()

        return {
            'success': True,
            'model_name': service.model_name,
            'dimension': service.dimension,
            'batch_size': service.batch_size,
            'gpu_enabled': service.use_gpu
        }
    except Exception as e:
        return {
            'success': False,
            'error': f"Failed to get model info: {e}"
        }