From 72aeaca54e8ccb65e3442fae99d99613e4c639ff Mon Sep 17 00:00:00 2001 From: Brad Skaggs Date: Thu, 17 Apr 2025 15:35:47 -0400 Subject: [PATCH 1/9] fix(amazonq) Previous and subsequent cells are used as context for completion in a Jupyter notebook --- ...-9b0e6490-39a8-445f-9d67-9d762de7421c.json | 4 + .../codewhisperer/util/editorContext.test.ts | 224 ++++++++++++++++++ .../src/codewhisperer/util/editorContext.ts | 118 ++++++++- 3 files changed, 343 insertions(+), 3 deletions(-) create mode 100644 packages/amazonq/.changes/next-release/Bug Fix-9b0e6490-39a8-445f-9d67-9d762de7421c.json diff --git a/packages/amazonq/.changes/next-release/Bug Fix-9b0e6490-39a8-445f-9d67-9d762de7421c.json b/packages/amazonq/.changes/next-release/Bug Fix-9b0e6490-39a8-445f-9d67-9d762de7421c.json new file mode 100644 index 00000000000..f17516bb8f4 --- /dev/null +++ b/packages/amazonq/.changes/next-release/Bug Fix-9b0e6490-39a8-445f-9d67-9d762de7421c.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "Previous and subsequent cells are used as context for completion in a Jupyter notebook" +} diff --git a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts index d5085e4db0c..d16a0735cad 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts @@ -6,6 +6,7 @@ import assert from 'assert' import * as codewhispererClient from 'aws-core-vscode/codewhisperer' import * as EditorContext from 'aws-core-vscode/codewhisperer' import { + createMockDocument, createMockTextEditor, createMockClientRequest, resetCodeWhispererGlobalVariables, @@ -15,6 +16,27 @@ import { } from 'aws-core-vscode/test' import { globals } from 'aws-core-vscode/shared' import { GenerateCompletionsRequest } from 'aws-core-vscode/codewhisperer' +import * as vscode from 'vscode' + +export function createNotebookCell( + document: vscode.TextDocument = createMockDocument('def example():\n return "test"'), + kind: vscode.NotebookCellKind = vscode.NotebookCellKind.Code, + notebook: vscode.NotebookDocument = {} as any, + index: number = 0, + outputs: vscode.NotebookCellOutput[] = [], + metadata: { readonly [key: string]: any } = {}, + executionSummary?: vscode.NotebookCellExecutionSummary +): vscode.NotebookCell { + return { + document, + kind, + notebook, + index, + outputs, + metadata, + executionSummary, + } +} describe('editorContext', function () { let telemetryEnabledDefault: boolean @@ -63,6 +85,44 @@ describe('editorContext', function () { } assert.deepStrictEqual(actual, expected) }) + + it('Should include context from other cells when in a notebook', async function () { + const cells: vscode.NotebookCellData[] = [ + new vscode.NotebookCellData(vscode.NotebookCellKind.Markup, 'Previous cell', 'python'), + new vscode.NotebookCellData( + vscode.NotebookCellKind.Code, + 'import numpy as np\nimport pandas as pd\n\ndef analyze_data(df):\n # Current cell with cursor here', + 'python' + ), + new vscode.NotebookCellData( + vscode.NotebookCellKind.Code, + '# Process the data\nresult = analyze_data(df)\nprint(result)', + 'python' + ), + ] + + const document = await vscode.workspace.openNotebookDocument( + 'jupyter-notebook', + new vscode.NotebookData(cells) + ) + const editor: any = { + document: document.cellAt(1).document, + selection: { active: new vscode.Position(4, 13) }, + } + + const actual = EditorContext.extractContextForCodeWhisperer(editor) + const expected: codewhispererClient.FileContext = { + filename: 'Untitled-1.py', + programmingLanguage: { + languageName: 'python', + }, + leftFileContent: + '# Previous cell\nimport numpy as np\nimport pandas as pd\n\ndef analyze_data(df):\n # Current', + rightFileContent: + ' cell with cursor here\n# Process the data\nresult = analyze_data(df)\nprint(result)\n', + } + assert.deepStrictEqual(actual, expected) + }) }) describe('getFileName', function () { @@ -115,6 +175,170 @@ describe('editorContext', function () { }) }) + describe('extractSingleCellContext', function () { + it('Should return cell text for python code cells when language is python', function () { + const mockCodeCell = createNotebookCell(createMockDocument('def example():\n return "test"')) + const result = EditorContext.extractSingleCellContext(mockCodeCell, 'python') + assert.strictEqual(result, 'def example():\n return "test"') + }) + + it('Should return java comments for python code cells when language is java', function () { + const mockCodeCell = createNotebookCell(createMockDocument('def example():\n return "test"')) + const result = EditorContext.extractSingleCellContext(mockCodeCell, 'java') + assert.strictEqual(result, '// def example():\n// return "test"') + }) + + it('Should return python comments for java code cells when language is python', function () { + const mockCodeCell = createNotebookCell(createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java')) + const result = EditorContext.extractSingleCellContext(mockCodeCell, 'python') + assert.strictEqual(result, '# println(1 + 1);') + }) + + it('Should add python comment prefixes for markdown cells when language is python', function () { + const mockMarkdownCell = createNotebookCell( + createMockDocument('# Heading\nThis is a markdown cell'), + vscode.NotebookCellKind.Markup + ) + const result = EditorContext.extractSingleCellContext(mockMarkdownCell, 'python') + assert.strictEqual(result, '# # Heading\n# This is a markdown cell') + }) + + it('Should add java comment prefixes for markdown cells when language is java', function () { + const mockMarkdownCell = createNotebookCell( + createMockDocument('# Heading\nThis is a markdown cell'), + vscode.NotebookCellKind.Markup + ) + const result = EditorContext.extractSingleCellContext(mockMarkdownCell, 'java') + assert.strictEqual(result, '// # Heading\n// This is a markdown cell') + }) + }) + + describe('extractPrefixCellsContext', function () { + it('Should extract content from cells in reverse order up to maxLength', function () { + const mockCells = [ + createNotebookCell(createMockDocument('First cell content')), + createNotebookCell(createMockDocument('Second cell content')), + createNotebookCell(createMockDocument('Third cell content')), + ] + + const result = EditorContext.extractPrefixCellsContext(mockCells, 100, 'python') + assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') + }) + + it('Should respect maxLength parameter', function () { + const mockCells = [ + createNotebookCell(createMockDocument('First')), + createNotebookCell(createMockDocument('Second')), + createNotebookCell(createMockDocument('Third')), + createNotebookCell(createMockDocument('Fourth')), + ] + + const result = EditorContext.extractPrefixCellsContext(mockCells, 15, 'python') + assert.strictEqual(result, 'd\nThird\nFourth\n') + }) + + it('Should handle empty cells array', function () { + const result = EditorContext.extractPrefixCellsContext([], 100, '') + assert.strictEqual(result, '') + }) + + it('Should add python comments to markdown cells', function () { + const mockCells = [ + createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), + createNotebookCell(createMockDocument('def example():\n return "test"')), + ] + const result = EditorContext.extractPrefixCellsContext(mockCells, 100, 'python') + assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') + }) + + it('Should add java comments to markdown and python cells when language is java', function () { + const mockCells = [ + createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), + createNotebookCell(createMockDocument('def example():\n return "test"')), + ] + const result = EditorContext.extractPrefixCellsContext(mockCells, 100, 'java') + assert.strictEqual(result, '// # Heading\n// This is markdown\n// def example():\n// return "test"\n') + }) + + it('Should handle code cells with different languages', function () { + const mockCells = [ + createNotebookCell( + createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java'), + vscode.NotebookCellKind.Code + ), + createNotebookCell(createMockDocument('def example():\n return "test"')), + ] + const result = EditorContext.extractPrefixCellsContext(mockCells, 100, 'python') + assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') + }) + }) + + describe('extractSuffixCellsContext', function () { + it('Should extract content from cells in order up to maxLength', function () { + const mockCells = [ + createNotebookCell(createMockDocument('First cell content')), + createNotebookCell(createMockDocument('Second cell content')), + createNotebookCell(createMockDocument('Third cell content')), + ] + + const result = EditorContext.extractSuffixCellsContext(mockCells, 100, 'python') + assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') + }) + + it('Should respect maxLength parameter', function () { + const mockCells = [ + createNotebookCell(createMockDocument('First')), + createNotebookCell(createMockDocument('Second')), + createNotebookCell(createMockDocument('Third')), + createNotebookCell(createMockDocument('Fourth')), + ] + + // Should only include first cell and part of second cell + const result = EditorContext.extractSuffixCellsContext(mockCells, 15, 'plaintext') + assert.strictEqual(result, 'First\nSecond\nTh') + }) + + it('Should handle empty cells array', function () { + const result = EditorContext.extractSuffixCellsContext([], 100, 'plaintext') + assert.strictEqual(result, '') + }) + + it('Should add python comments to markdown cells', function () { + const mockCells = [ + createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), + createNotebookCell(createMockDocument('def example():\n return "test"')), + ] + + const result = EditorContext.extractSuffixCellsContext(mockCells, 100, 'python') + assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') + }) + + it('Should add java comments to markdown cells', function () { + const mockCells = [ + createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), + createNotebookCell( + createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java'), + vscode.NotebookCellKind.Code + ), + ] + + const result = EditorContext.extractSuffixCellsContext(mockCells, 100, 'java') + assert.strictEqual(result, '// # Heading\n// This is markdown\nprintln(1 + 1);\n') + }) + + it('Should handle code cells with different languages', function () { + const mockCells = [ + createNotebookCell( + createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java'), + vscode.NotebookCellKind.Code + ), + createNotebookCell(createMockDocument('def example():\n return "test"')), + ] + const result = EditorContext.extractSuffixCellsContext(mockCells, 100, 'python') + assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') + }) + }) + describe('validateRequest', function () { it('Should return false if request filename.length is invalid', function () { const req = createMockClientRequest() diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts index 58301e176f6..f196b227f27 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -25,19 +25,95 @@ import { predictionTracker } from '../nextEditPrediction/activation' let tabSize: number = getTabSizeSetting() +const languageCommentChars: Record = { + python: '# ', + java: '// ', +} + +export function extractSingleCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string { + // Extract the text verbatim if the cell is code and the cell has the same language. + // Otherwise, add the correct comment string for the refeference language + const cellText = cell.document.getText() + if ( + cell.kind === vscode.NotebookCellKind.Markup || + (runtimeLanguageContext.normalizeLanguage(cell.document.languageId) ?? cell.document.languageId) !== + referenceLanguage + ) { + const commentPrefix = (referenceLanguage && languageCommentChars[referenceLanguage]) ?? '' + if (commentPrefix === '') { + return cellText + } + return cell.document + .getText() + .split('\n') + .map((line) => `${commentPrefix}${line}`) + .join('\n') + } + return cellText +} + +export function extractPrefixCellsContext( + cells: vscode.NotebookCell[], + maxLength: number, + referenceLanguage: string +): string { + const output: string[] = [] + for (let i = cells.length - 1; i >= 0; i--) { + let cellText = addNewlineIfMissing(extractSingleCellContext(cells[i], referenceLanguage)) + if (cellText.length > 0) { + if (cellText.length >= maxLength) { + output.unshift(cellText.substring(cellText.length - maxLength)) + break + } + output.unshift(cellText) + maxLength -= cellText.length + } + } + return output.join('') +} + +export function extractSuffixCellsContext( + cells: vscode.NotebookCell[], + maxLength: number, + referenceLanguage: string +): string { + const output: string[] = [] + for (let i = 0; i < cells.length; i++) { + let cellText = addNewlineIfMissing(extractSingleCellContext(cells[i], referenceLanguage)) + if (cellText.length > 0) { + if (!cellText.endsWith('\n')) { + cellText += '\n' + } + if (cellText.length >= maxLength) { + output.push(cellText.substring(0, maxLength)) + break + } + output.push(cellText) + maxLength -= cellText.length + } + } + return output.join('') +} + +export function addNewlineIfMissing(text: string): string { + if (text.length > 0 && !text.endsWith('\n')) { + text += '\n' + } + return text +} + export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codewhispererClient.FileContext { const document = editor.document const curPos = editor.selection.active const offset = document.offsetAt(curPos) - const caretLeftFileContext = editor.document.getText( + let caretLeftFileContext = editor.document.getText( new vscode.Range( document.positionAt(offset - CodeWhispererConstants.charactersLimit), document.positionAt(offset) ) ) - - const caretRightFileContext = editor.document.getText( + let caretRightFileContext = editor.document.getText( new vscode.Range( document.positionAt(offset), document.positionAt(offset + CodeWhispererConstants.charactersLimit) @@ -48,6 +124,42 @@ export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codew languageName = runtimeLanguageContext.normalizeLanguage(editor.document.languageId) ?? editor.document.languageId } + if (editor.document.uri.scheme === 'vscode-notebook-cell') { + // For notebook cells, first find the existing notebook with a cell that matches the current editor. + const notebook = vscode.workspace.notebookDocuments.find( + (nb) => + nb.notebookType === 'jupyter-notebook' && + nb.getCells().some((cell) => cell.document === editor.document) + ) + if (notebook) { + const allCells = notebook.getCells() + const cellIndex = allCells.findIndex((cell) => cell.document === editor.document) + + // Extract text from prior cells if there is enough room in left file context + if (caretLeftFileContext.length < CodeWhispererConstants.charactersLimit - 1) { + const leftCellsText = extractPrefixCellsContext( + allCells.slice(0, cellIndex), + CodeWhispererConstants.charactersLimit - (caretLeftFileContext.length + 1), + languageName + ) + if (leftCellsText.length > 0) { + caretLeftFileContext = addNewlineIfMissing(leftCellsText) + caretLeftFileContext + } + } + // Extract text from subsequent cells if there is enough room in right file context + if (caretRightFileContext.length < CodeWhispererConstants.charactersLimit - 1) { + const rightCellsText = extractSuffixCellsContext( + allCells.slice(cellIndex + 1), + CodeWhispererConstants.charactersLimit - (caretRightFileContext.length + 1), + languageName + ) + if (rightCellsText.length > 0) { + caretRightFileContext = addNewlineIfMissing(caretRightFileContext) + rightCellsText + } + } + } + } + return { filename: getFileRelativePath(editor), programmingLanguage: { From e95932b78f667766b0b5759155278948625793b9 Mon Sep 17 00:00:00 2001 From: Brad Skaggs Date: Fri, 18 Apr 2025 13:28:05 -0400 Subject: [PATCH 2/9] fix(amazonq) Previous and subsequent cells are used as context for completion in a Jupyter notebook --- .../codewhisperer/util/editorContext.test.ts | 117 +++++++++--------- .../src/codewhisperer/util/editorContext.ts | 56 ++++----- 2 files changed, 79 insertions(+), 94 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts index d16a0735cad..858ca4ddeb2 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts @@ -214,119 +214,114 @@ describe('editorContext', function () { }) describe('extractPrefixCellsContext', function () { - it('Should extract content from cells in reverse order up to maxLength', function () { + it('Should extract content from cells in reverse order up to maxLength from prefix cells', function () { const mockCells = [ createNotebookCell(createMockDocument('First cell content')), createNotebookCell(createMockDocument('Second cell content')), createNotebookCell(createMockDocument('Third cell content')), ] - const result = EditorContext.extractPrefixCellsContext(mockCells, 100, 'python') + const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', false) assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') }) - it('Should respect maxLength parameter', function () { + it('Should extract content from cells in reverse order up to maxLength from suffix cells', function () { + const mockCells = [ + createNotebookCell(createMockDocument('First cell content')), + createNotebookCell(createMockDocument('Second cell content')), + createNotebookCell(createMockDocument('Third cell content')), + ] + + const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', true) + assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') + }) + + it('Should respect maxLength parameter from prefix cells', function () { const mockCells = [ createNotebookCell(createMockDocument('First')), createNotebookCell(createMockDocument('Second')), createNotebookCell(createMockDocument('Third')), createNotebookCell(createMockDocument('Fourth')), ] - - const result = EditorContext.extractPrefixCellsContext(mockCells, 15, 'python') + // Should only include part of second cell and the last two cells + const result = EditorContext.extractCellsSliceContext(mockCells, 15, 'python', false) assert.strictEqual(result, 'd\nThird\nFourth\n') }) - it('Should handle empty cells array', function () { - const result = EditorContext.extractPrefixCellsContext([], 100, '') + it('Should respect maxLength parameter from suffix cells', function () { + const mockCells = [ + createNotebookCell(createMockDocument('First')), + createNotebookCell(createMockDocument('Second')), + createNotebookCell(createMockDocument('Third')), + createNotebookCell(createMockDocument('Fourth')), + ] + + // Should only include first cell and part of second cell + const result = EditorContext.extractCellsSliceContext(mockCells, 15, 'python', true) + assert.strictEqual(result, 'First\nSecond\nTh') + }) + + it('Should handle empty cells array from prefix cells', function () { + const result = EditorContext.extractCellsSliceContext([], 100, 'python', false) assert.strictEqual(result, '') }) - it('Should add python comments to markdown cells', function () { - const mockCells = [ - createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), - createNotebookCell(createMockDocument('def example():\n return "test"')), - ] - const result = EditorContext.extractPrefixCellsContext(mockCells, 100, 'python') - assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') + it('Should handle empty cells array from suffix cells', function () { + const result = EditorContext.extractCellsSliceContext([], 100, 'python', true) + assert.strictEqual(result, '') }) - it('Should add java comments to markdown and python cells when language is java', function () { + it('Should add python comments to markdown prefix cells', function () { const mockCells = [ createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), createNotebookCell(createMockDocument('def example():\n return "test"')), ] - const result = EditorContext.extractPrefixCellsContext(mockCells, 100, 'java') - assert.strictEqual(result, '// # Heading\n// This is markdown\n// def example():\n// return "test"\n') + const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', false) + assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') }) - it('Should handle code cells with different languages', function () { + it('Should add python comments to markdown suffix cells', function () { const mockCells = [ - createNotebookCell( - createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java'), - vscode.NotebookCellKind.Code - ), + createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), createNotebookCell(createMockDocument('def example():\n return "test"')), ] - const result = EditorContext.extractPrefixCellsContext(mockCells, 100, 'python') - assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') - }) - }) - describe('extractSuffixCellsContext', function () { - it('Should extract content from cells in order up to maxLength', function () { - const mockCells = [ - createNotebookCell(createMockDocument('First cell content')), - createNotebookCell(createMockDocument('Second cell content')), - createNotebookCell(createMockDocument('Third cell content')), - ] - - const result = EditorContext.extractSuffixCellsContext(mockCells, 100, 'python') - assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') + const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', true) + assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') }) - it('Should respect maxLength parameter', function () { + it('Should add java comments to markdown and python prefix cells when language is java', function () { const mockCells = [ - createNotebookCell(createMockDocument('First')), - createNotebookCell(createMockDocument('Second')), - createNotebookCell(createMockDocument('Third')), - createNotebookCell(createMockDocument('Fourth')), + createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), + createNotebookCell(createMockDocument('def example():\n return "test"')), ] - - // Should only include first cell and part of second cell - const result = EditorContext.extractSuffixCellsContext(mockCells, 15, 'plaintext') - assert.strictEqual(result, 'First\nSecond\nTh') - }) - - it('Should handle empty cells array', function () { - const result = EditorContext.extractSuffixCellsContext([], 100, 'plaintext') - assert.strictEqual(result, '') + const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'java', false) + assert.strictEqual(result, '// # Heading\n// This is markdown\n// def example():\n// return "test"\n') }) - it('Should add python comments to markdown cells', function () { + it('Should add java comments to markdown and python suffix cells when language is java', function () { const mockCells = [ createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), - createNotebookCell(createMockDocument('def example():\n return "test"')), + createNotebookCell(createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java')), ] - const result = EditorContext.extractSuffixCellsContext(mockCells, 100, 'python') - assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') + const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'java', true) + assert.strictEqual(result, '// # Heading\n// This is markdown\nprintln(1 + 1);\n') }) - it('Should add java comments to markdown cells', function () { + it('Should handle code prefix cells with different languages', function () { const mockCells = [ - createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), createNotebookCell( createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java'), vscode.NotebookCellKind.Code ), + createNotebookCell(createMockDocument('def example():\n return "test"')), ] - - const result = EditorContext.extractSuffixCellsContext(mockCells, 100, 'java') - assert.strictEqual(result, '// # Heading\n// This is markdown\nprintln(1 + 1);\n') + const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', false) + assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') }) - it('Should handle code cells with different languages', function () { + it('Should handle code suffix cells with different languages', function () { const mockCells = [ createNotebookCell( createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java'), @@ -334,7 +329,7 @@ describe('editorContext', function () { ), createNotebookCell(createMockDocument('def example():\n return "test"')), ] - const result = EditorContext.extractSuffixCellsContext(mockCells, 100, 'python') + const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', true) assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') }) }) diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts index f196b227f27..4240418dcb8 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -52,45 +52,33 @@ export function extractSingleCellContext(cell: vscode.NotebookCell, referenceLan return cellText } -export function extractPrefixCellsContext( +export function extractCellsSliceContext( cells: vscode.NotebookCell[], maxLength: number, - referenceLanguage: string + referenceLanguage: string, + fromStart: boolean ): string { - const output: string[] = [] - for (let i = cells.length - 1; i >= 0; i--) { - let cellText = addNewlineIfMissing(extractSingleCellContext(cells[i], referenceLanguage)) - if (cellText.length > 0) { - if (cellText.length >= maxLength) { - output.unshift(cellText.substring(cellText.length - maxLength)) - break - } - output.unshift(cellText) - maxLength -= cellText.length - } + let output: string[] = [] + if (!fromStart) { + cells = cells.reverse() } - return output.join('') -} - -export function extractSuffixCellsContext( - cells: vscode.NotebookCell[], - maxLength: number, - referenceLanguage: string -): string { - const output: string[] = [] - for (let i = 0; i < cells.length; i++) { - let cellText = addNewlineIfMissing(extractSingleCellContext(cells[i], referenceLanguage)) + cells.some((cell) => { + let cellText = addNewlineIfMissing(extractSingleCellContext(cell, referenceLanguage)) if (cellText.length > 0) { - if (!cellText.endsWith('\n')) { - cellText += '\n' - } if (cellText.length >= maxLength) { - output.push(cellText.substring(0, maxLength)) - break + if (fromStart) { + output.push(cellText.substring(0, maxLength)) + } else { + output.push(cellText.substring(cellText.length - maxLength)) + } + return true } output.push(cellText) maxLength -= cellText.length } + }) + if (!fromStart) { + output = output.reverse() } return output.join('') } @@ -137,10 +125,11 @@ export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codew // Extract text from prior cells if there is enough room in left file context if (caretLeftFileContext.length < CodeWhispererConstants.charactersLimit - 1) { - const leftCellsText = extractPrefixCellsContext( + const leftCellsText = extractCellsSliceContext( allCells.slice(0, cellIndex), CodeWhispererConstants.charactersLimit - (caretLeftFileContext.length + 1), - languageName + languageName, + true ) if (leftCellsText.length > 0) { caretLeftFileContext = addNewlineIfMissing(leftCellsText) + caretLeftFileContext @@ -148,10 +137,11 @@ export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codew } // Extract text from subsequent cells if there is enough room in right file context if (caretRightFileContext.length < CodeWhispererConstants.charactersLimit - 1) { - const rightCellsText = extractSuffixCellsContext( + const rightCellsText = extractCellsSliceContext( allCells.slice(cellIndex + 1), CodeWhispererConstants.charactersLimit - (caretRightFileContext.length + 1), - languageName + languageName, + false ) if (rightCellsText.length > 0) { caretRightFileContext = addNewlineIfMissing(caretRightFileContext) + rightCellsText From a8b5f1c04a436b079a4856017868ca3133080c27 Mon Sep 17 00:00:00 2001 From: Brad Skaggs Date: Fri, 18 Apr 2025 13:32:22 -0400 Subject: [PATCH 3/9] Fix test name to be extractCellsSliceContext --- .../amazonq/test/unit/codewhisperer/util/editorContext.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts index 858ca4ddeb2..76b716c2962 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts @@ -213,7 +213,7 @@ describe('editorContext', function () { }) }) - describe('extractPrefixCellsContext', function () { + describe('extractCellsSliceContext', function () { it('Should extract content from cells in reverse order up to maxLength from prefix cells', function () { const mockCells = [ createNotebookCell(createMockDocument('First cell content')), From b652541e9873714656a2fec1ee4ab31a23f4cdd2 Mon Sep 17 00:00:00 2001 From: Brad Skaggs <126105424+brdskggs@users.noreply.github.com> Date: Mon, 5 May 2025 10:56:56 -0400 Subject: [PATCH 4/9] Apply suggestions from code review Renaming `extractCellsSliceContext` to `getNotebookCellsSliceContext` Co-authored-by: Justin M. Keyes --- .../test/unit/codewhisperer/util/editorContext.test.ts | 4 ++-- packages/core/src/codewhisperer/util/editorContext.ts | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts index 76b716c2962..1d86b246475 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts @@ -86,7 +86,7 @@ describe('editorContext', function () { assert.deepStrictEqual(actual, expected) }) - it('Should include context from other cells when in a notebook', async function () { + it('in a notebook, includes context from other cells', async function () { const cells: vscode.NotebookCellData[] = [ new vscode.NotebookCellData(vscode.NotebookCellKind.Markup, 'Previous cell', 'python'), new vscode.NotebookCellData( @@ -175,7 +175,7 @@ describe('editorContext', function () { }) }) - describe('extractSingleCellContext', function () { + describe('getNotebookCellContext', function () { it('Should return cell text for python code cells when language is python', function () { const mockCodeCell = createNotebookCell(createMockDocument('def example():\n return "test"')) const result = EditorContext.extractSingleCellContext(mockCodeCell, 'python') diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts index 4240418dcb8..c665895b4ce 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -30,7 +30,7 @@ const languageCommentChars: Record = { java: '// ', } -export function extractSingleCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string { +export function getNotebookCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string { // Extract the text verbatim if the cell is code and the cell has the same language. // Otherwise, add the correct comment string for the refeference language const cellText = cell.document.getText() @@ -52,7 +52,7 @@ export function extractSingleCellContext(cell: vscode.NotebookCell, referenceLan return cellText } -export function extractCellsSliceContext( +export function getNotebookCellsSliceContext( cells: vscode.NotebookCell[], maxLength: number, referenceLanguage: string, From 80eb542070ff9f2ff11e6d67ab3741dfa8feaee7 Mon Sep 17 00:00:00 2001 From: Brad Skaggs Date: Wed, 7 May 2025 13:01:39 -0400 Subject: [PATCH 5/9] Complete function rename; refactored notebook logic out of main function. --- .../codewhisperer/util/editorContext.test.ts | 26 +++--- .../src/codewhisperer/util/editorContext.ts | 90 ++++++++++++------- 2 files changed, 69 insertions(+), 47 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts index 1d86b246475..527507ae094 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts @@ -213,7 +213,7 @@ describe('editorContext', function () { }) }) - describe('extractCellsSliceContext', function () { + describe('getNotebookCellsSliceContext', function () { it('Should extract content from cells in reverse order up to maxLength from prefix cells', function () { const mockCells = [ createNotebookCell(createMockDocument('First cell content')), @@ -221,7 +221,7 @@ describe('editorContext', function () { createNotebookCell(createMockDocument('Third cell content')), ] - const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', false) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false) assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') }) @@ -232,7 +232,7 @@ describe('editorContext', function () { createNotebookCell(createMockDocument('Third cell content')), ] - const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', true) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true) assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') }) @@ -244,7 +244,7 @@ describe('editorContext', function () { createNotebookCell(createMockDocument('Fourth')), ] // Should only include part of second cell and the last two cells - const result = EditorContext.extractCellsSliceContext(mockCells, 15, 'python', false) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 15, 'python', false) assert.strictEqual(result, 'd\nThird\nFourth\n') }) @@ -257,17 +257,17 @@ describe('editorContext', function () { ] // Should only include first cell and part of second cell - const result = EditorContext.extractCellsSliceContext(mockCells, 15, 'python', true) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 15, 'python', true) assert.strictEqual(result, 'First\nSecond\nTh') }) it('Should handle empty cells array from prefix cells', function () { - const result = EditorContext.extractCellsSliceContext([], 100, 'python', false) + const result = EditorContext.getNotebookCellsSliceContext([], 100, 'python', false) assert.strictEqual(result, '') }) it('Should handle empty cells array from suffix cells', function () { - const result = EditorContext.extractCellsSliceContext([], 100, 'python', true) + const result = EditorContext.getNotebookCellsSliceContext([], 100, 'python', true) assert.strictEqual(result, '') }) @@ -276,7 +276,7 @@ describe('editorContext', function () { createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), createNotebookCell(createMockDocument('def example():\n return "test"')), ] - const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', false) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false) assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') }) @@ -286,7 +286,7 @@ describe('editorContext', function () { createNotebookCell(createMockDocument('def example():\n return "test"')), ] - const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', true) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true) assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') }) @@ -295,7 +295,7 @@ describe('editorContext', function () { createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), createNotebookCell(createMockDocument('def example():\n return "test"')), ] - const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'java', false) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'java', false) assert.strictEqual(result, '// # Heading\n// This is markdown\n// def example():\n// return "test"\n') }) @@ -305,7 +305,7 @@ describe('editorContext', function () { createNotebookCell(createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java')), ] - const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'java', true) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'java', true) assert.strictEqual(result, '// # Heading\n// This is markdown\nprintln(1 + 1);\n') }) @@ -317,7 +317,7 @@ describe('editorContext', function () { ), createNotebookCell(createMockDocument('def example():\n return "test"')), ] - const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', false) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false) assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') }) @@ -329,7 +329,7 @@ describe('editorContext', function () { ), createNotebookCell(createMockDocument('def example():\n return "test"')), ] - const result = EditorContext.extractCellsSliceContext(mockCells, 100, 'python', true) + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true) assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') }) }) diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts index c665895b4ce..7a7ef020ef1 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -30,7 +30,52 @@ const languageCommentChars: Record = { java: '// ', } -export function getNotebookCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string { +function getEnclosingNotebook(editor: vscode.TextEditor): vscode.NotebookDocument | undefined { + // For notebook cells, find the existing notebook with a cell that matches the current editor. + return vscode.workspace.notebookDocuments.find( + (nb) => + nb.notebookType === 'jupyter-notebook' && nb.getCells().some((cell) => cell.document === editor.document) + ) +} + +export function extractNotebookContext( + notebook: vscode.NotebookDocument, + editor: vscode.TextEditor, + languageName: string, + caretLeftFileContext: string, + caretRightFileContext: string +) { + // Expand the context for a cell inside of a noteboo with whatever text fits from the preceding and subsequent cells + const allCells = notebook.getCells() + const cellIndex = allCells.findIndex((cell) => cell.document === editor.document) + // Extract text from prior cells if there is enough room in left file context + if (caretLeftFileContext.length < CodeWhispererConstants.charactersLimit - 1) { + const leftCellsText = getNotebookCellsSliceContext( + allCells.slice(0, cellIndex), + CodeWhispererConstants.charactersLimit - (caretLeftFileContext.length + 1), + languageName, + true + ) + if (leftCellsText.length > 0) { + caretLeftFileContext = addNewlineIfMissing(leftCellsText) + caretLeftFileContext + } + } + // Extract text from subsequent cells if there is enough room in right file context + if (caretRightFileContext.length < CodeWhispererConstants.charactersLimit - 1) { + const rightCellsText = getNotebookCellsSliceContext( + allCells.slice(cellIndex + 1), + CodeWhispererConstants.charactersLimit - (caretRightFileContext.length + 1), + languageName, + false + ) + if (rightCellsText.length > 0) { + caretRightFileContext = addNewlineIfMissing(caretRightFileContext) + rightCellsText + } + } + return { caretLeftFileContext, caretRightFileContext } +} + +export function extractSingleCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string { // Extract the text verbatim if the cell is code and the cell has the same language. // Otherwise, add the correct comment string for the refeference language const cellText = cell.document.getText() @@ -58,6 +103,8 @@ export function getNotebookCellsSliceContext( referenceLanguage: string, fromStart: boolean ): string { + // Extract context from array of notebook cells that fits inside `maxLength` characters, + // from either the start or the end of the array. let output: string[] = [] if (!fromStart) { cells = cells.reverse() @@ -113,40 +160,15 @@ export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codew runtimeLanguageContext.normalizeLanguage(editor.document.languageId) ?? editor.document.languageId } if (editor.document.uri.scheme === 'vscode-notebook-cell') { - // For notebook cells, first find the existing notebook with a cell that matches the current editor. - const notebook = vscode.workspace.notebookDocuments.find( - (nb) => - nb.notebookType === 'jupyter-notebook' && - nb.getCells().some((cell) => cell.document === editor.document) - ) + const notebook = getEnclosingNotebook(editor) if (notebook) { - const allCells = notebook.getCells() - const cellIndex = allCells.findIndex((cell) => cell.document === editor.document) - - // Extract text from prior cells if there is enough room in left file context - if (caretLeftFileContext.length < CodeWhispererConstants.charactersLimit - 1) { - const leftCellsText = extractCellsSliceContext( - allCells.slice(0, cellIndex), - CodeWhispererConstants.charactersLimit - (caretLeftFileContext.length + 1), - languageName, - true - ) - if (leftCellsText.length > 0) { - caretLeftFileContext = addNewlineIfMissing(leftCellsText) + caretLeftFileContext - } - } - // Extract text from subsequent cells if there is enough room in right file context - if (caretRightFileContext.length < CodeWhispererConstants.charactersLimit - 1) { - const rightCellsText = extractCellsSliceContext( - allCells.slice(cellIndex + 1), - CodeWhispererConstants.charactersLimit - (caretRightFileContext.length + 1), - languageName, - false - ) - if (rightCellsText.length > 0) { - caretRightFileContext = addNewlineIfMissing(caretRightFileContext) + rightCellsText - } - } + ;({ caretLeftFileContext, caretRightFileContext } = extractNotebookContext( + notebook, + editor, + languageName, + caretLeftFileContext, + caretRightFileContext + )) } } From f0f8b684cf9840ad70fee84e001b86a8d614a70d Mon Sep 17 00:00:00 2001 From: Brad Skaggs Date: Wed, 7 May 2025 13:20:47 -0400 Subject: [PATCH 6/9] Fix typo - stray semi-colon From b440e8ba421c2e3709e0bd3549e1ca081e4d059d Mon Sep 17 00:00:00 2001 From: Brad Skaggs Date: Wed, 7 May 2025 13:33:33 -0400 Subject: [PATCH 7/9] Finish function renaming --- .../test/unit/codewhisperer/util/editorContext.test.ts | 10 +++++----- packages/core/src/codewhisperer/util/editorContext.ts | 8 ++++---- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts index 527507ae094..f8265a4fa86 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts @@ -178,19 +178,19 @@ describe('editorContext', function () { describe('getNotebookCellContext', function () { it('Should return cell text for python code cells when language is python', function () { const mockCodeCell = createNotebookCell(createMockDocument('def example():\n return "test"')) - const result = EditorContext.extractSingleCellContext(mockCodeCell, 'python') + const result = EditorContext.getNotebookCellContext(mockCodeCell, 'python') assert.strictEqual(result, 'def example():\n return "test"') }) it('Should return java comments for python code cells when language is java', function () { const mockCodeCell = createNotebookCell(createMockDocument('def example():\n return "test"')) - const result = EditorContext.extractSingleCellContext(mockCodeCell, 'java') + const result = EditorContext.getNotebookCellContext(mockCodeCell, 'java') assert.strictEqual(result, '// def example():\n// return "test"') }) it('Should return python comments for java code cells when language is python', function () { const mockCodeCell = createNotebookCell(createMockDocument('println(1 + 1);', 'somefile.ipynb', 'java')) - const result = EditorContext.extractSingleCellContext(mockCodeCell, 'python') + const result = EditorContext.getNotebookCellContext(mockCodeCell, 'python') assert.strictEqual(result, '# println(1 + 1);') }) @@ -199,7 +199,7 @@ describe('editorContext', function () { createMockDocument('# Heading\nThis is a markdown cell'), vscode.NotebookCellKind.Markup ) - const result = EditorContext.extractSingleCellContext(mockMarkdownCell, 'python') + const result = EditorContext.getNotebookCellContext(mockMarkdownCell, 'python') assert.strictEqual(result, '# # Heading\n# This is a markdown cell') }) @@ -208,7 +208,7 @@ describe('editorContext', function () { createMockDocument('# Heading\nThis is a markdown cell'), vscode.NotebookCellKind.Markup ) - const result = EditorContext.extractSingleCellContext(mockMarkdownCell, 'java') + const result = EditorContext.getNotebookCellContext(mockMarkdownCell, 'java') assert.strictEqual(result, '// # Heading\n// This is a markdown cell') }) }) diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts index 7a7ef020ef1..f4437d166fd 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -38,7 +38,7 @@ function getEnclosingNotebook(editor: vscode.TextEditor): vscode.NotebookDocumen ) } -export function extractNotebookContext( +export function getNotebookContext( notebook: vscode.NotebookDocument, editor: vscode.TextEditor, languageName: string, @@ -75,7 +75,7 @@ export function extractNotebookContext( return { caretLeftFileContext, caretRightFileContext } } -export function extractSingleCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string { +export function getNotebookCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string { // Extract the text verbatim if the cell is code and the cell has the same language. // Otherwise, add the correct comment string for the refeference language const cellText = cell.document.getText() @@ -110,7 +110,7 @@ export function getNotebookCellsSliceContext( cells = cells.reverse() } cells.some((cell) => { - let cellText = addNewlineIfMissing(extractSingleCellContext(cell, referenceLanguage)) + let cellText = addNewlineIfMissing(getNotebookCellContext(cell, referenceLanguage)) if (cellText.length > 0) { if (cellText.length >= maxLength) { if (fromStart) { @@ -162,7 +162,7 @@ export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codew if (editor.document.uri.scheme === 'vscode-notebook-cell') { const notebook = getEnclosingNotebook(editor) if (notebook) { - ;({ caretLeftFileContext, caretRightFileContext } = extractNotebookContext( + ;({ caretLeftFileContext, caretRightFileContext } = getNotebookContext( notebook, editor, languageName, From 98a36dccfbfd756d1fe9717e2055d2a9c40329f2 Mon Sep 17 00:00:00 2001 From: Brad Skaggs Date: Wed, 7 May 2025 14:33:13 -0400 Subject: [PATCH 8/9] Lint fix --- packages/core/src/codewhisperer/util/editorContext.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts index f4437d166fd..9dbb8882c2c 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -110,7 +110,7 @@ export function getNotebookCellsSliceContext( cells = cells.reverse() } cells.some((cell) => { - let cellText = addNewlineIfMissing(getNotebookCellContext(cell, referenceLanguage)) + const cellText = addNewlineIfMissing(getNotebookCellContext(cell, referenceLanguage)) if (cellText.length > 0) { if (cellText.length >= maxLength) { if (fromStart) { From dbfe86b7aa70edeaa61cfa83e6fae1c023140fbf Mon Sep 17 00:00:00 2001 From: Brad Skaggs Date: Wed, 7 May 2025 20:42:25 -0400 Subject: [PATCH 9/9] Move language-to-comment mapping to runtimeLanguageContext --- .../util/runtimeLanguageContext.test.ts | 34 +++++++++++++ .../src/codewhisperer/util/editorContext.ts | 9 +--- .../util/runtimeLanguageContext.ts | 50 +++++++++++++++++++ 3 files changed, 86 insertions(+), 7 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/runtimeLanguageContext.test.ts b/packages/amazonq/test/unit/codewhisperer/util/runtimeLanguageContext.test.ts index 59c3771abb4..a5cc430a5a9 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/runtimeLanguageContext.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/runtimeLanguageContext.test.ts @@ -333,6 +333,40 @@ describe('runtimeLanguageContext', function () { } }) + describe('getSingleLineCommentPrefix', function () { + it('should return the correct comment prefix for supported languages', function () { + assert.strictEqual(languageContext.getSingleLineCommentPrefix('java'), '// ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('javascript'), '// ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('jsonc'), '// ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('kotlin'), '// ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('lua'), '-- ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('python'), '# ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('ruby'), '# ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('sql'), '-- ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('tf'), '# ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('typescript'), '// ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('vue'), '') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('yaml'), '# ') + }) + + it('should normalize language ID before getting comment prefix', function () { + assert.strictEqual(languageContext.getSingleLineCommentPrefix('hcl'), '# ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('javascriptreact'), '// ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('shellscript'), '# ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('typescriptreact'), '// ') + assert.strictEqual(languageContext.getSingleLineCommentPrefix('yml'), '# ') + }) + + it('should return empty string for unsupported languages', function () { + assert.strictEqual(languageContext.getSingleLineCommentPrefix('nonexistent'), '') + assert.strictEqual(languageContext.getSingleLineCommentPrefix(undefined), '') + }) + + it('should return empty string for plaintext', function () { + assert.strictEqual(languageContext.getSingleLineCommentPrefix('plaintext'), '') + }) + }) + // for now we will only jsx mapped to javascript, tsx mapped to typescript, all other language should remain the same describe('test covertCwsprRequest', function () { const leftFileContent = 'left' diff --git a/packages/core/src/codewhisperer/util/editorContext.ts b/packages/core/src/codewhisperer/util/editorContext.ts index 9dbb8882c2c..a3f787af6c6 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -25,11 +25,6 @@ import { predictionTracker } from '../nextEditPrediction/activation' let tabSize: number = getTabSizeSetting() -const languageCommentChars: Record = { - python: '# ', - java: '// ', -} - function getEnclosingNotebook(editor: vscode.TextEditor): vscode.NotebookDocument | undefined { // For notebook cells, find the existing notebook with a cell that matches the current editor. return vscode.workspace.notebookDocuments.find( @@ -77,14 +72,14 @@ export function getNotebookContext( export function getNotebookCellContext(cell: vscode.NotebookCell, referenceLanguage?: string): string { // Extract the text verbatim if the cell is code and the cell has the same language. - // Otherwise, add the correct comment string for the refeference language + // Otherwise, add the correct comment string for the reference language const cellText = cell.document.getText() if ( cell.kind === vscode.NotebookCellKind.Markup || (runtimeLanguageContext.normalizeLanguage(cell.document.languageId) ?? cell.document.languageId) !== referenceLanguage ) { - const commentPrefix = (referenceLanguage && languageCommentChars[referenceLanguage]) ?? '' + const commentPrefix = runtimeLanguageContext.getSingleLineCommentPrefix(referenceLanguage) if (commentPrefix === '') { return cellText } diff --git a/packages/core/src/codewhisperer/util/runtimeLanguageContext.ts b/packages/core/src/codewhisperer/util/runtimeLanguageContext.ts index 9a495cf5356..3a1403b453e 100644 --- a/packages/core/src/codewhisperer/util/runtimeLanguageContext.ts +++ b/packages/core/src/codewhisperer/util/runtimeLanguageContext.ts @@ -58,6 +58,13 @@ export class RuntimeLanguageContext { */ private supportedLanguageExtensionMap: ConstantMap + /** + * A map storing single-line comment prefixes for different languages + * Key: CodewhispererLanguage + * Value: Comment prefix string + */ + private languageSingleLineCommentPrefixMap: ConstantMap + constructor() { this.supportedLanguageMap = createConstantMap< CodeWhispererConstants.PlatformLanguageId | CodewhispererLanguage, @@ -146,6 +153,39 @@ export class RuntimeLanguageContext { psm1: 'powershell', r: 'r', }) + this.languageSingleLineCommentPrefixMap = createConstantMap({ + c: '// ', + cpp: '// ', + csharp: '// ', + dart: '// ', + go: '// ', + hcl: '# ', + java: '// ', + javascript: '// ', + json: '// ', + jsonc: '// ', + jsx: '// ', + kotlin: '// ', + lua: '-- ', + php: '// ', + plaintext: '', + powershell: '# ', + python: '# ', + r: '# ', + ruby: '# ', + rust: '// ', + scala: '// ', + shell: '# ', + sql: '-- ', + swift: '// ', + systemVerilog: '// ', + tf: '# ', + tsx: '// ', + typescript: '// ', + vue: '', // vue lacks a single-line comment prefix + yaml: '# ', + yml: '# ', + }) } /** @@ -159,6 +199,16 @@ export class RuntimeLanguageContext { return this.supportedLanguageMap.get(languageId) } + /** + * Get the comment prefix for a given language + * @param language The language to get comment prefix for + * @returns The comment prefix string, or empty string if not found + */ + public getSingleLineCommentPrefix(language?: string): string { + const normalizedLanguage = this.normalizeLanguage(language) + return normalizedLanguage ? (this.languageSingleLineCommentPrefixMap.get(normalizedLanguage) ?? '') : '' + } + /** * Normalize client side language id to service aware language id (service is not aware of jsx/tsx) * Only used when invoking CodeWhisperer service API, for client usage please use normalizeLanguage