preview added

src/extension.ts

@@ -1,297 +1,231 @@
-import * as vscode from "vscode";
-import axios from "axios";
+import * as vscode from 'vscode';
+import axios from 'axios';
 
-let VSConfig: vscode.WorkspaceConfiguration;
-let apiEndpoint: string;
-let apiAuthentication: string;
-let apiModel: string;
-let apiTemperature: number;
-let numPredict: number;
-let promptWindowSize: number;
-let completionKeys: string;
-let responsePreview: boolean | undefined;
-let responsePreviewMaxTokens: number;
-let responsePreviewDelay: number;
-let continueInline: boolean | undefined;
-let keepAlive: number | undefined;
-let topP: number | undefined;
+let config: {
+    apiEndpoint: string;
+    apiAuthentication: string;
+    apiModel: string;
+    apiTemperature: number;
+    numPredict: number;
+    promptWindowSize: number;
+    completionKeys: string[];
+    responsePreview: boolean;
+    responsePreviewMaxTokens: number;
+    responsePreviewDelay: number;
+    continueInline: boolean;
+    keepAlive: number;
+    topP: number;
+};
 
-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") || "fabelous-coder:latest";
-    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;
-    continueInline = VSConfig.get("continue inline");
-    apiTemperature = VSConfig.get("temperature") || 0.7;
-    keepAlive = VSConfig.get("keep alive") || 30;
-    topP = VSConfig.get("top p") || 1;
+let previewDecorationType: vscode.TextEditorDecorationType;
+
+function updateConfig() {
+    const vsConfig = vscode.workspace.getConfiguration('fabelous-autocoder');
+    config = {
+        apiEndpoint: vsConfig.get('endpoint') || 'http://localhost:11434/api/generate',
+        apiAuthentication: vsConfig.get('authentication') || '',
+        apiModel: vsConfig.get('model') || 'fabelous-coder:latest',
+        apiTemperature: vsConfig.get('temperature') || 0.7,
+        numPredict: vsConfig.get('max tokens predicted') || 1000,
+        promptWindowSize: vsConfig.get('prompt window size') || 2000,
+        completionKeys: (vsConfig.get('completion keys') as string || ' ').split(''),
+        responsePreview: vsConfig.get('response preview') || false,
+        responsePreviewMaxTokens: vsConfig.get('preview max tokens') || 50,
+        responsePreviewDelay: vsConfig.get('preview delay') || 0,
+        continueInline: vsConfig.get('continue inline') || false,
+        keepAlive: vsConfig.get('keep alive') || 30,
+        topP: vsConfig.get('top p') || 1,
+    };
 }
 
-updateVSConfig();
-vscode.workspace.onDidChangeConfiguration(updateVSConfig);
+function createPreviewDecorationType() {
+    previewDecorationType = vscode.window.createTextEditorDecorationType({
+        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;
-
-    for (let i = startLine; i <= endLine; i++) {
-        lines.push(document.lineAt(i).text);
-    }
-
-    return lines.join("\n");
+    return document.getText(new vscode.Range(startLine, 0, endLine, position.character));
 }
 
-
 function createFIMPrompt(prefix: string, language: string): string {
     return `<fim_prefix>${prefix}<fim_middle><fim_suffix>${language}\n`;
 }
 
-const previewDecorationType = vscode.window.createTextEditorDecorationType({
-    after: {
-        color: '#888888', // Grayed-out preview text
-        fontStyle: 'italic',
-    },
-    textDecoration: 'none; display: none;', // Hide the original text
-});
-
-async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationToken?: vscode.CancellationToken) {
-    const document = textEditor.document;
-    const position = textEditor.selection.active;
-    const contextLines = 2;
-    const startLine = Math.max(0, position.line - contextLines);
-    const context = getContextLines(document, position);
-    let isHandlingChange = false;
-    const fimPrompt = createFIMPrompt(context, document.languageId);
-
-    vscode.window.withProgress(
-        {
-            location: vscode.ProgressLocation.Notification,
-            title: "Fabelous Autocoder",
-            cancellable: true,
-        },
-        async (progress, progressCancellationToken) => {
-            try {
-                progress.report({ message: "Starting model..." });
-
-                let axiosCancelPost: () => void;
+async function generateCompletion(prompt: string, cancellationToken: vscode.CancellationToken): Promise<string> {
     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);
+        cancellationToken.onCancellationRequested(() => c('Request cancelled'));
     });
 
-                // Make the API request
-                const response = await axios.post(apiEndpoint, {
-                    model: apiModel,
-                    prompt: fimPrompt,
+    const response = await axios.post(config.apiEndpoint, {
+        model: config.apiModel,
+        prompt: prompt,
         stream: false,
         raw: true,
         options: {
-                        num_predict: numPredict,
-                        temperature: apiTemperature,
-                        stop: ["<fim_suffix>", "```"]
+            num_predict: config.numPredict,
+            temperature: config.apiTemperature,
+            stop: ['<fim_suffix>', '```'],
+            keep_alive: config.keepAlive,
+            top_p: config.topP,
         }
     }, {
         cancelToken: axiosCancelToken,
         headers: {
-                        'Authorization': apiAuthentication
+            'Authorization': config.apiAuthentication
         }
     });
 
-                progress.report({ message: "Generating..." });
-
-                let completionText = response.data.response;
-                completionText = completionText.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trim();
-
-                let previewInserted = false;
-                let originalContent: string;
-                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);
-            }
-        }
-    );
+    return response.data.response.replace(/<fim_middle>|<fim_suffix>|<fim_prefix>/g, '').trim();
 }
 
