From 14f39f556692dfa12422852cccd758a7db2ff8da Mon Sep 17 00:00:00 2001 From: David Ichim Date: Sun, 11 May 2025 16:02:30 +0300 Subject: [PATCH 1/6] fix(smart-apply): avoid extra response transforms when using the new smart apply The new Qwen based smart apply model performs an entire-file replace, so we can simply remove the leading newlines from the response. This simplifies the response transformation logic for this model since we no longer need todo extra formatting given the fact that we replace the entire file and it should be correctly formatted already --- vscode/src/edit/output/response-transformer.ts | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vscode/src/edit/output/response-transformer.ts b/vscode/src/edit/output/response-transformer.ts index 720461bf306f..ce17e78c401c 100644 --- a/vscode/src/edit/output/response-transformer.ts +++ b/vscode/src/edit/output/response-transformer.ts @@ -108,6 +108,15 @@ export function responseTransformer( isMessageInProgress: boolean ): string { const updatedText = extractSmartApplyCustomModelResponse(text, task) + + // if we use the new Qwen based smart apply which performs an entire-file + // replace it's enough for us to just remove the start and end newline + if ( + task.intent === 'smartApply' && + Object.values(SMART_APPLY_MODEL_IDENTIFIERS).includes(task.model) + ) { + return updatedText.replace(LEADING_SPACES_AND_NEW_LINES, '') + } const strippedText = stripText(updatedText, task) // Trim leading spaces From 4ae4237c6df965d25eecb2b656a6f52ffd483dcf Mon Sep 17 00:00:00 2001 From: David Ichim Date: Sun, 11 May 2025 23:12:16 +0300 Subject: [PATCH 2/6] fix(smart-apply): trim trailing whitespace for non-newline files The new Qwen based smart apply model performs an entire-file replace, but we should only trim the final newline if the original file ends with a newline. This commit updates the response transformation logic for the new smart apply model to conditionally trim trailing whitespace based on whether the original file ends with a newline character. This ensures that files without a final newline are not incorrectly modified. --- vscode/src/edit/output/response-transformer.ts | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/vscode/src/edit/output/response-transformer.ts b/vscode/src/edit/output/response-transformer.ts index ce17e78c401c..447aea33511d 100644 --- a/vscode/src/edit/output/response-transformer.ts +++ b/vscode/src/edit/output/response-transformer.ts @@ -107,6 +107,13 @@ export function responseTransformer( task: FixupTask, isMessageInProgress: boolean ): string { + if ( + task.intent === 'smartApply' && + Object.values(SMART_APPLY_MODEL_IDENTIFIERS).includes(task.model) && + isMessageInProgress + ) { + return text + } const updatedText = extractSmartApplyCustomModelResponse(text, task) // if we use the new Qwen based smart apply which performs an entire-file @@ -115,7 +122,7 @@ export function responseTransformer( task.intent === 'smartApply' && Object.values(SMART_APPLY_MODEL_IDENTIFIERS).includes(task.model) ) { - return updatedText.replace(LEADING_SPACES_AND_NEW_LINES, '') + return task.original.endsWith('\n') ? updatedText : updatedText.trimEnd() } const strippedText = stripText(updatedText, task) From fc8558664568a4ab2594c1b3d3ce6cba5ff03ec3 Mon Sep 17 00:00:00 2001 From: David Ichim Date: Wed, 14 May 2025 19:53:00 +0300 Subject: [PATCH 3/6] feat(smart-apply): improve newline handling and extraction for Qwen models This commit enhances the response transformation logic for the new Qwen-based smart apply model to provide more accurate and consistent results. Key changes: - Adds a `trimLLMNewlines` function to preserve or remove leading/trailing newlines based on the original file content. This ensures that files without a final newline are not incorrectly modified, and files with a newline retain it. - Skips processing for in-progress messages from smart apply custom models. - Simplifies the response transformation logic by consolidating newline handling and extraction into dedicated functions. - Improves test coverage to validate the newline preservation logic. --- .../edit/output/response-transformer.test.ts | 65 +++++++++++++++- .../src/edit/output/response-transformer.ts | 77 +++++++++++-------- 2 files changed, 109 insertions(+), 33 deletions(-) diff --git a/vscode/src/edit/output/response-transformer.test.ts b/vscode/src/edit/output/response-transformer.test.ts index 118f61704e6d..09e5b6820f5d 100644 --- a/vscode/src/edit/output/response-transformer.test.ts +++ b/vscode/src/edit/output/response-transformer.test.ts @@ -16,7 +16,12 @@ describe('Smart Apply Response Extraction', () => { const text = `<${SMART_APPLY_CUSTOM_PROMPT_TOPICS.FINAL_CODE}>const x = 1;` const task = createTask('smartApply', SMART_APPLY_MODEL_IDENTIFIERS.FireworksQwenCodeDefault) - const result = responseTransformer(text, task, true) + // When isMessageInProgress is true, the original text is returned without processing + const resultInProgress = responseTransformer(text, task, true) + expect(resultInProgress).toBe(text) + + // When isMessageInProgress is false, the text is processed + const result = responseTransformer(text, task, false) expect(result).toBe('const x = 1;') }) @@ -48,7 +53,12 @@ describe('Smart Apply Response Extraction', () => { const text = `<${SMART_APPLY_CUSTOM_PROMPT_TOPICS.FINAL_CODE}>outer <${SMART_APPLY_CUSTOM_PROMPT_TOPICS.FINAL_CODE}>inner content` const task = createTask('smartApply', SMART_APPLY_MODEL_IDENTIFIERS.FireworksQwenCodeDefault) - const result = responseTransformer(text, task, true) + // When isMessageInProgress is true, the original text is returned without processing + const resultInProgress = responseTransformer(text, task, true) + expect(resultInProgress).toBe(text) + + // When isMessageInProgress is false, the text is processed + const result = responseTransformer(text, task, false) expect(result).toBe( `outer <${SMART_APPLY_CUSTOM_PROMPT_TOPICS.FINAL_CODE}>inner content` ) @@ -58,7 +68,12 @@ describe('Smart Apply Response Extraction', () => { const text = `<${SMART_APPLY_CUSTOM_PROMPT_TOPICS.FINAL_CODE}>` const task = createTask('smartApply', SMART_APPLY_MODEL_IDENTIFIERS.FireworksQwenCodeDefault) - const result = responseTransformer(text, task, true) + // When isMessageInProgress is true, the original text is returned without processing + const resultInProgress = responseTransformer(text, task, true) + expect(resultInProgress).toBe(text) + + // When isMessageInProgress is false, the text is processed + const result = responseTransformer(text, task, false) expect(result).toBe('') }) @@ -91,6 +106,50 @@ describe('Smart Apply Response Extraction', () => { expect(result).toBe('const x = 1;\n') expect(result.endsWith('\n')).toBe(true) }) + + it('should preserve newlines based on original text', () => { + const text = `<${SMART_APPLY_CUSTOM_PROMPT_TOPICS.FINAL_CODE}>\nconst x = 1;\n` + + // Test 1: Original has no newlines, result should have no newlines + const task1 = { + ...createTask('smartApply', SMART_APPLY_MODEL_IDENTIFIERS.FireworksQwenCodeDefault), + original: 'const y = 2;', + } as any + const result1 = responseTransformer(text, task1, false) + expect(result1).toBe('const x = 1;') + expect(result1.startsWith('\n')).toBe(false) + expect(result1.endsWith('\n')).toBe(false) + + // Test 2: Original has starting newline, result should have starting newline + const task2 = { + ...createTask('smartApply', SMART_APPLY_MODEL_IDENTIFIERS.FireworksQwenCodeDefault), + original: '\nconst y = 2;', + } as any + const result2 = responseTransformer(text, task2, false) + expect(result2).toBe('\nconst x = 1;') + expect(result2.startsWith('\n')).toBe(true) + expect(result2.endsWith('\n')).toBe(false) + + // Test 3: Original has ending newline, result should have ending newline + const task3 = { + ...createTask('smartApply', SMART_APPLY_MODEL_IDENTIFIERS.FireworksQwenCodeDefault), + original: 'const y = 2;\n', + } as any + const result3 = responseTransformer(text, task3, false) + expect(result3).toBe('const x = 1;\n') + expect(result3.startsWith('\n')).toBe(false) + expect(result3.endsWith('\n')).toBe(true) + + // Test 4: Original has both newlines, result should have both newlines + const task4 = { + ...createTask('smartApply', SMART_APPLY_MODEL_IDENTIFIERS.FireworksQwenCodeDefault), + original: '\nconst y = 2;\n', + } as any + const result4 = responseTransformer(text, task4, false) + expect(result4).toBe('\nconst x = 1;\n') + expect(result4.startsWith('\n')).toBe(true) + expect(result4.endsWith('\n')).toBe(true) + }) }) describe('responseTransformer', () => { diff --git a/vscode/src/edit/output/response-transformer.ts b/vscode/src/edit/output/response-transformer.ts index 447aea33511d..3632074d8f76 100644 --- a/vscode/src/edit/output/response-transformer.ts +++ b/vscode/src/edit/output/response-transformer.ts @@ -38,6 +38,13 @@ const MARKDOWN_CODE_BLOCK_REGEX = new RegExp( const LEADING_SPACES_AND_NEW_LINES = /^\s*\n/ const LEADING_SPACES = /^[ ]+/ +/** + * Checks if the task is using a smart apply custom model + */ +function taskUsesSmartApplyCustomModel(task: FixupTask): boolean { + return Object.values(SMART_APPLY_MODEL_IDENTIFIERS).includes(task.model) +} + /** * Strips the text of any unnecessary content. * This includes: @@ -62,10 +69,7 @@ function stripText(text: string, task: FixupTask): string { } function extractSmartApplyCustomModelResponse(text: string, task: FixupTask): string { - if ( - task.intent !== 'smartApply' || - !Object.values(SMART_APPLY_MODEL_IDENTIFIERS).includes(task.model) - ) { + if (!taskUsesSmartApplyCustomModel(task)) { return text } @@ -86,15 +90,33 @@ function extractSmartApplyCustomModelResponse(text: string, task: FixupTask): st } /** -- * Regular expression to detect potential HTML entities. -- * Checks for named (&name;), decimal (&#digits;), or hex (&#xhex;) entities. -+ * Regular expression to detect the *few* entities we actually care about. -+ * We purposefully limit the named-entity part to the common escaping -+ * sequences that LLMs emit in source code: -+ * < > & " ' -+ * Everything else (e.g.  , ¤, ©, …) is ignored so that we -+ * don’t accidentally alter code like “¤t_value;”. - */ + * Preserves or removes newlines at the start and end of the text based on the original text. + * If the original text doesn't start/end with a newline, the corresponding newline in the updated text is removed. + */ +function trimLLMNewlines(text: string, original: string): string { + let result = text + + // Handle starting newline + if (result.startsWith('\n') && !original.startsWith('\n')) { + result = result.replace(/^\n/, '') + } + + // Handle ending newline + if (result.endsWith('\n') && !original.endsWith('\n')) { + result = result.replace(/\n$/, '') + } + + return result +} + +/** + * Regular expression to detect the *few* entities we actually care about. + * We purposefully limit the named-entity part to the common escaping + * sequences that LLMs emit in source code: + * < > & " ' + * Everything else (e.g.  , ¤, ©, …) is ignored so that we + * don't accidentally alter code like "¤t_value;". + */ const POTENTIAL_HTML_ENTITY_REGEX = /&(?:(?:lt|gt|amp|quot|apos)|#\d+|#x[0-9a-fA-F]+);/ /** @@ -107,24 +129,19 @@ export function responseTransformer( task: FixupTask, isMessageInProgress: boolean ): string { - if ( - task.intent === 'smartApply' && - Object.values(SMART_APPLY_MODEL_IDENTIFIERS).includes(task.model) && - isMessageInProgress - ) { - return text - } - const updatedText = extractSmartApplyCustomModelResponse(text, task) - - // if we use the new Qwen based smart apply which performs an entire-file - // replace it's enough for us to just remove the start and end newline - if ( - task.intent === 'smartApply' && - Object.values(SMART_APPLY_MODEL_IDENTIFIERS).includes(task.model) - ) { - return task.original.endsWith('\n') ? updatedText : updatedText.trimEnd() + // Skip processing for in-progress messages from smart apply custom models + if (taskUsesSmartApplyCustomModel(task)) { + if (isMessageInProgress) { + return text + } + + const updatedText = extractSmartApplyCustomModelResponse(text, task) + + // Preserve newlines only if they were in the original text + return trimLLMNewlines(updatedText, task.original) } - const strippedText = stripText(updatedText, task) + + const strippedText = stripText(text, task) // Trim leading spaces // - For `add` insertions, the LLM will attempt to continue the code from the position of the cursor, we handle the `insertionPoint` From 5a7e2fb6c5d0cf040babd8f1b56570481bcebe2e Mon Sep 17 00:00:00 2001 From: David Ichim Date: Sat, 17 May 2025 11:56:52 +0300 Subject: [PATCH 4/6] fix(smart-apply): improve smart apply insert mode and remove redundant newline test This commit addresses two issues related to the smart apply feature: - Fixes an issue where in-progress messages were not being skipped correctly when using insert mode with smart apply custom models. The condition to skip processing for in-progress messages is updated to include the insert mode. - Removes a redundant test case in `response-transformer.test.ts` that was asserting the addition of a newline character for smart apply without an empty selection range. This test is no longer necessary due to changes in how newlines are handled. --- .../edit/output/response-transformer.test.ts | 17 +---------------- vscode/src/edit/output/response-transformer.ts | 2 +- 2 files changed, 2 insertions(+), 17 deletions(-) diff --git a/vscode/src/edit/output/response-transformer.test.ts b/vscode/src/edit/output/response-transformer.test.ts index 09e5b6820f5d..75b1fdb9ce31 100644 --- a/vscode/src/edit/output/response-transformer.test.ts +++ b/vscode/src/edit/output/response-transformer.test.ts @@ -78,7 +78,7 @@ describe('Smart Apply Response Extraction', () => { }) it('should not add newline for smartApply with empty selection range', () => { - const text = `<${SMART_APPLY_CUSTOM_PROMPT_TOPICS.FINAL_CODE}>const x = 1;` + const text = 'const x = 1;' const task = { ...createTask('smartApply', SMART_APPLY_MODEL_IDENTIFIERS.FireworksQwenCodeDefault), mode: 'insert', @@ -92,21 +92,6 @@ describe('Smart Apply Response Extraction', () => { expect(result.endsWith('\n')).toBe(false) }) - it('should add newline for smartApply without empty selection range', () => { - const text = `<${SMART_APPLY_CUSTOM_PROMPT_TOPICS.FINAL_CODE}>const x = 1;` - const task = { - ...createTask('smartApply', SMART_APPLY_MODEL_IDENTIFIERS.FireworksQwenCodeDefault), - mode: 'insert', - selectionRange: { isEmpty: false }, - original: 'existing content', - fixupFile: { uri: {} as any }, - } as any - - const result = responseTransformer(text, task, false) - expect(result).toBe('const x = 1;\n') - expect(result.endsWith('\n')).toBe(true) - }) - it('should preserve newlines based on original text', () => { const text = `<${SMART_APPLY_CUSTOM_PROMPT_TOPICS.FINAL_CODE}>\nconst x = 1;\n` diff --git a/vscode/src/edit/output/response-transformer.ts b/vscode/src/edit/output/response-transformer.ts index 3632074d8f76..757e502c5955 100644 --- a/vscode/src/edit/output/response-transformer.ts +++ b/vscode/src/edit/output/response-transformer.ts @@ -131,7 +131,7 @@ export function responseTransformer( ): string { // Skip processing for in-progress messages from smart apply custom models if (taskUsesSmartApplyCustomModel(task)) { - if (isMessageInProgress) { + if (isMessageInProgress || task.mode === 'insert') { return text } From 9be3b888ff20bb0941f6eac70ae3d0440236dfc6 Mon Sep 17 00:00:00 2001 From: David Ichim Date: Sat, 17 May 2025 14:51:13 +0300 Subject: [PATCH 5/6] fix(smart-apply): improve entire file selection logic and add token limit validation This commit improves the logic for selecting the entire file in the `CustomModelSelectionProvider` and adds validation to prevent exceeding the context window's input token limit. - Modifies the `getSelectedText` method to prioritize `shouldAlwaysUseEntireFile` and refactors the token limit check to occur before the full file rewrite check. - Adds a test case to verify that an error is thrown when the token count exceeds the context window input limit. - Mocks the `vscode.Range` to avoid errors during testing. --- .../edit/adapters/smart-apply-custom.test.ts | 45 ++++++++++++++++++- .../smart-apply/selection/custom-model.ts | 7 +-- 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/vscode/src/edit/adapters/smart-apply-custom.test.ts b/vscode/src/edit/adapters/smart-apply-custom.test.ts index bf40a16f4175..f39c62bf7d43 100644 --- a/vscode/src/edit/adapters/smart-apply-custom.test.ts +++ b/vscode/src/edit/adapters/smart-apply-custom.test.ts @@ -1,4 +1,7 @@ -import { describe, expect, it } from 'vitest' +import { TokenCounterUtils, ps } from '@sourcegraph/cody-shared' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import * as vscode from 'vscode' +import { CustomModelSelectionProvider } from '../prompt/smart-apply/selection/custom-model' import type { ModelParametersInput } from './base' import { SmartApplyCustomModelParameterProvider } from './smart-apply-custom' @@ -105,3 +108,43 @@ describe('SmartApplyCustomModelParameterProvider', () => { }) }) }) + +describe('CustomModelSelectionProvider token count validation', () => { + beforeEach(() => { + vi.spyOn(TokenCounterUtils, 'countPromptString').mockImplementation(async () => 1500) + + vi.spyOn(vscode, 'Range').mockImplementation(() => ({}) as vscode.Range) + }) + + afterEach(() => { + vi.restoreAllMocks() + }) + + it('should throw an error when token count exceeds context window input limit', async () => { + const mockDocument = { + lineCount: 100, + uri: { fsPath: 'test.ts' }, + getText: vi.fn().mockReturnValue('mock document text'), + } as unknown as vscode.TextDocument + + // Mock token count to exceed context window input + const contextWindow = { input: 1000, output: 500 } + + const provider = new CustomModelSelectionProvider({ shouldAlwaysUseEntireFile: false }) + + await expect( + provider.getSelectedText({ + instruction: ps`test instruction`, + replacement: ps`test replacement`, + document: mockDocument, + model: 'gpt-4', + chatClient: {} as any, + contextWindow, + codyApiVersion: 1, + }) + ).rejects.toThrow("The amount of text in this document exceeds Cody's current capacity.") + + // Verify the token counter was called with the document text + expect(TokenCounterUtils.countPromptString).toHaveBeenCalled() + }) +}) diff --git a/vscode/src/edit/prompt/smart-apply/selection/custom-model.ts b/vscode/src/edit/prompt/smart-apply/selection/custom-model.ts index c537e3f4d84e..0c10478cd123 100644 --- a/vscode/src/edit/prompt/smart-apply/selection/custom-model.ts +++ b/vscode/src/edit/prompt/smart-apply/selection/custom-model.ts @@ -58,10 +58,6 @@ export class CustomModelSelectionProvider implements SmartApplySelectionProvider chatClient, contextWindow, }: SelectionPromptProviderArgs): Promise { - if (this.shouldAlwaysUseEntireFile) { - return 'ENTIRE_FILE' - } - const documentRange = new vscode.Range(0, 0, document.lineCount - 1, 0) const documentText = PromptString.fromDocumentText(document, documentRange) const tokenCount = await TokenCounterUtils.countPromptString(documentText) @@ -69,9 +65,10 @@ export class CustomModelSelectionProvider implements SmartApplySelectionProvider if (tokenCount > contextWindow.input) { throw new Error("The amount of text in this document exceeds Cody's current capacity.") } - if (tokenCount < FULL_FILE_REWRITE_TOKEN_TOKEN_LIMIT) { + if (tokenCount < FULL_FILE_REWRITE_TOKEN_TOKEN_LIMIT || this.shouldAlwaysUseEntireFile) { return 'ENTIRE_FILE' } + const { prefix, messages } = await this.getPrompt( instruction, replacement, From c9b654b77838e4235269d54b2faa94a271ebc2e6 Mon Sep 17 00:00:00 2001 From: David Ichim Date: Sat, 17 May 2025 22:32:00 +0300 Subject: [PATCH 6/6] change(smart-apply): improve token handling and add comprehensive tests for custom models This commit introduces several enhancements to the smart apply custom model feature, focusing on token handling and test coverage. - Adds tests to validate the behavior of `CustomModelSelectionProvider` under different token count scenarios, including exceeding the context window input limit, staying within limits but below a threshold, and always returning `ENTIRE_FILE` when `shouldAlwaysUseEntireFile` is true. - Updates the `trimLLMNewlines` function to handle both `\n` and `\r\n` newline characters, ensuring correct newline preservation across different operating systems. - Introduces a `SMART_APPLY_MODEL_SET` for efficient checking of smart apply custom models. --- .../edit/adapters/smart-apply-custom.test.ts | 90 +++++++++++++++++-- .../src/edit/output/response-transformer.ts | 14 +-- 2 files changed, 92 insertions(+), 12 deletions(-) diff --git a/vscode/src/edit/adapters/smart-apply-custom.test.ts b/vscode/src/edit/adapters/smart-apply-custom.test.ts index f39c62bf7d43..fa83b876c015 100644 --- a/vscode/src/edit/adapters/smart-apply-custom.test.ts +++ b/vscode/src/edit/adapters/smart-apply-custom.test.ts @@ -110,9 +110,11 @@ describe('SmartApplyCustomModelParameterProvider', () => { }) describe('CustomModelSelectionProvider token count validation', () => { - beforeEach(() => { - vi.spyOn(TokenCounterUtils, 'countPromptString').mockImplementation(async () => 1500) + const mockChatClient = { + getCompletion: vi.fn().mockResolvedValue('ENTIRE_FILE'), + } as any + beforeEach(() => { vi.spyOn(vscode, 'Range').mockImplementation(() => ({}) as vscode.Range) }) @@ -121,10 +123,14 @@ describe('CustomModelSelectionProvider token count validation', () => { }) it('should throw an error when token count exceeds context window input limit', async () => { + // Mock token count to exceed limit + vi.spyOn(TokenCounterUtils, 'countPromptString').mockImplementation(async () => 1500) + + const mockDocumentText = 'mock document text that exceeds token limit' const mockDocument = { lineCount: 100, uri: { fsPath: 'test.ts' }, - getText: vi.fn().mockReturnValue('mock document text'), + getText: vi.fn().mockReturnValue(mockDocumentText), } as unknown as vscode.TextDocument // Mock token count to exceed context window input @@ -144,7 +150,81 @@ describe('CustomModelSelectionProvider token count validation', () => { }) ).rejects.toThrow("The amount of text in this document exceeds Cody's current capacity.") - // Verify the token counter was called with the document text - expect(TokenCounterUtils.countPromptString).toHaveBeenCalled() + // Verify the token counter was called with the correct document text + expect(TokenCounterUtils.countPromptString).toHaveBeenCalledTimes(1) + const calledWith = vi.mocked(TokenCounterUtils.countPromptString).mock.calls[0][0] + expect(calledWith).toBeDefined() + expect(calledWith.toString()).toContain(mockDocumentText) + }) + + it('should return ENTIRE_FILE when token count is within limits and below threshold', async () => { + // Mock token count to be within limit but below threshold + vi.spyOn(TokenCounterUtils, 'countPromptString').mockImplementation(async () => 800) + + const mockDocumentText = 'mock document text within token limit' + const mockDocument = { + lineCount: 100, + uri: { fsPath: 'test.ts' }, + getText: vi.fn().mockReturnValue(mockDocumentText), + } as unknown as vscode.TextDocument + + const contextWindow = { input: 1000, output: 500 } + + const provider = new CustomModelSelectionProvider({ shouldAlwaysUseEntireFile: false }) + + const result = await provider.getSelectedText({ + instruction: ps`test instruction`, + replacement: ps`test replacement`, + document: mockDocument, + model: 'gpt-4', + chatClient: mockChatClient, + contextWindow, + codyApiVersion: 1, + }) + + // Verify result is ENTIRE_FILE when token count is below threshold + expect(result).toBe('ENTIRE_FILE') + + // Verify the token counter was called with the correct document text + expect(TokenCounterUtils.countPromptString).toHaveBeenCalledTimes(1) + const calledWith = vi.mocked(TokenCounterUtils.countPromptString).mock.calls[0][0] + expect(calledWith).toBeDefined() + expect(calledWith.toString()).toContain(mockDocumentText) + }) + + it('should always return ENTIRE_FILE when shouldAlwaysUseEntireFile is true regardless of token count', async () => { + // Mock token count to be above threshold but within limit + vi.spyOn(TokenCounterUtils, 'countPromptString').mockImplementation(async () => 15000) + + const mockDocumentText = 'mock document text above threshold' + const mockDocument = { + lineCount: 100, + uri: { fsPath: 'test.ts' }, + getText: vi.fn().mockReturnValue(mockDocumentText), + } as unknown as vscode.TextDocument + + const contextWindow = { input: 20000, output: 500 } // Large enough to not trigger error + + // Create provider instance with shouldAlwaysUseEntireFile set to true + const provider = new CustomModelSelectionProvider({ shouldAlwaysUseEntireFile: true }) + + const result = await provider.getSelectedText({ + instruction: ps`test instruction`, + replacement: ps`test replacement`, + document: mockDocument, + model: 'gpt-4', + chatClient: mockChatClient, + contextWindow, + codyApiVersion: 1, + }) + + // Verify result is ENTIRE_FILE when shouldAlwaysUseEntireFile is true + expect(result).toBe('ENTIRE_FILE') + + // Verify the token counter was called with the correct document text + expect(TokenCounterUtils.countPromptString).toHaveBeenCalledTimes(1) + const calledWith = vi.mocked(TokenCounterUtils.countPromptString).mock.calls[0][0] + expect(calledWith).toBeDefined() + expect(calledWith.toString()).toContain(mockDocumentText) }) }) diff --git a/vscode/src/edit/output/response-transformer.ts b/vscode/src/edit/output/response-transformer.ts index 757e502c5955..59732e52e774 100644 --- a/vscode/src/edit/output/response-transformer.ts +++ b/vscode/src/edit/output/response-transformer.ts @@ -38,11 +38,13 @@ const MARKDOWN_CODE_BLOCK_REGEX = new RegExp( const LEADING_SPACES_AND_NEW_LINES = /^\s*\n/ const LEADING_SPACES = /^[ ]+/ +const SMART_APPLY_MODEL_SET = new Set(Object.values(SMART_APPLY_MODEL_IDENTIFIERS)) + /** * Checks if the task is using a smart apply custom model */ function taskUsesSmartApplyCustomModel(task: FixupTask): boolean { - return Object.values(SMART_APPLY_MODEL_IDENTIFIERS).includes(task.model) + return SMART_APPLY_MODEL_SET.has(task.model) } /** @@ -97,13 +99,11 @@ function trimLLMNewlines(text: string, original: string): string { let result = text // Handle starting newline - if (result.startsWith('\n') && !original.startsWith('\n')) { - result = result.replace(/^\n/, '') + if (result.match(/^\r?\n/) && !original.match(/^\r?\n/)) { + result = result.replace(/^\r?\n/, '') } - - // Handle ending newline - if (result.endsWith('\n') && !original.endsWith('\n')) { - result = result.replace(/\n$/, '') + if (result.match(/\r?\n$/) && !original.match(/\r?\n$/)) { + result = result.replace(/\r?\n$/, '') } return result