From cf997dc2be676f29a46df22f89a096a0fcfbe3e4 Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sat, 19 Jul 2025 11:15:46 -0500 Subject: [PATCH] Added timeout to token providing --- package-lock.json | 2 +- package.json | 7 +- src/extension.ts | 195 +++++++----------- ...e.proposed.inlineCompletionsAdditions.d.ts | 130 ++++++++++++ 4 files changed, 214 insertions(+), 120 deletions(-) create mode 100644 vscode.proposed.inlineCompletionsAdditions.d.ts diff --git a/package-lock.json b/package-lock.json index 4f6bfb0..a83d43e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -25,7 +25,7 @@ "webpack-cli": "^6.0.1" }, "engines": { - "vscode": "^1.102.0" + "vscode": "^1.101.0" } }, "node_modules/@bcoe/v8-coverage": { diff --git a/package.json b/package.json index 4741386..651ba27 100644 --- a/package.json +++ b/package.json @@ -9,7 +9,9 @@ "categories": [ "Other" ], - "activationEvents": [], + "activationEvents": [ + "*" + ], "main": "./dist/extension.js", "contributes": { "commands": [ @@ -19,6 +21,9 @@ } ] }, + "enabledApiProposals": [ + "inlineCompletionsAdditions" + ], "scripts": { "vscode:prepublish": "npm run package", "compile": "webpack", diff --git a/src/extension.ts b/src/extension.ts index 884ee20..da73830 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -3,13 +3,16 @@ import * as vscode from 'vscode'; import { Ollama } from 'ollama/browser'; -const MODEL = 'deepseek-coder:1.3b'; +const MODEL = 'deepseek-coder:6.7b'; -const PREFIX_START = '' -const PREFIX_END = '' +const PREFIX_START = ''; +const PREFIX_END = ''; -const SUFFIX_START = '' -const SUFFIX_END = '' +const SUFFIX_START = ''; +const SUFFIX_END = ''; + +const MAX_TOKENS = 50; +const GENERATION_TIMEOUT = 200; const HOST = undefined; @@ -26,11 +29,11 @@ const getModelSupportsSuffix = async (model: string) => { // }) // model.capabilities.includes('suffix') - return false -} + return false; +}; const getPrompt = (document: vscode.TextDocument, position: vscode.Position) => { - const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character)) + const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character)); const messageHeader = `In an english code base with the file.\nfile:\nproject {PROJECT_NAME}\nfile {FILE_NAME}\nlanguage {LANG}\nFile:\n${PREFIX_START}\n`; @@ -39,146 +42,102 @@ const getPrompt = (document: vscode.TextDocument, position: vscode.Position) => .replace("{FILE_NAME}", document.fileName) .replace("{LANG}", document.languageId) + prefix; - return prompt -} + return prompt; +}; const getPromptWithSuffix = (document: vscode.TextDocument, position: vscode.Position) => { - const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character)) - const suffix = document.getText(new vscode.Range(position.line, position.character, document.lineCount - 1, document.lineAt(document.lineCount - 1).text.length)) + const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character)); + const suffix = document.getText(new vscode.Range(position.line, position.character, document.lineCount - 1, document.lineAt(document.lineCount - 1).text.length)); - const messageSuffix = `End of the file:\n${SUFFIX_START}\n${suffix}\n${SUFFIX_END}\n` - const messagePrefix = `Start of the file:\n${PREFIX_START}` + const messageSuffix = `End of the file:\n${SUFFIX_START}\n${suffix}\n${SUFFIX_END}\n`; + const messagePrefix = `Start of the file:\n${PREFIX_START}`; const messageHeader = `In an english code base with the file.\nfile:\nproject {PROJECT_NAME}\nfile {FILE_NAME}\nlanguage {LANG}\nThis is the end of and then the start of the file.` .replace("{PROJECT_NAME}", vscode.workspace.name || "Untitled") .replace("{FILE_NAME}", document.fileName) - .replace("{LANG}", document.languageId) + prefix; + .replace("{LANG}", document.languageId); - const prompt = `${messageHeader}\n${messageSuffix}\n${messagePrefix}\n`; + const prompt = `${messageHeader}\n${messageSuffix}\n${messagePrefix}\n${prefix}`; - return prompt -} + return prompt; +}; const getSuffix = (document: vscode.TextDocument, position: vscode.Position) => { - const suffix = document.getText(new vscode.Range(position.line, position.character, document.lineCount - 1, document.lineAt(document.lineCount - 1).text.length)) + const suffix = document.getText(new vscode.Range(position.line, position.character, document.lineCount - 1, document.lineAt(document.lineCount - 1).text.length)); - return suffix -} + return suffix; +}; -// This method is called when your extension is activated -// Your extension is activated the very first time the command is executed -export function activate(context: vscode.ExtensionContext) { +const tokenProvider = async ( + document: vscode.TextDocument, + position: vscode.Position, + context: vscode.InlineCompletionContext, + _token: vscode.CancellationToken, +) => { + const modelSupportsSuffix = await getModelSupportsSuffix(MODEL); + const prompt = modelSupportsSuffix ? getPrompt(document, position) : getPromptWithSuffix(document, position); + const suffix = modelSupportsSuffix ? getSuffix(document, position) : undefined; - // Use the console to output diagnostic information (console.log) and errors (console.error) - // This line of code will only be executed once when your extension is activated - console.log('Congratulations, your extension "ai-code" is now active!'); - - // The command has been defined in the package.json file - // Now provide the implementation of the command with registerCommand - // The commandId parameter must match the command field in package.json - const disposable = vscode.commands.registerCommand('ai-code.helloWorld', async () => { - // The code you place here will be executed every time your command is executed - // Display a message box to the user - // vscode.window.showInformationMessage('Hello world!'); - // try { - // vscode.window.showInformationMessage("asking ollama"); - // const response = await ollama.chat({ - // model: 'deepseek-coder:1.3b', - // messages: [{ role: 'user', content: 'Why is the sky blue?' }], - // stream: true, - // }); - // for await (const part of response) { - // vscode.window.showInformationMessage(part.message.content); - // } - // // vscode.window.showInformationMessage(response.message.content); - // } - // catch (err) { - // console.log(err) - // } + const response = await ollama.generate({ + model: MODEL, + prompt, + suffix, + raw: true, + stream: true, + options: { + num_predict: MAX_TOKENS, + stop: [PREFIX_END] + }, }); + return response; +}; + +export const activate = (context: vscode.ExtensionContext) => { + console.log('"ai-code" extensions loaded'); + const provider: vscode.InlineCompletionItemProvider = { - async provideInlineCompletionItems(document, position, _context, _token) { - console.log('provideInlineCompletionItems triggered'); - + async provideInlineCompletionItems(document, position, context, _token) { try { - const modelSupportsSuffix = await getModelSupportsSuffix(MODEL) - const prompt = modelSupportsSuffix ? getPrompt(document, position) : getPromptWithSuffix(document, position) - const suffix = modelSupportsSuffix ? undefined : getSuffix(document, position) + const response = await tokenProvider(document, position, context, _token); - const response = await ollama.generate({ - model: MODEL, - prompt, - suffix, - raw: true, - stream: true, - options: { - num_predict: 10, - stop: [PREFIX_END] - }, - }) - - const buffer = [] - for await (const part of response) { - process.stdout.write(part.response) - buffer.push(part.response) - } + const resultBuffer: string[] = await new Promise(async (resolve, reject) => { + const buffer: string[] = []; + const timeout = setTimeout(() => { + resolve(buffer); + }, GENERATION_TIMEOUT); + + try { + for await (const part of response) { + console.log(part.response); + buffer.push(part.response); + } + resolve(buffer); + } catch (err) { + reject(err); + } finally { + clearTimeout(timeout); + }; + }); + + const text = resultBuffer.join(''); return [ { - insertText: buffer.join(''), + insertText: text, range: new vscode.Range(position, position), } - ] - } catch (err) { - console.log(err) + ]; + } catch (err) { + console.log(err); } - return [] - - // const regexp = /\/\/ \[(.+?),(.+?)\)(.*?):(.*)/; - // if (position.line <= 0) { - // return; - // } - - // const result: vscode.InlineCompletionItem[] = [] - - // let offset = 1; - // while (offset > 0) { - // if (position.line - offset < 0) { - // break; - // } - - // const lineBefore = document.lineAt(position.line - offset).text; - // const matches = lineBefore.match(regexp); - // if (!matches) { - // break; - // } - // offset++; - - // const start = matches[1]; - // const startInt = parseInt(start, 10); - // const end = matches[2]; - // const endInt = - // end === '*' - // ? document.lineAt(position.line).text.length - // : parseInt(end, 10); - // const text = matches[4].replace(/\\n/g, '\n'); - - // result.push({ - // insertText: text, - // range: new vscode.Range(position.line, startInt, position.line, endInt), - // }); - // } - - // return result; + return []; }, }; vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, provider); - - context.subscriptions.push(disposable); -} +}; // This method is called when your extension is deactivated export function deactivate() {} diff --git a/vscode.proposed.inlineCompletionsAdditions.d.ts b/vscode.proposed.inlineCompletionsAdditions.d.ts new file mode 100644 index 0000000..8eb935c --- /dev/null +++ b/vscode.proposed.inlineCompletionsAdditions.d.ts @@ -0,0 +1,130 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +declare module 'vscode' { + + // https://github.com/microsoft/vscode/issues/124024 @hediet + + export namespace languages { + /** + * Registers an inline completion provider. + * + * Multiple providers can be registered for a language. In that case providers are asked in + * parallel and the results are merged. A failing provider (rejected promise or exception) will + * not cause a failure of the whole operation. + * + * @param selector A selector that defines the documents this provider is applicable to. + * @param provider An inline completion provider. + * @param metadata Metadata about the provider. + * @return A {@link Disposable} that unregisters this provider when being disposed. + */ + export function registerInlineCompletionItemProvider(selector: DocumentSelector, provider: InlineCompletionItemProvider, metadata: InlineCompletionItemProviderMetadata): Disposable; + } + + export interface InlineCompletionItem { + /** + * If set to `true`, unopened closing brackets are removed and unclosed opening brackets are closed. + * Defaults to `false`. + */ + completeBracketPairs?: boolean; + + warning?: InlineCompletionWarning; + + /** If set to `true`, this item is treated as inline edit. */ + isInlineEdit?: boolean; + + /** + * A range specifying when the edit can be shown based on the cursor position. + * If the cursor is within this range, the inline edit can be displayed. + */ + showRange?: Range; + + showInlineEditMenu?: boolean; + + action?: Command; + } + + export interface InlineCompletionWarning { + message: MarkdownString | string; + icon?: ThemeIcon; + } + + export interface InlineCompletionItemProviderMetadata { + /** + * Specifies a list of extension ids that this provider yields to if they return a result. + * If some inline completion provider registered by such an extension returns a result, this provider is not asked. + */ + yieldTo?: string[]; + + debounceDelayMs?: number; + + displayName?: string; + } + + export interface InlineCompletionItemProvider { + /** + * @param completionItem The completion item that was shown. + * @param updatedInsertText The actual insert text (after brackets were fixed). + */ + handleDidShowCompletionItem?(completionItem: InlineCompletionItem, updatedInsertText: string): void; + + /** + * @param completionItem The completion item that was rejected. + */ + handleDidRejectCompletionItem?(completionItem: InlineCompletionItem): void; + + /** + * Is called when an inline completion item was accepted partially. + * @param acceptedLength The length of the substring of the inline completion that was accepted already. + * @deprecated Use `handleDidPartiallyAcceptCompletionItem` with `PartialAcceptInfo` instead. + */ + handleDidPartiallyAcceptCompletionItem?(completionItem: InlineCompletionItem, acceptedLength: number): void; + + /** + * Is called when an inline completion item was accepted partially. + * @param info Additional info for the partial accepted trigger. + */ + handleDidPartiallyAcceptCompletionItem?(completionItem: InlineCompletionItem, info: PartialAcceptInfo): void; + + provideInlineEditsForRange?(document: TextDocument, range: Range, context: InlineCompletionContext, token: CancellationToken): ProviderResult; + + readonly debounceDelayMs?: number; + } + + export interface InlineCompletionContext { + readonly userPrompt?: string; + + readonly requestUuid?: string; + } + + export interface PartialAcceptInfo { + kind: PartialAcceptTriggerKind; + /** + * The length of the substring of the provided inline completion text that was accepted already. + */ + acceptedLength: number; + } + + export enum PartialAcceptTriggerKind { + Unknown = 0, + Word = 1, + Line = 2, + Suggest = 3, + } + + // When finalizing `commands`, make sure to add a corresponding constructor parameter. + export interface InlineCompletionList { + /** + * A list of commands associated with the inline completions of this list. + */ + commands?: Command[]; + + /** + * When set and the user types a suggestion without derivating from it, the inline suggestion is not updated. + * Defaults to false (might change). + */ + enableForwardStability?: boolean; + } +}