src/embedding_model.py

from __future__ import annotations
from typing import Optional, List
from src.config import settings
from src.logger import logger


class _STEmbeddings:
    """Minimal LangChain-compatible embeddings via sentence-transformers."""

    def __init__(self, model_name: str):
        from sentence_transformers import SentenceTransformer
        self.model_name = model_name
        self._model = SentenceTransformer(model_name)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self._model.encode(texts, show_progress_bar=False).tolist()

    def embed_query(self, text: str) -> List[float]:
        return self._model.encode([text], show_progress_bar=False)[0].tolist()

    # LangChain  / Chroma sometimes call this
    def __call__(self, text: str) -> List[float]:
        return self.embed_query(text)


def get_embeddings(model: Optional[str] = None):
    model_name = model or settings.embedding_model or "sentence-transformers/all-MiniLM-L6-v2"
    logger.info("Using SentenceTransformer embeddings", context={"model": model_name})
    return _STEmbeddings(model_name)