+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',
+            cancellable: true,
+        }, async (progress, progressCancellationToken) => {
+            progress.report({ message: 'Generating...' });
+            return await generateCompletion(fimPrompt, progressCancellationToken);
+        });
+
+        console.log('Completion generated:', completionText);
+
+        const completionManager = new CompletionManager(textEditor, position, completionText);
+        await completionManager.showPreview();
+
+        let isDisposed = false;
+
+        const dispose = () => {
+            if (!isDisposed) {
+                console.log('Disposing listeners');
+                disposable.dispose();
+                declineDisposable.dispose();
+                isDisposed = true;
+            }
+        };
+
+        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 (responsePreview) {
-        await new Promise(resolve => setTimeout(resolve, responsePreviewDelay * 1000));
+    if (config.responsePreview) {
+        await new Promise(resolve => setTimeout(resolve, config.responsePreviewDelay * 1000));
         if (cancellationToken.isCancellationRequested) {
             return [item];
         }
@@ -300,58 +234,36 @@ async function provideCompletionItems(document: vscode.TextDocument, position: v
         const fimPrompt = createFIMPrompt(context, document.languageId);
 
         try {
-            const response_preview = await axios.post(apiEndpoint, {
-                model: apiModel,
-                prompt: fimPrompt,
-                stream: false,
-                raw: true,
-                options: {
-                    num_predict: responsePreviewMaxTokens,
-                    temperature: apiTemperature,
-                    stop: ['<fim_suffix>', '\n', '```'],
-                    ...(keepAlive && { keep_alive: keepAlive }),
-                    ...(topP && { top_p: topP }),
+            const result = await generateCompletion(fimPrompt, cancellationToken);
+            const preview = (result as any).preview;
+            if (preview) {
+            item.detail = preview.split('\n')[0];
             }
-            }, {
-                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);
         }
     }
 
-    item.documentation = new vscode.MarkdownString('Press `Enter` to get an autocompletion from Fabelous Autocoder');
-    if (continueInline || !responsePreview) {
+    if (config.continueInline || !config.responsePreview) {
         item.command = {
             command: 'fabelous-autocoder.autocomplete',
             title: 'Fabelous Autocomplete',
-            arguments: [cancellationToken]
+            arguments: []
         };
     }
+
     return [item];
 }
 
-function activate(context: vscode.ExtensionContext) {
-    const completionProvider = vscode.languages.registerCompletionItemProvider("*", {
-        provideCompletionItems
-    },
-        ...completionKeys.split("")
+export function activate(context: vscode.ExtensionContext) {
+    updateConfig();
+    createPreviewDecorationType();
+
+    context.subscriptions.push(
+        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);
 }
 
-function deactivate() { }
-
-module.exports = {
-    activate,
-    deactivate,
-};
+export function deactivate() {}