From ed1e8f7240a28a7ec8d971d3970768e1d98b232d Mon Sep 17 00:00:00 2001 From: Tal Wertheimer Date: Wed, 13 Dec 2023 13:28:25 +0200 Subject: [PATCH 1/2] Enhance context gathering in Jupyter notebooks by including text from all cells Signed-off-by: Tal Wertheimer --- src/runCompletion.ts | 35 ++++++++++++++++++++++ src/vscode.proposed.inlineCompletions.d.ts | 10 +++++++ 2 files changed, 45 insertions(+) diff --git a/src/runCompletion.ts b/src/runCompletion.ts index 5fd329a877..e956037711 100644 --- a/src/runCompletion.ts +++ b/src/runCompletion.ts @@ -1,4 +1,5 @@ import { CancellationToken, Position, Range, TextDocument } from "vscode"; +import * as vscode from "vscode"; import { autocomplete, AutocompleteParams, @@ -14,6 +15,29 @@ import { import languages from "./globals/languages"; import { getSDKPath } from "./languages"; +function calculateContextForJupyterNotebook( + notebookEditor: vscode.window.NotebookEditor, + documentUri: vscode.Uri +): { before: string; after: string } { + const cells = notebookEditor.notebook.getCells(); + + const index = cells.findIndex( + (cell) => cell.document.uri.toString() === documentUri.toString() + ); + const before = cells + .slice(0, index) + .map((cell) => cell.document.getText()) + .join("\n"); + const after = cells + .slice(index + 1) + .map((cell) => cell.document.getText()) + .join("\n"); + return { + before, + after, + }; +} + export default async function runCompletion({ document, position, @@ -36,6 +60,7 @@ export default async function runCompletion({ const afterEndOffset = offset + CHAR_LIMIT; const beforeStart = document.positionAt(beforeStartOffset); const afterEnd = document.positionAt(afterEndOffset); + const requestData = { filename: getFileNameWithExtension(document), before: @@ -51,6 +76,16 @@ export default async function runCompletion({ indentation_size: getTabSize(), sdk_path: getSDKPath(document.languageId), }; + const notebookEditor = vscode.window.activeNotebookEditor; + + if (notebookEditor) { + const { before, after } = calculateContextForJupyterNotebook( + notebookEditor, + document.uri + ); + requestData.before = [before, requestData.before].join("\n"); + requestData.after = [requestData.after, after].join("\n"); + } const isEmptyLine = document.lineAt(position.line).text.trim().length === 0; diff --git a/src/vscode.proposed.inlineCompletions.d.ts b/src/vscode.proposed.inlineCompletions.d.ts index 438de75446..50439b5702 100644 --- a/src/vscode.proposed.inlineCompletions.d.ts +++ b/src/vscode.proposed.inlineCompletions.d.ts @@ -116,6 +116,16 @@ declare module "vscode" { * Be aware that this API will not ever be finalized. */ export namespace window { + export const activeNotebookEditor: NotebookEditor | undefined; + + class NotebookEditor { + notebook: { + getCells: () => { + document: TextDocument; + }[]; + }; + } + export function getInlineCompletionItemController< T extends InlineCompletionItem >(provider: InlineCompletionItemProvider): InlineCompletionController; From 928d9dd27efcb109800f41305e284ea0d590738e Mon Sep 17 00:00:00 2001 From: Tal Wertheimer Date: Thu, 14 Dec 2023 11:00:40 +0200 Subject: [PATCH 2/2] fix Signed-off-by: Tal Wertheimer --- src/runCompletion.ts | 37 ++++++++++++++++++++----------------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/src/runCompletion.ts b/src/runCompletion.ts index e956037711..b173d8e2da 100644 --- a/src/runCompletion.ts +++ b/src/runCompletion.ts @@ -1,5 +1,11 @@ -import { CancellationToken, Position, Range, TextDocument } from "vscode"; -import * as vscode from "vscode"; +import { + CancellationToken, + Position, + Range, + TextDocument, + window, + Uri, +} from "vscode"; import { autocomplete, AutocompleteParams, @@ -16,8 +22,8 @@ import languages from "./globals/languages"; import { getSDKPath } from "./languages"; function calculateContextForJupyterNotebook( - notebookEditor: vscode.window.NotebookEditor, - documentUri: vscode.Uri + notebookEditor: window.NotebookEditor, + documentUri: Uri ): { before: string; after: string } { const cells = notebookEditor.notebook.getCells(); @@ -33,8 +39,8 @@ function calculateContextForJupyterNotebook( .map((cell) => cell.document.getText()) .join("\n"); return { - before, - after, + before: before + "\n", + after: "\n" + after, }; } @@ -61,12 +67,19 @@ export default async function runCompletion({ const beforeStart = document.positionAt(beforeStartOffset); const afterEnd = document.positionAt(afterEndOffset); + const notebookEditor = window.activeNotebookEditor; + + const { before, after } = notebookEditor + ? calculateContextForJupyterNotebook(notebookEditor, document.uri) + : { before: "", after: "" }; + const requestData = { filename: getFileNameWithExtension(document), before: + before + document.getText(new Range(beforeStart, position)) + currentSuggestionText, - after: document.getText(new Range(position, afterEnd)), + after: document.getText(new Range(position, afterEnd)) + after, region_includes_beginning: beforeStartOffset === 0, region_includes_end: document.offsetAt(afterEnd) !== afterEndOffset, max_num_results: getMaxResults(), @@ -76,16 +89,6 @@ export default async function runCompletion({ indentation_size: getTabSize(), sdk_path: getSDKPath(document.languageId), }; - const notebookEditor = vscode.window.activeNotebookEditor; - - if (notebookEditor) { - const { before, after } = calculateContextForJupyterNotebook( - notebookEditor, - document.uri - ); - requestData.before = [before, requestData.before].join("\n"); - requestData.after = [requestData.after, after].join("\n"); - } const isEmptyLine = document.lineAt(position.line).text.trim().length === 0;