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

brdskggs · brdskggs · commit 5f29d4a796c8 · 2025-04-18T09:25:12.000-04:00
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
@@ -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
@@ -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
@@ -22,19 +22,95 @@ import { AuthUtil } from './authUtil'
 
 let tabSize: number = getTabSizeSetting()
 
+const languageCommentChars: Record<string, string> = {
+    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)
@@ -45,6 +121,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: {
