embedding_services.py

"""
Embedding service abstraction.

This implementation uses SentenceTransformers for local embedding generation,
which is appropriate for a Hugging Face Docker Space and does not require
remote API calls for embeddings.
"""

from __future__ import annotations

import logging
from typing import Iterable, List

from sentence_transformers import SentenceTransformer

from config import settings

logger = logging.getLogger(__name__)


class EmbeddingService:
    """Service responsible for generating vector embeddings."""

    def __init__(self, model_name: str | None = None, device: str | None = None) -> None:
        self.model_name = model_name or settings.embedding_model_name
        self.device = device or settings.embedding_device

        logger.info("Loading embedding model: %s on device: %s", self.model_name, self.device)
        self.model = SentenceTransformer(self.model_name, device=self.device)

    def embed_texts(self, texts: Iterable[str]) -> List[List[float]]:
        """
        Convert an iterable of texts into embeddings.

        Returns:
            List[List[float]]: One embedding vector per input text.
        """
        text_list = [text.strip() for text in texts if text and text.strip()]
        if not text_list:
            return []

        embeddings = self.model.encode(
            text_list,
            batch_size=settings.embedding_batch_size,
            show_progress_bar=False,
            normalize_embeddings=True,
            convert_to_numpy=True,
        )
        return embeddings.tolist()

    def embed_query(self, query_text: str) -> List[float]:
        """
        Convert a single query string into an embedding vector.
        """
        query = query_text.strip()
        if not query:
            raise ValueError("query_text cannot be empty")

        embedding = self.model.encode(
            query,
            show_progress_bar=False,
            normalize_embeddings=True,
            convert_to_numpy=True,
        )
        return embedding.tolist()