Added timeout to token providing

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": {

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",

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 = '<fileStart>'
+const PREFIX_START = '<fileStart>';
-const PREFIX_END = '</fileStart>'
+const PREFIX_END = '</fileStart>';
 
-const SUFFIX_START = '<fileEnd>'
+const SUFFIX_START = '<fileEnd>';
-const SUFFIX_END = '</fileEnd>'
+const SUFFIX_END = '</fileEnd>';
+
+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 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 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 messageSuffix = `End of the file:\n${SUFFIX_START}\n${suffix}\n${SUFFIX_END}\n`;
-	const messagePrefix = `Start of the file:\n${PREFIX_START}`
+	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
+const tokenProvider = async (
-// Your extension is activated the very first time the command is executed
+	document: vscode.TextDocument,
-export function activate(context: vscode.ExtensionContext) {
+	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)
+	const response = await ollama.generate({
-	// This line of code will only be executed once when your extension is activated
+		model: MODEL,
-	console.log('Congratulations, your extension "ai-code" is now active!');
+		prompt,
-
+		suffix,
-	// The command has been defined in the package.json file
+		raw: true,
-	// Now provide the implementation of the command with registerCommand
+		stream: true,
-	// The commandId parameter must match the command field in package.json
+		options: {
-	const disposable = vscode.commands.registerCommand('ai-code.helloWorld', async () => {
+			num_predict: MAX_TOKENS,
-		// The code you place here will be executed every time your command is executed
+			stop: [PREFIX_END]
-		// 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)
-		// }
 	});
 
+	return response;
+};
+
+export const activate = (context: vscode.ExtensionContext) => {
+	console.log('"ai-code" extensions loaded');
+
 	const provider: vscode.InlineCompletionItemProvider = {
-		async provideInlineCompletionItems(document, position, _context, _token) {
+		async provideInlineCompletionItems(document, position, context, _token) {
-			console.log('provideInlineCompletionItems triggered');
-
 			try {
-				const modelSupportsSuffix = await getModelSupportsSuffix(MODEL)
+				const response = await tokenProvider(document, position, context, _token);
-				const prompt = modelSupportsSuffix ? getPrompt(document, position) : getPromptWithSuffix(document, position)
-				const suffix = modelSupportsSuffix ? undefined : getSuffix(document, position)
 
-				const response = await ollama.generate({
+				const resultBuffer: string[] = await new Promise(async (resolve, reject) => {
-					model: MODEL,
+					const buffer: string[] = [];
-					prompt,
+					const timeout = setTimeout(() => {
-					suffix,
+						resolve(buffer);
-					raw: true,
+					}, GENERATION_TIMEOUT);
-					stream: true,
+
-					options: {
+					try {
-						num_predict: 10,
+						for await (const part of response) {
-						stop: [PREFIX_END]
+							console.log(part.response);
-					},
+							buffer.push(part.response);
-				})
+						}
-	
+						resolve(buffer);
-				const buffer = []
+					} catch (err) {
-				for await (const part of response) {
+						reject(err);
-					process.stdout.write(part.response)
+					} finally {
-					buffer.push(part.response)
+						clearTimeout(timeout);
-				}
+					};
+				});
+
+				const text = resultBuffer.join('');
 
 				return [
 					{
-						insertText: buffer.join(''),
+						insertText: text,
 						range: new vscode.Range(position, position),
 					}
-				]
+				];
-			} catch (err) { 
+			} catch (err) {
-				console.log(err)
+				console.log(err);
 			}
 
-			return []
+			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;
 		},
 	};
 
 	vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, provider);
-
+};
-	context.subscriptions.push(disposable);
-}
 
 // This method is called when your extension is deactivated
 export function deactivate() {}

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<InlineCompletionItem[] | InlineCompletionList>;
+
+		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;
+	}
+}