serving/app/services/embedding.py

"""Embedding service port and adapters.

The ``EmbeddingService`` protocol abstracts the embedding provider so the rest
of the application depends on a capability, not on a deployment choice.  Two
adapters are provided:

* ``TEIEmbeddingService`` — calls HuggingFace Text Embeddings Inference over
  HTTP.  This is the production path.
* ``MockEmbeddingService`` — returns deterministic vectors of the configured
  dimension.  Used in tests; not a fake implementation, just a test double.
"""

from __future__ import annotations

import hashlib
import math
from typing import Protocol, runtime_checkable

import httpx

from app.config import Settings, get_settings
from app.exceptions import EmbeddingServiceUnavailableError
from app.logging import get_logger

logger = get_logger(__name__)

# ---------------------------------------------------------------------------
# Port
# ---------------------------------------------------------------------------


@runtime_checkable
class EmbeddingService(Protocol):
    """Capability: embed text into dense vectors."""

    @property
    def model_name(self) -> str:
        """The model identifier used for provenance tracking."""
        ...

    @property
    def model_version(self) -> str:
        """A version tag for the embedder, stored on every chunk row."""
        ...

    @property
    def dimension(self) -> int:
        """Dimensionality of the vectors this service produces."""
        ...

    async def embed_texts(self, texts: list[str]) -> list[list[float]]:
        """Embed a batch of passage texts.

        Args:
            texts: The texts to embed.

        Returns:
            A list of embedding vectors, one per input text.

        Raises:
            EmbeddingServiceUnavailableError: If the service is unreachable.
        """
        ...

    async def embed_query(self, query: str) -> list[float]:
        """Embed a single search query.

        Some models use different prefixes for queries vs passages.  This
        method handles that distinction.

        Args:
            query: The search query text.

        Returns:
            The embedding vector.

        Raises:
            EmbeddingServiceUnavailableError: If the service is unreachable.
        """
        ...


# ---------------------------------------------------------------------------
# TEI adapter (production)
# ---------------------------------------------------------------------------

_TEI_EMBED_PATH = "/embed"
_TEI_TIMEOUT_SECONDS = 30.0


class TEIEmbeddingService:
    """Calls HuggingFace Text Embeddings Inference over HTTP.

    TEI serves models like ``bge-m3`` and exposes a simple JSON API.  This
    adapter batches according to ``settings.embedding_batch_size`` and
    retries once on transient failures.

    Attributes:
        _endpoint: Base URL of the TEI instance.
        _model: Model name for provenance.
        _dim: Expected vector dimension.
        _batch_size: Max texts per HTTP request.
    """

    def __init__(self, settings: Settings | None = None) -> None:
        """Initialize from settings.

        Args:
            settings: Optional configuration override.
        """
        cfg = settings or get_settings()
        self._endpoint = cfg.embedding_endpoint.rstrip("/")
        self._model = cfg.embedding_model
        self._dim = cfg.embedding_dim
        self._batch_size = cfg.embedding_batch_size
        self._version = "1"

    @property
    def model_name(self) -> str:
        """Return the configured model identifier."""
        return self._model

    @property
    def model_version(self) -> str:
        """Return the embedding version tag."""
        return self._version

    @property
    def dimension(self) -> int:
        """Return the configured embedding dimension."""
        return self._dim

    async def embed_texts(self, texts: list[str]) -> list[list[float]]:
        """Embed passage texts via TEI, batched.

        Args:
            texts: The texts to embed.

        Returns:
            A list of embedding vectors.

        Raises:
            EmbeddingServiceUnavailableError: On HTTP or connection error.
        """
        if not texts:
            return []

        all_embeddings: list[list[float]] = []
        for i in range(0, len(texts), self._batch_size):
            batch = texts[i : i + self._batch_size]
            embeddings = await self._call_tei(batch)
            all_embeddings.extend(embeddings)
        return all_embeddings

    async def embed_query(self, query: str) -> list[float]:
        """Embed a search query via TEI.

        Args:
            query: The search query.

        Returns:
            The embedding vector.

        Raises:
            EmbeddingServiceUnavailableError: On HTTP or connection error.
        """
        results = await self._call_tei([query])
        return results[0]

    async def _call_tei(self, inputs: list[str]) -> list[list[float]]:
        """Make a single TEI embed request.

        Args:
            inputs: Batch of texts to embed.

        Returns:
            The embedding vectors.

        Raises:
            EmbeddingServiceUnavailableError: On any failure.
        """
        url = f"{self._endpoint}{_TEI_EMBED_PATH}"
        payload = {"inputs": inputs}
        try:
            async with httpx.AsyncClient(timeout=_TEI_TIMEOUT_SECONDS) as client:
                response = await client.post(url, json=payload)
                response.raise_for_status()
                data: list[list[float]] = response.json()
                return data
        except (httpx.HTTPError, httpx.ConnectError, Exception) as exc:
            logger.error(
                "embedding_service_error",
                endpoint=self._endpoint,
                error=str(exc)[:200],
            )
            raise EmbeddingServiceUnavailableError(
                f"Embedding service at {self._endpoint} is unavailable."
            ) from exc


