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
N. Taylor Mullen
parent
5ec254253f
commit
3217576743
|
|
@ -148,7 +148,7 @@ function createToolRegistry(config: Config): ToolRegistry {
|
|||
new ReadFileTool(targetDir),
|
||||
new GrepTool(targetDir),
|
||||
new GlobTool(targetDir),
|
||||
new EditTool(targetDir),
|
||||
new EditTool(config),
|
||||
new WriteFileTool(targetDir),
|
||||
new WebFetchTool(), // Note: WebFetchTool takes no arguments
|
||||
new ReadManyFilesTool(targetDir),
|
||||
|
|
|
|||
|
|
@ -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 result = await this.client.models.generateContent({
|
||||
model: 'gemini-2.0-flash',
|
||||
config: {
|
||||
const requestConfig = {
|
||||
...this.generateContentConfig,
|
||||
...config,
|
||||
};
|
||||
const result = await this.client.models.generateContent({
|
||||
model,
|
||||
config: {
|
||||
...requestConfig,
|
||||
systemInstruction: getCoreSystemPrompt(),
|
||||
responseSchema: schema,
|
||||
responseMimeType: 'application/json',
|
||||
|
|
|
|||
|
|
@ -19,6 +19,9 @@ import { SchemaValidator } from '../utils/schemaValidator.js';
|
|||
import { makeRelative, shortenPath } from '../utils/paths.js';
|
||||
import { isNodeError } from '../utils/errors.js';
|
||||
import { ReadFileTool } from './read-file.js';
|
||||
import { GeminiClient } from '../core/client.js';
|
||||
import { Config } from '../config/config.js';
|
||||
import { countOccurrences, ensureCorrectEdit } from '../utils/editCorrector.js';
|
||||
|
||||
/**
|
||||
* Parameters for the Edit tool
|
||||
|
|
@ -38,11 +41,6 @@ export interface EditToolParams {
|
|||
* The text to replace it with
|
||||
*/
|
||||
new_string: string;
|
||||
|
||||
/**
|
||||
* The expected number of replacements to perform (optional, defaults to 1)
|
||||
*/
|
||||
expected_replacements?: number;
|
||||
}
|
||||
|
||||
interface CalculatedEdit {
|
||||
|
|
@ -59,23 +57,25 @@ interface CalculatedEdit {
|
|||
export class EditTool extends BaseTool<EditToolParams, ToolResult> {
|
||||
static readonly Name = 'replace';
|
||||
private shouldAlwaysEdit = false;
|
||||
private readonly rootDirectory: string;
|
||||
private readonly client: GeminiClient;
|
||||
|
||||
/**
|
||||
* Creates a new instance of the EditLogic
|
||||
* @param rootDirectory Root directory to ground this tool in.
|
||||
*/
|
||||
constructor(private readonly rootDirectory: string) {
|
||||
constructor(config: Config) {
|
||||
super(
|
||||
EditTool.Name,
|
||||
'Edit',
|
||||
`Replaces a single, unique occurrence of text within a file. This tool requires providing significant context around the change to ensure uniqueness and precise targeting. Always use the ${ReadFileTool} tool to examine the file's current content before attempting a text replacement.
|
||||
|
||||
Expectation for parameters:
|
||||
1. 'file_path' MUST be an absolute path; otherwise an error will be thrown.
|
||||
2. 'old_string' MUST be the exact literal text to replace (including all whitespace, indentation, newlines, and surrounding code etc.).
|
||||
3. 'new_string' MUST be the exact literal text to replace 'old_string' with (also including all whitespace, indentation, newlines, and surrounding code etc.).
|
||||
4. NEVER escape 'old_string' or 'new_string', JSON encoding will handle that automatically.
|
||||
**Important:** If ANY of the above are not satisfied the tool will fail.`,
|
||||
1. \`file_path\` MUST be an absolute path; otherwise an error will be thrown.
|
||||
2. \`old_string\` MUST be the exact literal text to replace (including all whitespace, indentation, newlines, and surrounding code etc.).
|
||||
3. \`new_string\` MUST be the exact literal text to replace \`old_string\` with (also including all whitespace, indentation, newlines, and surrounding code etc.). Ensure the resulting code is correct and idiomatic.
|
||||
4. NEVER escape \`old_string\` or \`new_string\`, that would break the exact literal text requirement.
|
||||
**Important:** If ANY of the above are not satisfied, the tool will fail. CRITICAL for \`old_string\`: Must uniquely identify the single instance to change. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string matches multiple locations, or does not match exactly, the tool will fail.`,
|
||||
{
|
||||
properties: {
|
||||
file_path: {
|
||||
|
|
@ -85,7 +85,7 @@ Expectation for parameters:
|
|||
},
|
||||
old_string: {
|
||||
description:
|
||||
'The exact literal text to replace, preferably unescaped. The tool will attempt to match this string as-is. CRITICAL: Must uniquely identify the single instance to change. Include at least 3-5 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string matches multiple locations, is escaped in a way that is not present in the original file, or does not match exactly, the tool will fail.',
|
||||
'The exact literal text to replace, preferably unescaped. CRITICAL: Must uniquely identify the single instance to change. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string is not the exact literal text (i.e. you escaped it), matches multiple locations, or does not match exactly, the tool will fail.',
|
||||
type: 'string',
|
||||
},
|
||||
new_string: {
|
||||
|
|
@ -98,7 +98,8 @@ Expectation for parameters:
|
|||
type: 'object',
|
||||
},
|
||||
);
|
||||
this.rootDirectory = path.resolve(rootDirectory);
|
||||
this.rootDirectory = path.resolve(config.getTargetDir());
|
||||
this.client = new GeminiClient(config);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -142,13 +143,6 @@ Expectation for parameters:
|
|||
return `File path must be within the root directory (${this.rootDirectory}): ${params.file_path}`;
|
||||
}
|
||||
|
||||
if (
|
||||
params.expected_replacements !== undefined &&
|
||||
params.expected_replacements < 0
|
||||
) {
|
||||
return 'Expected replacements must be a non-negative number';
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
|
|
@ -158,11 +152,8 @@ Expectation for parameters:
|
|||
* @returns An object describing the potential edit outcome
|
||||
* @throws File system errors if reading the file fails unexpectedly (e.g., permissions)
|
||||
*/
|
||||
private calculateEdit(params: EditToolParams): CalculatedEdit {
|
||||
const expectedReplacements =
|
||||
params.expected_replacements === undefined
|
||||
? 1
|
||||
: params.expected_replacements;
|
||||
private async calculateEdit(params: EditToolParams): Promise<CalculatedEdit> {
|
||||
const expectedReplacements = 1;
|
||||
let currentContent: string | null = null;
|
||||
let fileExists = false;
|
||||
let isNewFile = false;
|
||||
|
|
@ -194,22 +185,22 @@ Expectation for parameters:
|
|||
};
|
||||
} else if (currentContent !== null) {
|
||||
// Editing an existing file
|
||||
occurrences = this.countOccurrences(currentContent, params.old_string);
|
||||
occurrences = countOccurrences(currentContent, params.old_string);
|
||||
|
||||
if (params.old_string === '') {
|
||||
// Error: Trying to create a file that already exists
|
||||
error = {
|
||||
display: `File already exists. Use a non-empty old_string to edit.`,
|
||||
display: `Failed to edit. Attempted to create a file that already exists.`,
|
||||
raw: `File already exists, cannot create: ${params.file_path}`,
|
||||
};
|
||||
} else if (occurrences === 0) {
|
||||
error = {
|
||||
display: `No edits made. The exact text in old_string was not found. Check whitespace, indentation, and context. Use ReadFile tool to verify. `,
|
||||
raw: `Failed to edit, 0 occurrences found for old_string in ${params.file_path}`,
|
||||
display: `Failed to edit, could not find the string to replace.`,
|
||||
raw: `Failed to edit, 0 occurrences found for old_string in ${params.file_path}. No edits made. The exact text in old_string was not found. Ensure you're not escaping content incorrectly and check whitespace, indentation, and context. Use ReadFile tool to verify.`,
|
||||
};
|
||||
} else if (occurrences !== expectedReplacements) {
|
||||
error = {
|
||||
display: `Failed to edit, expected ${expectedReplacements} occurrence(s) but found ${occurrences}. Make old_string more specific with more context.`,
|
||||
display: `Failed to edit, expected ${expectedReplacements} occurrence(s) but found ${occurrences}.`,
|
||||
raw: `Failed to edit, Expected ${expectedReplacements} occurrences but found ${occurrences} for old_string in file: ${params.file_path}`,
|
||||
};
|
||||
} else {
|
||||
|
|
@ -223,7 +214,7 @@ Expectation for parameters:
|
|||
} else {
|
||||
// Should not happen if fileExists and no exception was thrown, but defensively:
|
||||
error = {
|
||||
display: `Failed to read content of existing file.`,
|
||||
display: `Failed to read content of file.`,
|
||||
raw: `Failed to read content of existing file: ${params.file_path}`,
|
||||
};
|
||||
}
|
||||
|
|
@ -273,15 +264,14 @@ Expectation for parameters:
|
|||
} else if (!fileExists) {
|
||||
return false;
|
||||
} else if (currentContent !== null) {
|
||||
const occurrences = this.countOccurrences(
|
||||
currentContent,
|
||||
params.old_string,
|
||||
);
|
||||
const expectedReplacements =
|
||||
params.expected_replacements === undefined
|
||||
? 1
|
||||
: params.expected_replacements;
|
||||
if (occurrences === 0 || occurrences !== expectedReplacements) {
|
||||
// Use the correctEdit utility to potentially correct params and get occurrences
|
||||
const { params: correctedParams, occurrences: correctedOccurrences } =
|
||||
await ensureCorrectEdit(currentContent, params, this.client);
|
||||
|
||||
params.old_string = correctedParams.old_string;
|
||||
params.new_string = correctedParams.new_string;
|
||||
|
||||
if (correctedOccurrences === 0 || correctedOccurrences !== 1) {
|
||||
return false;
|
||||
}
|
||||
newContent = this.replaceAll(
|
||||
|
|
@ -347,7 +337,7 @@ Expectation for parameters:
|
|||
|
||||
let editData: CalculatedEdit;
|
||||
try {
|
||||
editData = this.calculateEdit(params);
|
||||
editData = await this.calculateEdit(params);
|
||||
} catch (error) {
|
||||
const errorMsg = error instanceof Error ? error.message : String(error);
|
||||
return {
|
||||
|
|
@ -402,22 +392,6 @@ Expectation for parameters:
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Counts occurrences of a substring in a string
|
||||
*/
|
||||
private countOccurrences(str: string, substr: string): number {
|
||||
if (substr === '') {
|
||||
return 0;
|
||||
}
|
||||
let count = 0;
|
||||
let pos = str.indexOf(substr);
|
||||
while (pos !== -1) {
|
||||
count++;
|
||||
pos = str.indexOf(substr, pos + 1); // Ensure overlap is not counted if substr repeats
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
/**
|
||||
* Replaces all occurrences of a substring in a string
|
||||
*/
|
||||
|
|
|
|||
|
|
@ -0,0 +1,117 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import {
|
||||
countOccurrences,
|
||||
unescapeStringForGeminiBug,
|
||||
} from './editCorrector.js';
|
||||
|
||||
describe('editCorrector', () => {
|
||||
describe('countOccurrences', () => {
|
||||
it('should return 0 for empty string', () => {
|
||||
expect(countOccurrences('', 'a')).toBe(0);
|
||||
});
|
||||
|
||||
it('should return 0 for empty substring', () => {
|
||||
expect(countOccurrences('abc', '')).toBe(0);
|
||||
});
|
||||
|
||||
it('should return 0 if substring is not found', () => {
|
||||
expect(countOccurrences('abc', 'd')).toBe(0);
|
||||
});
|
||||
|
||||
it('should return 1 if substring is found once', () => {
|
||||
expect(countOccurrences('abc', 'b')).toBe(1);
|
||||
});
|
||||
|
||||
it('should return correct count for multiple occurrences', () => {
|
||||
expect(countOccurrences('ababa', 'a')).toBe(3);
|
||||
expect(countOccurrences('ababab', 'ab')).toBe(3);
|
||||
});
|
||||
|
||||
it('should count non-overlapping occurrences', () => {
|
||||
expect(countOccurrences('aaaaa', 'aa')).toBe(2); // Non-overlapping: aa_aa_
|
||||
expect(countOccurrences('ababab', 'aba')).toBe(1); // Non-overlapping: aba_ab -> 1
|
||||
});
|
||||
|
||||
it('should correctly count occurrences when substring is longer', () => {
|
||||
expect(countOccurrences('abc', 'abcdef')).toBe(0);
|
||||
});
|
||||
|
||||
it('should be case sensitive', () => {
|
||||
expect(countOccurrences('abcABC', 'a')).toBe(1);
|
||||
expect(countOccurrences('abcABC', 'A')).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('unescapeStringForGeminiBug', () => {
|
||||
it('should unescape common sequences', () => {
|
||||
expect(unescapeStringForGeminiBug('\\n')).toBe('\n');
|
||||
expect(unescapeStringForGeminiBug('\\t')).toBe('\t');
|
||||
expect(unescapeStringForGeminiBug("\\'")).toBe("'");
|
||||
expect(unescapeStringForGeminiBug('\\"')).toBe('"');
|
||||
expect(unescapeStringForGeminiBug('\\`')).toBe('`');
|
||||
});
|
||||
|
||||
it('should handle multiple escaped sequences', () => {
|
||||
expect(unescapeStringForGeminiBug('Hello\\nWorld\\tTest')).toBe(
|
||||
'Hello\nWorld\tTest',
|
||||
);
|
||||
});
|
||||
|
||||
it('should not alter already correct sequences', () => {
|
||||
expect(unescapeStringForGeminiBug('\n')).toBe('\n');
|
||||
expect(unescapeStringForGeminiBug('Correct string')).toBe(
|
||||
'Correct string',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle mixed correct and incorrect sequences', () => {
|
||||
expect(unescapeStringForGeminiBug('\\nCorrect\t\\`')).toBe(
|
||||
'\nCorrect\t`',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle backslash followed by actual newline character', () => {
|
||||
expect(unescapeStringForGeminiBug('\\\n')).toBe('\n');
|
||||
expect(unescapeStringForGeminiBug('First line\\\nSecond line')).toBe(
|
||||
'First line\nSecond line',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle multiple backslashes before an escapable character', () => {
|
||||
expect(unescapeStringForGeminiBug('\\\\n')).toBe('\n'); // \\n -> \n
|
||||
expect(unescapeStringForGeminiBug('\\\\\\t')).toBe('\t'); // \\\t -> \t
|
||||
expect(unescapeStringForGeminiBug('\\\\\\\\`')).toBe('`'); // \\\\` -> `
|
||||
});
|
||||
|
||||
it('should return empty string for empty input', () => {
|
||||
expect(unescapeStringForGeminiBug('')).toBe('');
|
||||
});
|
||||
|
||||
it('should not alter strings with no targeted escape sequences', () => {
|
||||
expect(unescapeStringForGeminiBug('abc def')).toBe('abc def');
|
||||
// \\F and \\S are not targeted escapes, so they should remain as \\F and \\S
|
||||
expect(unescapeStringForGeminiBug('C:\\Folder\\File')).toBe(
|
||||
'C:\\Folder\\File',
|
||||
);
|
||||
});
|
||||
|
||||
it('should correctly process strings with some targeted escapes', () => {
|
||||
// \\U is not targeted, \\n is.
|
||||
expect(unescapeStringForGeminiBug('C:\\Users\\name')).toBe(
|
||||
'C:\\Users\name',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle complex cases with mixed slashes and characters', () => {
|
||||
expect(
|
||||
unescapeStringForGeminiBug('\\\\\\nLine1\\\nLine2\\tTab\\\\`Tick\\"'),
|
||||
).toBe('\nLine1\nLine2\tTab`Tick"');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,292 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
Content,
|
||||
GenerateContentConfig,
|
||||
SchemaUnion,
|
||||
Type,
|
||||
} from '@google/genai';
|
||||
import { GeminiClient } from '../core/client.js';
|
||||
import { EditToolParams } from '../tools/edit.js';
|
||||
|
||||
const EditModel = 'gemini-2.5-flash-preview-04-17';
|
||||
const EditConfig: GenerateContentConfig = {
|
||||
thinkingConfig: {
|
||||
thinkingBudget: 0,
|
||||
},
|
||||
};
|
||||
|
||||
/**
|
||||
* Counts occurrences of a substring in a string
|
||||
*/
|
||||
export function countOccurrences(str: string, substr: string): number {
|
||||
if (substr === '') {
|
||||
return 0;
|
||||
}
|
||||
let count = 0;
|
||||
let pos = str.indexOf(substr);
|
||||
while (pos !== -1) {
|
||||
count++;
|
||||
pos = str.indexOf(substr, pos + substr.length); // Start search after the current match
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
/**
|
||||
* Attempts to correct edit parameters if the original old_string is not found.
|
||||
* It tries unescaping, and then LLM-based correction.
|
||||
*
|
||||
* @param currentContent The current content of the file.
|
||||
* @param params The original EditToolParams.
|
||||
* @param client The GeminiClient for LLM calls.
|
||||
* @returns A promise resolving to an object containing the (potentially corrected) EditToolParams and the final occurrences count.
|
||||
*/
|
||||
export async function ensureCorrectEdit(
|
||||
currentContent: string,
|
||||
originalParams: EditToolParams,
|
||||
client: GeminiClient,
|
||||
): Promise<CorrectedEditResult> {
|
||||
let occurrences = countOccurrences(currentContent, originalParams.old_string);
|
||||
const currentParams = { ...originalParams };
|
||||
|
||||
if (occurrences === 1) {
|
||||
return { params: currentParams, occurrences };
|
||||
}
|
||||
|
||||
const unescapedOldString = unescapeStringForGeminiBug(
|
||||
currentParams.old_string,
|
||||
);
|
||||
occurrences = countOccurrences(currentContent, unescapedOldString);
|
||||
|
||||
if (occurrences === 1) {
|
||||
currentParams.old_string = unescapedOldString;
|
||||
currentParams.new_string = unescapeStringForGeminiBug(
|
||||
currentParams.new_string,
|
||||
);
|
||||
} else if (occurrences === 0) {
|
||||
const llmCorrectedOldString = await correctOldStringMismatch(
|
||||
client,
|
||||
currentContent,
|
||||
unescapedOldString,
|
||||
);
|
||||
occurrences = countOccurrences(currentContent, llmCorrectedOldString);
|
||||
|
||||
if (occurrences === 1) {
|
||||
const llmCorrectedNewString = await correctNewString(
|
||||
client,
|
||||
unescapedOldString,
|
||||
llmCorrectedOldString,
|
||||
currentParams.new_string,
|
||||
);
|
||||
currentParams.old_string = llmCorrectedOldString;
|
||||
currentParams.new_string = llmCorrectedNewString;
|
||||
} else {
|
||||
// If LLM correction also results in 0 or >1 occurrences,
|
||||
// return the original params and 0 occurrences,
|
||||
// letting the caller handle the "still not found" case.
|
||||
return { params: originalParams, occurrences: 0 };
|
||||
}
|
||||
} else {
|
||||
// If unescaping resulted in >1 occurrences, return original params and that count.
|
||||
return { params: originalParams, occurrences };
|
||||
}
|
||||
|
||||
return { params: currentParams, occurrences };
|
||||
}
|
||||
|
||||
/**
|
||||
* Attempts to correct potential formatting/escaping issues in a snippet using an LLM call.
|
||||
*/
|
||||
async function correctOldStringMismatch(
|
||||
geminiClient: GeminiClient,
|
||||
fileContent: string,
|
||||
problematicSnippet: string,
|
||||
): Promise<string> {
|
||||
const prompt = `
|
||||
Context: A process needs to find an exact literal, unique match for a specific text snippet within a file's content. The provided snippet failed to match exactly. This is most likely because it has been overly escaped.
|
||||
|
||||
Task: Analyze the provided file content and the problematic target snippet. Identify the segment in the file content that the snippet was *most likely* intended to match. Output the *exact*, literal text of that segment from the file content. Focus *only* on removing extra escape characters and correcting formatting, whitespace, or minor differences to achieve a PERFECT literal match. The output must be the exact literal text as it appears in the file.
|
||||
|
||||
Problematic target snippet:
|
||||
\`\`\`
|
||||
${problematicSnippet}
|
||||
\`\`\`
|
||||
|
||||
File Content:
|
||||
\`\`\`
|
||||
${fileContent}
|
||||
\`\`\`
|
||||
|
||||
For example, if the problematic target snippet was "\\\\\\nconst greeting = \`Hello \\\\\`\${name}\\\\\`\`;" and the file content had content that looked like "\nconst greeting = \`Hello ${'\\`'}\${name}${'\\`'}\`;", then corrected_target_snippet should likely be "\nconst greeting = \`Hello ${'\\`'}\${name}${'\\`'}\`;" to fix the incorrect escaping to match the original file content.
|
||||
If the differences are only in whitespace or formatting, apply similar whitespace/formatting changes to the corrected_target_snippet.
|
||||
|
||||
Return ONLY the corrected target snippet in the specified JSON format with the key 'corrected_target_snippet'. If no clear, unique match can be found, return an empty string for 'corrected_target_snippet'.
|
||||
`.trim();
|
||||
|
||||
const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }];
|
||||
|
||||
try {
|
||||
const result = await geminiClient.generateJson(
|
||||
contents,
|
||||
OLD_STRING_CORRECTION_SCHEMA,
|
||||
EditModel,
|
||||
EditConfig,
|
||||
);
|
||||
|
||||
if (
|
||||
result &&
|
||||
typeof result.corrected_target_snippet === 'string' &&
|
||||
result.corrected_target_snippet.length > 0
|
||||
) {
|
||||
return result.corrected_target_snippet;
|
||||
} else {
|
||||
return problematicSnippet;
|
||||
}
|
||||
} catch (error) {
|
||||
console.error(
|
||||
'Error during LLM call for old string snippet correction:',
|
||||
error,
|
||||
);
|
||||
return problematicSnippet;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Adjusts the new_string to align with a corrected old_string, maintaining the original intent.
|
||||
*/
|
||||
async function correctNewString(
|
||||
geminiClient: GeminiClient,
|
||||
originalOldString: string,
|
||||
correctedOldString: string,
|
||||
originalNewString: string,
|
||||
): Promise<string> {
|
||||
if (originalOldString === correctedOldString) {
|
||||
return originalNewString;
|
||||
}
|
||||
|
||||
const prompt = `
|
||||
Context: A text replacement operation was planned. The original text to be replaced (original_old_string) was slightly different from the actual text in the file (corrected_old_string). The original_old_string has now been corrected to match the file content.
|
||||
We now need to adjust the replacement text (original_new_string) so that it makes sense as a replacement for the corrected_old_string, while preserving the original intent of the change.
|
||||
|
||||
original_old_string (what was initially intended to be found):
|
||||
\`\`\`
|
||||
${originalOldString}
|
||||
\`\`\`
|
||||
|
||||
corrected_old_string (what was actually found in the file and will be replaced):
|
||||
\`\`\`
|
||||
${correctedOldString}
|
||||
\`\`\`
|
||||
|
||||
original_new_string (what was intended to replace original_old_string):
|
||||
\`\`\`
|
||||
${originalNewString}
|
||||
\`\`\`
|
||||
|
||||
Task: Based on the differences between original_old_string and corrected_old_string, and the content of original_new_string, generate a corrected_new_string. This corrected_new_string should be what original_new_string would have been if it was designed to replace corrected_old_string directly, while maintaining the spirit of the original transformation.
|
||||
|
||||
For example, if original_old_string was "\\\\\\nconst greeting = \`Hello \\\\\`\${name}\\\\\`\`;" and corrected_old_string is "\nconst greeting = \`Hello ${'\\`'}\${name}${'\\`'}\`;", and original_new_string was "\\\\\\nconst greeting = \`Hello \\\\\`\${name} \${lastName}\\\\\`\`;", then corrected_new_string should likely be "\nconst greeting = \`Hello ${'\\`'}\${name} \${lastName}${'\\`'}\`;" to fix the incorrect escaping.
|
||||
If the differences are only in whitespace or formatting, apply similar whitespace/formatting changes to the corrected_new_string.
|
||||
|
||||
Return ONLY the corrected string in the specified JSON format with the key 'corrected_new_string'. If no adjustment is deemed necessary or possible, return the original_new_string.
|
||||
`.trim();
|
||||
|
||||
const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }];
|
||||
|
||||
try {
|
||||
const result = await geminiClient.generateJson(
|
||||
contents,
|
||||
NEW_STRING_CORRECTION_SCHEMA,
|
||||
EditModel,
|
||||
EditConfig,
|
||||
);
|
||||
|
||||
if (
|
||||
result &&
|
||||
typeof result.corrected_new_string === 'string' &&
|
||||
result.corrected_new_string.length > 0
|
||||
) {
|
||||
return result.corrected_new_string;
|
||||
} else {
|
||||
return originalNewString;
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Error during LLM call for new_string correction:', error);
|
||||
return originalNewString;
|
||||
}
|
||||
}
|
||||
|
||||
export interface CorrectedEditResult {
|
||||
params: EditToolParams;
|
||||
occurrences: number;
|
||||
}
|
||||
|
||||
// Define the expected JSON schema for the LLM response for old_string correction
|
||||
const OLD_STRING_CORRECTION_SCHEMA: SchemaUnion = {
|
||||
type: Type.OBJECT,
|
||||
properties: {
|
||||
corrected_target_snippet: {
|
||||
type: Type.STRING,
|
||||
description:
|
||||
'The corrected version of the target snippet that exactly and uniquely matches a segment within the provided file content.',
|
||||
},
|
||||
},
|
||||
required: ['corrected_target_snippet'],
|
||||
};
|
||||
|
||||
// Define the expected JSON schema for the new_string correction LLM response
|
||||
const NEW_STRING_CORRECTION_SCHEMA: SchemaUnion = {
|
||||
type: Type.OBJECT,
|
||||
properties: {
|
||||
corrected_new_string: {
|
||||
type: Type.STRING,
|
||||
description:
|
||||
'The original_new_string adjusted to be a suitable replacement for the corrected_old_string, while maintaining the original intent of the change.',
|
||||
},
|
||||
},
|
||||
required: ['corrected_new_string'],
|
||||
};
|
||||
|
||||
/**
|
||||
* Unescapes a string that might have been overly escaped by an LLM.
|
||||
*/
|
||||
export function unescapeStringForGeminiBug(inputString: string): string {
|
||||
// Regex explanation:
|
||||
// \\+ : Matches one or more literal backslash characters.
|
||||
// (n|t|r|'|"|`|\n) : This is a capturing group. It matches one of the following:
|
||||
// n, t, r, ', ", ` : These match the literal characters 'n', 't', 'r', single quote, double quote, or backtick.
|
||||
// This handles cases like "\\n", "\\\\`", etc.
|
||||
// \n : This matches an actual newline character. This handles cases where the input
|
||||
// string might have something like "\\\n" (a literal backslash followed by a newline).
|
||||
// g : Global flag, to replace all occurrences.
|
||||
|
||||
return inputString.replace(/\\+(n|t|r|'|"|`|\n)/g, (match, capturedChar) => {
|
||||
// 'match' is the entire erroneous sequence, e.g., if the input (in memory) was "\\\\`", match is "\\\\`".
|
||||
// 'capturedChar' is the character that determines the true meaning, e.g., '`'.
|
||||
|
||||
switch (capturedChar) {
|
||||
case 'n':
|
||||
return '\n'; // Correctly escaped: \n (newline character)
|
||||
case 't':
|
||||
return '\t'; // Correctly escaped: \t (tab character)
|
||||
case 'r':
|
||||
return '\r'; // Correctly escaped: \r (carriage return character)
|
||||
case "'":
|
||||
return "'"; // Correctly escaped: ' (apostrophe character)
|
||||
case '"':
|
||||
return '"'; // Correctly escaped: " (quotation mark character)
|
||||
case '`':
|
||||
return '`'; // Correctly escaped: ` (backtick character)
|
||||
case '\n': // This handles when 'capturedChar' is an actual newline
|
||||
return '\n'; // Replace the whole erroneous sequence (e.g., "\\\n" in memory) with a clean newline
|
||||
default:
|
||||
// This fallback should ideally not be reached if the regex captures correctly.
|
||||
// It would return the original matched sequence if an unexpected character was captured.
|
||||
return match;
|
||||
}
|
||||
});
|
||||
}
|
||||
Loading…
Reference in New Issue