add option to disable markup and cell outputs

zhutony · Feb 7, 2024 · 539db98 · 539db98
1 parent c36cdff
commit 539db98
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 7 deletions.
diff --git a/package.json b/package.json
@@ -36,6 +36,21 @@
       {
         "title": "Llama coder",
         "properties": {
+          "notebook.includeMarkup": {
+            "type": "boolean",
+            "default": true,
+            "description": "Include markup cell types in prompt"
+          },
+          "notebook.includeCellOutputs": {
+            "type": "boolean",
+            "default": false,
+            "description": "Include Cell previous output results in the prompt"
+          },
+          "notebook.cellOutputLimit": {
+            "type": "number",
+            "default": 256,
+            "description": "truncate cell output result if exceeds this limit"
+          },
           "inference.endpoint": {
             "type": "string",
             "default": "",
@@ -121,7 +136,7 @@
     "pretest": "yarn run compile && yarn run lint",
     "lint": "eslint src --ext ts",
     "test": "jest",
-    "package": "vsce package"
+    "package": "npx @vscode/vsce package"
   },
   "devDependencies": {
     "@types/jest": "^29.5.10",

diff --git a/src/config.ts b/src/config.ts
@@ -45,6 +45,20 @@ class Config {
         };
     }
 
+    // Notebook
+    get notebook() {
+        let config = vscode.workspace.getConfiguration('notebook');
+
+        let includeMarkup = config.get('includeMarkup') as boolean;
+        let includeCellOutputs = config.get('includeCellOutputs') as boolean;
+        let cellOutputLimit = config.get('cellOutputLimit') as number;
+        return {
+            includeMarkup,
+            includeCellOutputs,
+            cellOutputLimit,
+        };
+    }
+
     get #config() {
         return vscode.workspace.getConfiguration('inference');
     };

diff --git a/src/prompts/preparePrompt.ts b/src/prompts/preparePrompt.ts
@@ -2,6 +2,7 @@ import vscode from 'vscode';
 import { detectLanguage } from './processors/detectLanguage';
 import { fileHeaders } from './processors/fileHeaders';
 import { languages } from './processors/languages';
+import { config } from '../config';
 
 var decoder = new TextDecoder("utf8");
 
@@ -13,12 +14,13 @@ function getNotebookDocument(document: vscode.TextDocument): vscode.NotebookDocu
 export async function preparePrompt(document: vscode.TextDocument, position: vscode.Position, context: vscode.InlineCompletionContext) {
 
     // Load document text
-    console.log(document);
     let text = document.getText();
     let offset = document.offsetAt(position);
     let prefix = text.slice(0, offset);
     let suffix: string = text.slice(offset);
 
+    let notebookConfig = config.notebook;
+
     // If this is a notebook, add the surrounding cells to the prefix and suffix
     let notebookDocument = getNotebookDocument(document);
     let language = detectLanguage(document.uri.fsPath, document.languageId);
@@ -43,21 +45,26 @@ export async function preparePrompt(document: vscode.TextDocument, position: vsc
 
             // add the markdown cell output to the prompt as a comment
             if (cell.kind === vscode.NotebookCellKind.Markup && commentStart) {
-                for (const line of cell.document.getText().split('\n')) {
-                    out += `\n${commentStart}${line}`;
+                if (notebookConfig.includeMarkup) {
+                    for (const line of cell.document.getText().split('\n')) {
+                        out += `\n${commentStart}${line}`;
+                    }
                 }
             } else {
                 out += cell.document.getText();
             }
 
             // if there is any outputs add them to the prompt as a comment
-            if (cell.kind === vscode.NotebookCellKind.Code && commentStart) {
-                console.log(cell.outputs);
+            const addCellOutputs = notebookConfig.includeCellOutputs
+                                    && beforeCurrentCell
+                                    && cell.kind === vscode.NotebookCellKind.Code
+                                    && commentStart;
+            if (addCellOutputs) {
                 let cellOutputs = cell.outputs
                     .map(x => x.items
                                 .filter(x => x.mime === 'text/plain')
                                 .map(x => decoder.decode(x.data))
-                                .map(x => x.slice(0, 256).split('\n'))) // limit to 256 characters
+                                .map(x => x.slice(0, notebookConfig.cellOutputLimit).split('\n')))
                     .flat(3);
 
                 if (cellOutputs.length > 0) {