ZerdaZerda

GeminiEmbedder

On this page

embeddings.gemini_embedder.GeminiEmbedder(BaseEmbedder) — Google Gemini Embeddings API client, using the new google-genai SDK.

[!WARNING] SDK choice matters This class imports from google import genai — the new SDK. It explicitly is not compatible with the older, deprecated google-generativeai package; if that's the one installed, the ImportError message raised here says so directly.

Purpose

Wraps client.models.embed_content() with task-type support (RETRIEVAL_QUERY, RETRIEVAL_DOCUMENT, SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING), Matryoshka-style output-dimension reduction (output_dimensionality), retry logic, and usage tracking.

Constructor

Python · 17 lines
GeminiEmbedder(
    model_name: str = "gemini-embedding-001",
    api_key: Optional[str] = None,
    task_type: Optional[str] = None,
    output_dimensionality: Optional[int] = None,
    device: Optional[str] = None,
    normalize_embeddings: bool = True,
    batch_size: int = 50,
    max_length: Optional[int] = None,
    cache_embeddings: bool = True,
    show_progress: bool = False,
    timeout: int = 60,
    max_retries: int = 3,
    retry_delay: float = 2.0,
    track_usage: bool = True,
    **kwargs,
)

Parameters

NameTypeDefaultNotes
model_namestr"gemini-embedding-001"See MODEL_SPECS below; unknown names fall back to a generic 768-dim/50-batch spec with status: "unknown"
api_keyOptional[str]NoneFalls back to GOOGLE_API_KEY, then GEMINI_API_KEY; raises ValueError if neither is set
task_typeOptional[str]NoneMust be a key of TASK_TYPES, or a warning is logged (the value is still stored, just without recognized-task benefits)
output_dimensionalityOptional[int]NoneOnly honored if model_spec['output_dimensionality_support'] is True and the value is within [min_output_dimensionality, default_dimensions]; otherwise ignored with a warning
batch_sizeint50Automatically capped to model_spec['max_batch_size'] (100 for all listed models) if exceeded
cache_embeddingsboolTruePlain dict cache, keyed by text plus a suffix encoding task_type/output_dimensionality
max_retriesint3
retry_delayfloat2.0Base for exponential backoff
track_usageboolTrueEnables self.usage_stats = GeminiUsageStats()

MODEL_SPECS (class attribute)

ModelDefault dimsMax tokensOutput-dim reduction?Status
gemini-embedding-00130728192Yes, recommended sizes [768, 1536, 3072]GA
text-embedding-0047682048Yesdeprecated (2026-01-14)
embedding-0017682048Nodeprecated (2025-08-14)

TASK_TYPES (class attribute)

RETRIEVAL_QUERY, RETRIEVAL_DOCUMENT, SEMANTIC_SIMILARITY, CLASSIFICATION, CLUSTERING — each with Arabic/English descriptions and a use-case string. Duplicated at module scope in embeddings/__init__.py as __task_types_gemini__.

Raises

  • ValueError if no API key is resolvable.
  • ImportError if google-genai isn't installed (with an explicit warning about the deprecated google-generativeai package).
  • ValueError from _validate_api_key() if the API responds with an invalid-key or model-not-found error (other errors during validation are only logged as warnings, not raised).

encode()

Python · 8 lines
encode(
    texts: Union[str, List[str]],
    show_progress_bar: bool = False,
    convert_to_numpy: bool = True,
    task_type: Optional[str] = None,
    batch_size: Optional[int] = None,
    **kwargs,
) -> np.ndarray

Behavior

  1. If task_type is passed, it temporarily overrides self.task_type for the duration of the call (restored in a finally block).
  2. Wraps single string in a list; raises ValueError if empty after filtering blank entries.
  3. Iterates in batch_size-sized chunks; per text, checks the cache using a key that includes the active task_type and output_dimensionality (so the same text cached under different task types produces separate cache entries — correctly avoiding cross-task-type cache collisions).
  4. Non-cached texts go through _process_batch(), which builds a types.EmbedContentConfig (if task_type or output_dimensionality is set) and calls client.models.embed_content(). Retries on rate-limit (429/quota) and timeout errors with exponential backoff; other APIErrors are not retried.
  5. Validates the returned embedding count matches input count (raises ValueError on mismatch) and warns (does not raise) on a dimension mismatch against the expected dimension.
  6. Stacks batches, L2-normalizes if normalize_embeddings=True, updates stats.

Returns

np.ndarray of shape (n, embedding_dim) — always 2D, no single-string unwrap.

Raises

  • ValueError for empty input or an embedding-count mismatch from the API.
  • RuntimeError if all retries in _process_batch are exhausted.
  • Re-raises non-retryable APIErrors and any other unexpected exception.

Diagnostics

get_usage_stats() -> Dict[str, Any], clear_cache(), get_model_info()

Standard introspection; get_model_info() includes the full model_spec, current task_type, and (if tracking) GeminiUsageStats.get_summary().

GeminiUsageStats (dataclass)

Fields: total_requests, total_texts, total_characters, start_time, errors, cache_hits. .add_request(text_count, char_count) and .add_error() mutate counters; .get_summary() derives avg_texts_per_request, avg_chars_per_text, and success_rate.

Context manager

__enter__ returns self; __exit__ logs a session summary (requests, texts, characters, elapsed time, success rate) if track_usage=True.

Examples

Python · 10 lines
from embeddings import GeminiEmbedder

embedder = zerda.GeminiEmbedder(api_key="your-key", task_type="RETRIEVAL_DOCUMENT")
vectors = embedder.encode(["first document", "second document"])

# Dimension reduction via MRL
embedder = GeminiEmbedder(api_key="your-key", output_dimensionality=768)

# Overriding task type per-call without changing the instance default
query_vec = embedder.encode("what is the capital of Egypt?", task_type="RETRIEVAL_QUERY")

Related pages