jan-leila/ai-code
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Added timeout to token providing

Browse source
This commit is contained in:
Leyla Becker 2025-07-19 11:15:46 -05:00
parent afd26b271c
commit cf997dc2be
4 changed files with 214 additions and 120 deletions
Download patch file Download diff file Expand all files Collapse all files

2
package-lock.json generated
View file

@ -25,7 +25,7 @@
"webpack-cli": "^6.0.1"
},
"engines": {
"vscode": "^1.102.0"
"vscode": "^1.101.0"
}
},
"node_modules/@bcoe/v8-coverage": {

7
package.json
View file

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

191
src/extension.ts
View file

@ -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_END = '</fileStart>'
const PREFIX_START = '<fileStart>';
const PREFIX_END = '</fileStart>';
const SUFFIX_START = '<fileEnd>'
const SUFFIX_END = '</fileEnd>'
const SUFFIX_START = '<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 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 resultBuffer: string[] = await new Promise(async (resolve, reject) => {
const buffer: string[] = [];
const timeout = setTimeout(() => {
resolve(buffer);
}, GENERATION_TIMEOUT);
const buffer = []
for await (const part of response) {
process.stdout.write(part.response)
buffer.push(part.response)
}
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)
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() {}

130
vscode.proposed.inlineCompletionsAdditions.d.ts vendored Normal file
View file

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