fix(amazonq): support inline completion in notebook

package.json

@@ -74,6 +74,7 @@
         "webpack-merge": "^5.10.0"
     },
     "dependencies": {
+        "@aws/language-server-runtimes": "^0.2.125",
         "@types/node": "^22.7.5",
         "jaro-winkler": "^0.2.8",
         "vscode-nls": "^5.2.0",

packages/amazonq/.changes/next-release/Bug Fix-f80b74d6-d598-4d64-ae80-a7da47fecc32.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Enable inline completion in Jupyter Notebook"
+}

packages/amazonq/src/app/inline/notebookUtil.ts

@@ -0,0 +1,98 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+
+import { CodeWhispererConstants, runtimeLanguageContext } from 'aws-core-vscode/codewhisperer'
+import { InlineCompletionWithReferencesParams } from '@aws/language-server-runtimes/server-interface'
+
+function getEnclosingNotebook(document: vscode.TextDocument): vscode.NotebookDocument | undefined {
+    // For notebook cells, find the existing notebook with a cell that matches the current document.
+    return vscode.workspace.notebookDocuments.find(
+        (nb) => nb.notebookType === 'jupyter-notebook' && nb.getCells().some((cell) => cell.document === document)
+    )
+}
+
+export function getNotebookContext(
+    notebook: vscode.NotebookDocument,
+    document: vscode.TextDocument,
+    position: vscode.Position
+) {
+    // 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 === document)
+    let caretLeftFileContext = ''
+    let caretRightFileContext = ''
+
+    if (cellIndex >= 0 && cellIndex < allCells.length) {
+        // Add content from previous cells
+        for (let i = 0; i < cellIndex; i++) {
+            caretLeftFileContext += convertCellContent(allCells[i]) + '\n'
+        }
+
+        // Add content from current cell up to cursor
+        caretLeftFileContext += allCells[cellIndex].document.getText(
+            new vscode.Range(new vscode.Position(0, 0), position)
+        )
+
+        // Add content from cursor to end of current cell
+        caretRightFileContext = allCells[cellIndex].document.getText(
+            new vscode.Range(
+                position,
+                allCells[cellIndex].document.positionAt(allCells[cellIndex].document.getText().length)
+            )
+        )
+
+        // Add content from following cells
+        for (let i = cellIndex + 1; i < allCells.length; i++) {
+            caretRightFileContext += '\n' + convertCellContent(allCells[i])
+        }
+    }
+    caretLeftFileContext = caretLeftFileContext.slice(-CodeWhispererConstants.charactersLimit)
+    caretRightFileContext = caretRightFileContext.slice(0, CodeWhispererConstants.charactersLimit)
+    return { caretLeftFileContext, caretRightFileContext }
+}
+
+// Convert the markup cells into code with comments
+export function convertCellContent(cell: vscode.NotebookCell) {
+    const cellText = cell.document.getText()
+    if (cell.kind === vscode.NotebookCellKind.Markup) {
+        const commentPrefix = runtimeLanguageContext.getSingleLineCommentPrefix(
+            runtimeLanguageContext.normalizeLanguage(cell.document.languageId) ?? cell.document.languageId
+        )
+        if (commentPrefix === '') {
+            return cellText
+        }
+        return cell.document
+            .getText()
+            .split('\n')
+            .map((line) => `${commentPrefix}${line}`)
+            .join('\n')
+    }
+    return cellText
+}
+
+export function extractFileContextInNotebooks(
+    document: vscode.TextDocument,
+    position: vscode.Position
+): InlineCompletionWithReferencesParams['fileContextOverride'] | undefined {
+    let caretLeftFileContext = ''
+    let caretRightFileContext = ''
+    const languageName = runtimeLanguageContext.normalizeLanguage(document.languageId) ?? document.languageId
+    if (document.uri.scheme === 'vscode-notebook-cell') {
+        const notebook = getEnclosingNotebook(document)
+        if (notebook) {
+            ;({ caretLeftFileContext, caretRightFileContext } = getNotebookContext(notebook, document, position))
+            return {
+                leftFileContent: caretLeftFileContext,
+                rightFileContent: caretRightFileContext,
+                filename: document.fileName,
+                fileUri: document.uri.toString(),
+                programmingLanguage: languageName,
+            }
+        }
+    }
+    return undefined
+}

