Merge commit '1e2431af2477f14b924f13c123f657aa7402f33d'

Author: ex3ndr

package.json

@@ -37,6 +37,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": "",
@@ -114,6 +129,14 @@
             "default": 256,
             "description": "Max number of new tokens to be generated.",
             "order": 7
+          },
+          "inference.delay": {
+            "type": "number",
+            "default": 250,
+            "description": "Completion request delay in milliseconds (0 - no delay, -1 - no completions).",
+            "order": 8,
+            "minimum": -1,
+            "maximum": 5000
           }
         }
       }
@@ -126,7 +149,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",

src/config.ts

@@ -35,13 +35,30 @@ class Config {
             }
         }
 
+        let delay = config.get('delay') as number;
+
         return {
             endpoint,
             maxLines,
             maxTokens,
             temperature,
             modelName,
-            modelFormat
+            modelFormat,
+            delay
+        };
+    }
+
+    // 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,
         };
     }
 

src/prompts/filter.ts

@@ -1,7 +1,7 @@
 import type vscode from 'vscode';
 
 export function isSupported(doc: vscode.TextDocument) {
-    return doc.uri.scheme === 'file';
+    return doc.uri.scheme === 'file' || doc.uri.scheme === 'vscode-notebook-cell';
 }
 
 export function isNotNeeded(doc: vscode.TextDocument, position: vscode.Position, context: vscode.InlineCompletionContext): boolean {

src/prompts/preparePrompt.ts

@@ -2,6 +2,14 @@ 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");
+
+function getNotebookDocument(document: vscode.TextDocument): vscode.NotebookDocument | undefined  {
+    return  vscode.workspace.notebookDocuments
+        .find(x => x.uri.path === document.uri.path);
+}
 
 export async function preparePrompt(document: vscode.TextDocument, position: vscode.Position, context: vscode.InlineCompletionContext) {
 
@@ -11,6 +19,75 @@ export async function preparePrompt(document: vscode.TextDocument, position: vsc
     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);
+    let commentStart: string | undefined = undefined;
+    if (language) {
+        commentStart = languages[language].comment?.start;
+    }
+
+    if (notebookDocument) {
+        let beforeCurrentCell = true;
+
+        let prefixCells = "";
+        let suffixCells = "";
+
+        notebookDocument.getCells().forEach((cell) => {
+            let out = "";
+
+            if (cell.document.uri.fragment === document.uri.fragment) {
+                beforeCurrentCell = false; // switch to suffix mode
+                return;
+            }
+            
+            // add the markdown cell output to the prompt as a comment
+            if (cell.kind === vscode.NotebookCellKind.Markup && commentStart) {
+                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
+            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, notebookConfig.cellOutputLimit).split('\n')))
+                    .flat(3);
+                
+                if (cellOutputs.length > 0) {
+                    out += `\n${commentStart}Output:`;
+                    for (const line of cellOutputs) {
+                        out += `\n${commentStart}${line}`;
+                    }
+                }
+            }
+
+            // update the prefix/suffix
+            if (beforeCurrentCell) {
+                prefixCells += out;
+            } else {
+                suffixCells += out;
+            }
+
+        });
+
+        prefix = prefixCells + prefix;
+        suffix = suffix + suffixCells;
+    }
+
     // Trim suffix
     // If suffix is too small it is safe to assume that it could be ignored which would allow us to use
     // more powerful completition instead of in middle one
@@ -22,7 +99,6 @@ export async function preparePrompt(document: vscode.TextDocument, position: vsc
     // NOTE: Most networks don't have a concept of filenames and expected language, but we expect that some files in training set has something in title that 
     //       would indicate filename and language
     // NOTE: If we can't detect language, we could ignore this since the number of languages that need detection is limited
-    let language = detectLanguage(document.uri.fsPath, document.languageId);
     if (language) {
         prefix = fileHeaders(prefix, document.uri.fsPath, languages[language]);
     }

src/prompts/processors/languages.ts

@@ -111,7 +111,7 @@ export const languages: { [key in Language]: LanguageDescriptor } = {
     },
     python: {
         name: 'Python',
-        extensions: ['.py'],
+        extensions: ['.py', 'ipynb'],
         comment: { start: '#' }
     },
     c: {

src/prompts/processors/models.ts

@@ -19,7 +19,7 @@ export function adaptPrompt(args: { format: ModelFormat, prefix: string, suffix:
     }
 
     // Stable code FIM
-    if (args.format === 'stable-code')  {
+    if (args.format === 'stable-code') {
         return {
             prompt: `<fim_prefix>${args.prefix}<fim_suffix>${args.suffix}<fim_middle>`,
             stop: [`<|endoftext|>`]
@@ -28,7 +28,7 @@ export function adaptPrompt(args: { format: ModelFormat, prefix: string, suffix:
 
     // Codellama FIM
     return {
-        prompt: `<PRE> ${args.prefix} <SUF>${args.suffix} <MID>`,
-        stop: [`<PRE>`, `<SUF>`, `<MID>`, `<END>`]
+        prompt: `<PRE> ${args.prefix} <SUF> ${args.suffix} <MID>`,
+        stop: [`<END>`, `<EOD>`, `<EOT>`]
     };
 }

src/prompts/provider.ts

@@ -20,7 +20,21 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider {
         this.context = context;
     }
 
+    async delayCompletion(delay: number, token: vscode.CancellationToken): Promise<boolean> {
+        if (config.inference.delay < 0) {
+            return false;
+        }
+        await new Promise(p => setTimeout(p, delay));
+        if (token.isCancellationRequested) {
+            return false;
+        }
+        return true;
+    }
+
     async provideInlineCompletionItems(document: vscode.TextDocument, position: vscode.Position, context: vscode.InlineCompletionContext, token: vscode.CancellationToken): Promise<vscode.InlineCompletionItem[] | vscode.InlineCompletionList | undefined | null> {
+        if (!await this.delayCompletion(config.inference.delay, token)) {
+            return;
+        }
 
         try {
 
@@ -66,15 +80,6 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider {
 
                     // Config
                     let inferenceConfig = config.inference;
-                    // let config = vscode.workspace.getConfiguration('inference');
-                    // let endpoint = config.get('endpoint') as string;
-                    // let model = config.get('model') as string;
-                    // let maxLines = config.get('maxLines') as number;
-                    // let maxTokens = config.get('maxTokens') as number;
-                    // let temperature = config.get('temperature') as number;
-                    // if (endpoint.endsWith('/')) {
-                    //     endpoint = endpoint.slice(0, endpoint.length - 1);
-                    // }
 
                     // Update status
                     this.statusbar.text = `$(sync~spin) Llama Coder`;