diff --git a/src/extension.ts b/src/extension.ts index 0b3c6b4..1291669 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -1,269 +1,179 @@ -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 `${prefix}${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 generateCompletion(prompt: string, cancellationToken: vscode.CancellationToken): Promise { + const axiosCancelToken = new axios.CancelToken((c) => { + cancellationToken.onCancellationRequested(() => c('Request cancelled')); + }); - -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; - 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); - }); - - // Make the API request - const response = await axios.post(apiEndpoint, { - model: apiModel, - prompt: fimPrompt, - stream: false, - raw: true, - options: { - num_predict: numPredict, - temperature: apiTemperature, - stop: ["", "```"] - } - }, { - cancelToken: axiosCancelToken, - headers: { - 'Authorization': apiAuthentication - } - }); - - progress.report({ message: "Generating..." }); - - let completionText = response.data.response; - completionText = completionText.replace(/||/g, '').trim(); - - let previewInserted = false; - let originalContent: string; - let previewStartLine: number; - let previewEndLine: number; - - const storeAndInsertPreview = async () => { - const currentLine = position.line; - previewStartLine = currentLine; - const previewLines = completionText.split('\n'); - previewEndLine = previewStartLine + previewLines.length; - - // Ensure the previewEndLine doesn't exceed the document's line count - const documentEndLine = document.lineCount - 1; - previewEndLine = Math.min(previewEndLine, documentEndLine + 1); - - // Store original content - const fullRange = new vscode.Range(previewStartLine, 0, previewEndLine, document.lineAt(previewEndLine - 1).text.length); - originalContent = document.getText(fullRange); - - // Prepare the new content - const newContent = previewLines.join('\n'); - - const edit = new vscode.WorkspaceEdit(); - edit.replace(document.uri, fullRange, newContent); - await vscode.workspace.applyEdit(edit); - - // Highlight only the new lines - const previewRanges: vscode.DecorationOptions[] = previewLines.map((line: string, index: number) => ({ - range: new vscode.Range(previewStartLine + index, 0, previewStartLine + index, line.length), - renderOptions: { - after: { - contentText: line, - } - } - })); - 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; - - 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; - } - }); - - 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) { - await acceptPreview(textEditor, textEditor.document, startLine, textEditor.selection.active, completionText); - await vscode.commands.executeCommand('default:type', { text: '\n' }); - break; - } - - if (change.text === '' && change.rangeLength === 1 && changeStartLine >= previewStartLine) { - await restoreOriginalContent(); - break; - } - } - } finally { - isHandlingChange = false; - } - }); - - const restoreOriginalContent = async () => { - if (!previewInserted) return; - - const fullRange = new vscode.Range(previewStartLine, 0, previewEndLine, 0); - 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 - }; - - const acceptPreview = async (textEditor: vscode.TextEditor, document: vscode.TextDocument, startLine: number, position: vscode.Position, completionText: string) => { - textEditor.setDecorations(previewDecorationType, []); - const edit = new vscode.WorkspaceEdit(); - - const previewLines = completionText.split('\n'); - const endLine = startLine + previewLines.length; - - // Ensure that endLine does not exceed document bounds - const documentEndLine = document.lineCount - 1; - const finalEndLine = Math.min(endLine, documentEndLine + 1); - - // Prepare the new content - const newContent = previewLines.join('\n'); - - // Replace the range with the new content - const replaceRange = new vscode.Range(startLine, 0, finalEndLine, document.lineAt(finalEndLine - 1).text.length); - edit.replace(document.uri, replaceRange, newContent.trim()); - - await vscode.workspace.applyEdit(edit); - await document.save(); - - disposable.dispose(); // Cancel listener when preview is accepted - previewInserted = false; - }; - - await storeAndInsertPreview(); - } catch (err: any) { - vscode.window.showErrorMessage( - "Fabelous Autocoder encountered an error: " + err.message - ); - console.log(err); - } + const response = await axios.post(config.apiEndpoint, { + model: config.apiModel, + prompt: prompt, + stream: false, + raw: true, + options: { + num_predict: config.numPredict + 100, // Generate extra lines for preview + temperature: config.apiTemperature, + stop: ['', '```'], + keep_alive: config.keepAlive, + top_p: config.topP, } - ); + }, { + cancelToken: axiosCancelToken, + headers: { + 'Authorization': config.apiAuthentication + } + }); + + const fullCompletion = response.data.response.replace(/||/g, '').trim(); + const completionLines = fullCompletion.split('\n'); + const usedContextLines = prompt.split('\n').length; + + // Remove used context and take extra lines for preview + return completionLines.slice(usedContextLines, usedContextLines + config.numPredict).join('\n'); } +class CompletionManager { + private textEditor: vscode.TextEditor; + private document: vscode.TextDocument; + private startPosition: vscode.Position; + private completionText: string; + private previewLines: string[]; + constructor(textEditor: vscode.TextEditor, startPosition: vscode.Position, completionText: string) { + this.textEditor = textEditor; + this.document = textEditor.document; + this.startPosition = startPosition; + this.completionText = completionText; + this.previewLines = completionText.split('\n'); + } + + public async showPreview() { + const previewRanges: vscode.DecorationOptions[] = this.previewLines.map((line, index) => ({ + range: new vscode.Range(this.startPosition.translate(index, 0), this.startPosition.translate(index, 0)), + renderOptions: { + after: { + contentText: line, + } + } + })); + this.textEditor.setDecorations(previewDecorationType, previewRanges); + } + + public async acceptCompletion() { + const edit = new vscode.WorkspaceEdit(); + edit.insert(this.document.uri, this.startPosition, this.completionText); + await vscode.workspace.applyEdit(edit); + this.clearPreview(); + } + + public clearPreview() { + this.textEditor.setDecorations(previewDecorationType, []); + } +} + +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); + }); + + const completionManager = new CompletionManager(textEditor, position, completionText); + await completionManager.showPreview(); + + const disposable = vscode.window.onDidChangeTextEditorSelection(async (event) => { + if (event.textEditor !== textEditor) return; + + if (event.kind === vscode.TextEditorSelectionChangeKind.Keyboard) { + await completionManager.acceptCompletion(); + disposable.dispose(); + } + }); + } catch (err: any) { + vscode.window.showErrorMessage(`Fabelous Autocoder encountered an error: ${err.message}`); + console.error(err); + } 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]; } @@ -272,58 +182,33 @@ 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: ['', '\n', '```'], - ...(keepAlive && { keep_alive: keepAlive }), - ...(topP && { top_p: topP }), - } - }, { - cancelToken: new axios.CancelToken((c) => { - cancellationToken.onCancellationRequested(() => c("Autocompletion request terminated by completion cancel")); - }) - }); + const previewText = await generateCompletion(fimPrompt, cancellationToken); + item.detail = previewText.split('\n')[0]; // Show first line as preview } 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() {}