updadet code

package.json

@@ -2,7 +2,7 @@
 	"name": "fabelous-autocoder",
 	"displayName": "Fabelous Autocoder",
 	"description": "A simple to use Ollama autocompletion engine with options exposed and streaming functionality",
-	"version": "0.0.4",
+	"version": "0.1.0",
 	"icon": "icon.png",
 	"publisher": "fabel",
 	"license": "CC BY-ND 4.0",

src/extension.ts

@@ -1,5 +1,3 @@
-// Original script was GPT4 but it has been deeply Ship of Theseused. 
-
 import * as vscode from "vscode";
 import axios from "axios";
 
@@ -19,19 +17,18 @@ let continueInline: boolean | undefined;
 let keepAlive: number | undefined;
 let topP: number | undefined;
 
-
 function updateVSConfig() {
     VSConfig = vscode.workspace.getConfiguration("fabelous-autocoder");
     apiEndpoint = VSConfig.get("endpoint") || "http://localhost:11434/api/generate";
     apiAuthentication = VSConfig.get("authentication") || "";
-	apiModel = VSConfig.get("model") || "openhermes2.5-mistral:7b-q4_K_M"; // The model I tested with
+    apiModel = VSConfig.get("model") || "fabelous-coder:latest"; // Updated to use FIM model
     apiMessageHeader = VSConfig.get("message header") || "";
     numPredict = VSConfig.get("max tokens predicted") || 1000;
     promptWindowSize = VSConfig.get("prompt window size") || 2000;
     completionKeys = VSConfig.get("completion keys") || " ";
     responsePreview = VSConfig.get("response preview");
     responsePreviewMaxTokens = VSConfig.get("preview max tokens") || 50;
-	responsePreviewDelay = VSConfig.get("preview delay") || 0; // Must be || 0 instead of || [default] because of truthy
+    responsePreviewDelay = VSConfig.get("preview delay") || 0;
     continueInline = VSConfig.get("continue inline");
     apiTemperature = VSConfig.get("temperature") || 0.7;
     keepAlive = VSConfig.get("keep alive") || 30;
@@ -39,9 +36,8 @@ function updateVSConfig() {
 }
 
 updateVSConfig();
-// No need for restart for any of these settings
 vscode.workspace.onDidChangeConfiguration(updateVSConfig);
-// Give model additional information
+
 function messageHeaderSub(document: vscode.TextDocument) {
     const sub = apiMessageHeader
         .replace("{LANG}", document.languageId)
@@ -49,23 +45,45 @@ function messageHeaderSub(document: vscode.TextDocument) {
         .replace("{PROJECT_NAME}", vscode.workspace.name || "Untitled");
     return sub;
 }
-// internal function for autocomplete, not directly exposed
+
+function getContextLines(document: vscode.TextDocument, position: vscode.Position): string {
+    const lines = [];
+    const lineCount = document.lineCount;
+    
+    // Get more context for FIM
+    const startLine = Math.max(0, position.line - 10);
+    const endLine = Math.min(lineCount - 1, position.line + 10);
+    
+    for (let i = startLine; i <= endLine; i++) {
+        lines.push(document.lineAt(i).text);
+    }
+    
+    return lines.join("\n");
+}
+
 async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationToken?: vscode.CancellationToken) {
     const document = textEditor.document;
     const position = textEditor.selection.active;
 
-    // Get the current prompt
+    // Get the current context
-    let prompt = document.getText(new vscode.Range(document.lineAt(0).range.start, position));
+    const context = getContextLines(document, position);
-    prompt = prompt.substring(Math.max(0, prompt.length - promptWindowSize), prompt.length);
 
-    // Replace {Prompt} with the extracted text from the document
+    // Split the context into prefix and suffix for FIM
-    const sub = messageHeaderSub(document).replace("{Prompt}", prompt);
+    const lines = context.split("\n");
+    const currentLineIndex = position.line - Math.max(0, position.line - 10);
+    const prefix = lines.slice(0, currentLineIndex + 1).join("\n");
+    const suffix = lines.slice(currentLineIndex + 1).join("\n");
+
+    // Create FIM prompt
+    const fimPrompt = `<fim_prefix>${prefix}<fim_suffix>${suffix}<fim_middle>`;
+
+    // Replace {Prompt} with the FIM prompt
+    const sub = messageHeaderSub(document).replace("{PROMPT}", fimPrompt);
 
-    // Show a progress message
     vscode.window.withProgress(
         {
             location: vscode.ProgressLocation.Notification,
-            title: "Ollama Autocoder",
+            title: "Fabelous Autocoder",
             cancellable: true,
         },
         async (progress, progressCancellationToken) => {
@@ -84,16 +102,15 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo
                     vscode.workspace.onDidCloseTextDocument(cancelPost);
                 });
 
-                // Make a request to the ollama.ai REST API
                 const response = await axios.post(apiEndpoint, {
-                    model: apiModel, // Change this to the model you want to use
+                    model: apiModel,
-                    prompt: sub, // Use the modified sub string with the replaced prompt
+                    prompt: sub,
                     stream: true,
                     raw: true,
                     options: {
                         num_predict: numPredict,
                         temperature: apiTemperature,
-                        stop: ["```"]
+                        stop: ["<fim_suffix>", "```"]
                     }
                 }, {
                     cancelToken: axiosCancelToken,
@@ -101,30 +118,28 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo
                     headers: {
                         'Authorization': apiAuthentication
                     }
-				}
+                });
-				);
+
-				//tracker
                 let currentPosition = position;
