app/embeddings.py

"""Local multilingual embeddings via fastembed (ONNX, no torch).

Default model: sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2
(384-dim, multilingual incl. ES/PT/EN). If an e5 model is configured, the
"query: " / "passage: " prefixes it expects are applied automatically.
Inference runs in a worker thread; the model is loaded lazily and cached.
"""

from __future__ import annotations

import math

import anyio

from app.config import get_settings

_model = None


def _get_model():
    global _model
    if _model is None:
        from fastembed import TextEmbedding

        _model = TextEmbedding(
            model_name=get_settings().embedding_model, cache_dir=".fastembed_cache"
        )
    return _model


def _embed_sync(texts: list[str]) -> list[list[float]]:
    model = _get_model()
    return [[float(x) for x in vec] for vec in model.embed(texts)]


def _apply_prefix(kind: str, texts: list[str]) -> list[str]:
    # e5 family expects "query: " / "passage: " prefixes; others must not.
    if "e5" in get_settings().embedding_model.lower():
        return [f"{kind}: {t}" for t in texts]
    return texts


async def embed_texts(texts: list[str], *, kind: str = "passage") -> list[list[float]]:
    if not texts:
        return []
    return await anyio.to_thread.run_sync(_embed_sync, _apply_prefix(kind, texts))


async def embed_query(text: str) -> list[float]:
    result = await embed_texts([text], kind="query")
    return result[0]


def cosine(a: list[float], b: list[float]) -> float:
    if not a or not b:
        return 0.0
    dot = sum(x * y for x, y in zip(a, b, strict=False))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)