gemini_embed_pool

Shared Gemini API key pool for rate-limit distribution.

All embedding calls use next_gemini_embed_key() and all flash-lite generation calls use next_gemini_flash_key() for round-robin key selection. Both accessors draw from the same underlying key pool but maintain independent rotation cycles so they don’t interfere with each other.

Supports GEMINI_EMBED_KEY_POOL env var (comma-separated keys) with fallback to the default hardcoded pool.

Embedding fallback cascade

When the free pool starts 429-ing (consecutive failures exceed PAID_KEY_FALLBACK_THRESHOLD), the embed call sites cascade in this order:

1. Free-tier pool (round-robin via next_gemini_embed_key()).

2. OpenRouter (openrouter_embed_batch() / ..._sync) — engaged via set_openrouter_only() so subsequent calls go straight here.

3. Paid tier-3 Gemini key (gemini_embed_paid_fallback() / ..._sync) — used as the absolute last resort when both the free pool and OpenRouter have failed. Its key is returned by get_paid_fallback_key().

Every credential this module uses – the free-tier pool, the paid key, and the OpenRouter key – is resolved through the Redis-backed system_secret_store (names SECRET_EMBED_POOL, SECRET_PAID_KEY, SECRET_OPENROUTER_KEY). None is hardcoded here. Provision them with python -m scripts.system_secrets_cli status, and call reload_pool() once at service startup – module import deliberately does no I/O, so the pool is empty until that call lands.

When openrouter_only is pinned, callers must still fall through to the paid-key helper if OpenRouter fails on a specific call before raising.

OpenRouter embed fallback (openrouter_embed_batch / openrouter_embed_batch_sync) retries transient network errors and HTTP 429 / 5xx with exponential backoff. Tune with env: OPENROUTER_EMBED_MAX_ATTEMPTS (default 24), OPENROUTER_EMBED_RETRY_BASE_DELAY, OPENROUTER_EMBED_RETRY_MAX_DELAY. Credential/account failures (HTTP 401/402/403) open a short per-key, process-wide cooldown so concurrent batches stop hammering a key that cannot currently serve embeddings and can fall through to the paid Gemini tier immediately. Tune it with OPENROUTER_EMBED_NONRETRIABLE_COOLDOWN_SECONDS (default 1800).

The paid Gemini last-resort path also retries transient network errors and HTTP 408/429/500/502/503/504 with capped exponential backoff and jitter. Tune it with GEMINI_PAID_EMBED_MAX_ATTEMPTS (default 8), GEMINI_PAID_EMBED_RETRY_BASE_DELAY (default 2), and GEMINI_PAID_EMBED_RETRY_MAX_DELAY (default 60).

Daily quota tracking

Quotas are tracked per model class ("embed" vs "generate"). Embedding RPD and generation RPD are separate quotas on the same key, so a key spent for embeddings can still serve flash-lite generation and vice versa.

Each key’s usage is tracked in Redis and two in-memory spent sets. Keys that receive a daily-quota 429 (PerDay in the quotaId) are excluded from rotation for the relevant model class until midnight Pacific Time. A background probe every 2 hours detects keys exhausted by external usage.

async gemini_embed_pool.aclose_shared_http_clients()

Close shared clients owned by the running loop (shutdown hook).

Return type:

None

gemini_embed_pool.get_paid_fallback_key()

Return the paid tier-3 Gemini API key used as the absolute last-resort fallback.

Resolves gemini.paid through system_secret_store.get_secret_sync(): the GEMINI_EMBED_PAID_KEY env override wins, then the shared TTL cache (warmed by reload_pool() at startup), then Redis, then the gemini_api_key entry in config.yaml. The store suppresses its blocking Redis read when a running event loop is detected, so this stays safe on the async hot path.

Called within this module by _next_active_key(), gemini_embed_paid_fallback(), and gemini_embed_paid_fallback_sync(), and externally by gemini_kg_bulk_client.py and classifiers/build_tool_index.py for their own paid-key fallbacks.

Returns:

The paid key, or None when it has not been provisioned.

Return type:

str | None

async gemini_embed_pool.reload_pool()

Load the base key pool from the secret store, merge donated keys, rebuild cycles.

