fix: improve rate limiting for gemini-embedding-001 high-dimensional model

- Add Gemini-specific constants for enhanced rate limiting
- Implement model-aware batch size reduction (25k tokens vs 100k)
- Add longer retry delays for high-dimensional models (1000ms vs 500ms)
- Increase max retries for rate-limited requests (5 vs 3)
- Add custom batching logic for models with 3000+ dimensions
- Update tests to cover new model-aware behavior
- Fix dimension documentation for gemini-embedding-001 (3072 not 2048)

Fixes #5774

src/services/code-index/constants/index.ts

@@ -28,3 +28,6 @@ export const BATCH_PROCESSING_CONCURRENCY = 10
 
 /**Gemini Embedder */
 export const GEMINI_MAX_ITEM_TOKENS = 2048
+export const GEMINI_MAX_BATCH_TOKENS = 25000 // Reduced batch size for high-dimensional models
+export const GEMINI_INITIAL_RETRY_DELAY_MS = 1000 // Longer initial delay for Gemini
+export const GEMINI_MAX_BATCH_RETRIES = 5 // More retries for rate-limited requests

src/services/code-index/embedders/__tests__/gemini.spec.ts

@@ -6,6 +6,17 @@ import { OpenAICompatibleEmbedder } from "../openai-compatible"
 // Mock the OpenAICompatibleEmbedder
 vitest.mock("../openai-compatible")
 
