used trie to store auto complete results

src/extension.ts

@@ -2,15 +2,12 @@
 // Import the module and reference it with the alias vscode in your code below
 import * as vscode from 'vscode';
 import { Ollama } from 'ollama/browser';
-import { trieInsert, trieLookup, TrieNode } from './trie';
+import { flattenTrie, trieInsert, trieLookup, TrieNode } from './trie';
 
 const MODEL = 'deepseek-coder:6.7b';
 
 const PREFIX_START = '<prefixStart>';
-const PREFIX_END = '<prefixEnd>';
-
-const SUFFIX_START = '<suffixStart>';
-const SUFFIX_END = '<suffixEnd>';
+const PREFIX_ENDS = ['<prefixEnd>', '</prefixEnd>'];
 
 const MAX_TOKENS = 50;
 const GENERATION_TIMEOUT = 200;
@@ -80,6 +77,14 @@ const tokenProvider = async (
 
 	const suffix = modelSupportsSuffix ? getSuffix(document, position) : undefined;
 
+	console.log(JSON.stringify(trieRoot));
+
+	const result = trieRootLookup(prefix);
+
+	if (result !== null) {
+		return flattenTrie(result);
+	}
+
 	const response = await ollama.generate({
 		model: MODEL,
 		prompt,
@@ -88,21 +93,10 @@ const tokenProvider = async (
 		stream: true,
 		options: {
 			num_predict: MAX_TOKENS,
-			stop: [PREFIX_END]
+			stop: PREFIX_ENDS,
 		},
 	});
 
-	return response;
-};
-
-export const activate = (context: vscode.ExtensionContext) => {
-	console.log('"ai-code" extensions loaded');
-
-	const provider: vscode.InlineCompletionItemProvider = {
-		async provideInlineCompletionItems(document, position, context, token) {
-			try {
-				const response = await tokenProvider(document, position, context, token);
-
 	const resultBuffer: string[] = await new Promise(async (resolve, reject) => {
 		const buffer: string[] = [];
 		const timeout = setTimeout(() => {
@@ -111,8 +105,8 @@ export const activate = (context: vscode.ExtensionContext) => {
 
 		try {
 			for await (const part of response) {
-							// process.stdout.write(part.response);
 				buffer.push(part.response);
+				trieRootInsert(prefix + buffer.join(''));
 			}
 			resolve(buffer);
 		} catch (err) {
@@ -123,14 +117,23 @@ export const activate = (context: vscode.ExtensionContext) => {
 		};
 	});
 
-				const text = resultBuffer.join('');
-
 	return [
-					{
+		resultBuffer.join('')
+	];
+};
+
+export const activate = (context: vscode.ExtensionContext) => {
+	console.log('"ai-code" extensions loaded');
+
+	const provider: vscode.InlineCompletionItemProvider = {
+		async provideInlineCompletionItems(document, position, context, token) {
+			try {
+				const completions = await tokenProvider(document, position, context, token);
+
+				return completions.map((text) => ({
 					insertText: text,
 					range: new vscode.Range(position, position),
-					}
-				];
+				}));
 			} catch (err) {
 				console.log(err);
 			}

src/trie.ts

@@ -132,3 +132,16 @@ export const trieLookup = (node: TrieNode, text: string): TrieNode | null => {
 	const child = node.children[childKey];
 	return child === undefined ? null : trieLookup(child, childText);
 };
+
+
+export const flattenTrie = (node: TrieNode): string[] => {
+	if (node.isLeaf) {
+		return [
+			node.value
+		];
+	}
+
+	return Object.entries(node.children).flatMap(([key, child]) => {
+		return flattenTrie(child).map((value) => node.value + key + value);
+	});
+};