packages/amazonq/src/app/inline/recommendationService.ts

@@ -25,6 +25,7 @@ import { DocumentEventListener } from './documentEventListener'
 import { getOpenFilesInWindow } from 'aws-core-vscode/utils'
 import { asyncCallWithTimeout } from '../../util/timeoutUtil'
 import { EditSuggestionState } from './editSuggestionState'
+import { extractFileContextInNotebooks } from './notebookUtil'
 
 export interface GetAllRecommendationsOptions {
     emitTelemetry?: boolean
@@ -97,6 +98,9 @@ export class RecommendationService {
         if (options.editsStreakToken) {
             request = { ...request, partialResultToken: options.editsStreakToken }
         }
+        if (document.uri.scheme === 'vscode-notebook-cell') {
+            request.fileContextOverride = extractFileContextInNotebooks(document, position)
+        }
         const requestStartTime = performance.now()
         const statusBar = CodeWhispererStatusBarManager.instance
 

packages/amazonq/test/unit/app/inline/notebookUtil.test.ts

@@ -0,0 +1,87 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import * as assert from 'assert'
+import { createMockDocument } from 'aws-core-vscode/test'
+import { convertCellContent, getNotebookContext } from '../../../../src/app/inline/notebookUtil'
+import { CodeWhispererConstants } from 'aws-core-vscode/codewhisperer'
+
+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('Notebook Util', function () {
+    describe('convertCellContent', function () {
+        it('should return code cell content as-is', function () {
+            const codeCell = createNotebookCell(
+                createMockDocument('def example():\n    return "test"'),
+                vscode.NotebookCellKind.Code
+            )
+            const result = convertCellContent(codeCell)
+            assert.strictEqual(result, 'def example():\n    return "test"')
+        })
+
+        it('should convert markdown cell content to comments for Python', function () {
+            const markdownCell = createNotebookCell(
+                createMockDocument('# Heading\nSome text'),
+                vscode.NotebookCellKind.Markup
+            )
+            const result = convertCellContent(markdownCell)
+            assert.strictEqual(result, '# # Heading\n# Some text')
+        })
+    })
+
+    describe('getNotebookContext', function () {
+        it('should combine context from multiple cells', function () {
+            const currentDoc = createMockDocument('cell2 content', 'b.ipynb')
+            const notebook = {
+                getCells: () => [
+                    createNotebookCell(createMockDocument('cell1 content', 'a.ipynb'), vscode.NotebookCellKind.Code),
+                    createNotebookCell(currentDoc, vscode.NotebookCellKind.Code),
+                    createNotebookCell(createMockDocument('cell3 content', 'c.ipynb'), vscode.NotebookCellKind.Code),
+                ],
+            } as vscode.NotebookDocument
+
+            const position = new vscode.Position(0, 5)
+
+            const { caretLeftFileContext, caretRightFileContext } = getNotebookContext(notebook, currentDoc, position)
+
+            assert.strictEqual(caretLeftFileContext, 'cell1 content\ncell2')
+            assert.strictEqual(caretRightFileContext, ' content\ncell3 content')
+        })
+
+        it('should respect character limits', function () {
+            const longContent = 'a'.repeat(10000)
+            const notebook = {
+                getCells: () => [createNotebookCell(createMockDocument(longContent), vscode.NotebookCellKind.Code)],
+            } as vscode.NotebookDocument
+
+            const currentDoc = createMockDocument(longContent)
+            const position = new vscode.Position(0, 5000)
+
+            const { caretLeftFileContext, caretRightFileContext } = getNotebookContext(notebook, currentDoc, position)
+
+            assert.ok(caretLeftFileContext.length <= CodeWhispererConstants.charactersLimit)
+            assert.ok(caretRightFileContext.length <= CodeWhispererConstants.charactersLimit)
+        })
+    })
+})

packages/core/src/shared/utilities/index.ts

@@ -9,3 +9,4 @@ export * from './functionUtils'
 export * as messageUtils from './messages'
 export * as CommentUtils from './commentUtils'
 export * from './editorUtilities'
+export * from './tsUtils'