microdao-daarion/services/rag-service/app/embedding.py

"""
Embedding service for RAG
Uses SentenceTransformers via Haystack
"""

import logging
from typing import Optional, List, Dict, Any

from sentence_transformers import SentenceTransformer

from app.core.config import settings

logger = logging.getLogger(__name__)


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

    def run(self, texts: Optional[List[str]] = None, text: Optional[str] = None) -> Dict[str, Any]:
        if texts is not None:
            embeddings = self.model.encode(texts, convert_to_numpy=True).tolist()
            return {"embeddings": embeddings}
        if text is not None:
            embedding = self.model.encode([text], convert_to_numpy=True).tolist()
            return {"embedding": embedding}
        return {"embeddings": []}


# Global embedder instance
_text_embedder: Optional[SimpleEmbedder] = None


def get_text_embedder() -> SimpleEmbedder:
    global _text_embedder
    if _text_embedder is not None:
        return _text_embedder
    logger.info(f"Loading embedding model: {settings.EMBED_MODEL_NAME}")
    logger.info(f"Device: {settings.EMBED_DEVICE}")
    _text_embedder = SimpleEmbedder(settings.EMBED_MODEL_NAME, settings.EMBED_DEVICE)
    logger.info("Text embedder initialized successfully")
    return _text_embedder


def reset_embedder():
    global _text_embedder
    _text_embedder = None