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

Author: brdskggs

packages/amazonq/test/unit/codewhisperer/util/editorContext.test.ts

@@ -214,127 +214,122 @@ 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'),
                     vscode.NotebookCellKind.Code
                 ),
                 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')
         })
     })

packages/core/src/codewhisperer/util/editorContext.ts

@@ -51,45 +51,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('')
 }
@@ -136,21 +124,23 @@ 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
                 }
             }
             // 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