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

Author: brdskggs

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/core/src/codewhisperer/util/editorContext.ts

@@ -22,19 +22,87 @@ import { AuthUtil } from './authUtil'
 
 let tabSize: number = getTabSizeSetting()
 
+const languageCommentChars: Record<string, string> = {
+    python: '# ',
+    java: '// ',
+}
+
+export function extractSingleCellContext(cell: vscode.NotebookCell, commentPrefix: string): string {
+    const cellText = cell.document.getText()
+    if (cell.kind === vscode.NotebookCellKind.Markup) {
+        // Converting markdown to a comment block is language dependent.
+        // If a commentPrefix is specified, add it to each line of markdown.
+        // Otherwise, just return the text as is.
+        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,
+    commentPrefix: string
+): string {
+    const output: string[] = []
+    for (let i = cells.length - 1; i >= 0; i--) {
+        let cellText = extractSingleCellContext(cells[i], commentPrefix)
+        if (cellText.length > 0) {
+            if (!cellText.endsWith('\n')) {
+                cellText += '\n'
+            }
+            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,
+    commentPrefix: string
+): string {
+    const output: string[] = []
+    for (let i = 0; i < cells.length; i++) {
+        let cellText = extractSingleCellContext(cells[i], commentPrefix)
+        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 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 +113,45 @@ 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)
+
+            // Add appropriate comment string at beginning of each line if language is recognized
+            const commentPrefix = languageCommentChars[languageName] ?? ''
+
+            // 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),
+                    commentPrefix
+                )
+                if (leftCellsText.length > 0) {
+                    caretLeftFileContext = leftCellsText + '\n' + 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),
+                    commentPrefix
+                )
+                if (rightCellsText.length > 0) {
+                    caretRightFileContext = caretRightFileContext + '\n' + rightCellsText
+                }
+            }
+        }
+    }
+
     return {
         filename: getFileRelativePath(editor),
         programmingLanguage: {