feat: Enhance replace tool reliability with multi-stage edit correction

This commit significantly improves the `replace` tool's robustness by introducing a multi-stage correction mechanism. This directly addresses challenges with LLM-generated tool inputs, particularly the over-escaping of strings sometimes observed with Gemini models, and other minor discrepancies that previously led to failed edits.

The correction process is as follows:
1.  **Targeted Unescaping:** The system first applies a specialized unescaping function to the `old_string` and `new_string` to counteract common LLM-induced escaping patterns.
2.  **LLM-Powered Discrepancy Resolution:** If a unique match for the `old_string` is still not found, the system leverages a Gemini model (`gemini-2.5-flash-preview-04-17`) to:
    *   Identify the most probable intended `old_string` in the file by intelligently correcting minor formatting or escaping differences.
    *   Adjust the `new_string` to correspond with any corrections made to the `old_string`, maintaining the original edit's intent.

This enhancement makes the `replace` tool more resilient and effective, leading to a higher success rate for automated code modifications. The `expected_replacements` parameter has been removed as the tool now focuses on finding a single, unique, and correctable match. The tool's description and error reporting have been updated to reflect these new capabilities.

Fixes https://b.corp.google.com/issues/416933027
This commit is contained in:
Taylor Mullen
2025-05-12 23:23:24 -07:00
committed by N. Taylor Mullen
parent 5ec254253f
commit 3217576743
5 changed files with 449 additions and 60 deletions

View File

@@ -193,12 +193,18 @@ export class GeminiClient {
async generateJson(
contents: Content[],
schema: SchemaUnion,
model: string = 'gemini-2.0-flash',
config: GenerateContentConfig = {},
): Promise<Record<string, unknown>> {
try {
const requestConfig = {
...this.generateContentConfig,
...config,
};
const result = await this.client.models.generateContent({
model: 'gemini-2.0-flash',
model,
config: {
...this.generateContentConfig,
...requestConfig,
systemInstruction: getCoreSystemPrompt(),
responseSchema: schema,
responseMimeType: 'application/json',