From b00c82b1e2df145c578b2bfb5ab38c31cf651ba5 Mon Sep 17 00:00:00 2001 From: Nathan Hedge <23344786+10Nates@users.noreply.github.com> Date: Sat, 27 Jan 2024 21:22:14 -0600 Subject: [PATCH] More reliable user input detect + undo --- package.json | 2 +- src/extension.ts | 161 +++++++++++++++++++++++++++-------------------- 2 files changed, 94 insertions(+), 69 deletions(-) diff --git a/package.json b/package.json index 97da3e7..d921e63 100644 --- a/package.json +++ b/package.json @@ -2,7 +2,7 @@ "name": "ollama-autocoder", "displayName": "Ollama Autocoder", "description": "A simple to use Ollama autocompletion engine with options exposed and streaming functionality", - "version": "0.0.6", + "version": "0.0.7", "icon": "icon.png", "publisher": "10nates", "license": "MIT", diff --git a/src/extension.ts b/src/extension.ts index ff59419..03f7734 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -38,6 +38,18 @@ updateVSConfig(); // No need for restart for any of these settings vscode.workspace.onDidChangeConfiguration(updateVSConfig); +// internal function to handle undo functionality +async function handleUndo(document: vscode.TextDocument, startingPosition: vscode.Position, endingPosition: vscode.Position, optons?: vscode.MessageOptions) { + vscode.window.showInformationMessage("Ollama autocompletion was cancelled. Would you like it to undo its changes?", "Keep", "Undo").then(r => { + if (r === "Undo") { + const edit = new vscode.WorkspaceEdit(); + const range = new vscode.Range(startingPosition, endingPosition); + edit.delete(document.uri, range); + vscode.workspace.applyEdit(edit); + } + }); +} + // internal function for autocomplete, not directly exposed async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationToken?: vscode.CancellationToken) { const document = textEditor.document; @@ -58,10 +70,16 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo try { progress.report({ message: "Starting model..." }); + //tracker + // let oldPosition = position; + let currentPosition = position; + // let lastToken = ""; + let axiosCancelPost: () => void; const axiosCancelToken = new axios.CancelToken((c) => { const cancelPost = function () { c("Autocompletion request terminated by user cancel"); + handleUndo(document, position, currentPosition); }; axiosCancelPost = cancelPost; if (cancellationToken) cancellationToken.onCancellationRequested(cancelPost); @@ -85,27 +103,31 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo } ); - //tracker - let oldPosition = position; - let currentPosition = position; - let lastToken = ""; - - response.data.on('data', async (d: Uint8Array) => { progress.report({ message: "Generating..." }); // Check for user input (cancel) - if (lastToken != "") { - const lastInput = document.getText(new vscode.Range(oldPosition, textEditor.selection.active)); - if (lastInput !== lastToken) { - axiosCancelPost(); // cancel axios => cancel finished promise => close notification - return; - } + if (currentPosition != textEditor.selection.start) { + console.log(currentPosition, textEditor.selection.start); + axiosCancelPost(); // cancel axios => cancel finished promise => close notification + return; } + // if (lastToken != "") { + // const lastInput = document.getText(new vscode.Range(oldPosition, textEditor.selection.active)); + // if (lastInput !== lastToken) { + // console.log(lastInput, lastToken); + // axiosCancelPost(); // cancel axios => cancel finished promise => close notification + // return; + // } + // } // Get a completion from the response const completion: string = JSON.parse(d.toString()).response; - lastToken = completion; + // lastToken = completion; + + if (completion === "") { + return; + } //complete edit for token const edit = new vscode.WorkspaceEdit(); @@ -126,7 +148,7 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo newPosition, newPosition ); - oldPosition = currentPosition; + // oldPosition = currentPosition; currentPosition = newPosition; // completion bar @@ -161,64 +183,67 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo ); } +// Completion item provider callback for activate +async function provideCompletionItems(document: vscode.TextDocument, position: vscode.Position, cancellationToken: vscode.CancellationToken) { + + // Create a completion item + const item = new vscode.CompletionItem("Autocomplete with Ollama"); + + // Set the insert text to a placeholder + item.insertText = new vscode.SnippetString('${1:}'); + + // Wait before initializing Ollama to reduce compute usage + if (responsePreview) await new Promise(resolve => setTimeout(resolve, responsePreviewDelay * 1000)); + if (cancellationToken.isCancellationRequested) { + return [ item ]; + } + + // Set the label & inset text to a shortened, non-stream response + if (responsePreview) { + let prompt = document.getText(new vscode.Range(document.lineAt(0).range.start, position)); + prompt = prompt.substring(Math.max(0, prompt.length - promptWindowSize), prompt.length); + const response_preview = await axios.post(apiEndpoint, { + model: apiModel, // Change this to the model you want to use + prompt: prompt, + stream: false, + system: apiSystemMessage, + raw: rawInput, + options: { + num_predict: responsePreviewMaxTokens, // reduced compute max + stop: ['\n'] + } + }, { + cancelToken: new axios.CancelToken((c) => { + const cancelPost = function () { + c("Autocompletion request terminated by completion cancel"); + }; + cancellationToken.onCancellationRequested(cancelPost); + }) + }); + + if (response_preview.data.response.trim() != "") { // default if empty + item.label = response_preview.data.response.trimStart(); // tended to add whitespace at the beginning + item.insertText = response_preview.data.response.trimStart(); + } + } + + // Set the documentation to a message + item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Ollama'); + // Set the command to trigger the completion + if (continueInline || !responsePreview) item.command = { + command: 'ollama-autocoder.autocomplete', + title: 'Autocomplete with Ollama', + arguments: [cancellationToken] + }; + // Return the completion item + return [item]; +} + // This method is called when extension is activated function activate(context: vscode.ExtensionContext) { // Register a completion provider for JavaScript files const completionProvider = vscode.languages.registerCompletionItemProvider("*", { - async provideCompletionItems(document, position, cancellationToken) { - - // Create a completion item - const item = new vscode.CompletionItem("Autocomplete with Ollama"); - - // Set the insert text to a placeholder - item.insertText = new vscode.SnippetString('${1:}'); - - // Wait before initializing Ollama to reduce compute usage - if (responsePreview) await new Promise(resolve => setTimeout(resolve, responsePreviewDelay * 1000)); - if (cancellationToken.isCancellationRequested) { - return [ item ]; - } - - // Set the label & inset text to a shortened, non-stream response - if (responsePreview) { - let prompt = document.getText(new vscode.Range(document.lineAt(0).range.start, position)); - prompt = prompt.substring(Math.max(0, prompt.length - promptWindowSize), prompt.length); - const response_preview = await axios.post(apiEndpoint, { - model: apiModel, // Change this to the model you want to use - prompt: prompt, - stream: false, - system: apiSystemMessage, - raw: rawInput, - options: { - num_predict: responsePreviewMaxTokens, // reduced compute max - stop: ['\n'] - } - }, { - cancelToken: new axios.CancelToken((c) => { - const cancelPost = function () { - c("Autocompletion request terminated by completion cancel"); - }; - cancellationToken.onCancellationRequested(cancelPost); - }) - }); - - if (response_preview.data.response.trim() != "") { // default if empty - item.label = response_preview.data.response.trimStart(); // tended to add whitespace at the beginning - item.insertText = response_preview.data.response.trimStart(); - } - } - - // Set the documentation to a message - item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Ollama'); - // Set the command to trigger the completion - if (continueInline || !responsePreview) item.command = { - command: 'ollama-autocoder.autocomplete', - title: 'Autocomplete with Ollama', - arguments: [cancellationToken] - }; - // Return the completion item - return [item]; - }, + provideCompletionItems }, ...completionKeys.split("") );