From 1a98cfe1e15a71abc452b1bddfec441644bc8776 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Wed, 14 Aug 2024 07:31:21 +0200 Subject: [PATCH] version2 --- src/extension.ts | 126 +++++++++++++++++++++-------------------------- 1 file changed, 55 insertions(+), 71 deletions(-) diff --git a/src/extension.ts b/src/extension.ts index 6977072..69ac704 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -5,7 +5,6 @@ let VSConfig: vscode.WorkspaceConfiguration; let apiEndpoint: string; let apiAuthentication: string; let apiModel: string; -let apiMessageHeader: string; let apiTemperature: number; let numPredict: number; let promptWindowSize: number; @@ -21,8 +20,7 @@ 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"; // Updated to use FIM model - apiMessageHeader = VSConfig.get("message header") || ""; + 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") || " "; @@ -38,21 +36,12 @@ function updateVSConfig() { updateVSConfig(); vscode.workspace.onDidChangeConfiguration(updateVSConfig); -function messageHeaderSub(document: vscode.TextDocument) { - const sub = apiMessageHeader - .replace("{LANG}", document.languageId) - .replace("{FILE_NAME}", document.fileName) - .replace("{PROJECT_NAME}", vscode.workspace.name || "Untitled"); - return sub; -} - 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); + const startLine = Math.max(0, position.line - promptWindowSize / 2); + const endLine = Math.min(lineCount - 1, position.line + promptWindowSize / 2); for (let i = startLine; i <= endLine; i++) { lines.push(document.lineAt(i).text); @@ -61,24 +50,21 @@ function getContextLines(document: vscode.TextDocument, position: vscode.Positio return lines.join("\n"); } +function createFIMPrompt(prefix: string, suffix: string, language: string): string { + return `${prefix}${suffix}\`\`\`${language}\n`; +} + async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationToken?: vscode.CancellationToken) { const document = textEditor.document; const position = textEditor.selection.active; - // Get the current context const context = getContextLines(document, position); - - // Split the context into prefix and suffix for FIM const lines = context.split("\n"); - const currentLineIndex = position.line - Math.max(0, position.line - 10); + const currentLineIndex = position.line - Math.max(0, position.line - promptWindowSize / 2); const prefix = lines.slice(0, currentLineIndex + 1).join("\n"); const suffix = lines.slice(currentLineIndex + 1).join("\n"); - // Create FIM prompt - const fimPrompt = `${prefix}${suffix}`; - - // Replace {Prompt} with the FIM prompt - const sub = messageHeaderSub(document).replace("{PROMPT}", fimPrompt); + const fimPrompt = createFIMPrompt(prefix, suffix, document.languageId); vscode.window.withProgress( { @@ -92,19 +78,17 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo let axiosCancelPost: () => void; const axiosCancelToken = new axios.CancelToken((c) => { - const cancelPost = function () { + axiosCancelPost = () => { c("Autocompletion request terminated by user cancel"); }; - - axiosCancelPost = cancelPost; - if (cancellationToken) cancellationToken.onCancellationRequested(cancelPost); - progressCancellationToken.onCancellationRequested(cancelPost); - vscode.workspace.onDidCloseTextDocument(cancelPost); + if (cancellationToken) cancellationToken.onCancellationRequested(axiosCancelPost); + progressCancellationToken.onCancellationRequested(axiosCancelPost); + vscode.workspace.onDidCloseTextDocument(axiosCancelPost); }); const response = await axios.post(apiEndpoint, { model: apiModel, - prompt: sub, + prompt: fimPrompt, stream: true, raw: true, options: { @@ -124,7 +108,7 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo let completionText = ""; response.data.on('data', async (d: Uint8Array) => { progress.report({ message: "Generating..." }); - if (currentPosition.line != textEditor.selection.end.line || currentPosition.character != textEditor.selection.end.character) { + if (currentPosition.line !== textEditor.selection.end.line || currentPosition.character !== textEditor.selection.end.character) { axiosCancelPost(); return; } @@ -145,10 +129,7 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo currentPosition.line + completionLines.length - 1, (completionLines.length > 1 ? 0 : currentPosition.character) + completionLines[completionLines.length - 1].length ); - const newSelection = new vscode.Selection( - position, - newPosition - ); + const newSelection = new vscode.Selection(position, newPosition); currentPosition = newPosition; progress.report({ message: "Generating...", increment: 1 / (numPredict / 100) }); @@ -166,7 +147,6 @@ async function autocompleteCommand(textEditor: vscode.TextEditor, cancellationTo }); await finished; - // Remove any remaining FIM tokens from the completion completionText = completionText.replace(/||/g, ''); const finalEdit = new vscode.WorkspaceEdit(); finalEdit.replace(document.uri, new vscode.Range(position, currentPosition), completionText); @@ -186,52 +166,56 @@ async function provideCompletionItems(document: vscode.TextDocument, position: v const item = new vscode.CompletionItem("Fabelous autocompletion"); item.insertText = new vscode.SnippetString('${1:}'); - if (responsePreview) await new Promise(resolve => setTimeout(resolve, responsePreviewDelay * 1000)); - if (cancellationToken.isCancellationRequested) { - return [ item ]; - } - if (responsePreview) { + await new Promise(resolve => setTimeout(resolve, responsePreviewDelay * 1000)); + if (cancellationToken.isCancellationRequested) { + return [ item ]; + } + const context = getContextLines(document, position); const lines = context.split("\n"); - const currentLineIndex = position.line - Math.max(0, position.line - 10); + const currentLineIndex = position.line - Math.max(0, position.line - promptWindowSize / 2); const prefix = lines.slice(0, currentLineIndex + 1).join("\n"); const suffix = lines.slice(currentLineIndex + 1).join("\n"); - const fimPrompt = `${prefix}${suffix}`; + const fimPrompt = createFIMPrompt(prefix, suffix, document.languageId); - const response_preview = await axios.post(apiEndpoint, { - model: apiModel, - prompt: messageHeaderSub(document) + fimPrompt, - stream: false, - raw: true, - options: { - num_predict: responsePreviewMaxTokens, - temperature: apiTemperature, - stop: ['', '\n', '```'], - ...keepAlive && { keep_alive: keepAlive }, - ...topP && { top_p: topP }, + 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")); + }) + }); + + if (response_preview.data.response.trim() !== "") { + const previewText = response_preview.data.response.replace(/||/g, '').trimStart(); + item.label = previewText; + item.insertText = previewText; } - }, { - cancelToken: new axios.CancelToken((c) => { - const cancelPost = function () { - c("Autocompletion request terminated by completion cancel"); - }; - cancellationToken.onCancellationRequested(cancelPost); - }) - }); - if (response_preview.data.response.trim() != "") { - const previewText = response_preview.data.response.replace(/||/g, '').trimStart(); - item.label = previewText; - item.insertText = previewText; + } catch (error) { + console.error("Error fetching preview:", error); } } 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: [cancellationToken] - }; + if (continueInline || !responsePreview) { + item.command = { + command: 'fabelous-autocoder.autocomplete', + title: 'Fabelous Autocomplete', + arguments: [cancellationToken] + }; + } return [item]; }