diff --git a/packages/server/src/utils/LruCache.ts b/packages/server/src/utils/LruCache.ts new file mode 100644 index 00000000..076828c4 --- /dev/null +++ b/packages/server/src/utils/LruCache.ts @@ -0,0 +1,41 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export class LruCache { + private cache: Map; + private maxSize: number; + + constructor(maxSize: number) { + this.cache = new Map(); + this.maxSize = maxSize; + } + + get(key: K): V | undefined { + const value = this.cache.get(key); + if (value) { + // Move to end to mark as recently used + this.cache.delete(key); + this.cache.set(key, value); + } + return value; + } + + set(key: K, value: V): void { + if (this.cache.has(key)) { + this.cache.delete(key); + } else if (this.cache.size >= this.maxSize) { + const firstKey = this.cache.keys().next().value; + if (firstKey !== undefined) { + this.cache.delete(firstKey); + } + } + this.cache.set(key, value); + } + + clear(): void { + this.cache.clear(); + } +} diff --git a/packages/server/src/utils/editCorrector.ts b/packages/server/src/utils/editCorrector.ts index a17de6d4..91aa85a3 100644 --- a/packages/server/src/utils/editCorrector.ts +++ b/packages/server/src/utils/editCorrector.ts @@ -12,6 +12,7 @@ import { } from '@google/genai'; import { GeminiClient } from '../core/client.js'; import { EditToolParams } from '../tools/edit.js'; +import { LruCache } from './LruCache.js'; const EditModel = 'gemini-2.5-flash-preview-04-17'; const EditConfig: GenerateContentConfig = { @@ -20,6 +21,13 @@ const EditConfig: GenerateContentConfig = { }, }; +const MAX_CACHE_SIZE = 50; + +// Cache for ensureCorrectEdit results +const editCorrectionCache = new LruCache( + MAX_CACHE_SIZE, +); + /** * Defines the structure of the parameters within CorrectedEditResult */ @@ -53,6 +61,12 @@ export async function ensureCorrectEdit( originalParams: EditToolParams, // This is the EditToolParams from edit.ts, without \'corrected\' client: GeminiClient, ): Promise { + const cacheKey = `${currentContent}---${originalParams.old_string}---${originalParams.new_string}`; + const cachedResult = editCorrectionCache.get(cacheKey); + if (cachedResult) { + return cachedResult; + } + let finalNewString = originalParams.new_string; const newStringPotentiallyEscaped = unescapeStringForGeminiBug(originalParams.new_string) !== @@ -74,6 +88,7 @@ export async function ensureCorrectEdit( params: { ...originalParams }, occurrences, }; + editCorrectionCache.set(cacheKey, result); return result; } else { // occurrences is 0 or some other unexpected state initially @@ -124,6 +139,7 @@ export async function ensureCorrectEdit( params: { ...originalParams }, occurrences: 0, // Explicitly 0 as LLM failed }; + editCorrectionCache.set(cacheKey, result); return result; } } else { @@ -132,6 +148,7 @@ export async function ensureCorrectEdit( params: { ...originalParams }, occurrences, // This will be > 1 }; + editCorrectionCache.set(cacheKey, result); return result; } } @@ -153,6 +170,7 @@ export async function ensureCorrectEdit( }, occurrences: countOccurrences(currentContent, finalOldString), // Recalculate occurrences with the final old_string }; + editCorrectionCache.set(cacheKey, result); return result; } @@ -449,3 +467,7 @@ export function countOccurrences(str: string, substr: string): number { } return count; } + +export function resetEditCorrectorCaches_TEST_ONLY() { + editCorrectionCache.clear(); +}