src/embeddings.py

"""Unified embedding interface supporting local, OpenAI, and Gemini backends."""

import os
import time

import numpy as np
from dotenv import load_dotenv

load_dotenv()

LOCAL_MODEL_NAME = "all-MiniLM-L6-v2"
OPENAI_MODEL_NAME = "text-embedding-3-small"
AZURE_OPENAI_MODEL_NAME = "text-embedding-3-large"
GEMINI_MODEL_NAME = "gemini-embedding-2"
BATCH_SIZE = 128
GEMINI_BATCH_SIZE = 100  # stay within Gemini rate limits


def get_provider() -> str:
    return os.environ.get("EMBEDDING_PROVIDER", "local").lower()


def embeddings_filename(provider: str | None = None, multi: bool = False) -> str:
    p = provider or get_provider()
    suffix = ".multi" if multi else ""
    return f"embeddings.{p}{suffix}.npy"


def default_min_score(provider: str | None = None) -> float:
    p = provider or get_provider()
    if p in ("openai", "azure", "gemini"):
        return 0.45
    return 0.35


class Embedder:
    def __init__(self):
        self.provider = get_provider()
        if self.provider == "openai":
            from openai import OpenAI

            self._client = OpenAI()
            self.dim = 1536
        elif self.provider == "azure":
            from openai import AzureOpenAI

            self._client = AzureOpenAI(
                api_key=os.environ["AZURE_OPENAI_API_KEY"],
                azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
                api_version=os.environ.get("AZURE_OPENAI_API_VERSION", "2024-09-01-preview"),
            )
            self._azure_deployment = os.environ.get("AZURE_OPENAI_DEPLOYMENT", AZURE_OPENAI_MODEL_NAME)
            self.dim = 3072
        elif self.provider == "gemini":
            from google import genai

            self._client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
            self.dim = 1536
        else:
            from sentence_transformers import SentenceTransformer

            self._model = SentenceTransformer(LOCAL_MODEL_NAME)
            self.dim = self._model.get_embedding_dimension()

    def embed_batch(self, texts: list[str], show_progress: bool = True) -> np.ndarray:
        if self.provider == "openai":
            return self._embed_batch_openai(texts)
        if self.provider == "azure":
            return self._embed_batch_azure(texts)
        if self.provider == "gemini":
            return self._embed_batch_gemini(texts, show_progress)
        return self._embed_batch_local(texts, show_progress)

    def embed_query(self, query: str) -> np.ndarray:
        if self.provider == "openai":
            resp = self._client.embeddings.create(
                model=OPENAI_MODEL_NAME, input=[query]
            )
            vec = np.array(resp.data[0].embedding, dtype=np.float32)
        elif self.provider == "azure":
            resp = self._client.embeddings.create(
                model=self._azure_deployment, input=[query]
            )
            vec = np.array(resp.data[0].embedding, dtype=np.float32)
        elif self.provider == "gemini":
            resp = self._client.models.embed_content(
                model=GEMINI_MODEL_NAME,
                contents=query,
                config={"output_dimensionality": 1536},
            )
            vec = np.array(resp.embeddings[0].values, dtype=np.float32)
        else:
            vec = self._model.encode(query, convert_to_numpy=True).astype(np.float32)
        return vec / np.linalg.norm(vec)

    def _embed_batch_azure(self, texts: list[str]) -> np.ndarray:
        all_embeddings: list[list[float]] = []
        total = len(texts)
        for i in range(0, total, BATCH_SIZE):
            batch = texts[i : i + BATCH_SIZE]
            n = i // BATCH_SIZE + 1
            total_batches = (total + BATCH_SIZE - 1) // BATCH_SIZE
            print(f"  Embedding batch {n}/{total_batches} ({len(batch)} texts) ...")
            resp = self._client.embeddings.create(
                model=self._azure_deployment, input=batch
            )
            all_embeddings.extend([d.embedding for d in resp.data])
        return np.array(all_embeddings, dtype=np.float32)

    def _embed_batch_openai(self, texts: list[str]) -> np.ndarray:
        all_embeddings: list[list[float]] = []
        total = len(texts)
        for i in range(0, total, BATCH_SIZE):
            batch = texts[i : i + BATCH_SIZE]
            n = i // BATCH_SIZE + 1
            total_batches = (total + BATCH_SIZE - 1) // BATCH_SIZE
            print(f"  Embedding batch {n}/{total_batches} ({len(batch)} texts) ...")
            resp = self._client.embeddings.create(
                model=OPENAI_MODEL_NAME, input=batch
            )
            all_embeddings.extend([d.embedding for d in resp.data])
        return np.array(all_embeddings, dtype=np.float32)

    def _embed_batch_gemini(self, texts: list[str], show_progress: bool) -> np.ndarray:
        # The Gemini SDK returns one aggregated embedding for list inputs,
        # so we embed one text at a time and rate-limit per GEMINI_BATCH_SIZE.
        all_embeddings: list[list[float]] = []
        total = len(texts)
        for idx, text in enumerate(texts):
            if show_progress and idx % GEMINI_BATCH_SIZE == 0:
                batch_n = idx // GEMINI_BATCH_SIZE + 1
                total_batches = (total + GEMINI_BATCH_SIZE - 1) // GEMINI_BATCH_SIZE
                print(f"  Embedding batch {batch_n}/{total_batches} ({min(GEMINI_BATCH_SIZE, total - idx)} texts) ...")
            resp = self._client.models.embed_content(
                model=GEMINI_MODEL_NAME,
                contents=text,
                config={"output_dimensionality": 1536},
            )
            all_embeddings.append(resp.embeddings[0].values)
            # Brief pause every GEMINI_BATCH_SIZE texts to respect rate limits
            if (idx + 1) % GEMINI_BATCH_SIZE == 0 and (idx + 1) < total:
                time.sleep(1)
        return np.array(all_embeddings, dtype=np.float32)

    def _embed_batch_local(self, texts: list[str], show_progress: bool) -> np.ndarray:
        return self._model.encode(
            texts,
            batch_size=BATCH_SIZE,
            show_progress_bar=show_progress,
            convert_to_numpy=True,
        ).astype(np.float32)