DEV2-564 Enhance context gathering in Jupyter notebooks by including text from…

src/runCompletion.ts

@@ -1,4 +1,11 @@
-import { CancellationToken, Position, Range, TextDocument } from "vscode";
+import {
+  CancellationToken,
+  Position,
+  Range,
+  TextDocument,
+  window,
+  Uri,
+} from "vscode";
 import {
   autocomplete,
   AutocompleteParams,
@@ -14,6 +21,29 @@
 import languages from "./globals/languages";
 import { getSDKPath } from "./languages";
 
+function calculateContextForJupyterNotebook(
+  notebookEditor: window.NotebookEditor,
+  documentUri: 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: before + "\n",
+    after: "\n" + after,
+  };
+}
+
 export default async function runCompletion({
   document,
   position,
@@ -36,12 +66,20 @@
   const afterEndOffset = offset + CHAR_LIMIT;
   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(),

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<T>): InlineCompletionController<T>;