mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-05 08:10:14 +00:00
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:
parent
88c4261829
commit
5f2771224f
4 changed files with 55 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
)
|
||||
})
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue