From 1634d5fcca29e7c64d37f99e17e42d303e12062d Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Mon, 15 Sep 2025 20:46:41 -0400 Subject: [PATCH] refactor(core): Use BaseLlmClient for utility LLM calls in edit corrector (#8443) --- packages/core/src/tools/edit.test.ts | 9 +- packages/core/src/tools/edit.ts | 1 + packages/core/src/tools/write-file.test.ts | 28 +++-- packages/core/src/tools/write-file.ts | 3 +- packages/core/src/utils/editCorrector.test.ts | 100 ++++++------------ packages/core/src/utils/editCorrector.ts | 91 ++++++++++------ 6 files changed, 120 insertions(+), 112 deletions(-) diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index cd5d559842..d2ca62d929 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -58,6 +58,7 @@ describe('EditTool', () => { let rootDir: string; let mockConfig: Config; let geminiClient: any; + let baseLlmClient: any; beforeEach(() => { vi.restoreAllMocks(); @@ -69,8 +70,13 @@ describe('EditTool', () => { generateJson: mockGenerateJson, // mockGenerateJson is already defined and hoisted }; + baseLlmClient = { + generateJson: vi.fn(), + }; + mockConfig = { getGeminiClient: vi.fn().mockReturnValue(geminiClient), + getBaseLlmClient: vi.fn().mockReturnValue(baseLlmClient), getTargetDir: () => rootDir, getApprovalMode: vi.fn(), setApprovalMode: vi.fn(), @@ -424,11 +430,12 @@ describe('EditTool', () => { // Set a specific mock for this test case let mockCalled = false; mockEnsureCorrectEdit.mockImplementationOnce( - async (_, content, p, client) => { + async (_, content, p, client, baseClient) => { mockCalled = true; expect(content).toBe(originalContent); expect(p).toBe(params); expect(client).toBe(geminiClient); + expect(baseClient).toBe(baseLlmClient); return { params: { file_path: filePath, diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index 24f9332042..cf7029355e 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -164,6 +164,7 @@ class EditToolInvocation implements ToolInvocation { currentContent, params, this.config.getGeminiClient(), + this.config.getBaseLlmClient(), abortSignal, ); finalOldString = correctedEdit.params.old_string; diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index 5693978971..7e67134165 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -26,6 +26,7 @@ import path from 'node:path'; import fs from 'node:fs'; import os from 'node:os'; import { GeminiClient } from '../core/client.js'; +import type { BaseLlmClient } from '../core/baseLlmClient.js'; import type { CorrectedEditResult } from '../utils/editCorrector.js'; import { ensureCorrectEdit, @@ -47,6 +48,7 @@ vi.mock('../ide/ide-client.js', () => ({ }, })); let mockGeminiClientInstance: Mocked; +let mockBaseLlmClientInstance: Mocked; const mockEnsureCorrectEdit = vi.fn(); const mockEnsureCorrectFileContent = vi.fn(); const mockIdeClient = { @@ -70,6 +72,7 @@ const mockConfigInternal = { getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT), setApprovalMode: vi.fn(), getGeminiClient: vi.fn(), // Initialize as a plain mock function + getBaseLlmClient: vi.fn(), // Initialize as a plain mock function getFileSystemService: () => fsService, getIdeMode: vi.fn(() => false), getWorkspaceContext: () => createMockWorkspaceContext(rootDir), @@ -123,15 +126,23 @@ describe('WriteFileTool', () => { ) as Mocked; vi.mocked(GeminiClient).mockImplementation(() => mockGeminiClientInstance); + // Setup BaseLlmClient mock + mockBaseLlmClientInstance = { + generateJson: vi.fn(), + } as unknown as Mocked; + vi.mocked(ensureCorrectEdit).mockImplementation(mockEnsureCorrectEdit); vi.mocked(ensureCorrectFileContent).mockImplementation( mockEnsureCorrectFileContent, ); - // Now that mockGeminiClientInstance is initialized, set the mock implementation for getGeminiClient + // Now that mock instances are initialized, set the mock implementations for config getters mockConfigInternal.getGeminiClient.mockReturnValue( mockGeminiClientInstance, ); + mockConfigInternal.getBaseLlmClient.mockReturnValue( + mockBaseLlmClientInstance, + ); tool = new WriteFileTool(mockConfig); @@ -148,7 +159,8 @@ describe('WriteFileTool', () => { _currentContent: string, params: EditToolParams, _client: GeminiClient, - signal?: AbortSignal, // Make AbortSignal optional to match usage + _baseClient: BaseLlmClient, + signal?: AbortSignal, ): Promise => { if (signal?.aborted) { return Promise.reject(new Error('Aborted')); @@ -162,10 +174,9 @@ describe('WriteFileTool', () => { mockEnsureCorrectFileContent.mockImplementation( async ( content: string, - _client: GeminiClient, + _baseClient: BaseLlmClient, signal?: AbortSignal, ): Promise => { - // Make AbortSignal optional if (signal?.aborted) { return Promise.reject(new Error('Aborted')); } @@ -263,7 +274,7 @@ describe('WriteFileTool', () => { expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith( proposedContent, - mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockEnsureCorrectEdit).not.toHaveBeenCalled(); @@ -307,6 +318,7 @@ describe('WriteFileTool', () => { file_path: filePath, }, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled(); @@ -383,7 +395,7 @@ describe('WriteFileTool', () => { expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith( proposedContent, - mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(confirmation).toEqual( @@ -433,6 +445,7 @@ describe('WriteFileTool', () => { file_path: filePath, }, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(confirmation).toEqual( @@ -599,7 +612,7 @@ describe('WriteFileTool', () => { expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith( proposedContent, - mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(result.llmContent).toMatch( @@ -663,6 +676,7 @@ describe('WriteFileTool', () => { file_path: filePath, }, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(result.llmContent).toMatch(/Successfully overwrote file/); diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index 3253fc2d6e..c4cc3b56c8 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -121,6 +121,7 @@ export async function getCorrectedFileContent( file_path: filePath, }, config.getGeminiClient(), + config.getBaseLlmClient(), abortSignal, ); correctedContent = correctedParams.new_string; @@ -128,7 +129,7 @@ export async function getCorrectedFileContent( // This implies new file (ENOENT) correctedContent = await ensureCorrectFileContent( proposedContent, - config.getGeminiClient(), + config.getBaseLlmClient(), abortSignal, ); } diff --git a/packages/core/src/utils/editCorrector.test.ts b/packages/core/src/utils/editCorrector.test.ts index 7f76bed4ef..ab2c5e68f3 100644 --- a/packages/core/src/utils/editCorrector.test.ts +++ b/packages/core/src/utils/editCorrector.test.ts @@ -9,6 +9,7 @@ import type { Mock } from 'vitest'; import { vi, describe, it, expect, beforeEach, type Mocked } from 'vitest'; import * as fs from 'node:fs'; import { EditTool } from '../tools/edit.js'; +import type { BaseLlmClient } from '../core/baseLlmClient.js'; // MOCKS let callCount = 0; @@ -28,10 +29,9 @@ vi.mock('../core/client.js', () => ({ this: any, _config: Config, ) { - this.generateJson = (...params: any[]) => mockGenerateJson(...params); // Corrected: use mockGenerateJson - this.startChat = (...params: any[]) => mockStartChat(...params); // Corrected: use mockStartChat + this.startChat = (...params: any[]) => mockStartChat(...params); this.sendMessageStream = (...params: any[]) => - mockSendMessageStream(...params); // Corrected: use mockSendMessageStream + mockSendMessageStream(...params); return this; }), })); @@ -154,6 +154,7 @@ describe('editCorrector', () => { describe('ensureCorrectEdit', () => { let mockGeminiClientInstance: Mocked; + let mockBaseLlmClientInstance: Mocked; let mockToolRegistry: Mocked; let mockConfigInstance: Config; const abortSignal = new AbortController().signal; @@ -233,6 +234,9 @@ describe('editCorrector', () => { mockConfigInstance, ) as Mocked; mockGeminiClientInstance.getHistory = vi.fn().mockResolvedValue([]); + mockBaseLlmClientInstance = { + generateJson: mockGenerateJson, + } as unknown as Mocked; resetEditCorrectorCaches_TEST_ONLY(); }); @@ -252,6 +256,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); @@ -271,6 +276,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(0); @@ -293,6 +299,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); @@ -312,6 +319,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(0); @@ -335,6 +343,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); @@ -354,6 +363,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(0); @@ -373,6 +383,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(0); @@ -397,6 +408,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); @@ -420,6 +432,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(2); @@ -441,6 +454,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); @@ -464,6 +478,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); @@ -486,6 +501,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); @@ -505,6 +521,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(0); @@ -529,6 +546,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(2); @@ -585,6 +603,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); @@ -595,65 +614,10 @@ describe('editCorrector', () => { }); describe('ensureCorrectFileContent', () => { - let mockGeminiClientInstance: Mocked; - let mockToolRegistry: Mocked; - let mockConfigInstance: Config; + let mockBaseLlmClientInstance: Mocked; const abortSignal = new AbortController().signal; beforeEach(() => { - mockToolRegistry = new ToolRegistry({} as Config) as Mocked; - const configParams = { - apiKey: 'test-api-key', - model: 'test-model', - sandbox: false as boolean | string, - targetDir: '/test', - debugMode: false, - question: undefined as string | undefined, - fullContext: false, - coreTools: undefined as string[] | undefined, - toolDiscoveryCommand: undefined as string | undefined, - toolCallCommand: undefined as string | undefined, - mcpServerCommand: undefined as string | undefined, - mcpServers: undefined as Record | undefined, - userAgent: 'test-agent', - userMemory: '', - geminiMdFileCount: 0, - alwaysSkipModificationConfirmation: false, - }; - mockConfigInstance = { - ...configParams, - getApiKey: vi.fn(() => configParams.apiKey), - getModel: vi.fn(() => configParams.model), - getSandbox: vi.fn(() => configParams.sandbox), - getTargetDir: vi.fn(() => configParams.targetDir), - getToolRegistry: vi.fn(() => mockToolRegistry), - getDebugMode: vi.fn(() => configParams.debugMode), - getQuestion: vi.fn(() => configParams.question), - getFullContext: vi.fn(() => configParams.fullContext), - getCoreTools: vi.fn(() => configParams.coreTools), - getToolDiscoveryCommand: vi.fn(() => configParams.toolDiscoveryCommand), - getToolCallCommand: vi.fn(() => configParams.toolCallCommand), - getMcpServerCommand: vi.fn(() => configParams.mcpServerCommand), - getMcpServers: vi.fn(() => configParams.mcpServers), - getUserAgent: vi.fn(() => configParams.userAgent), - getUserMemory: vi.fn(() => configParams.userMemory), - setUserMemory: vi.fn((mem: string) => { - configParams.userMemory = mem; - }), - getGeminiMdFileCount: vi.fn(() => configParams.geminiMdFileCount), - setGeminiMdFileCount: vi.fn((count: number) => { - configParams.geminiMdFileCount = count; - }), - getAlwaysSkipModificationConfirmation: vi.fn( - () => configParams.alwaysSkipModificationConfirmation, - ), - setAlwaysSkipModificationConfirmation: vi.fn((skip: boolean) => { - configParams.alwaysSkipModificationConfirmation = skip; - }), - getQuotaErrorOccurred: vi.fn().mockReturnValue(false), - setQuotaErrorOccurred: vi.fn(), - } as unknown as Config; - callCount = 0; mockResponses.length = 0; mockGenerateJson = vi @@ -667,12 +631,10 @@ describe('editCorrector', () => { if (response === undefined) return Promise.resolve({}); return Promise.resolve(response); }); - mockStartChat = vi.fn(); - mockSendMessageStream = vi.fn(); - mockGeminiClientInstance = new GeminiClient( - mockConfigInstance, - ) as Mocked; + mockBaseLlmClientInstance = { + generateJson: mockGenerateJson, + } as unknown as Mocked; resetEditCorrectorCaches_TEST_ONLY(); }); @@ -680,7 +642,7 @@ describe('editCorrector', () => { const content = 'This is normal content without escaping issues'; const result = await ensureCorrectFileContent( content, - mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); expect(result).toBe(content); @@ -696,7 +658,7 @@ describe('editCorrector', () => { const result = await ensureCorrectFileContent( content, - mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); @@ -716,7 +678,7 @@ describe('editCorrector', () => { const result = await ensureCorrectFileContent( content, - mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); @@ -731,7 +693,7 @@ describe('editCorrector', () => { const result = await ensureCorrectFileContent( content, - mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); @@ -751,7 +713,7 @@ describe('editCorrector', () => { const result = await ensureCorrectFileContent( content, - mockGeminiClientInstance, + mockBaseLlmClientInstance, abortSignal, ); diff --git a/packages/core/src/utils/editCorrector.ts b/packages/core/src/utils/editCorrector.ts index 81c1aec7c6..809973e79f 100644 --- a/packages/core/src/utils/editCorrector.ts +++ b/packages/core/src/utils/editCorrector.ts @@ -6,6 +6,7 @@ import type { Content, GenerateContentConfig } from '@google/genai'; import type { GeminiClient } from '../core/client.js'; +import type { BaseLlmClient } from '../core/baseLlmClient.js'; import type { EditToolParams } from '../tools/edit.js'; import { EditTool } from '../tools/edit.js'; import { WriteFileTool } from '../tools/write-file.js'; @@ -19,14 +20,27 @@ import { isFunctionCall, } from '../utils/messageInspectors.js'; import * as fs from 'node:fs'; +import { promptIdContext } from './promptIdContext.js'; -const EditModel = DEFAULT_GEMINI_FLASH_LITE_MODEL; -const EditConfig: GenerateContentConfig = { +const EDIT_MODEL = DEFAULT_GEMINI_FLASH_LITE_MODEL; +const EDIT_CONFIG: GenerateContentConfig = { thinkingConfig: { thinkingBudget: 0, }, }; +const CODE_CORRECTION_SYSTEM_PROMPT = ` +You are an expert code-editing assistant. Your task is to analyze a failed edit attempt and provide a corrected version of the text snippets. +The correction should be as minimal as possible, staying very close to the original. +Focus ONLY on fixing issues like whitespace, indentation, line endings, or incorrect escaping. +Do NOT invent a completely new edit. Your job is to fix the provided parameters to make the edit succeed. +Return ONLY the corrected snippet in the specified JSON format. +`.trim(); + +function getPromptId(): string { + return promptIdContext.getStore() ?? `edit-corrector-${Date.now()}`; +} + const MAX_CACHE_SIZE = 50; // Cache for ensureCorrectEdit results @@ -159,7 +173,8 @@ export async function ensureCorrectEdit( filePath: string, currentContent: string, originalParams: EditToolParams, // This is the EditToolParams from edit.ts, without \'corrected\' - client: GeminiClient, + geminiClient: GeminiClient, + baseLlmClient: BaseLlmClient, abortSignal: AbortSignal, ): Promise { const cacheKey = `${currentContent}---${originalParams.old_string}---${originalParams.new_string}`; @@ -181,7 +196,7 @@ export async function ensureCorrectEdit( if (occurrences === expectedReplacements) { if (newStringPotentiallyEscaped) { finalNewString = await correctNewStringEscaping( - client, + baseLlmClient, finalOldString, originalParams.new_string, abortSignal, @@ -228,7 +243,7 @@ export async function ensureCorrectEdit( finalOldString = unescapedOldStringAttempt; if (newStringPotentiallyEscaped) { finalNewString = await correctNewString( - client, + baseLlmClient, originalParams.old_string, // original old unescapedOldStringAttempt, // corrected old originalParams.new_string, // original new (which is potentially escaped) @@ -242,7 +257,7 @@ export async function ensureCorrectEdit( // our system has done const lastEditedByUsTime = await findLastEditTimestamp( filePath, - client, + geminiClient, ); // Add a 1-second buffer to account for timing inaccuracies. If the file @@ -265,7 +280,7 @@ export async function ensureCorrectEdit( } const llmCorrectedOldString = await correctOldStringMismatch( - client, + baseLlmClient, currentContent, unescapedOldStringAttempt, abortSignal, @@ -284,7 +299,7 @@ export async function ensureCorrectEdit( originalParams.new_string, ); finalNewString = await correctNewString( - client, + baseLlmClient, originalParams.old_string, // original old llmCorrectedOldString, // corrected old baseNewStringForLLMCorrection, // base new for correction @@ -335,7 +350,7 @@ export async function ensureCorrectEdit( export async function ensureCorrectFileContent( content: string, - client: GeminiClient, + baseLlmClient: BaseLlmClient, abortSignal: AbortSignal, ): Promise { const cachedResult = fileContentCorrectionCache.get(content); @@ -352,7 +367,7 @@ export async function ensureCorrectFileContent( const correctedContent = await correctStringEscaping( content, - client, + baseLlmClient, abortSignal, ); fileContentCorrectionCache.set(content, correctedContent); @@ -373,7 +388,7 @@ const OLD_STRING_CORRECTION_SCHEMA: Record = { }; export async function correctOldStringMismatch( - geminiClient: GeminiClient, + baseLlmClient: BaseLlmClient, fileContent: string, problematicSnippet: string, abortSignal: AbortSignal, @@ -402,13 +417,15 @@ Return ONLY the corrected target snippet in the specified JSON format with the k const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }]; try { - const result = await geminiClient.generateJson( + const result = await baseLlmClient.generateJson({ contents, - OLD_STRING_CORRECTION_SCHEMA, + schema: OLD_STRING_CORRECTION_SCHEMA, abortSignal, - EditModel, - EditConfig, - ); + model: EDIT_MODEL, + config: EDIT_CONFIG, + systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT, + promptId: getPromptId(), + }); if ( result && @@ -450,7 +467,7 @@ const NEW_STRING_CORRECTION_SCHEMA: Record = { * Adjusts the new_string to align with a corrected old_string, maintaining the original intent. */ export async function correctNewString( - geminiClient: GeminiClient, + baseLlmClient: BaseLlmClient, originalOldString: string, correctedOldString: string, originalNewString: string, @@ -490,13 +507,15 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }]; try { - const result = await geminiClient.generateJson( + const result = await baseLlmClient.generateJson({ contents, - NEW_STRING_CORRECTION_SCHEMA, + schema: NEW_STRING_CORRECTION_SCHEMA, abortSignal, - EditModel, - EditConfig, - ); + model: EDIT_MODEL, + config: EDIT_CONFIG, + systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT, + promptId: getPromptId(), + }); if ( result && @@ -530,7 +549,7 @@ const CORRECT_NEW_STRING_ESCAPING_SCHEMA: Record = { }; export async function correctNewStringEscaping( - geminiClient: GeminiClient, + baseLlmClient: BaseLlmClient, oldString: string, potentiallyProblematicNewString: string, abortSignal: AbortSignal, @@ -559,13 +578,15 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }]; try { - const result = await geminiClient.generateJson( + const result = await baseLlmClient.generateJson({ contents, - CORRECT_NEW_STRING_ESCAPING_SCHEMA, + schema: CORRECT_NEW_STRING_ESCAPING_SCHEMA, abortSignal, - EditModel, - EditConfig, - ); + model: EDIT_MODEL, + config: EDIT_CONFIG, + systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT, + promptId: getPromptId(), + }); if ( result && @@ -603,7 +624,7 @@ const CORRECT_STRING_ESCAPING_SCHEMA: Record = { export async function correctStringEscaping( potentiallyProblematicString: string, - client: GeminiClient, + baseLlmClient: BaseLlmClient, abortSignal: AbortSignal, ): Promise { const prompt = ` @@ -625,13 +646,15 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }]; try { - const result = await client.generateJson( + const result = await baseLlmClient.generateJson({ contents, - CORRECT_STRING_ESCAPING_SCHEMA, + schema: CORRECT_STRING_ESCAPING_SCHEMA, abortSignal, - EditModel, - EditConfig, - ); + model: EDIT_MODEL, + config: EDIT_CONFIG, + systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT, + promptId: getPromptId(), + }); if ( result &&