Fabelous-Autocoder/src/extension.ts

import * as vscode from 'vscode';
import axios from 'axios';

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;
};

let previewDecorationType: vscode.TextEditorDecorationType;
let activeCompletionManager: CompletionManager | null = null;

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,
    };
}

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 startLine = Math.max(0, position.line - 1);
    const endLine = position.line;
    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`;
}

async function generateCompletion(prompt: string, cancellationToken: vscode.CancellationToken): Promise<string> {
    const axiosCancelToken = new axios.CancelToken((c) => {
        cancellationToken.onCancellationRequested(() => c('Request cancelled'));
    });

    const response = await axios.post(config.apiEndpoint, {
        model: config.apiModel,
        prompt: prompt,
        stream: false,
        raw: true,
        options: {
            num_predict: config.numPredict,
            temperature: config.apiTemperature,
            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();
}

class CompletionManager {
    private textEditor: vscode.TextEditor;
    private document: vscode.TextDocument;
    private startPosition: vscode.Position;
    private completionText: string;
    private insertedLineCount: number = 0; // Track the number of inserted lines

    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() {
        if (!previewDecorationType) {
            createPreviewDecorationType();
        }
    
        const completionLines = this.completionText.split('\n').length;
    
        // Adjust the start position to line after the original start position
        const adjustedStartPosition = this.startPosition.translate(0, 0);
    
        // Step 1: Insert blank lines to make space for the preview
        const edit = new vscode.WorkspaceEdit();
        const linePadding = '\n'.repeat(completionLines + 1); // Include extra line break for visual separation
        edit.insert(this.document.uri, adjustedStartPosition, linePadding);
        await vscode.workspace.applyEdit(edit);
    
        this.insertedLineCount = completionLines + 1;
    
        // Step 2: Apply decorations
        const previewRanges: vscode.DecorationOptions[] = this.completionText.split('\n').map((line, index) => {
            const lineNumber = adjustedStartPosition.line + index + 1; // Start preview one line later
            return {
                range: new vscode.Range(
                    new vscode.Position(lineNumber, 0),
                    new vscode.Position(lineNumber, 0)
                ),
                renderOptions: {
                    after: {
                        contentText: line,
                        color: '#888888',
                        fontStyle: 'italic',
                    },
                },
            };
        });
    
        this.textEditor.setDecorations(previewDecorationType, previewRanges);
    }
    

    public async acceptCompletion() {
        const edit = new vscode.WorkspaceEdit();
        const completionLines = this.completionText.split('\n');
        const numberOfLines = completionLines.length;
        
        // Ensure the start position is never negative
        const safeStartPosition = new vscode.Position(Math.max(0, this.startPosition.line - 1), 0);

        // Prepare the range to replace
        const rangeToReplace = new vscode.Range(
            safeStartPosition,
            this.startPosition.translate(numberOfLines, 0)
        );

        // Construct the content to insert
        const contentToInsert = (safeStartPosition.line === 0 ? '' : '\n') + this.completionText + '\n';
        edit.replace(this.document.uri, rangeToReplace, contentToInsert);
        
        await vscode.workspace.applyEdit(edit);
        
        // Clear the preview decorations
        this.clearPreview();

        // Set activeCompletionManager to null
        activeCompletionManager = null;

        // Calculate the new cursor position from the inserted content 
        const lastCompletionLine = completionLines[completionLines.length - 1];
        const newPosition = new vscode.Position(
            this.startPosition.line + numberOfLines - 1, 
            lastCompletionLine.length
        );

        // Set the new cursor position
        this.textEditor.selection = new vscode.Selection(newPosition, newPosition);
    }

    

    public clearPreview() {
        this.textEditor.setDecorations(previewDecorationType, []);  // Remove all preview decorations
    }

    public async declineCompletion() {
        this.clearPreview();  // Clear the preview decorations
    
        try {
            const document = this.textEditor.document;
            const currentPosition = this.textEditor.selection.active;
    
            // Calculate the range of lines to remove
            const startLine = this.startPosition.line + 1;
            const endLine = currentPosition.line;
            if (endLine > startLine) {
                const workspaceEdit = new vscode.WorkspaceEdit();
    
                // Create a range from start of startLine to end of endLine
                const range = new vscode.Range(
                    new vscode.Position(startLine, 0),
                    new vscode.Position(endLine, document.lineAt(endLine).text.length)
                );
    
                // Delete the range
                workspaceEdit.delete(document.uri, range);
    
                // Apply the edit
                await vscode.workspace.applyEdit(workspaceEdit);
    
                // Move the cursor back to the original position
                this.textEditor.selection = new vscode.Selection(this.startPosition, this.startPosition);
    
                console.log(`Lines ${startLine + 1} to ${endLine + 1} removed successfully`);
                activeCompletionManager = null;
            } else {
                console.log('No lines to remove');
            }
        } catch (error) {
            console.error('Error declining completion:', error);
            vscode.window.showErrorMessage(`Error removing lines: ${error}`);
        }
    }
}

 
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();
        activeCompletionManager = completionManager;


    } catch (err: any) {
        console.error('Error in autocompleteCommand:', err);
        vscode.window.showErrorMessage(`Fabelous Autocoder encountered an error: ${err.message}`);
    } finally {
        cancellationTokenSource.dispose();
    }
}

async function handleTab() {
    if (activeCompletionManager) {
        await activeCompletionManager.acceptCompletion();
    } else {
        await vscode.commands.executeCommand('tab');
    }
}

async function handleBackspace() {
    if (activeCompletionManager) {
        await activeCompletionManager.declineCompletion();
    } else {
        await vscode.commands.executeCommand('deleteLeft');
    }
}


export function activate(context: vscode.ExtensionContext) {
    updateConfig();
    createPreviewDecorationType();

    context.subscriptions.push(
        vscode.workspace.onDidChangeConfiguration(updateConfig),
        vscode.commands.registerTextEditorCommand('fabelous-autocoder.autocomplete', autocompleteCommand),
        vscode.commands.registerCommand('fabelous-autocoder.handleTab', handleTab),
        vscode.commands.registerCommand('fabelous-autocoder.handleBackspace', handleBackspace)
    );
}

export function deactivate() {}