services/embedding.py

"""
Embedding Service
=================
Generates dense vector embeddings using sentence-transformers (100% free, runs locally).
Model: all-MiniLM-L6-v2 (384-dim, fast, great quality for enterprise text)

Singleton pattern: model loads once on first call and stays in memory.
Thread-safe for concurrent FastAPI requests.
"""
from __future__ import annotations

import asyncio
import threading
from typing import List, Optional

import numpy as np
import structlog

logger = structlog.get_logger(__name__)


class EmbeddingService:
    """
    Wraps sentence-transformers SentenceTransformer for async-compatible embedding.
    Falls back gracefully if the model is not yet downloaded (first run).
    """

    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model_name = model_name
        self._model = None
        self._lock = threading.Lock()
        self._dim: Optional[int] = None

    def _load_model(self):
        """Lazy-load the embedding model (downloads on first use)."""
        if self._model is None:
            with self._lock:
                if self._model is None:
                    logger.info(
                        "Loading embedding model",
                        model=self.model_name,
                        note="Downloading on first run, ~90MB",
                    )
                    from sentence_transformers import SentenceTransformer
                    self._model = SentenceTransformer(self.model_name)
                    # Warm up with a test sentence
                    test = self._model.encode(["test"], show_progress_bar=False)
                    self._dim = test.shape[1]
                    logger.info(
                        "Embedding model ready",
                        model=self.model_name,
                        dim=self._dim,
                    )

    @property
    def dim(self) -> int:
        """Return embedding dimensions (loaded lazily)."""
        if self._dim is None:
            self._load_model()
        return self._dim  # type: ignore

    def encode_sync(
        self,
        texts: List[str],
        batch_size: int = 64,
        normalize: bool = True,
    ) -> np.ndarray:
        """
        Synchronous encoding — for use in background tasks.

        Args:
            texts:      List of strings to embed
            batch_size: Batch size for encoding (larger = faster on GPU)
            normalize:  L2-normalize embeddings (recommended for cosine similarity)

        Returns:
            numpy array of shape (len(texts), dim)
        """
        self._load_model()
        embeddings = self._model.encode(  # type: ignore
            texts,
            batch_size=batch_size,
            normalize_embeddings=normalize,
            show_progress_bar=False,
        )
        return embeddings

    async def encode(
        self,
        texts: List[str],
        batch_size: int = 64,
        normalize: bool = True,
    ) -> np.ndarray:
        """
        Async encoding — runs encoding in a thread pool to avoid blocking the event loop.

        Args:
            texts:      List of strings to embed
            batch_size: Batch size for encoding
            normalize:  L2-normalize embeddings

        Returns:
            numpy array of shape (len(texts), dim)
        """
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None, lambda: self.encode_sync(texts, batch_size, normalize)
        )

    async def encode_query(self, query: str) -> np.ndarray:
        """Encode a single query string and return 1D array."""
        result = await self.encode([query])
        return result[0]

    async def encode_documents(self, texts: List[str]) -> np.ndarray:
        """
        Encode a list of document chunks. Adds a prefix per MTEB best practice.
        """
        prefixed = [f"passage: {t}" for t in texts]
        return await self.encode(prefixed)


# ── Singleton instance ─────────────────────────────────────────
embedding_service = EmbeddingService(model_name="all-MiniLM-L6-v2")