This is the authoritative pool load. The import-time seed is env-only (see _seed_key_pool_from_env()), so every service must call this once at startup – after init_quota_tracking() – for the Redis-backed pool to take effect. Also warms the secret store’s TTL cache for get_paid_fallback_key() and get_openrouter_api_key(), which is what keeps those synchronous accessors off the network on the hot path.

Returns the new total pool size. Safe to call multiple times.

Return type:

int

gemini_embed_pool.init_quota_tracking(redis_client)

Provide the async Redis client for daily quota persistence.

Must be called once at service startup before any embedding calls are made. Call reload_pool() afterwards (from an async context) to merge donated keys.

Return type:

None

Parameters:

redis_client (Any)

gemini_embed_pool.next_gemini_embed_key()

Thread-safe round-robin selection from the Gemini embedding key pool.

Skips keys that have been marked as daily-spent for embeddings.

Return type:

str

gemini_embed_pool.next_gemini_flash_key()

Thread-safe round-robin selection for Gemini flash-lite generation calls.

Uses the same key pool as embeddings but an independent cycle so embed and generation rotations don’t interfere. Skips keys that have been marked as daily-spent for generation.

Return type:

str

gemini_embed_pool.is_daily_quota_429(resp)

Return True if resp is a 429 caused by a daily (RPD) quota limit.

Parses the structured QuotaFailure violations in the response body and looks for a quotaId containing PerDay.

Return type:

bool

Parameters:

resp (httpx.Response)

gemini_embed_pool.is_daily_quota_429_for_model(resp)

Identify which model class a daily-quota 429 belongs to.

Returns "embed" or "generate" based on the quotaDimensions.model field, or None if the response is not a daily 429.

Return type:

Optional[Literal['embed', 'generate']]

Parameters:

resp (httpx.Response)

async gemini_embed_pool.record_key_usage(api_key)

Increment the daily request counter for api_key in Redis (best-effort).

Bumps gemini_key_daily_usage:<suffix>:count and (re)sets its TTL to the next midnight Pacific via a non-transactional pipeline, giving the diagnostics in get_pool_status() a rough per-key call count that self-expires daily. No-ops when the async Redis client is unwired, and swallows any Redis error rather than disrupting the embed call.

Called within this module by probe_all_keys() and gemini_embed_paid_fallback(), and externally by the embed transport in openrouter_client/transport.py after each Gemini POST.

Parameters:

api_key (str) – The Gemini key whose usage to record.

Return type:

None

async gemini_embed_pool.mark_key_daily_spent(api_key, model_class='embed')

Mark api_key as daily-spent for model_class.

Updates both the in-memory set and the Redis flag.

Return type:

None

Parameters:

• api_key (str)

• model_class (Literal['embed', 'generate'])

async gemini_embed_pool.sync_spent_keys_from_redis()

Rebuild both in-memory daily-spent sets from the per-key flags in Redis.

Lets a worker pick up exhaustion decisions made by other processes: for every pooled key it reads gemini_key_daily_usage:<suffix>:spent:embed and ...:spent:generate (plus the legacy suffix-less ...:spent key, which counts for both classes) and replaces _spent_keys_embed / _spent_keys_generate accordingly under _spent_keys_lock. Calls _maybe_reset_day() first so a quota-day rollover is honored, and no-ops when the async Redis client is unwired; Redis errors are swallowed.

Called within this module at the start of probe_all_keys(); no external callers were found.

Return type:

None

async gemini_embed_pool.get_pool_status()

Return a snapshot of pool health (active vs spent keys, usage) for diagnostics.

Calls _maybe_reset_day() to honor a quota-day rollover, then reports the total key count, active and daily-spent counts split per model class (embed vs generate), and the current quota day. When the async Redis client is wired it also reads each key’s gemini_key_daily_usage:<suffix>:count and includes a ...<suffix>-keyed usage_counts map; Redis failures are tolerated and simply omit those counts.

Called externally by the donate_embed_key tool (tools/donate_embed_key.py) to surface pool status to users.

Returns:

Pool health fields including total_keys, embed_active / embed_spent, generate_active / generate_spent, quota_day, and (when Redis is available) usage_counts.

Return type:

dict[str, Any]

async gemini_embed_pool.probe_all_keys()

