src/extension.ts

@@ -1,231 +1,297 @@
-import * as vscode from 'vscode';
+import * as vscode from "vscode";
-import axios from 'axios';
+import axios from "axios";
 
-let config: {
+let VSConfig: vscode.WorkspaceConfiguration;
-    apiEndpoint: string;
+let apiEndpoint: string;
-    apiAuthentication: string;
+let apiAuthentication: string;
-    apiModel: string;
+let apiModel: string;
-    apiTemperature: number;
+let apiTemperature: number;
-    numPredict: number;
+let numPredict: number;
-    promptWindowSize: number;
+let promptWindowSize: number;
-    completionKeys: string[];
+let completionKeys: string;
-    responsePreview: boolean;
+let responsePreview: boolean | undefined;
-    responsePreviewMaxTokens: number;
+let responsePreviewMaxTokens: number;
-    responsePreviewDelay: number;
+let responsePreviewDelay: number;
-    continueInline: boolean;
+let continueInline: boolean | undefined;
-    keepAlive: number;
+let keepAlive: number | undefined;
-    topP: number;
+let topP: number | undefined;
-};
 
-let previewDecorationType: vscode.TextEditorDecorationType;
+function updateVSConfig() {
-
+    VSConfig = vscode.workspace.getConfiguration("fabelous-autocoder");
-function updateConfig() {
+    apiEndpoint = VSConfig.get("endpoint") || "http://localhost:11434/api/generate";
-    const vsConfig = vscode.workspace.getConfiguration('fabelous-autocoder');
+    apiAuthentication = VSConfig.get("authentication") || "";
-    config = {
+    apiModel = VSConfig.get("model") || "fabelous-coder:latest";
-        apiEndpoint: vsConfig.get('endpoint') || 'http://localhost:11434/api/generate',
+    numPredict = VSConfig.get("max tokens predicted") || 1000;
-        apiAuthentication: vsConfig.get('authentication') || '',
+    promptWindowSize = VSConfig.get("prompt window size") || 2000;
-        apiModel: vsConfig.get('model') || 'fabelous-coder:latest',
+    completionKeys = VSConfig.get("completion keys") || " ";
-        apiTemperature: vsConfig.get('temperature') || 0.7,
+    responsePreview = VSConfig.get("response preview");
-        numPredict: vsConfig.get('max tokens predicted') || 1000,
+    responsePreviewMaxTokens = VSConfig.get("preview max tokens") || 50;
-        promptWindowSize: vsConfig.get('prompt window size') || 2000,
+    responsePreviewDelay = VSConfig.get("preview delay") || 0;
-        completionKeys: (vsConfig.get('completion keys') as string || ' ').split(''),
+    continueInline = VSConfig.get("continue inline");
-        responsePreview: vsConfig.get('response preview') || false,
+    apiTemperature = VSConfig.get("temperature") || 0.7;
-        responsePreviewMaxTokens: vsConfig.get('preview max tokens') || 50,
+    keepAlive = VSConfig.get("keep alive") || 30;
-        responsePreviewDelay: vsConfig.get('preview delay') || 0,
+    topP = VSConfig.get("top p") || 1;
-        continueInline: vsConfig.get('continue inline') || false,
-        keepAlive: vsConfig.get('keep alive') || 30,
-        topP: vsConfig.get('top p') || 1,
-    };
 }
 
-function createPreviewDecorationType() {
+updateVSConfig();
-    previewDecorationType = vscode.window.createTextEditorDecorationType({
+vscode.workspace.onDidChangeConfiguration(updateVSConfig);
-        after: {
-            color: '#888888',
-            fontStyle: 'italic',
-        },
-        textDecoration: 'none; display: none;',
-    });
-}
 
 function getContextLines(document: vscode.TextDocument, position: vscode.Position): string {
+    const lines = [];
     const startLine = Math.max(0, position.line - 1);
     const endLine = position.line;
-    return document.getText(new vscode.Range(startLine, 0, endLine, position.character));
+
+    for (let i = startLine; i <= endLine; i++) {
+        lines.push(document.lineAt(i).text);
+    }
+
+    return lines.join("\n");
 }
 
+
 function createFIMPrompt(prefix: string, language: string): string {
     return `<fim_prefix>${prefix}<fim_middle><fim_suffix>${language}\n`;
 }
 
-async function generateCompletion(prompt: string, cancellationToken: vscode.CancellationToken): Promise<string> {
+const previewDecorationType = vscode.window.createTextEditorDecorationType({
-    const axiosCancelToken = new axios.CancelToken((c) => {
+    after: {
-        cancellationToken.onCancellationRequested(() => c('Request cancelled'));
+        color: '#888888', // Grayed-out preview text
-    });
+        fontStyle: 'italic',
+    },
+    textDecoration: 'none; display: none;', // Hide the original text
+});
 
-    const response = await axios.post(config.apiEndpoint, {
+async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationToken?: vscode.CancellationToken) {
-        model: config.apiModel,
+    const document = textEditor.document;
-        prompt: prompt,
+    const position = textEditor.selection.active;
-        stream: false,
+    const contextLines = 2;
-        raw: true,
+    const startLine = Math.max(0, position.line - contextLines);
-        options: {
+    const context = getContextLines(document, position);
-            num_predict: config.numPredict,
+    let isHandlingChange = false;
-            temperature: config.apiTemperature,
+    const fimPrompt = createFIMPrompt(context, document.languageId);
-            stop: ['<fim_suffix>', '```'],
-            keep_alive: config.keepAlive,
-            top_p: config.topP,
-        }
-    }, {
-        cancelToken: axiosCancelToken,
-        headers: {
-            'Authorization': config.apiAuthentication
-        }
-    });
 
-    return response.data.response.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trim();
+    vscode.window.withProgress(
-}
+        {
-
-class CompletionManager {
-    private textEditor: vscode.TextEditor;
-    private document: vscode.TextDocument;
-    private startPosition: vscode.Position;
-    private completionText: string;
-
-    constructor(textEditor: vscode.TextEditor, startPosition: vscode.Position, completionText: string) {
-        this.textEditor = textEditor;
-        this.document = textEditor.document;
-        this.startPosition = startPosition;
-        this.completionText = completionText;
-    }
-
-    public async showPreview() {
-        const completionLines = this.completionText.split('\n');
-        const previewLines = [
-            '',  // Empty line before
-            ...completionLines,
-            ''   // Empty line after
-        ];
-        
-        const previewRanges: vscode.DecorationOptions[] = previewLines.map((line, index) => {
-            const lineNumber = Math.max(0, this.startPosition.line + index - 1);
-            return {
-                range: new vscode.Range(
-                    new vscode.Position(lineNumber, 0),
-                    new vscode.Position(lineNumber, Number.MAX_VALUE)
-                ),
-                renderOptions: {
-                    after: {
-                        contentText: line,
-                    }
-                }
-            };
-        });
-        this.textEditor.setDecorations(previewDecorationType, previewRanges);
-    }
-    
-
-    public async acceptCompletion() {
-        const edit = new vscode.WorkspaceEdit();
-        const startLine = Math.max(0, this.startPosition.line - 1);
-        const range = new vscode.Range(
-            new vscode.Position(startLine, 0),
-            this.startPosition.translate(0, Number.MAX_VALUE)
-        );
-        edit.replace(this.document.uri, range, this.completionText);
-        await vscode.workspace.applyEdit(edit);
-        this.clearPreview();
-    }
-    
-    public clearPreview() {
-        this.textEditor.setDecorations(previewDecorationType, []);
-    }
-    public declineCompletion() {
-        this.clearPreview();
-    }
-}
-
-async function autocompleteCommand(textEditor: vscode.TextEditor, edit: vscode.TextEditorEdit, ...args: any[]) {
-    const cancellationTokenSource = new vscode.CancellationTokenSource();
-    const cancellationToken = cancellationTokenSource.token;
-
-    try {
-        const document = textEditor.document;
-        const position = textEditor.selection.active;
-        const context = getContextLines(document, position);
-        const fimPrompt = createFIMPrompt(context, document.languageId);
-
-        const completionText = await vscode.window.withProgress({
             location: vscode.ProgressLocation.Notification,
-            title: 'Fabelous Autocoder',
+            title: "Fabelous Autocoder",
             cancellable: true,
-        }, async (progress, progressCancellationToken) => {
+        },
-            progress.report({ message: 'Generating...' });
+        async (progress, progressCancellationToken) => {
-            return await generateCompletion(fimPrompt, progressCancellationToken);
+            try {
-        });
+                progress.report({ message: "Starting model..." });
 
-        console.log('Completion generated:', completionText);
+                let axiosCancelPost: () => void;
+                const axiosCancelToken = new axios.CancelToken((c) => {
+                    axiosCancelPost = () => {
+                        c("Autocompletion request terminated by user cancel");
+                    };
+                    if (cancellationToken) cancellationToken.onCancellationRequested(axiosCancelPost);
+                    progressCancellationToken.onCancellationRequested(axiosCancelPost);
+                    vscode.workspace.onDidCloseTextDocument(axiosCancelPost);
+                });
 
-        const completionManager = new CompletionManager(textEditor, position, completionText);
+                // Make the API request
-        await completionManager.showPreview();
+                const response = await axios.post(apiEndpoint, {
+                    model: apiModel,
+                    prompt: fimPrompt,
+                    stream: false,
+                    raw: true,
+                    options: {
+                        num_predict: numPredict,
+                        temperature: apiTemperature,
+                        stop: ["<fim_suffix>", "```"]
+                    }
+                }, {
+                    cancelToken: axiosCancelToken,
+                    headers: {
+                        'Authorization': apiAuthentication
+                    }
+                });
 
-        let isDisposed = false;
+                progress.report({ message: "Generating..." });
 
-        const dispose = () => {
+                let completionText = response.data.response;
-            if (!isDisposed) {
+                completionText = completionText.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trim();
-                console.log('Disposing listeners');
+
-                disposable.dispose();
+                let previewInserted = false;
-                declineDisposable.dispose();
+                let originalContent: string;
-                isDisposed = true;
+                let previewStartLine: number;
+
+
+                const storeAndInsertPreview = async () => {
+                    previewStartLine = startLine;
+                    const endLine = document.lineCount - 1;
+                    const endCharacter = document.lineAt(endLine).text.length;
+                    const fullRange = new vscode.Range(startLine, 0, endLine, endCharacter);
+                    originalContent = document.getText(fullRange);
+                
+                    const previewContent = completionText + '\n'.repeat(1);
+                    const edit = new vscode.WorkspaceEdit();
+                    edit.replace(document.uri, fullRange, previewContent); // Overwrite the content
+                    await vscode.workspace.applyEdit(edit);
+                
+                    // Split the preview content into lines
+                    const previewLines = previewContent.split('\n');
+                    
+                
+                    // Set decorations on the newly inserted lines
+                    const previewRanges: vscode.DecorationOptions[] = [];
+                    for (let i = 0; i < previewLines.length; i++) {
+                        const range = new vscode.Range(previewStartLine + i, 0, previewStartLine + i, previewLines[i].length);
+                        previewRanges.push({
+                            range,
+                            renderOptions: {
+                                after: {
+                                    contentText: previewLines[i],
+                                }
+                            }
+                        });
+                    }
+                    textEditor.setDecorations(previewDecorationType, previewRanges);
+                    previewInserted = true;
+                };
+                const disposable = vscode.window.onDidChangeTextEditorSelection(async (event) => {
+                    const textEditor = vscode.window.activeTextEditor;
+                    if (!textEditor || !previewInserted || isHandlingChange) {
+                        return;
+                    }
+
+                    isHandlingChange = true;
+
+                    try {
+                        const activeSelection = textEditor.selection;
+                        const changeStartLine = activeSelection.active.line;
+
+                        // Detect Tab key press by checking the active selection
+                        if (event.kind === vscode.TextEditorSelectionChangeKind.Keyboard && changeStartLine >= previewStartLine) {
+                            const changeText = textEditor.document.getText(activeSelection);
+
+                            if (changeText === '') {
+                                // Tab key (empty selection) -> Accept the preview
+                                await acceptPreview(textEditor, textEditor.document, startLine, activeSelection.active, completionText);
+                            }
+                        }
+                    } finally {
+                        isHandlingChange = false;
+                    }
+                });
+
+                // Handles Enter key separately
+                vscode.workspace.onDidChangeTextDocument(async (event) => {
+                    const textEditor = vscode.window.activeTextEditor;
+                    if (!textEditor || event.document.uri.toString() !== textEditor.document.uri.toString() || !previewInserted || isHandlingChange) {
+                        return;
+                    }
+
+                    isHandlingChange = true;
+
+                    try {
+                        for (const change of event.contentChanges) {
+                            const changeStartLine = change.range.start.line;
+
+                            if (change.text.includes('\n') && changeStartLine >= previewStartLine) {
+                                // Accept the preview and move to the next line
+                                await acceptPreview(textEditor, textEditor.document, startLine, textEditor.selection.active, completionText);
+                                await vscode.commands.executeCommand('default:type', { text: '\n' });
+                                break;
+                            }
+
+                            // Handle Backspace
+                            if (change.text === '' && change.rangeLength === 1 && changeStartLine >= previewStartLine) {
+                                // Discard the preview if Backspace is pressed
+                                await restoreOriginalContent();
+                                break;
+                            }
+                        }
+                    } finally {
+                        isHandlingChange = false;
+                    }
+                });
+
+                vscode.window.onDidChangeTextEditorSelection(async (event) => {
+                    const textEditor = vscode.window.activeTextEditor;
+                    if (!textEditor || !previewInserted || isHandlingChange) {
+                        return;
+                    }
+
+                    isHandlingChange = true;
+
+                    try {
+                        // Handle arrow keys or any other navigation keys
+                        const currentSelection = event.selections[0];
+                        const { document } = textEditor;
+
+                        // Detect unwanted acceptance from simple navigation
+                        if (currentSelection.start.line < previewStartLine) {
+                            await restoreOriginalContent();
+                        }
+                    } finally {
+                        isHandlingChange = false;
+                    }
+                });
+
+
+
+                
+                // Restore original content if Backspace is pressed (decline the preview)
+                const restoreOriginalContent = async () => {
+                    if (!previewInserted) return;
+                
+                    const endLine = document.lineCount - 1;
+                    const fullRange = new vscode.Range(previewStartLine, 0, endLine, document.lineAt(endLine).text.length);
+                    const edit = new vscode.WorkspaceEdit();
+                
+                    edit.replace(document.uri, fullRange, originalContent);
+                    await vscode.workspace.applyEdit(edit);
+                
+                    textEditor.setDecorations(previewDecorationType, []);
+                    previewInserted = false;
+                    disposable.dispose(); // Cancel listener when preview is discarded
+                };
+                
+                // Accept the preview when Tab is pressed
+                const acceptPreview = async (textEditor: vscode.TextEditor, document: vscode.TextDocument, startLine: number, position: vscode.Position, completionText: string) => {
+                    textEditor.setDecorations(previewDecorationType, []);
+                    const edit = new vscode.WorkspaceEdit();
+                    
+                    // Adjust the insertion logic to avoid duplicate newlines
+                    const insertText = completionText;
+                    
+                    // Replace the range from the start of the context to the current position
+                    const replaceRange = new vscode.Range(startLine, 0, position.line, position.character);
+                    edit.replace(document.uri, replaceRange, insertText);
+                    
+                    await vscode.workspace.applyEdit(edit);
+                    await document.save();
+                    
+                    disposable.dispose(); // Cancel listener when preview is accepted
+                    previewInserted = false;
+                };
+                
+                // Call this function to initiate the preview
+                await storeAndInsertPreview();
+                
+
+                                
+
+            } catch (err: any) {
+                vscode.window.showErrorMessage(
+                    "Fabelous Autocoder encountered an error: " + err.message
+                );
+                console.log(err);
             }
-        };
+        }
-
+    );
-        const disposable = vscode.Disposable.from(
-            vscode.window.onDidChangeTextEditorSelection(async (event) => {
-                if (event.textEditor !== textEditor) return;
-
-                if (event.kind === vscode.TextEditorSelectionChangeKind.Keyboard) {
-                    console.log('Accepting completion');
-                    await completionManager.acceptCompletion();
-                    dispose();
-                }
-            }),
-            vscode.window.onDidChangeActiveTextEditor(() => {
-                console.log('Active editor changed, clearing preview');
-                completionManager.clearPreview();
-                dispose();
-            }),
-            vscode.workspace.onDidChangeTextDocument((event) => {
-                if (event.document === document) {
-                    console.log('Document changed, clearing preview');
-                    completionManager.clearPreview();
-                    dispose();
-                }
-            })
-        );
-
-        const declineDisposable = vscode.commands.registerCommand('type', async (args) => {
-            if (args.text === '\b') { // Backspace key
-                console.log('Declining completion');
-                completionManager.declineCompletion();
-                dispose();
-            }
-        });
-
-    } catch (err: any) {
-        console.error('Error in autocompleteCommand:', err);
-        vscode.window.showErrorMessage(`Fabelous Autocoder encountered an error: ${err.message}`);
-    } finally {
-        cancellationTokenSource.dispose();
-    }
 }
 
+
+
 async function provideCompletionItems(document: vscode.TextDocument, position: vscode.Position, cancellationToken: vscode.CancellationToken) {
-    const item = new vscode.CompletionItem('Fabelous autocompletion');
+    const item = new vscode.CompletionItem("Fabelous autocompletion");
     item.insertText = new vscode.SnippetString('${1:}');
-    item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Fabelous Autocoder');
 
-    if (config.responsePreview) {
+    if (responsePreview) {
-        await new Promise(resolve => setTimeout(resolve, config.responsePreviewDelay * 1000));
+        await new Promise(resolve => setTimeout(resolve, responsePreviewDelay * 1000));
         if (cancellationToken.isCancellationRequested) {
             return [item];
         }
@@ -234,36 +300,58 @@ async function provideCompletionItems(document: vscode.TextDocument, position: v
         const fimPrompt = createFIMPrompt(context, document.languageId);
 
         try {
-            const result = await generateCompletion(fimPrompt, cancellationToken);
+            const response_preview = await axios.post(apiEndpoint, {
-            const preview = (result as any).preview;
+                model: apiModel,
-            if (preview) {
+                prompt: fimPrompt,
-            item.detail = preview.split('\n')[0];
+                stream: false,
-            }
+                raw: true,
+                options: {
+                    num_predict: responsePreviewMaxTokens,
+                    temperature: apiTemperature,
+                    stop: ['<fim_suffix>', '\n', '```'],
+                    ...(keepAlive && { keep_alive: keepAlive }),
+                    ...(topP && { top_p: topP }),
+                }
+            }, {
+                cancelToken: new axios.CancelToken((c) => {
+                    cancellationToken.onCancellationRequested(() => c("Autocompletion request terminated by completion cancel"));
+                })
+            });
         } catch (error) {
-            console.error('Error fetching preview:', error);
+            console.error("Error fetching preview:", error);
         }
     }
 
-    if (config.continueInline || !config.responsePreview) {
+    item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Fabelous Autocoder');
+    if (continueInline || !responsePreview) {
         item.command = {
             command: 'fabelous-autocoder.autocomplete',
             title: 'Fabelous Autocomplete',
-            arguments: []
+            arguments: [cancellationToken]
         };
     }
-
     return [item];
 }
 
-export function activate(context: vscode.ExtensionContext) {
+function activate(context: vscode.ExtensionContext) {
-    updateConfig();
+    const completionProvider = vscode.languages.registerCompletionItemProvider("*", {
-    createPreviewDecorationType();
+        provideCompletionItems
-
+    },
-    context.subscriptions.push(
+        ...completionKeys.split("")
-        vscode.workspace.onDidChangeConfiguration(updateConfig),
-        vscode.languages.registerCompletionItemProvider('*', { provideCompletionItems }, ...config.completionKeys),
-        vscode.commands.registerTextEditorCommand('fabelous-autocoder.autocomplete', autocompleteCommand)
     );
+    const externalAutocompleteCommand = vscode.commands.registerTextEditorCommand(
+        "fabelous-autocoder.autocomplete",
+        (textEditor, _, cancellationToken?) => {
+            autocompleteCommand(textEditor, cancellationToken);
+        }
+    );
+    context.subscriptions.push(completionProvider);
+    context.subscriptions.push(externalAutocompleteCommand);
 }
 
-export function deactivate() {}
+function deactivate() { }
+
+module.exports = {
+    activate,
+    deactivate,
+};