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..f8265a4fa86 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('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( + 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,165 @@ 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.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.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.getNotebookCellContext(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.getNotebookCellContext(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.getNotebookCellContext(mockMarkdownCell, 'java') + assert.strictEqual(result, '// # Heading\n// This is a markdown cell') + }) + }) + + 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')), + createNotebookCell(createMockDocument('Second cell content')), + createNotebookCell(createMockDocument('Third cell content')), + ] + + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', false) + assert.strictEqual(result, 'First cell content\nSecond cell content\nThird cell content\n') + }) + + 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.getNotebookCellsSliceContext(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')), + ] + // Should only include part of second cell and the last two cells + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 15, 'python', false) + assert.strictEqual(result, 'd\nThird\nFourth\n') + }) + + 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.getNotebookCellsSliceContext(mockCells, 15, 'python', true) + assert.strictEqual(result, 'First\nSecond\nTh') + }) + + it('Should handle empty cells array from prefix cells', function () { + const result = EditorContext.getNotebookCellsSliceContext([], 100, 'python', false) + assert.strictEqual(result, '') + }) + + it('Should handle empty cells array from suffix cells', function () { + const result = EditorContext.getNotebookCellsSliceContext([], 100, 'python', true) + assert.strictEqual(result, '') + }) + + 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.getNotebookCellsSliceContext(mockCells, 100, 'python', false) + assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') + }) + + it('Should add python comments to markdown suffix cells', function () { + const mockCells = [ + createNotebookCell(createMockDocument('# Heading\nThis is markdown'), vscode.NotebookCellKind.Markup), + createNotebookCell(createMockDocument('def example():\n return "test"')), + ] + + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'python', true) + assert.strictEqual(result, '# # Heading\n# This is markdown\ndef example():\n return "test"\n') + }) + + it('Should add java comments to markdown and python prefix 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.getNotebookCellsSliceContext(mockCells, 100, 'java', false) + assert.strictEqual(result, '// # Heading\n// This is markdown\n// def example():\n// return "test"\n') + }) + + 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('println(1 + 1);', 'somefile.ipynb', 'java')), + ] + + const result = EditorContext.getNotebookCellsSliceContext(mockCells, 100, 'java', true) + assert.strictEqual(result, '// # Heading\n// This is markdown\nprintln(1 + 1);\n') + }) + + it('Should handle code prefix 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.getNotebookCellsSliceContext(mockCells, 100, 'python', false) + assert.strictEqual(result, '# println(1 + 1);\ndef example():\n return "test"\n') + }) + + it('Should handle code suffix 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.getNotebookCellsSliceContext(mockCells, 100, 'python', true) + 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/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 58301e176f6..a3f787af6c6 100644 --- a/packages/core/src/codewhisperer/util/editorContext.ts +++ b/packages/core/src/codewhisperer/util/editorContext.ts @@ -25,19 +25,125 @@ import { predictionTracker } from '../nextEditPrediction/activation' let tabSize: number = getTabSizeSetting() +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 getNotebookContext( + 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 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 reference language + const cellText = cell.document.getText() + if ( + cell.kind === vscode.NotebookCellKind.Markup || + (runtimeLanguageContext.normalizeLanguage(cell.document.languageId) ?? cell.document.languageId) !== + referenceLanguage + ) { + const commentPrefix = runtimeLanguageContext.getSingleLineCommentPrefix(referenceLanguage) + if (commentPrefix === '') { + return cellText + } + return cell.document + .getText() + .split('\n') + .map((line) => `${commentPrefix}${line}`) + .join('\n') + } + return cellText +} + +export function getNotebookCellsSliceContext( + cells: vscode.NotebookCell[], + maxLength: number, + 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() + } + cells.some((cell) => { + const cellText = addNewlineIfMissing(getNotebookCellContext(cell, referenceLanguage)) + if (cellText.length > 0) { + if (cellText.length >= maxLength) { + 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('') +} + +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 +154,19 @@ export function extractContextForCodeWhisperer(editor: vscode.TextEditor): codew languageName = runtimeLanguageContext.normalizeLanguage(editor.document.languageId) ?? editor.document.languageId } + if (editor.document.uri.scheme === 'vscode-notebook-cell') { + const notebook = getEnclosingNotebook(editor) + if (notebook) { + ;({ caretLeftFileContext, caretRightFileContext } = getNotebookContext( + notebook, + editor, + languageName, + caretLeftFileContext, + caretRightFileContext + )) + } + } + return { filename: getFileRelativePath(editor), programmingLanguage: { 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