Probe each non-spent key for both embedding and generation daily exhaustion.

Sends a minimal request to each endpoint and marks the key spent for the relevant model class on a daily 429.

Return type:

None

exception gemini_embed_pool.OpenRouterEmbedParseError

Bases: RuntimeError

OpenRouter returned HTTP 200 but the body was not the expected {"data": [{"index": int, "embedding": [..]}...]} shape.

Treated as non-retriable by the outer retry loops in openrouter_client and rag_system.openrouter_embeddings, since the underlying provider is returning a malformed success payload (typically an upstream error surfaced as a 200) and immediate retries will hit the same issue.

exception gemini_embed_pool.OpenRouterEmbedHTTPError(status_code, body)

Bases: RuntimeError

A non-success response from the OpenRouter embeddings endpoint.

Parameters:

• status_code (int)

• body (str)

exception gemini_embed_pool.OpenRouterEmbedCircuitOpenError(retry_after_seconds, status_code)

Bases: RuntimeError

OpenRouter embeds are cooling down after an account-level failure.

Parameters:

• retry_after_seconds (float)

• status_code (int | None)

gemini_embed_pool.get_openrouter_api_key()

Return the OpenRouter API key for the embed fallback tier.

Prefers the legacy API_KEY env var when it is the only thing set, then resolves openrouter.embed through system_secret_store.get_secret_sync(), which itself checks OPENROUTER_API_KEY, the TTL cache warmed by reload_pool(), Redis, and finally the api_key entry in config.yaml. The result becomes the Bearer token for the embeddings endpoint. No key is baked into source, so an unprovisioned store yields None and the caller raises rather than silently authenticating with a leaked credential.

Called within this module by openrouter_embed_batch() and openrouter_embed_batch_sync() when no explicit api_key is passed.

Returns:

The resolved OpenRouter key, or None when unprovisioned.

Return type:

str | None

async gemini_embed_pool.openrouter_embed_batch(texts, *, model='google/gemini-embedding-001', api_key=None, dimensions=3072)

Embed texts via the OpenRouter /embeddings endpoint (async).

Retries on transient network failures and retriable HTTP codes (429, 5xx) until success or OPENROUTER_EMBED_MAX_ATTEMPTS is exhausted.

Empty or whitespace-only strings are not sent to the API; those positions receive zero vectors of length dimensions.

Return type:

list[list[float]]

Parameters:

• texts (list[str])

• model (str)

• api_key (str | None)

• dimensions (int)

gemini_embed_pool.openrouter_embed_batch_sync(texts, *, model='google/gemini-embedding-001', api_key=None, dimensions=3072)

Embed texts via the OpenRouter /embeddings endpoint (sync).

Synchronous version for callers that cannot use await (e.g. ChromaDB embedding functions). Same retry policy as openrouter_embed_batch().

Empty or whitespace-only strings are not sent to the API; those positions receive zero vectors of length dimensions.

Return type:

list[list[float]]

Parameters:

• texts (list[str])

• model (str)

• api_key (str | None)

• dimensions (int)

async gemini_embed_pool.gemini_embed_paid_fallback(texts, *, model='google/gemini-embedding-001', dimensions=3072, task_type=None)

Last-resort embedding via the paid tier-3 Gemini key (async).

Retries provider-pressure failures before returning control to the caller. Transient network errors and HTTP 408/429/500/502/503/504 retry with capped exponential backoff and jitter. Credential/protocol failures and daily-quota exhaustion fail immediately. Key usage is recorded in Redis when wired.

Empty/whitespace-only inputs receive zero vectors of length dimensions.

Raises RuntimeError when no paid key is configured, when the call fails non-2xx, or when the paid key 429s.

Return type:

list[list[float]]

Parameters:

• texts (list[str])

• model (str)

• dimensions (int)

• task_type (str | None)

gemini_embed_pool.gemini_embed_paid_fallback_sync(texts, *, model='google/gemini-embedding-001', dimensions=3072, task_type=None)

Sync mirror of gemini_embed_paid_fallback().

Used by the ChromaDB sync embedding paths. Daily-quota tracking in Redis is best-effort (no-op when the sync Redis client is unavailable), since the paid key is only hit on the slow last-resort path.

Return type:

