fix(vscode): get text from other cells of current notebook as prompt. (…

…TabbyML#904)
ben1009 · Nov 27, 2023 · 7d8bc3d · 7d8bc3d
1 parent 0d79f9f
commit 7d8bc3d
Showing 1 changed file with 54 additions and 3 deletions.
diff --git a/clients/vscode/src/TabbyCompletionProvider.ts b/clients/vscode/src/TabbyCompletionProvider.ts
@@ -7,7 +7,10 @@ import {
   Position,
   Range,
   TextDocument,
+  NotebookDocument,
+  NotebookRange,
   env,
+  window,
   workspace,
 } from "vscode";
 import { EventEmitter } from "events";
@@ -62,11 +65,13 @@ export class TabbyCompletionProvider extends EventEmitter implements InlineCompl
       return null;
     }
 
+    const additionalContext = this.buildAdditionalContext(document);
+
     const request: CompletionRequest = {
       filepath: document.uri.fsPath,
       language: document.languageId, // https://code.visualstudio.com/docs/languages/identifiers
-      text: document.getText(),
-      position: document.offsetAt(position),
+      text: additionalContext.prefix + document.getText() + additionalContext.suffix,
+      position: additionalContext.prefix.length + document.offsetAt(position),
       clipboard: await env.clipboard.readText(),
       manually: context.triggerKind === InlineCompletionTriggerKind.Invoke,
     };
@@ -102,7 +107,10 @@ export class TabbyCompletionProvider extends EventEmitter implements InlineCompl
         return [
           new InlineCompletionItem(
             choice.text,
-            new Range(document.positionAt(choice.replaceRange.start), document.positionAt(choice.replaceRange.end)),
+            new Range(
+              document.positionAt(choice.replaceRange.start - additionalContext.prefix.length),
+              document.positionAt(choice.replaceRange.end - additionalContext.prefix.length),
+            ),
             {
               title: "",
               command: "tabby.applyCallback",
@@ -167,4 +175,47 @@ export class TabbyCompletionProvider extends EventEmitter implements InlineCompl
       this.emit("triggerModeUpdated");
     }
   }
+
+  private buildAdditionalContext(document: TextDocument): { prefix: string; suffix: string } {
+    if (
+      document.uri.scheme === "vscode-notebook-cell" &&
+      window.activeNotebookEditor?.notebook.uri.path === document.uri.path
+    ) {
+      // Add all the cells in the notebook as context
+      const notebook = window.activeNotebookEditor.notebook;
+      const current = window.activeNotebookEditor.selection.start;
+      const prefix = this.buildNotebookContext(notebook, new NotebookRange(0, current), document.languageId);
+      const suffix = this.buildNotebookContext(
+        notebook,
+        new NotebookRange(current + 1, notebook.cellCount),
+        document.languageId,
+      );
+      return { prefix, suffix };
+    }
+    return { prefix: "", suffix: "" };
+  }
+
+  private notebookLanguageComments: { [languageId: string]: (code: string) => string } = {
+    markdown: (code) => "```\n" + code + "\n```",
+    python: (code) =>
+      code
+        .split("\n")
+        .map((l) => "# " + l)
+        .join("\n"),
+  };
+
+  private buildNotebookContext(notebook: NotebookDocument, range: NotebookRange, languageId: string): string {
+    return notebook
+      .getCells(range)
+      .map((cell) => {
+        if (cell.document.languageId === languageId) {
+          return cell.document.getText() + "\n\n";
+        } else if (Object.keys(this.notebookLanguageComments).includes(languageId)) {
+          return this.notebookLanguageComments[languageId](cell.document.getText()) + "\n\n";
+        } else {
+          return "";
+        }
+      })
+      .join("");
+  }
 }