# ---------------------------------------------------------------------------
# Mock adapter (tests)
# ---------------------------------------------------------------------------


class MockEmbeddingService:
    """Returns deterministic vectors for testing.

    Vectors are derived from a hash of the input text so the same text always
    produces the same vector, enabling assertions in tests.  This is *not* a
    fake implementation — it makes no attempt at semantic similarity.
    """

    def __init__(self, dimension: int = 1024) -> None:
        """Initialize with a target dimension.

        Args:
            dimension: Vector dimensionality to produce.
        """
        self._dim = dimension

    @property
    def model_name(self) -> str:
        """Return a test model identifier."""
        return "mock-embedder"

    @property
    def model_version(self) -> str:
        """Return a test version tag."""
        return "test-v1"

    @property
    def dimension(self) -> int:
        """Return the configured dimension."""
        return self._dim

    async def embed_texts(self, texts: list[str]) -> list[list[float]]:
        """Return deterministic vectors for each text.

        Args:
            texts: The texts to embed.

        Returns:
            Deterministic vectors derived from text hashes.
        """
        return [self._deterministic_vector(t) for t in texts]

    async def embed_query(self, query: str) -> list[float]:
        """Return a deterministic vector for the query.

        Args:
            query: The search query.

        Returns:
            A deterministic vector.
        """
        return self._deterministic_vector(query)

    def _deterministic_vector(self, text: str) -> list[float]:
        """Produce a unit-length vector deterministically from text.

        Args:
            text: Input text.

        Returns:
            A normalized vector of ``self._dim`` dimensions.
        """
        digest = hashlib.sha256(text.encode("utf-8")).digest()
        # Expand the 32-byte hash to fill the dimension by repeating
        raw = list(digest) * ((self._dim // len(digest)) + 1)
        raw = raw[: self._dim]
        # Map bytes to floats in [-1, 1] and normalize
        vec = [(b - 128) / 128.0 for b in raw]
        norm = math.sqrt(sum(v * v for v in vec)) or 1.0
        return [v / norm for v in vec]


# ---------------------------------------------------------------------------
# Singleton management
# ---------------------------------------------------------------------------

_embedding_service: EmbeddingService | None = None


def get_embedding_service() -> EmbeddingService:
    """Return the process-wide embedding service instance.

    Returns a ``MockEmbeddingService`` when no TEI endpoint is configured
    (i.e. in test or when the endpoint is empty).

    Returns:
        The embedding service singleton.
    """
    global _embedding_service  # noqa: PLW0603
    if _embedding_service is not None:
        return _embedding_service

    settings = get_settings()
    if settings.embedding_endpoint and settings.environment != "test":
        _embedding_service = TEIEmbeddingService(settings)
        logger.info(
            "embedding_service_initialized",
            adapter="TEI",
            endpoint=settings.embedding_endpoint,
            model=settings.embedding_model,
        )
    else:
        _embedding_service = MockEmbeddingService(dimension=settings.embedding_dim)
        logger.info("embedding_service_initialized", adapter="Mock")
    return _embedding_service


def set_embedding_service(service: EmbeddingService) -> None:
    """Override the embedding service singleton (for tests).

    Args:
        service: The service to install.
    """
    global _embedding_service  # noqa: PLW0603
    _embedding_service = service


def reset_embedding_service() -> None:
    """Clear the embedding service singleton (for tests)."""
    global _embedding_service  # noqa: PLW0603
    _embedding_service = None