From ed1e8f7240a28a7ec8d971d3970768e1d98b232d Mon Sep 17 00:00:00 2001
From: Tal Wertheimer <tal@Tals-MacBook-Pro-2.local>
Date: Wed, 13 Dec 2023 13:28:25 +0200
Subject: [PATCH] Enhance context gathering in Jupyter notebooks by including
 text from all cells

Signed-off-by: Tal Wertheimer <tal@Tals-MacBook-Pro-2.local>
---
 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<T>): InlineCompletionController<T>;