From 59537fde51b9b8b09dc27b19352e979e407d57fc Mon Sep 17 00:00:00 2001
From: Leyla Becker <git@jan-leila.com>
Date: Mon, 4 Aug 2025 18:14:40 -0500
Subject: [PATCH] moved predictTokens into its own function

---
 src/autoComplete.ts | 44 ++++++++++++++++++++++++++++++--------------
 1 file changed, 30 insertions(+), 14 deletions(-)

diff --git a/src/autoComplete.ts b/src/autoComplete.ts
index e750865..a7d3ef8 100644
--- a/src/autoComplete.ts
+++ b/src/autoComplete.ts
@@ -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) {