fix: implement model-specific rate limiting for gemini-embedding-001

- Add Gemini-specific constants for gemini-embedding-001:
  - Reduced batch token limit (20,000 vs 100,000)
  - Longer retry delays (2000ms vs 500ms)
  - Smaller batch size limit (10 items)
- Update OpenAICompatibleEmbedder to accept configurable rate limiting parameters
- Add inter-batch delays for models with stricter rate limits
- Update GeminiEmbedder to use model-specific configuration
- Fix test expectations to match new constructor signature

Fixes #5713: gemini-embedding-001 quota limit issues during indexing
This commit is contained in:
Roo Code 2025-07-15 00:41:43 +00:00
parent 88c4261829
commit 5f2771224f
4 changed files with 55 additions and 4 deletions

View file

@ -28,3 +28,6 @@ export const BATCH_PROCESSING_CONCURRENCY = 10
/**Gemini Embedder */
export const GEMINI_MAX_ITEM_TOKENS = 2048
export const GEMINI_EMBEDDING_001_MAX_BATCH_TOKENS = 20000 // Reduced batch size for gemini-embedding-001
export const GEMINI_EMBEDDING_001_RETRY_DELAY_MS = 2000 // Longer delay for gemini-embedding-001
export const GEMINI_EMBEDDING_001_MAX_BATCH_SIZE = 10 // Smaller batch size for gemini-embedding-001

View file

@ -38,6 +38,9 @@ describe("GeminiEmbedder", () => {
apiKey,
"gemini-embedding-001",
2048,
20000, // GEMINI_EMBEDDING_001_MAX_BATCH_TOKENS
2000, // GEMINI_EMBEDDING_001_RETRY_DELAY_MS
10, // GEMINI_EMBEDDING_001_MAX_BATCH_SIZE
)
})
@ -55,6 +58,9 @@ describe("GeminiEmbedder", () => {
apiKey,
"text-embedding-004",
2048,
100000, // MAX_BATCH_TOKENS (default for text-embedding-004)
500, // INITIAL_RETRY_DELAY_MS (default for text-embedding-004)
undefined, // maxBatchSize (undefined for text-embedding-004)
)
})

View file

@ -1,6 +1,13 @@
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_EMBEDDING_001_MAX_BATCH_TOKENS,
GEMINI_EMBEDDING_001_RETRY_DELAY_MS,
GEMINI_EMBEDDING_001_MAX_BATCH_SIZE,
MAX_BATCH_TOKENS,
INITIAL_RETRY_DELAY_MS,
} from "../constants"
import { t } from "../../../i18n"
import { TelemetryEventName } from "@roo-code/types"
import { TelemetryService } from "@roo-code/telemetry"
@ -32,12 +39,21 @@ export class GeminiEmbedder implements IEmbedder {
// Use provided model or default
this.modelId = modelId || GeminiEmbedder.DEFAULT_MODEL
// Get model-specific configuration for gemini-embedding-001
const isGeminiEmbedding001 = this.modelId === "gemini-embedding-001"
const maxBatchTokens = isGeminiEmbedding001 ? GEMINI_EMBEDDING_001_MAX_BATCH_TOKENS : MAX_BATCH_TOKENS
const retryDelayMs = isGeminiEmbedding001 ? GEMINI_EMBEDDING_001_RETRY_DELAY_MS : INITIAL_RETRY_DELAY_MS
const maxBatchSize = isGeminiEmbedding001 ? GEMINI_EMBEDDING_001_MAX_BATCH_SIZE : undefined
// Create an OpenAI Compatible embedder with Gemini's configuration
this.openAICompatibleEmbedder = new OpenAICompatibleEmbedder(
GeminiEmbedder.GEMINI_BASE_URL,
apiKey,
this.modelId,
GEMINI_MAX_ITEM_TOKENS,
maxBatchTokens,
retryDelayMs,
maxBatchSize,
)
}

View file

@ -37,6 +37,9 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
private readonly apiKey: string
private readonly isFullUrl: boolean
private readonly maxItemTokens: number
private readonly maxBatchTokens: number
private readonly retryDelayMs: number
private readonly maxBatchSize?: number
/**
* Creates a new OpenAI Compatible embedder
@ -44,8 +47,19 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
* @param apiKey The API key for authentication
* @param modelId Optional model identifier (defaults to "text-embedding-3-small")
* @param maxItemTokens Optional maximum tokens per item (defaults to MAX_ITEM_TOKENS)
* @param maxBatchTokens Optional maximum tokens per batch (defaults to MAX_BATCH_TOKENS)
* @param retryDelayMs Optional initial retry delay in milliseconds (defaults to INITIAL_DELAY_MS)
* @param maxBatchSize Optional maximum number of items per batch
*/
constructor(baseUrl: string, apiKey: string, modelId?: string, maxItemTokens?: number) {
constructor(
baseUrl: string,
apiKey: string,
modelId?: string,
maxItemTokens?: number,
maxBatchTokens?: number,
retryDelayMs?: number,
maxBatchSize?: number,
) {
if (!baseUrl) {
throw new Error(t("embeddings:validation.baseUrlRequired"))
}
@ -63,6 +77,9 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
// Cache the URL type check for performance
this.isFullUrl = this.isFullEndpointUrl(baseUrl)
this.maxItemTokens = maxItemTokens || MAX_ITEM_TOKENS
this.maxBatchTokens = maxBatchTokens || MAX_BATCH_TOKENS
this.retryDelayMs = retryDelayMs || INITIAL_DELAY_MS
this.maxBatchSize = maxBatchSize
}
/**
@ -124,7 +141,10 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
continue
}
if (currentBatchTokens + itemTokens <= MAX_BATCH_TOKENS) {
if (
currentBatchTokens + itemTokens <= this.maxBatchTokens &&
(!this.maxBatchSize || currentBatch.length < this.maxBatchSize)
) {
currentBatch.push(text)
currentBatchTokens += itemTokens
processedIndices.push(i)
@ -143,6 +163,12 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
allEmbeddings.push(...batchResult.embeddings)
usage.promptTokens += batchResult.usage.promptTokens
usage.totalTokens += batchResult.usage.totalTokens
// Add delay between batches if there are more batches to process
// This helps with rate limiting, especially for gemini-embedding-001
if (remainingTexts.length > 0 && this.retryDelayMs > INITIAL_DELAY_MS) {
await new Promise((resolve) => setTimeout(resolve, this.retryDelayMs / 4))
}
}
}
@ -299,7 +325,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
// Check if it's a rate limit error
const httpError = error as HttpError
if (httpError?.status === 429 && hasMoreAttempts) {
const delayMs = INITIAL_DELAY_MS * Math.pow(2, attempts)
const delayMs = this.retryDelayMs * Math.pow(2, attempts)
console.warn(
t("embeddings:rateLimitRetry", {
delayMs,