list[list[float]]

Parameters:

• texts (list[str])

• model (str)

• dimensions (int)

• task_type (str | None)

gemini_embed_pool.is_openrouter_only()

Return the in-memory OpenRouter-only flag without touching Redis.

A cheap, synchronous read of the process-local _openrouter_only state set by set_openrouter_only() / cleared by clear_openrouter_only(). Unlike check_openrouter_only(), it never consults Redis. It does enforce the process-local TTL, but cannot observe activations from other processes.

No in-repo callers were found outside this module; it serves as a public no-I/O accessor for the circuit-breaker state.

Returns:

True if OpenRouter-only mode is currently flagged in this process.

Return type:

bool

async gemini_embed_pool.check_openrouter_only()

Return whether OpenRouter-only mode is active.

When Redis is wired (init_quota_tracking), reads the TTL key so the mode expires when Redis lapses or is cleared. When Redis is not wired, the same TTL is enforced against a process-local monotonic deadline.

Return type:

bool

gemini_embed_pool.check_openrouter_only_sync()

Sync variant: check Redis before each embedding call.

Used by SyncOpenRouterEmbeddings (ChromaDB path). Lazy-creates a sync Redis client from config on first use.

Return type:

bool

async gemini_embed_pool.set_openrouter_only()

Activate fallback-only mode for embedding calls.

Sets both the in-memory flag and a Redis key with a TTL of _OPENROUTER_ONLY_TTL (30 minutes).

Return type:

None

gemini_embed_pool.set_openrouter_only_sync()

Sync mirror of set_openrouter_only().

Activates the process-local fallback-only flag immediately and persists the legacy embed:openrouter_only TTL key when sync Redis is available.

Return type:

None

async gemini_embed_pool.clear_openrouter_only()

Deactivate OpenRouter-only mode, restoring the free Gemini pool as primary.

Clears the in-memory _openrouter_only flag and deletes the embed:openrouter_only Redis key (when the async client is wired) so the embed cascade resumes using the free pool first. Acts as the manual / auto counterpart to set_openrouter_only(); Redis errors are swallowed.

Called within this module by _auto_clear_openrouter_only_on_parse_error(), and externally by the embedding-refresh jobs classifiers/update_tool_embeddings.py and classifiers/update_changed_tool_embeddings.py at the end of a run.

Return type:

None

gemini_embed_pool.clear_openrouter_only_sync()

Sync variant: deactivate OpenRouter-only mode.

Mirrors clear_openrouter_only() for code paths that cannot await, using the same sync Redis client that check_openrouter_only_sync() initializes on first use.

Return type:

None

async gemini_embed_pool.embed_batch_via_gemini(texts, model='google/gemini-embedding-001', *, chunk_size=50)

Embed a batch of texts via the native Gemini API using the shared key pool.

Returns one embedding vector per input text, in order. Empty or whitespace-only texts are replaced with zero vectors. Retries on transient errors.

Cascade on failure (after PAID_KEY_FALLBACK_THRESHOLD consecutive non-daily 429s on the free pool):

1. OpenRouter via openrouter_embed_batch() (and pin openrouter_only for the next 30 minutes).

2. Paid tier-3 Gemini key via gemini_embed_paid_fallback() (last resort).

When openrouter_only is already pinned, OpenRouter is tried first; if it fails on a specific batch, the paid Gemini key is tried before raising.

Return type:

list[list[float]]

Parameters:

• texts (list[str])

• model (str)

• chunk_size (int)

async gemini_embed_pool.batch_check_keys_usage(redis_client, api_keys)

Fetch the daily usage counts for many API keys in one pipelined Redis MGET.

Reads spent:api_key:<suffix> for every supplied key (keyed by the last 8 characters for privacy) in a single round trip and maps each full key to its integer count, defaulting missing entries to 0. Note this uses a different key namespace from the gemini_key_daily_usage:* counters written by record_key_usage(), so it reflects a separately maintained tally.

No production callers were found in the repo; it is currently exercised by tests/test_context_whitelisting.py.

Parameters:

• redis_client (Any) – An async Redis client supporting mget.

• api_keys (list[str]) – The full API keys to look up.

Returns:

A mapping from each input API key to its usage count.

Return type:

dict[str, int]