From ffd384fbe44c020705924f04ed6d09c0f9f6238c Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sat, 19 Jul 2025 18:57:11 -0500 Subject: [PATCH] created trie --- package.json | 7 -- src/extension.ts | 46 ++++++------- src/test/extension.test.ts | 15 ----- src/trie.ts | 134 +++++++++++++++++++++++++++++++++++++ 4 files changed, 157 insertions(+), 45 deletions(-) delete mode 100644 src/test/extension.test.ts create mode 100644 src/trie.ts diff --git a/package.json b/package.json index 651ba27..ab8165a 100644 --- a/package.json +++ b/package.json @@ -15,15 +15,8 @@ "main": "./dist/extension.js", "contributes": { "commands": [ - { - "command": "ai-code.helloWorld", - "title": "Hello World" - } ] }, - "enabledApiProposals": [ - "inlineCompletionsAdditions" - ], "scripts": { "vscode:prepublish": "npm run package", "compile": "webpack", diff --git a/src/extension.ts b/src/extension.ts index 90ddc82..f66ff54 100644 --- a/src/extension.ts +++ b/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); }; }); diff --git a/src/test/extension.test.ts b/src/test/extension.test.ts deleted file mode 100644 index 4ca0ab4..0000000 --- a/src/test/extension.test.ts +++ /dev/null @@ -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)); - }); -}); diff --git a/src/trie.ts b/src/trie.ts new file mode 100644 index 0000000..3834544 --- /dev/null +++ b/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); +};