+// Mock the embedding models module
+vitest.mock("../../../../shared/embeddingModels", () => ({
+	getModelDimension: vitest.fn((provider: string, modelId: string) => {
+		if (provider === "gemini") {
+			if (modelId === "gemini-embedding-001") return 3072
+			if (modelId === "text-embedding-004") return 768
+		}
+		return undefined
+	}),
+}))
+
 // Mock TelemetryService
 vitest.mock("@roo-code/telemetry", () => ({
 	TelemetryService: {
@@ -88,15 +99,37 @@ describe("GeminiEmbedder", () => {
 				MockedOpenAICompatibleEmbedder.prototype.createEmbeddings = mockCreateEmbeddings
 			})
 
-			it("should use instance model when no model parameter provided", async () => {
+			it("should use standard implementation for low-dimensional models (text-embedding-004)", async () => {
 				// Arrange
-				embedder = new GeminiEmbedder("test-api-key")
+				embedder = new GeminiEmbedder("test-api-key", "text-embedding-004")
 				const texts = ["test text 1", "test text 2"]
 				const mockResponse = {
 					embeddings: [
 						[0.1, 0.2],
 						[0.3, 0.4],
 					],
+					usage: { promptTokens: 10, totalTokens: 15 },
+				}
+				mockCreateEmbeddings.mockResolvedValue(mockResponse)
+
+				// Act
+				const result = await embedder.createEmbeddings(texts)
+
+				// Assert
+				expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "text-embedding-004")
+				expect(result).toEqual(mockResponse)
+			})
+
+			it("should use custom rate limiting for high-dimensional models (gemini-embedding-001)", async () => {
+				// Arrange
+				embedder = new GeminiEmbedder("test-api-key", "gemini-embedding-001")
+				const texts = ["test text 1", "test text 2"]
+				const mockResponse = {
+					embeddings: [
+						[0.1, 0.2],
+						[0.3, 0.4],
+					],
+					usage: { promptTokens: 10, totalTokens: 15 },
 				}
 				mockCreateEmbeddings.mockResolvedValue(mockResponse)
 
@@ -117,6 +150,7 @@ describe("GeminiEmbedder", () => {
 						[0.1, 0.2],
 						[0.3, 0.4],
 					],
+					usage: { promptTokens: 10, totalTokens: 15 },
 				}
 				mockCreateEmbeddings.mockResolvedValue(mockResponse)
 
@@ -128,9 +162,28 @@ describe("GeminiEmbedder", () => {
 				expect(result).toEqual(mockResponse)
 			})
 
-			it("should handle errors from OpenAICompatibleEmbedder", async () => {
+			it("should use custom rate limiting implementation for high-dimensional models", async () => {
 				// Arrange
-				embedder = new GeminiEmbedder("test-api-key")
+				embedder = new GeminiEmbedder("test-api-key", "gemini-embedding-001")
+				const texts = ["test text"]
+				const mockResponse = {
+					embeddings: [[0.1, 0.2]],
+					usage: { promptTokens: 5, totalTokens: 8 },
+				}
+				mockCreateEmbeddings.mockResolvedValue(mockResponse)
+
+				// Act
+				const result = await embedder.createEmbeddings(texts)
+
+				// Assert
+				// For high-dimensional models, the custom implementation should be used
+				expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "gemini-embedding-001")
+				expect(result).toEqual(mockResponse)
+			})
+
+			it("should handle errors from OpenAICompatibleEmbedder for low-dimensional models", async () => {
+				// Arrange
+				embedder = new GeminiEmbedder("test-api-key", "text-embedding-004")
 				const texts = ["test text"]
 				const error = new Error("Embedding failed")
 				mockCreateEmbeddings.mockRejectedValue(error)
@@ -138,6 +191,38 @@ describe("GeminiEmbedder", () => {
 				// Act & Assert
 				await expect(embedder.createEmbeddings(texts)).rejects.toThrow("Embedding failed")
 			})
+
+			it("should handle errors from OpenAICompatibleEmbedder for high-dimensional models", async () => {
+				// Arrange
+				embedder = new GeminiEmbedder("test-api-key", "gemini-embedding-001")
+				const texts = ["test text"]
+				const error = new Error("Embedding failed")
+				mockCreateEmbeddings.mockRejectedValue(error)
+
+				// Act & Assert
+				// For high-dimensional models, errors go through custom retry logic
+				await expect(embedder.createEmbeddings(texts)).rejects.toThrow("failedMaxAttempts")
+			})
+
+			it("should handle missing usage data gracefully", async () => {
+				// Arrange
+				embedder = new GeminiEmbedder("test-api-key", "gemini-embedding-001")
+				const texts = ["test text"]
+				const mockResponse = {
+					embeddings: [[0.1, 0.2]],
+					usage: undefined, // Missing usage data
+				}
+				mockCreateEmbeddings.mockResolvedValue(mockResponse)
+
+				// Act
+				const result = await embedder.createEmbeddings(texts)
+
+				// Assert
+				expect(result).toEqual({
+					embeddings: [[0.1, 0.2]],
+					usage: { promptTokens: 0, totalTokens: 0 },
+				})
+			})
 		})
 	})
 

src/services/code-index/embedders/gemini.ts

@@ -1,6 +1,12 @@
 import { OpenAICompatibleEmbedder } from "./openai-compatible"
 import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces/embedder"
-import { GEMINI_MAX_ITEM_TOKENS } from "../constants"
+import {
+	GEMINI_MAX_ITEM_TOKENS,
+	GEMINI_MAX_BATCH_TOKENS,
+	GEMINI_INITIAL_RETRY_DELAY_MS,
+	GEMINI_MAX_BATCH_RETRIES,
+} from "../constants"
+import { getModelDimension } from "../../../shared/embeddingModels"
 import { t } from "../../../i18n"
 import { TelemetryEventName } from "@roo-code/types"
 import { TelemetryService } from "@roo-code/telemetry"
@@ -11,7 +17,7 @@ import { TelemetryService } from "@roo-code/telemetry"
  *
  * Supported models:
  * - text-embedding-004 (dimension: 768)
- * - gemini-embedding-001 (dimension: 2048)
+ * - gemini-embedding-001 (dimension: 3072)
  */
 export class GeminiEmbedder implements IEmbedder {
 	private readonly openAICompatibleEmbedder: OpenAICompatibleEmbedder
@@ -42,16 +48,26 @@ export class GeminiEmbedder implements IEmbedder {
 	}
 
 	/**
-	 * Creates embeddings for the given texts using Gemini's embedding API
+	 * Creates embeddings for the given texts using Gemini's embedding API with model-aware rate limiting
 	 * @param texts Array of text strings to embed
 	 * @param model Optional model identifier (uses constructor model if not provided)
 	 * @returns Promise resolving to embedding response
 	 */
 	async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
 		try {
-			// Use the provided model or fall back to the instance's model
 			const modelToUse = model || this.modelId
-			return await this.openAICompatibleEmbedder.createEmbeddings(texts, modelToUse)
+
+			// Check if this is a high-dimensional model that needs special handling
+			const modelDimension = getModelDimension("gemini", modelToUse)
+			const isHighDimensionalModel = modelDimension && modelDimension >= 3000
+
+			if (isHighDimensionalModel) {
+				// Use custom batching and rate limiting for high-dimensional models
+				return await this._createEmbeddingsWithCustomRateLimiting(texts, modelToUse)
+			} else {
+				// Use standard implementation for lower-dimensional models
+				return await this.openAICompatibleEmbedder.createEmbeddings(texts, modelToUse)
+			}
 		} catch (error) {
 			TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
 				error: error instanceof Error ? error.message : String(error),
@@ -62,6 +78,121 @@ export class GeminiEmbedder implements IEmbedder {
 		}
 	}
 
+	/**
+	 * Custom embedding implementation with enhanced rate limiting for high-dimensional models
+	 * @param texts Array of text strings to embed
+	 * @param model Model identifier to use
+	 * @returns Promise resolving to embedding response
+	 */
+	private async _createEmbeddingsWithCustomRateLimiting(texts: string[], model: string): Promise<EmbeddingResponse> {
+		const allEmbeddings: number[][] = []
+		const usage = { promptTokens: 0, totalTokens: 0 }
+		const remainingTexts = [...texts]
+
+		while (remainingTexts.length > 0) {
+			const currentBatch: string[] = []
+			let currentBatchTokens = 0
+			const processedIndices: number[] = []
+
+			// Use smaller batch sizes for high-dimensional models
+			for (let i = 0; i < remainingTexts.length; i++) {
+				const text = remainingTexts[i]
+				const itemTokens = Math.ceil(text.length / 4)
+
+				if (itemTokens > GEMINI_MAX_ITEM_TOKENS) {
+					console.warn(
+						t("embeddings:textExceedsTokenLimit", {
+							index: i,
+							itemTokens,
+							maxTokens: GEMINI_MAX_ITEM_TOKENS,
+						}),
+					)
+					processedIndices.push(i)
+					continue
+				}
+
+				if (currentBatchTokens + itemTokens <= GEMINI_MAX_BATCH_TOKENS) {
+					currentBatch.push(text)
+					currentBatchTokens += itemTokens
+					processedIndices.push(i)
+				} else {
+					break
+				}
+			}
+
+			// Remove processed items from remainingTexts (in reverse order to maintain correct indices)
+			for (let i = processedIndices.length - 1; i >= 0; i--) {
+				remainingTexts.splice(processedIndices[i], 1)
+			}
+
+			if (currentBatch.length > 0) {
+				const batchResult = await this._embedBatchWithGeminiRateLimiting(currentBatch, model)
+				allEmbeddings.push(...batchResult.embeddings)
+				usage.promptTokens += batchResult.usage.promptTokens
+				usage.totalTokens += batchResult.usage.totalTokens
+			}
+		}
+
+		return { embeddings: allEmbeddings, usage }
+	}
+
+	/**
+	 * Helper method to handle batch embedding with Gemini-specific retries and rate limiting
+	 * @param batchTexts Array of texts to embed in this batch
+	 * @param model Model identifier to use
+	 * @returns Promise resolving to embeddings and usage statistics
+	 */
+	private async _embedBatchWithGeminiRateLimiting(
+		batchTexts: string[],
+		model: string,
+	): Promise<{ embeddings: number[][]; usage: { promptTokens: number; totalTokens: number } }> {
+		for (let attempts = 0; attempts < GEMINI_MAX_BATCH_RETRIES; attempts++) {
+			try {
+				// Delegate to the underlying OpenAI Compatible embedder for the actual API call
+				const response = await this.openAICompatibleEmbedder.createEmbeddings(batchTexts, model)
+				return {
+					embeddings: response.embeddings,
+					usage: response.usage || { promptTokens: 0, totalTokens: 0 },
+				}
+			} catch (error) {
+				TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
+					error: error instanceof Error ? error.message : String(error),
+					stack: error instanceof Error ? error.stack : undefined,
+					location: "GeminiEmbedder:_embedBatchWithGeminiRateLimiting",
+					attempt: attempts + 1,
+				})
+
+				const hasMoreAttempts = attempts < GEMINI_MAX_BATCH_RETRIES - 1
+
+				// Check if it's a rate limit error (429)
+				const httpError = error as any
+				if (httpError?.status === 429 && hasMoreAttempts) {
+					// Use longer delays for Gemini high-dimensional models
+					const delayMs = GEMINI_INITIAL_RETRY_DELAY_MS * Math.pow(2, attempts)
+					console.warn(
+						t("embeddings:rateLimitRetry", {
+							delayMs,
+							attempt: attempts + 1,
+							maxRetries: GEMINI_MAX_BATCH_RETRIES,
+						}),
+					)
+					await new Promise((resolve) => setTimeout(resolve, delayMs))
+					continue
+				}
+
+				// Log the error for debugging
+				console.error(`Gemini embedder error (attempt ${attempts + 1}/${GEMINI_MAX_BATCH_RETRIES}):`, error)
+
+				// If it's the last attempt or not a rate limit error, throw the error
+				if (!hasMoreAttempts) {
+					throw new Error(t("embeddings:failedMaxAttempts", { attempts: GEMINI_MAX_BATCH_RETRIES }))
+				}
+			}
+		}
+
+		throw new Error(t("embeddings:failedMaxAttempts", { attempts: GEMINI_MAX_BATCH_RETRIES }))
+	}
+
 	/**
 	 * Validates the Gemini embedder configuration by delegating to the underlying OpenAI-compatible embedder
 	 * @returns Promise resolving to validation result with success status and optional error message