cosma-backend 0.1.0__py3-none-any.whl

cosma_backend/embedder/embedder.py

@@ -0,0 +1,637 @@
+#!/usr/bin/env python
+"""
+@File    :   embedder.py
+@Time    :   2025/07/14
+@Author  :
+@Version :   1.0
+@Contact :
+@License :
+@Desc    :   Embedding generation for semantic search
+"""
+
+from datetime import datetime, timezone
+import os
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Any
+
+# Import AI libraries
+import litellm
+import numpy as np
+from backend.models import File
+import logging
+from backend.logging import sm
+from backend.models.status import ProcessingStatus
+
+# Configure structured logger
+logger = logging.getLogger(__name__)
+
+
+class EmbedderError(Exception):
+    """Base exception for embedder errors."""
+
+
+class EmbeddingProviderError(EmbedderError):
+    """Exception for embedding provider-specific errors."""
+
+
+class BaseEmbedder(ABC):
+    """Abstract base class for text embedders."""
+
+    def __init__(self, model_name: str, dimensions: int) -> None:
+        """
+        Initialize embedder with model specifications.
+
+        Args:
+            model_name: Name of the embedding model
+            dimensions: Dimension of the output embeddings
+        """
+        self.model_name = model_name
+        self.dimensions = dimensions
+        logger.info(sm("Initializing embedder", model=model_name, dimensions=dimensions))
+
+    @abstractmethod
+    def embed_text(self, text: str | list[str]) -> np.ndarray:
+        """
+        Generate embeddings for text input.
+
+        Args:
+            text: Single text or list of texts to embed
+
+        Returns:
+            Numpy array of embeddings (single vector or matrix)
+        """
+    
+    async def embed_text_async(self, text: str | list[str]) -> np.ndarray:
+        """
+        Async version of embed_text that runs in a thread pool.
+        
+        Args:
+            text: Single text or list of texts to embed
+
+        Returns:
+            Numpy array of embeddings (single vector or matrix)
+        """
+        # Check if this is an OnlineEmbedder with direct async support
+        if hasattr(self, '_embed_text_async'):
+            return await self._embed_text_async(text)
+        
+        # Fallback to thread pool for legacy implementations
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, self.embed_text, text)
+
+    @abstractmethod
+    def is_available(self) -> bool:
+        """Check if this embedder is available for use."""
+
+    def _validate_text(self, text: str | list[str]) -> list[str]:
+        """Validate and normalize text input."""
+        texts = [text] if isinstance(text, str) else text
+
+        # Filter out empty texts
+        valid_texts = [t for t in texts if t and t.strip()]
+
+        if not valid_texts:
+            msg = "No valid text provided for embedding"
+            raise ValueError(msg)
+
+        return valid_texts
+
+
+class OnlineEmbedder(BaseEmbedder):
+    """Embedder using online models via LiteLLM (OpenAI API)."""
+
+    def __init__(self, model: str | None = None, api_key: str | None = None, dimensions: int | None = None) -> None:
+        """
+        Initialize online embedder.
+
+        Args:
+            model: Model name (default: text-embedding-3-small)
+            api_key: API key (default from env)
+            dimensions: Embedding dimensions (default: 512 for efficiency)
+        """
+        # Default to text-embedding-3-small with 512 dimensions for efficiency
+        self.model = model or os.getenv("EMBEDDING_MODEL", "text-embedding-3-small")
+        self.configured_dimensions = dimensions or int(os.getenv("EMBEDDING_DIMENSIONS", "512"))
+
+        # Initialize base class
+        super().__init__(model_name=self.model, dimensions=self.configured_dimensions)
+
+        # Set API key if provided
+        if api_key:
+            os.environ["OPENAI_API_KEY"] = api_key
+
+        # Validate model compatibility
+        if self.model == "text-embedding-3-small":
+            if not (512 <= self.configured_dimensions <= 1536):
+                msg = f"Dimensions must be between 512 and 1536 for {self.model}"
+                raise ValueError(msg)
+
+        logger.info(sm("Online embedder initialized",
+                   model=self.model,
+                   dimensions=self.configured_dimensions))
+
+    def is_available(self) -> bool:
+        """Check if online embedder is available."""
+        return bool(os.getenv("OPENAI_API_KEY"))
+
+    def embed_text(self, text: str | list[str]) -> np.ndarray:
+        """Generate embeddings using online model."""
+        texts = self._validate_text(text)
+
+        logger.debug(sm("Generating embeddings",
+                    model=self.model,
+                    num_texts=len(texts),
+                    dimensions=self.configured_dimensions))
+
+        try:
+            # Call litellm embedding endpoint
+            response = litellm.embedding(
+                model=self.model,
+                input=texts,
+                dimensions=self.configured_dimensions,  # Only for models that support it
+                timeout=60,
+                max_retries=2
+            )
+
+            # Extract embeddings from response
+            embeddings = []
+            for item in response.data:
+                embeddings.append(item["embedding"])
+
+            # Convert to numpy array
+            embeddings_array = np.array(embeddings, dtype=np.float32)
+
+            # Return single vector if input was single text
+            if isinstance(text, str):
+                return embeddings_array[0]
+
+            return embeddings_array
+
+        except Exception as e:
+            error_msg = f"Online embedding generation failed: {e!s}"
+            logger.exception(sm(error_msg, model=self.model))
+            raise EmbeddingProviderError(error_msg)
+    
+    async def _embed_text_async(self, text: str | list[str]) -> np.ndarray:
+        """Truly async embedding generation using litellm async API."""
+        texts = self._validate_text(text)
+
+        logger.debug(sm("Generating embeddings async",
+                    model=self.model,
+                    num_texts=len(texts),
+                    dimensions=self.configured_dimensions))
+
+        try:
+            # Call async litellm embedding endpoint
+            response = await litellm.aembedding(
+                model=self.model,
+                input=texts,
+                dimensions=self.configured_dimensions,
+                timeout=60,
+                max_retries=2
+            )
+
+            # Extract embeddings from response
+            embeddings = []
+            for item in response.data:
+                embeddings.append(item["embedding"])
+
+            # Convert to numpy array
+            embeddings_array = np.array(embeddings, dtype=np.float32)
+
+            # Return single vector if input was single text
+            if isinstance(text, str):
+                return embeddings_array[0]
+
+            return embeddings_array
+
+        except Exception as e:
+            error_msg = f"Async online embedding generation failed: {e!s}"
+            logger.exception(sm(error_msg, model=self.model))
+            raise EmbeddingProviderError(error_msg)
+
+
+class LocalEmbedder(BaseEmbedder):
+    """Embedder using local models via sentence-transformers."""
+
+    def __init__(self, model_name: str | None = None, dimensions: int | None = None) -> None:
+        """
+        Initialize local embedder.
+
+        Args:
+            model_name: Model name (default: Qwen/Qwen3-Embedding-0.6B)
+            dimensions: Embedding dimensions (default: 768 for efficiency)
+        """
+        # Import sentence-transformers lazily
+        try:
+            from sentence_transformers import SentenceTransformer
+            self.sentence_transformers_available = True
+        except ImportError:
+            self.sentence_transformers_available = False
+            logger.warning(sm("sentence-transformers not installed, local embeddings unavailable"))
+
+        # Default to Qwen model with 768 dimensions for efficiency
+        self.model_name = model_name or os.getenv("LOCAL_EMBEDDING_MODEL", "Qwen/Qwen3-Embedding-0.6B")
+        self.configured_dimensions = dimensions or int(os.getenv("LOCAL_EMBEDDING_DIMENSIONS", "768"))
+
+        # Initialize base class
+        super().__init__(model_name=self.model_name, dimensions=self.configured_dimensions)
+
+        # Validate model compatibility
+        if "Qwen3-Embedding" in self.model_name:
+            if not (32 <= self.configured_dimensions <= 1024):
+                msg = f"Dimensions must be between 32 and 1024 for {self.model_name}"
+                raise ValueError(msg)
+
+        # Initialize model if available
+        self.model = None
+        if self.sentence_transformers_available:
+            try:
+                logger.info(sm("Loading local embedding model",
+                           model=self.model_name,
+                           dimensions=self.configured_dimensions))
+                # This will automatically download from HuggingFace if model doesn't exist locally
+                self.model = SentenceTransformer(self.model_name)
+                logger.info(sm("Local embedder initialized",
+                           model=self.model_name,
+                           dimensions=self.configured_dimensions))
+            except Exception as e:
+                logger.exception(sm("Failed to load local model",
+                            model=self.model_name,
+                            error=str(e)))
+                logger.info(sm("Model will be downloaded from HuggingFace on first use"))
+
+    def is_available(self) -> bool:
+        """Check if local embedder is available."""
+        return self.sentence_transformers_available and self.model is not None
+
+    def embed_text(self, text: str | list[str]) -> np.ndarray:
+        """Generate embeddings using local model."""
+        if not self.is_available():
+            msg = "Local embedder not available"
+            raise EmbeddingProviderError(msg)
+
+        texts = self._validate_text(text)
+
+        logger.debug(sm("Generating local embeddings",
+                    model=self.model_name,
+                    num_texts=len(texts),
+                    dimensions=self.configured_dimensions))
+
+        try:
+            # Generate embeddings
+            embeddings = self.model.encode(
+                texts,
+                normalize_embeddings=True,
+                show_progress_bar=False
+            )
+
+            # Truncate to configured dimensions if needed
+            if embeddings.shape[1] > self.configured_dimensions:
+                embeddings = embeddings[:, :self.configured_dimensions]
+
+            # Convert to float32
+            embeddings = embeddings.astype(np.float32)
+
+            # Return single vector if input was single text
+            if isinstance(text, str):
+                return embeddings[0]
+
+            return embeddings
+
+        except Exception as e:
+            error_msg = f"Local embedding generation failed: {e!s}"
+            logger.exception(sm(error_msg, model=self.model_name))
+            raise EmbeddingProviderError(error_msg)
+
+
+class AutoEmbedder:
+    """
+    Automatic embedder that selects the best available provider.
+
+    Tries providers in order of preference:
+    1. User-specified provider
+    2. Online models (OpenAI)
+    3. Local models (fallback)
+    """
+
+    def __init__(self, preferred_provider: str | None = None) -> None:
+        """
+        Initialize auto embedder.
+
+        Args:
+            preferred_provider: Preferred provider ('local', 'online')
+        """
+        self.preferred_provider = preferred_provider or os.getenv("EMBEDDING_PROVIDER", "local")
+        logger.debug(sm("Preferred provider", og=preferred_provider, provider=self.preferred_provider))
+        self.embedders = {}
+
+        logger.info(sm("AutoEmbedder initializing",
+                   preferred_provider=self.preferred_provider))
+        
+        # Eagerly initialize models based on preferred provider
+        self._eagerly_initialize_models()
+
+    def _eagerly_initialize_models(self) -> None:
+        """Initialize embedding models based on provider preference - eager for local, lazy for online."""
+        logger.info(sm("Initializing embedding models"))
+        
+        if self.preferred_provider == "local":
+            # Eagerly initialize local embedder for local preference
+            logger.info(sm("Eagerly loading local embedding models"))
+            local_embedder = self._get_local_embedder()
+            if local_embedder:
+                logger.info(sm("Local embedder ready", 
+                           model=local_embedder.model_name,
+                           dimensions=local_embedder.dimensions))
+            else:
+                logger.warning(sm("Local embedder failed to initialize"))
+                
+            # Check online availability but don't initialize (lazy loading)
+            logger.info(sm("Checking online embedding provider availability (lazy loading)"))
+            online_available = self._check_online_availability()
+            if online_available:
+                logger.info(sm("Online embedder available as fallback (will load on first use)"))
+            else:
+                logger.warning(sm("Online embedder not available - check API keys"))
+        else:
+            # For online preference, check availability but don't initialize (lazy loading)
+            logger.info(sm("Checking online embedding provider availability (lazy loading)"))
+            online_available = self._check_online_availability()
+            if online_available:
+                logger.info(sm("Online embedder ready (will load on first use)", 
+                           provider="online"))
+            else:
+                logger.warning(sm("Online embedder not available - check API keys"))
+                
+            # Skip local model initialization when user explicitly chose online
+            logger.info(sm("Skipping local embedding model initialization (online provider preferred)"))
+            logger.info(sm("To use local models as fallback, set EMBEDDING_PROVIDER=local"))
+        
+        # Summary of initialization strategy
+        if self.preferred_provider == "local":
+            logger.info(sm("AutoEmbedder configured: LOCAL models preloaded, ONLINE models lazy-loaded"))
+        else:
+            logger.info(sm("AutoEmbedder configured: ONLINE models only (LOCAL models skipped)"))
+
+        logger.info(sm("AutoEmbedder initialization complete",
+                   preferred_provider=self.preferred_provider,
+                   strategy="eager_local_lazy_online" if self.preferred_provider == "local" else "online_only"))
+
+    def _check_online_availability(self) -> bool:
+        """Check if online embedder is available without initializing it."""
+        try:
+            return bool(os.getenv("OPENAI_API_KEY"))
+        except Exception:
+            return False
+
+    def _check_local_availability(self) -> bool:
+        """Check if local embedder is available without initializing it."""
+        try:
+            import importlib.util
+            return importlib.util.find_spec("sentence_transformers") is not None
+        except Exception:
+            return False
+
+    def _get_online_embedder(self) -> OnlineEmbedder | None:
+        """Get or create online embedder if available."""
+        if "online" not in self.embedders:
+            try:
+                embedder = OnlineEmbedder()
+                if embedder.is_available():
+                    self.embedders["online"] = embedder
+                    logger.info(sm("Online embedder available"))
+                else:
+                    logger.debug(sm("Online embedder not available"))
+                    return None
+            except Exception as e:
+                logger.debug(sm("Failed to create online embedder", error=str(e)))
+                return None
+
+        return self.embedders.get("online")
+
+    def _get_local_embedder(self) -> LocalEmbedder | None:
+        """Get or create local embedder if available."""
+        if "local" not in self.embedders:
+            try:
+                embedder = LocalEmbedder()
+                if embedder.is_available():
+                    self.embedders["local"] = embedder
+                    logger.info(sm("Local embedder available"))
+                else:
+                    logger.debug(sm("Local embedder not available"))
+                    return None
+            except Exception as e:
+                logger.warning(sm("Failed to create local embedder", error=str(e)))
+                return None
+
+        return self.embedders.get("local")
+
+    def embed_text(self, text: str | list[str]) -> np.ndarray:
+        """
+        Generate embeddings using the best available provider.
+
+        Args:
+            text: Text or list of texts to embed
+
+        Returns:
+            Numpy array of embeddings
+
+        Raises:
+            EmbedderError: If no embedders are available or all fail
+        """
+        providers = []
+
+        # Build provider list based on preference
+        if self.preferred_provider == "online":
+            # When online is explicitly preferred, don't initialize local models as fallback
+            providers = [self._get_online_embedder()]
+        elif self.preferred_provider == "local":
+            providers = [self._get_local_embedder(), self._get_online_embedder()]
+        else:  # default: local first
+            providers = [self._get_local_embedder(), self._get_online_embedder()]
+            
+        logger.debug(sm("All available providers", providers=providers))
+
+        # Try each provider
+        for embedder in providers:
+            if embedder:
+                try:
+                    logger.info(sm("Attempting embedding generation",
+                               provider=type(embedder).__name__))
+                    return embedder.embed_text(text)
+                except Exception as e:
+                    logger.warning(sm("Embedder failed, trying next provider",
+                                 provider=type(embedder).__name__,
+                                 error=str(e)))
+                    continue
+
+        error_msg = "All embedding providers failed or are unavailable"
+        logger.error(sm(error_msg, preferred_provider=self.preferred_provider))
+        raise EmbedderError(error_msg)
+    
+    async def embed_text_async(self, text: str | list[str]) -> np.ndarray:
+        """
+        Async version of embed_text with fallback providers.
+
+        Args:
+            text: Text or list of texts to embed
+
+        Returns:
+            Numpy array of embeddings
+
+        Raises:
+            EmbedderError: If no embedders are available or all fail
+        """
+        providers = []
+
+        # Build provider list based on preference
+        if self.preferred_provider == "online":
+            providers = [self._get_online_embedder()]
+        elif self.preferred_provider == "local":
+            providers = [self._get_local_embedder(), self._get_online_embedder()]
+        else:  # default: local first
+            providers = [self._get_local_embedder(), self._get_online_embedder()]
+
+        # Try each provider
+        for embedder in providers:
+            if embedder:
+                try:
+                    logger.info(sm("Attempting async embedding generation",
+                               provider=type(embedder).__name__))
+                    return await embedder.embed_text_async(text)
+                except Exception as e:
+                    logger.warning(sm("Async embedder failed, trying next provider",
+                                 provider=type(embedder).__name__,
+                                 error=str(e)))
+                    continue
+
+        error_msg = "All embedding providers failed or are unavailable"
+        logger.error(sm(error_msg, preferred_provider=self.preferred_provider))
+        raise EmbedderError(error_msg)
+
+    def get_model_info(self) -> dict[str, Any]:
+        """Get information about the currently available model."""
+        if self.preferred_provider == "online" and self._get_online_embedder():
+            embedder = self._get_online_embedder()
+            return {
+                "provider": "online",
+                "model": embedder.model_name,
+                "dimensions": embedder.dimensions
+            }
+        if self.preferred_provider == "local" and self._get_local_embedder():
+            embedder = self._get_local_embedder()
+            return {
+                "provider": "local",
+                "model": embedder.model_name,
+                "dimensions": embedder.dimensions
+            }
+        # Auto mode - return first available (respecting online-only preference)
+        if self._get_online_embedder():
+            embedder = self._get_online_embedder()
+            return {
+                "provider": "online",
+                "model": embedder.model_name,
+                "dimensions": embedder.dimensions
+            }
+        # Only try local if not explicitly using online-only
+        if self.preferred_provider != "online" and self._get_local_embedder():
+            embedder = self._get_local_embedder()
+            return {
+                "provider": "local",
+                "model": embedder.model_name,
+                "dimensions": embedder.dimensions
+            }
+
+        return {
+            "provider": None,
+            "model": None,
+            "dimensions": None
+        }
+
+    def get_available_providers(self) -> list[str]:
+        """Get list of available providers (respects online-only preference)."""
+        providers = []
+
+        if self._get_online_embedder():
+            providers.append("online")
+
+        # Only check local if not explicitly using online-only
+        if self.preferred_provider != "online" and self._get_local_embedder():
+            providers.append("local")
+
+        return providers
+        
+    def _prepare_embedding_text(self, file: File) -> str:
+        """
+        Prepare text for embedding generation.
+    
+        Args:
+            file: File metadata to prepare text from
+    
+        Returns:
+            Text prepared for embedding
+        """
+        parts = []
+    
+        # Add title
+        if file.title:
+            parts.append(f"Title: {file.title}")
+    
+        # Add summary
+        if file.summary:
+            parts.append(f"Summary: {file.summary}")
+    
+        # Add keywords
+        if file.keywords:
+            parts.append(f"Keywords: {', '.join(file.keywords)}")
+    
+        # Add content (truncated)
+        # if file.content:
+        #     content = file.content[:1000]  # Limit content length
+        #     parts.append(f"Content: {content}")
+    
+        return " ".join(parts)
+        
+    async def embed(self, file: File):
+        embedding_text = self._prepare_embedding_text(file)
+        embedding = await self.embed_text_async(embedding_text)
+        
+        model_info = self.get_model_info()
+        
+        file.embedding = embedding
+        file.embedding_model = model_info["model"]
+        file.embedding_dimensions = model_info["dimensions"]
+        file.embedded_at = datetime.now(timezone.utc)
+        
+        file.status = ProcessingStatus.COMPLETE
+
+
+# Convenience functions for easier usage
+def generate_embedding(text: str | list[str],
+                      provider: str | None = None) -> np.ndarray:
+    """
+    Convenience function to generate embeddings.
+
+    Args:
+        text: Text or list of texts to embed
+        provider: Preferred embedding provider
+
+    Returns:
+        Numpy array of embeddings
+    """
+    embedder = AutoEmbedder(preferred_provider=provider)
+    return embedder.embed_text(text)
+
+
+def get_available_embedders() -> list[str]:
+    """Get list of available embedding providers."""
+    embedder = AutoEmbedder()
+    return embedder.get_available_providers()
+
+
+def get_embedder_info() -> dict[str, Any]:
+    """Get information about the current embedder configuration."""
+    embedder = AutoEmbedder()
+    return embedder.get_model_info()