Enhance context gathering in Jupyter notebooks by including text from…

… all cells

Signed-off-by: Tal Wertheimer <[email protected]>

src/runCompletion.ts

@@ -1,4 +1,5 @@
-import { CancellationToken, Position, Range, TextDocument } from "vscode";
+import { CancellationToken, Position, Range, TextDocument} from "vscode";
+import * as vscode from "vscode"; 
 import {
   autocomplete,
   AutocompleteParams,
@@ -31,12 +32,16 @@ export default async function runCompletion({
     timeout?: number;
   };
 }): Promise<AutocompleteResult | null | undefined> {
+
+
+
   const offset = document.offsetAt(position);
   const beforeStartOffset = Math.max(0, offset - CHAR_LIMIT);
   const afterEndOffset = offset + CHAR_LIMIT;
   const beforeStart = document.positionAt(beforeStartOffset);
   const afterEnd = document.positionAt(afterEndOffset);
-  const requestData = {
+
+  let requestData = {
     filename: getFileNameWithExtension(document),
     before:
       document.getText(new Range(beforeStart, position)) +
@@ -52,8 +57,27 @@ export default async function runCompletion({
     sdk_path: getSDKPath(document.languageId),
   };
 
+  // If the 'document' represents a cell within a Jupyter notebook, it is important to 
+  // gather text from all cells to establish the full context. This approach ensures 
+  // comprehensive understanding, not limited to the content of only the current cell.
+  const notebookEditor = vscode.window.activeNotebookEditor;
+  if (notebookEditor) {
+    const cells = notebookEditor.notebook.getCells();
+
+      const index = cells.findIndex((cell) => cell.document.uri.toString() === document.uri.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");
+
+
+      requestData.before = [before, requestData.before].join("\n");
+      requestData.after = [requestData.after + after].join("\n");
+
+  }
+
+
   const isEmptyLine = document.lineAt(position.line).text.trim().length === 0;
 
+
   const result = await autocomplete(
     requestData,
     isEmptyLine ? INLINE_REQUEST_TIMEOUT : timeout