Ensure edit correction isn't re-done after confirm.
- Edit corretion leans on LLM-isms to ensure we properly fix poorly escaped content. Beacues of this we need to ensure that we don't re-run edit correction in many cases. - To ensure this an `LruCache` has been added to capture intermediate steps of edit correction to avoid re-computations. - Max cache size is 50 currently. This means a user can have a muti-confirmation flow of 25 items without recomputing anything (assuming they all break edit correction). - Laid some groundwork for future testing. Part of https://github.com/google-gemini/gemini-cli/issues/484
This commit is contained in:
parent
c181fc1cf3
commit
1a5fe16b22
|
@ -0,0 +1,41 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
export class LruCache<K, V> {
|
||||
private cache: Map<K, V>;
|
||||
private maxSize: number;
|
||||
|
||||
constructor(maxSize: number) {
|
||||
this.cache = new Map<K, V>();
|
||||
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();
|
||||
}
|
||||
}
|
|
@ -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<string, CorrectedEditResult>(
|
||||
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<CorrectedEditResult> {
|
||||
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();
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue