app/services/embedding_service.py

from dataclasses import dataclass

import google.genai as genai
import google.genai.local_tokenizer as local_tokenizer

GEMINI_EMBEDDING_MAX_TOKENS = 8192
# The SDK does not currently expose LocalTokenizer support for `gemini-embedding-2`,
# so keep chunk sizing on a supported Gemini-family tokenizer until Google ships one.
GEMINI_LOCAL_TOKENIZER_MODEL = "gemini-2.0-flash"


class GeminiTokenizerAdapter:
    model_max_length = GEMINI_EMBEDDING_MAX_TOKENS

    def __init__(self, model_name: str = GEMINI_LOCAL_TOKENIZER_MODEL):
        self._local_tokenizer = local_tokenizer.LocalTokenizer(model_name=model_name)
        self._tokenizer = self._local_tokenizer._tokenizer

    def __call__(
        self,
        text: str,
        *,
        add_special_tokens: bool = True,
        truncation: bool = False,
    ) -> dict[str, list[int]]:
        token_ids = list(self._tokenizer.encode(text or "", out_type=int))
        if truncation:
            token_ids = token_ids[: self.model_max_length]
        return {"input_ids": token_ids}

    def decode(
        self,
        token_ids: list[int],
        *,
        skip_special_tokens: bool = True,
    ) -> str:
        del skip_special_tokens
        return self._tokenizer.decode(token_ids)


@dataclass
class GeminiEmbeddingModel:
    tokenizer: GeminiTokenizerAdapter
    max_seq_length: int = GEMINI_EMBEDDING_MAX_TOKENS


class EmbeddingService:
    def __init__(
        self,
        model_name: str,
        api_key: str,
        *,
        client: genai.Client | None = None,
        tokenizer_model_name: str = GEMINI_LOCAL_TOKENIZER_MODEL,
    ):
        if not api_key:
            raise ValueError("GEMINI_EMBEDDING_API_KEY (or GEMINI_API_KEY) is required")

        self.api_key = api_key
        self.model_name = model_name
        self.client = client or genai.Client(api_key=api_key)
        self.model = GeminiEmbeddingModel(
            tokenizer=GeminiTokenizerAdapter(model_name=tokenizer_model_name)
        )

    def embed(self, text: str) -> list[float]:
        try:
            response = self.client.models.embed_content(
                model=self.model_name,
                contents=text,
            )
        except Exception as exc:
            raise RuntimeError(
                f"Gemini embedding request failed for model {self.model_name}"
            ) from exc

        embeddings = getattr(response, "embeddings", None) or []
        if len(embeddings) != 1:
            raise RuntimeError(
                f"Gemini embedding response returned {len(embeddings)} embeddings"
            )

        values = getattr(embeddings[0], "values", None)
        if values is None:
            raise RuntimeError("Gemini embedding response did not include values")

        return list(values)