created trie

src/extension.ts

@@ -2,6 +2,7 @@
 // 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';
 
 const MODEL = 'deepseek-coder:6.7b';
 
@@ -32,9 +33,7 @@ const getModelSupportsSuffix = async (model: string) => {
 	return false;
 };
 
-const getPrompt = (document: vscode.TextDocument, position: vscode.Position) => {
-	const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character));
-
+const getPrompt = (document: vscode.TextDocument, position: vscode.Position, prefix: string) => {
 	const messageHeader = `In an english code base with the file.\nfile:\nproject {PROJECT_NAME}\nfile {FILE_NAME}\nlanguage {LANG}`
 		.replace("{PROJECT_NAME}", vscode.workspace.name || "Untitled")
 		.replace("{FILE_NAME}", document.fileName)
@@ -48,37 +47,37 @@ const getPrompt = (document: vscode.TextDocument, position: vscode.Position) =>
 	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 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}\n.`
-		.replace("{PROJECT_NAME}", vscode.workspace.name || "Untitled")
-		.replace("{FILE_NAME}", document.fileName)
-		.replace("{LANG}", document.languageId);
-
-	const prompt = `${messageHeader}\n${messageSuffix}\n${messagePrefix}\n${prefix}`;
-
-	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));
 
 	return suffix;
 };
 
+let trieRoot: TrieNode = {
+	isLeaf: false,
+	value: '',
+	children: {},
+};
+
+const trieRootInsert = (text: string) => {
+	trieRoot = trieInsert(trieRoot, text);
+};
+
+const trieRootLookup = (text: string) => {
+	return trieLookup(trieRoot, text);
+};
+
 const tokenProvider = async (
 	document: vscode.TextDocument,
 	position: vscode.Position,
-	context: vscode.InlineCompletionContext,
+	_context: vscode.InlineCompletionContext,
 	_token: vscode.CancellationToken,
 ) => {
+	const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character));
+	
 	const modelSupportsSuffix = await getModelSupportsSuffix(MODEL);
-	const prompt = modelSupportsSuffix ? getPrompt(document, position) : getPromptWithSuffix(document, position);
+	const prompt = getPrompt(document, position, prefix);
+
 	const suffix = modelSupportsSuffix ? getSuffix(document, position) : undefined;
 
 	const response = await ollama.generate({
@@ -112,13 +111,14 @@ export const activate = (context: vscode.ExtensionContext) => {
 
 					try {
 						for await (const part of response) {
-							console.log(part.response);
+							// process.stdout.write(part.response);
 							buffer.push(part.response);
 						}
 						resolve(buffer);
 					} catch (err) {
 						reject(err);
 					} finally {
+						response.abort();
 						clearTimeout(timeout);
 					};
 				});

src/test/extension.test.ts

@@ -1,15 +0,0 @@
-import * as assert from 'assert';
-
-// You can import and use all API from the 'vscode' module
-// as well as import your extension to test it
-import * as vscode from 'vscode';
-// import * as myExtension from '../../extension';
-
-suite('Extension Test Suite', () => {
-	vscode.window.showInformationMessage('Start all tests.');
-
-	test('Sample test', () => {
-		assert.strictEqual(-1, [1, 2, 3].indexOf(5));
-		assert.strictEqual(-1, [1, 2, 3].indexOf(0));
-	});
-});

src/trie.ts

@@ -0,0 +1,134 @@
+
+interface TrieLeaf {
+	isLeaf: true
+	value: string
+	children?: never
+}
+
+interface TrieBranch {
+	isLeaf: false
+	value: string
+	children: { [key: string]: TrieNode }
+}
+
+export type TrieNode = TrieLeaf | TrieBranch
+
+/**
+ * Creates a new TrieNode based on node that has text added to it
+ * @param node node that we are basing the update on
+ * @param text text that is being added to the node
+ * @returns a new node with text added to it
+ */
+export const trieInsert = (node: TrieNode, text: string): TrieNode => {
+	// TODO: mutate node to add text to it
+	for (let index = 0; index < text.length; index++) {
+		// If the inserted text is longer then the nodes text update the node with the new text
+		if (index >= node.value.length) {
+			// If the current node is a leaf we can just replace it with a larger leaf
+			if (node.isLeaf) {
+				const newLeaf: TrieLeaf = {
+					isLeaf: true,
+					value: text,
+				};
+				return newLeaf;
+			}
+
+			// If the current node is a branch then we need add the remaining text to one of its children
+			const childKey = text[index];
+			const childText = text.substring(index + 1);
+			const child = node.children[childKey];
+
+			const newBranch: TrieBranch = {
+				isLeaf: false,
+				value: node.value,
+				children: {
+					...node.children,
+					[childKey]: child === undefined ? {
+						isLeaf: true,
+						value: childText,
+					} : trieInsert(child, childText)
+				},
+			};
+			return newBranch;
+		}
+
+		// If our inserted text does not match the node then we need to split the node
+		if (node.value[index] !== text[index]) {
+			// If the node is a leaf we need to split it into a branch
+			if (node.isLeaf) {
+				const newBranch: TrieBranch = {
+					isLeaf: false,
+					value: text.substring(0, index),
+					children: {
+						[text[index]]: {
+							isLeaf: true,
+							value: text.substring(index + 1),
+						},
+						[node.value[index]]: {
+							isLeaf: true,
+							value: node.value.substring(index + 1),
+						},
+					},
+				};
+				return newBranch;
+			}
+
+			// If the node is a branch then we need to create a new leaf on it
+			const newBranch: TrieBranch = {
+				isLeaf: false,
+				value: text.substring(0, index),
+				children: {
+					[text[index]]: {
+						isLeaf: true,
+						value: text.substring(index + 1),
+					},
+					[node.value[index]]: {
+						isLeaf: false,
+						value: node.value.substring(index + 1),
+						children: node.children,
+					},
+				},
+			};
+			return newBranch;
+		}
+	}
+
+	return node;
+};
+
+export const trieLookup = (node: TrieNode, text: string): TrieNode | null => {
+	for (let index = 0; index < node.value.length; index++) {
+		// If our node still has text left but we have no more search query return a node starting at where we ran out of characters
+		if (index >= text.length) {
+			if (node.isLeaf) {
+				const newNode: TrieLeaf = {
+					isLeaf: true,
+					value: node.value.substring(index),
+				};
+				return newNode;
+			}
+			const newNode: TrieBranch = {
+				isLeaf: false,
+				value: node.value.substring(index),
+				children: node.children,
+			};
+			return newNode;
+		}
+
+		// If we have a difference then there is no match
+		if (node.value[index] !== text[index]) {
+			return null;
+		}
+	}
+
+	// If we get past the end of the node and it is a leaf then there is no match
+	if (node.isLeaf) {
+		return null;
+	}
+
+	// Continue matching on the child node
+	const childKey = text[node.value.length];
+	const childText = text.substring(node.value.length + 1);
+	const child = node.children[childKey];
+	return child === undefined ? null : trieLookup(child, childText);
+};