app/services/embeddings.py

from typing import List
from functools import lru_cache

import numpy as np
from sentence_transformers import SentenceTransformer


class EmbeddingService:
    def __init__(self, model_name: str) -> None:
        self.model = SentenceTransformer(model_name)

    def encode(self, texts: List[str]) -> np.ndarray:
        vectors = self.model.encode(
            texts,
            normalize_embeddings=True,
            show_progress_bar=False,
            convert_to_numpy=True,
        )
        return np.asarray(vectors, dtype=np.float32)

    @lru_cache(maxsize=2048)
    def encode_query_cached(self, text: str) -> bytes:
        vec = self.model.encode(
            [text],
            normalize_embeddings=True,
            show_progress_bar=False,
            convert_to_numpy=True,
        )
        arr = np.asarray(vec, dtype=np.float32)
        return arr.tobytes()

    def encode_query(self, text: str) -> np.ndarray:
        buf = self.encode_query_cached(text)
        return np.frombuffer(buf, dtype=np.float32).reshape(1, -1)