+                let completionText = "";
                 response.data.on('data', async (d: Uint8Array) => {
                     progress.report({ message: "Generating..." });
-					// Check for user input (cancel)
                     if (currentPosition.line != textEditor.selection.end.line || currentPosition.character != textEditor.selection.end.character) {
-						axiosCancelPost(); // cancel axios => cancel finished promise => close notification
+                        axiosCancelPost();
                         return;
                     }
-					// Get a completion from the response
                     const completion: string = JSON.parse(d.toString()).response;
-					// lastToken = completion;
-					
                     
                     if (completion === "") {
                         return;
                     }
-					//complete edit for token
+
+                    completionText += completion;
+                    
                     const edit = new vscode.WorkspaceEdit();
                     edit.insert(document.uri, currentPosition, completion);
                     await vscode.workspace.applyEdit(edit);
-					// Move the cursor to the end of the completion
+
                     const completionLines = completion.split("\n");
                     const newPosition = new vscode.Position(
                         currentPosition.line + completionLines.length - 1,
@@ -135,24 +150,29 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo
                         newPosition
                     );
                     currentPosition = newPosition;
-					// completion bar
+
                     progress.report({ message: "Generating...", increment: 1 / (numPredict / 100) });
-					// move cursor
                     textEditor.selection = newSelection;
                 });
-				// Keep cancel window available
+
                 const finished = new Promise((resolve) => {
                     response.data.on('end', () => {
                         progress.report({ message: "Fabelous completion finished." });
                         resolve(true);
                     });
-					axiosCancelToken.promise.finally(() => { // prevent notification from freezing on user input cancel
+                    axiosCancelToken.promise.finally(() => {
                         resolve(false);
                     });
                 });
                 await finished;
+
+                // Remove any remaining FIM tokens from the completion
+                completionText = completionText.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '');
+                const finalEdit = new vscode.WorkspaceEdit();
+                finalEdit.replace(document.uri, new vscode.Range(position, currentPosition), completionText);
+                await vscode.workspace.applyEdit(finalEdit);
+
             } catch (err: any) {
-				// Show an error message
                 vscode.window.showErrorMessage(
                     "Fabelous Autocoder encountered an error: " + err.message
                 );
@@ -161,30 +181,33 @@ 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("Fabelous autocompletion");
-	// 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));
+        const context = getContextLines(document, position);
-		prompt = prompt.substring(Math.max(0, prompt.length - promptWindowSize), prompt.length);
+        const lines = context.split("\n");
+        const currentLineIndex = position.line - Math.max(0, position.line - 10);
+        const prefix = lines.slice(0, currentLineIndex + 1).join("\n");
+        const suffix = lines.slice(currentLineIndex + 1).join("\n");
+        const fimPrompt = `<fim_prefix>${prefix}<fim_suffix>${suffix}<fim_middle>`;
+
         const response_preview = await axios.post(apiEndpoint, {
-			model: apiModel, // Change this to the model you want to use
+            model: apiModel,
-			prompt: messageHeaderSub(document) + prompt,
+            prompt: messageHeaderSub(document) + fimPrompt,
             stream: false,
             raw: true,
             options: {
-				num_predict: responsePreviewMaxTokens, // reduced compute max
+                num_predict: responsePreviewMaxTokens,
                 temperature: apiTemperature,
-				stop: ['\n', '```'],
+                stop: ['<fim_suffix>', '\n', '```'],
                 ...keepAlive && { keep_alive: keepAlive },
                 ...topP && { top_p: topP },
             }
@@ -196,45 +219,38 @@ async function provideCompletionItems(document: vscode.TextDocument, position: v
                 cancellationToken.onCancellationRequested(cancelPost);
             })
         });
-		if (response_preview.data.response.trim() != "") { // default if empty
+        if (response_preview.data.response.trim() != "") {
-			item.label = response_preview.data.response.trimStart(); // tended to add whitespace at the beginning
+            const previewText = response_preview.data.response.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trimStart();
-			item.insertText = response_preview.data.response.trimStart();
+            item.label = previewText;
+            item.insertText = previewText;
         }
     }
-	// Set the documentation to a message
+
-	item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Ollama');
+    item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Fabelous Autocoder');
-	// Set the command to trigger the completion
     if (continueInline || !responsePreview) item.command = {
         command: 'fabelous-autocoder.autocomplete',
         title: 'Fabelous Autocomplete',
         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("*", {
         provideCompletionItems
     },
         ...completionKeys.split("")
     );
-	// Register a command for getting a completion from Ollama through command/keybind
     const externalAutocompleteCommand = vscode.commands.registerTextEditorCommand(
         "fabelous-autocoder.autocomplete",
         (textEditor, _, cancellationToken?) => {
-			// no cancellation token from here, but there is one from completionProvider
             autocompleteCommand(textEditor, cancellationToken);
         }
     );
-	// Add the commands & completion provider to the context
     context.subscriptions.push(completionProvider);
     context.subscriptions.push(externalAutocompleteCommand);
 }
 
-// This method is called when extension is deactivated
-// eslint-disable-next-line @typescript-eslint/no-empty-function
 function deactivate() { }
 
 module.exports = {