feat: GeminiEmbedding rate-limit handling (#2237)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/nice-zebras-taste.md

@@ -0,0 +1,6 @@
+---
+"@llamaindex/google": minor
+"@llamaindex/examples": patch
+---
+
+respect Gemini's requests-per-minute rate limit with waits

examples/models/gemini/embedding_ratelimits.ts

@@ -0,0 +1,20 @@
+import { GEMINI_EMBEDDING_MODEL, GeminiEmbedding } from "@llamaindex/google";
+
+const requests_per_minute_limit = 3000; // cf. https://ai.google.dev/gemini-api/docs/rate-limits
+
+async function main() {
+  if (!process.env.GOOGLE_API_KEY) {
+    throw new Error("Please set the GOOGLE_API_KEY environment variable.");
+  }
+  const embedModel = new GeminiEmbedding({
+    model: GEMINI_EMBEDDING_MODEL.EMBEDDING_001,
+  });
+  const texts = Array.from(
+    { length: requests_per_minute_limit + 1000 },
+    (_, i) => `text ${i}`,
+  );
+  const embeddings = await embedModel.getTextEmbeddingsBatch(texts);
+  console.log(`\nWe have ${embeddings.length} embeddings`);
+}
+
+main().catch(console.error);

packages/providers/google/src/GeminiEmbedding.ts

@@ -43,12 +43,59 @@ export class GeminiEmbedding extends BaseEmbedding {
     this.embedBatchSize = opts?.embedBatchSize ?? DEFAULT_EMBED_BATCH_SIZE;
   }
 
+  // Add a retry wrapper for embedContent to handle rate limits (5s wait, up to 20 tries)
+  private async embedWithRetry(args: {
+    model: string;
+    contents: string | string[];
+  }) {
+    const MAX_TRIES = 20;
+    const DELAY_MS = 5000;
+
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const isRateLimitError = (err: any) => {
+      if (!err) return false;
+      return (
+        err.status === 429 ||
+        err.message.indexOf('"status":"RESOURCE_EXHAUSTED"') !== -1
+      );
+    };
+    /* error looks like this
+      {"error":{"code":429,"message":"You exceeded your current quota, please check your plan and billing details. For more information on this error, head to: https://ai.google.dev/gemini-api/docs/rate-limits. To monitor your current usage, head to: https://ai.dev/usage?tab=rate-limit. \n* Quota exceeded for metric: generativelanguage.googleapis.com/embed_content_paid_tier_requests, limit: 0","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.Help","links":[{"description":"Learn more about Gemini API quotas","url":"https://ai.google.dev/gemini-api/docs/rate-limits"}]},{"@type":"type.googleapis.com/google.rpc.QuotaFailure","violations":[{"quotaMetric":"generativelanguage.googleapis.com/embed_content_paid_tier_requests","quotaId":"EmbedContentPerMinutePerProjectPerUserPerModel-PaidTier"}]}]}}
+    */
+
+    const sleep = (ms: number) => new Promise((res) => setTimeout(res, ms));
+
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    let lastErr: any = null;
+    for (let attempt = 1; attempt <= MAX_TRIES; attempt++) {
+      try {
+        return await this.ai.models.embedContent(args);
+      } catch (err) {
+        lastErr = err;
+        if (isRateLimitError(err) && attempt < MAX_TRIES) {
+          await sleep(DELAY_MS);
+          console.debug(
+            `Gemini Embeddings rate limit error encountered. Retrying attempt ${attempt}...`,
+          );
+          continue;
+        }
+        throw err;
+      }
+    }
+    // If we exit loop unexpectedly, throw the last error.
+    throw lastErr;
+  }
+
   getTextEmbeddings = async (texts: string[]) => {
-    const result = await this.ai.models.embedContent({
+    const result = await this.embedWithRetry({
       model: this.model,
       contents: texts,
     });
-    return result.embeddings?.map((embedding) => embedding.values ?? []) ?? [];
+
+    return (
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      result.embeddings?.map((embedding: any) => embedding.values ?? []) ?? []
+    );
   };
 
   async getTextEmbeddingsBatch(
@@ -64,7 +111,7 @@ export class GeminiEmbedding extends BaseEmbedding {
   }
 
   async getTextEmbedding(text: string): Promise<number[]> {
-    const result = await this.ai.models.embedContent({
+    const result = await this.embedWithRetry({
       model: this.model,
       contents: text,
     });