From 9ecd77e23b47c8e2b27213f9e849f066911d947c Mon Sep 17 00:00:00 2001 From: Roo Code Date: Wed, 16 Jul 2025 08:56:17 +0000 Subject: [PATCH] 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 | 3 + .../embedders/__tests__/gemini.spec.ts | 93 +++++++++++- src/services/code-index/embedders/gemini.ts | 141 +++++++++++++++++- 3 files changed, 228 insertions(+), 9 deletions(-) diff --git a/src/services/code-index/constants/index.ts b/src/services/code-index/constants/index.ts index c2567f5635..cada8ea156 100644 --- a/src/services/code-index/constants/index.ts +++ b/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 diff --git a/src/services/code-index/embedders/__tests__/gemini.spec.ts b/src/services/code-index/embedders/__tests__/gemini.spec.ts index d41a4dc1e9..b79b230c1d 100644 --- a/src/services/code-index/embedders/__tests__/gemini.spec.ts +++ b/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 }, + }) + }) }) }) diff --git a/src/services/code-index/embedders/gemini.ts b/src/services/code-index/embedders/gemini.ts index 7e795875c9..05493e46be 100644 --- a/src/services/code-index/embedders/gemini.ts +++ b/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 { 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 { + 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