Added timeout to token providing
This commit is contained in:
parent
afd26b271c
commit
cf997dc2be
4 changed files with 214 additions and 120 deletions
195
src/extension.ts
195
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_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 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() {}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue