moved predictTokens into its own function

This commit is contained in:
Leyla Becker 2025-08-04 18:14:40 -05:00
parent fbc588368e
commit 59537fde51

View file

@ -66,15 +66,13 @@ const trieRootPrune = (text: string) => {
return triePrune(trieRoot, text);
};
const tokenProvider = async (
const predictTokens = async (
extension: ExtensionState,
document: vscode.TextDocument,
position: vscode.Position,
_context: vscode.InlineCompletionContext,
token: vscode.CancellationToken,
prefix: string,
) => {
const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character));
const model = extension.configuration.inlineCompletion.model;
const modelSupportsSuffix = await getModelSupportsSuffix(extension, model);
@ -82,16 +80,6 @@ const tokenProvider = async (
const suffix = modelSupportsSuffix ? getSuffix(extension, document, position) : undefined;
const result = trieRootLookup(prefix);
if (result !== null) {
return flattenTrie(result);
}
if (token.isCancellationRequested) {
return [];
}
const response = await extension.ollama.generate({
model,
prompt,
@ -140,6 +128,34 @@ const tokenProvider = async (
];
};
const tokenProvider = async (
extension: ExtensionState,
document: vscode.TextDocument,
position: vscode.Position,
_context: vscode.InlineCompletionContext,
token: vscode.CancellationToken,
) => {
const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character));
const result = trieRootLookup(prefix);
if (result !== null) {
return flattenTrie(result);
}
if (token.isCancellationRequested) {
return [];
}
return predictTokens(
extension,
document,
position,
token,
prefix,
);
};
export const getAutoCompleteProvider = (extension: ExtensionState) => {
const provider: vscode.InlineCompletionItemProvider = {
async provideInlineCompletionItems(document, position, context, token) {