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';
|
} from '@google/genai';
|
||||||
import { GeminiClient } from '../core/client.js';
|
import { GeminiClient } from '../core/client.js';
|
||||||
import { EditToolParams } from '../tools/edit.js';
|
import { EditToolParams } from '../tools/edit.js';
|
||||||
|
import { LruCache } from './LruCache.js';
|
||||||
|
|
||||||
const EditModel = 'gemini-2.5-flash-preview-04-17';
|
const EditModel = 'gemini-2.5-flash-preview-04-17';
|
||||||
const EditConfig: GenerateContentConfig = {
|
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
|
* 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\'
|
originalParams: EditToolParams, // This is the EditToolParams from edit.ts, without \'corrected\'
|
||||||
client: GeminiClient,
|
client: GeminiClient,
|
||||||
): Promise<CorrectedEditResult> {
|
): 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;
|
let finalNewString = originalParams.new_string;
|
||||||
const newStringPotentiallyEscaped =
|
const newStringPotentiallyEscaped =
|
||||||
unescapeStringForGeminiBug(originalParams.new_string) !==
|
unescapeStringForGeminiBug(originalParams.new_string) !==
|
||||||
|
@ -74,6 +88,7 @@ export async function ensureCorrectEdit(
|
||||||
params: { ...originalParams },
|
params: { ...originalParams },
|
||||||
occurrences,
|
occurrences,
|
||||||
};
|
};
|
||||||
|
editCorrectionCache.set(cacheKey, result);
|
||||||
return result;
|
return result;
|
||||||
} else {
|
} else {
|
||||||
// occurrences is 0 or some other unexpected state initially
|
// occurrences is 0 or some other unexpected state initially
|
||||||
|
@ -124,6 +139,7 @@ export async function ensureCorrectEdit(
|
||||||
params: { ...originalParams },
|
params: { ...originalParams },
|
||||||
occurrences: 0, // Explicitly 0 as LLM failed
|
occurrences: 0, // Explicitly 0 as LLM failed
|
||||||
};
|
};
|
||||||
|
editCorrectionCache.set(cacheKey, result);
|
||||||
return result;
|
return result;
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
|
@ -132,6 +148,7 @@ export async function ensureCorrectEdit(
|
||||||
params: { ...originalParams },
|
params: { ...originalParams },
|
||||||
occurrences, // This will be > 1
|
occurrences, // This will be > 1
|
||||||
};
|
};
|
||||||
|
editCorrectionCache.set(cacheKey, result);
|
||||||
return result;
|
return result;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -153,6 +170,7 @@ export async function ensureCorrectEdit(
|
||||||
},
|
},
|
||||||
occurrences: countOccurrences(currentContent, finalOldString), // Recalculate occurrences with the final old_string
|
occurrences: countOccurrences(currentContent, finalOldString), // Recalculate occurrences with the final old_string
|
||||||
};
|
};
|
||||||
|
editCorrectionCache.set(cacheKey, result);
|
||||||
return result;
|
return result;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -449,3 +467,7 @@ export function countOccurrences(str: string, substr: string): number {
|
||||||
}
|
}
|
||||||
return count;
|
return count;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export function resetEditCorrectorCaches_TEST_ONLY() {
|
||||||
|
editCorrectionCache.clear();
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue