fix(amazonq): Previous and subsequent cells are used as context for completion in a Jupyter notebook

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"
+}

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()

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'