backend/app/services/embedding_service.py

from sentence_transformers import SentenceTransformer
from app.config import settings
from typing import List
import logging

logger = logging.getLogger(__name__)

# e5 models require task-specific prefixes for best performance
_QUERY_PREFIX = "query: "
_PASSAGE_PREFIX = "passage: "


class EmbeddingService:
    def __init__(self):
        logger.info(f"Loading embedding model: {settings.embedding_model}")
        self._model = SentenceTransformer(settings.embedding_model)
        logger.info("Embedding model ready")

    def embed_query(self, text: str) -> List[float]:
        prefixed = _QUERY_PREFIX + text
        return self._model.encode(prefixed, normalize_embeddings=True).tolist()

    def embed_passages(self, texts: List[str]) -> List[List[float]]:
        prefixed = [_PASSAGE_PREFIX + t for t in texts]
        return self._model.encode(prefixed, normalize_embeddings=True).tolist()

    @property
    def dimension(self) -> int:
        return self._model.get_sentence_embedding_dimension()


_instance: EmbeddingService | None = None


def get_embedding_service() -> EmbeddingService:
    global _instance
    if _instance is None:
        _instance = EmbeddingService()
    return _instance