from src.labdaps.config import EMBEDDING_MODEL, EMBED_QUERY_PREFIX, EMBED_PASSAGE_PREFIX


class Embedder:
    def __init__(self):
        from sentence_transformers import SentenceTransformer
        print(f"[INFO] Carregando modelo de embeddings: {EMBEDDING_MODEL}")
        self._model = SentenceTransformer(EMBEDDING_MODEL)
        print("[INFO] Modelo carregado.")

    def embed_passages(self, texts: list[str], batch_size: int = 64) -> list[list[float]]:
        prefixed = [f"{EMBED_PASSAGE_PREFIX}{t}" for t in texts]
        vecs = self._model.encode(
            prefixed,
            batch_size=batch_size,
            normalize_embeddings=True,
            show_progress_bar=True,
        )
        return vecs.tolist()

    def embed_query(self, query: str) -> list[float]:
        vec = self._model.encode(
            f"{EMBED_QUERY_PREFIX}{query}",
            normalize_embeddings=True,
        )
        return vec.tolist()