From ffd384fbe44c020705924f04ed6d09c0f9f6238c Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sat, 19 Jul 2025 18:57:11 -0500 Subject: [PATCH 01/10] 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); +}; From 55c72ff98ca1cb348dd9131fd577c01251a225a1 Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sat, 19 Jul 2025 19:38:41 -0500 Subject: [PATCH 02/10] used trie to store auto complete results --- src/extension.ts | 79 +++++++++++++++++++++++++----------------------- src/trie.ts | 13 ++++++++ 2 files changed, 54 insertions(+), 38 deletions(-) diff --git a/src/extension.ts b/src/extension.ts index f66ff54..02cfdb1 100644 --- a/src/extension.ts +++ b/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 = ''; -const PREFIX_END = ''; - -const SUFFIX_START = ''; -const SUFFIX_END = ''; +const PREFIX_ENDS = ['', '']; const MAX_TOKENS = 50; const GENERATION_TIMEOUT = 200; @@ -74,12 +71,20 @@ const tokenProvider = async ( _token: vscode.CancellationToken, ) => { const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character)); - + const modelSupportsSuffix = await getModelSupportsSuffix(MODEL); const prompt = getPrompt(document, position, prefix); 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,11 +93,33 @@ const tokenProvider = async ( stream: true, options: { num_predict: MAX_TOKENS, - stop: [PREFIX_END] + stop: PREFIX_ENDS, }, }); - return response; + const resultBuffer: string[] = await new Promise(async (resolve, reject) => { + const buffer: string[] = []; + const timeout = setTimeout(() => { + resolve(buffer); + }, GENERATION_TIMEOUT); + + try { + for await (const part of response) { + buffer.push(part.response); + trieRootInsert(prefix + buffer.join('')); + } + resolve(buffer); + } catch (err) { + reject(err); + } finally { + response.abort(); + clearTimeout(timeout); + }; + }); + + return [ + resultBuffer.join('') + ]; }; export const activate = (context: vscode.ExtensionContext) => { @@ -101,36 +128,12 @@ export const activate = (context: vscode.ExtensionContext) => { const provider: vscode.InlineCompletionItemProvider = { async provideInlineCompletionItems(document, position, context, token) { try { - const response = await tokenProvider(document, position, context, token); + const completions = await tokenProvider(document, position, context, token); - const resultBuffer: string[] = await new Promise(async (resolve, reject) => { - const buffer: string[] = []; - const timeout = setTimeout(() => { - resolve(buffer); - }, GENERATION_TIMEOUT); - - try { - for await (const part of response) { - // process.stdout.write(part.response); - buffer.push(part.response); - } - resolve(buffer); - } catch (err) { - reject(err); - } finally { - response.abort(); - clearTimeout(timeout); - }; - }); - - const text = resultBuffer.join(''); - - return [ - { - insertText: text, - range: new vscode.Range(position, position), - } - ]; + return completions.map((text) => ({ + insertText: text, + range: new vscode.Range(position, position), + })); } catch (err) { console.log(err); } @@ -143,4 +146,4 @@ export const activate = (context: vscode.ExtensionContext) => { }; // This method is called when your extension is deactivated -export function deactivate() {} +export function deactivate() { } diff --git a/src/trie.ts b/src/trie.ts index 3834544..678df44 100644 --- a/src/trie.ts +++ b/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); + }); +}; From 02e65c5700e409ee02f718cf39cb49d75c749590 Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sat, 19 Jul 2025 19:45:24 -0500 Subject: [PATCH 03/10] added back suffix support --- src/extension.ts | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/src/extension.ts b/src/extension.ts index 02cfdb1..62afc80 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -7,7 +7,10 @@ import { flattenTrie, trieInsert, trieLookup, TrieNode } from './trie'; const MODEL = 'deepseek-coder:6.7b'; const PREFIX_START = ''; -const PREFIX_ENDS = ['', '']; +const PREFIX_ENDS = ['', '', '', '']; + +const SUFFIX_START = ''; +const SUFFIX_END = ''; const MAX_TOKENS = 50; const GENERATION_TIMEOUT = 200; @@ -44,6 +47,22 @@ const getPrompt = (document: vscode.TextDocument, position: vscode.Position, pre return prompt; }; +const getPromptWithSuffix = (document: vscode.TextDocument, position: vscode.Position, prefix: string) => { + 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)); @@ -73,12 +92,10 @@ const tokenProvider = async ( const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character)); const modelSupportsSuffix = await getModelSupportsSuffix(MODEL); - const prompt = getPrompt(document, position, prefix); + const prompt = modelSupportsSuffix ? getPrompt(document, position, prefix) : getPromptWithSuffix(document, position, prefix); const suffix = modelSupportsSuffix ? getSuffix(document, position) : undefined; - console.log(JSON.stringify(trieRoot)); - const result = trieRootLookup(prefix); if (result !== null) { From f93d24dc39c618916dc93742fc405f2f04d9356f Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sat, 19 Jul 2025 19:50:34 -0500 Subject: [PATCH 04/10] canceled ollama requests on token canceled --- src/extension.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/extension.ts b/src/extension.ts index 62afc80..31e681e 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -87,7 +87,7 @@ const tokenProvider = async ( document: vscode.TextDocument, position: vscode.Position, _context: vscode.InlineCompletionContext, - _token: vscode.CancellationToken, + token: vscode.CancellationToken, ) => { const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character)); @@ -114,6 +114,10 @@ const tokenProvider = async ( }, }); + token.onCancellationRequested(() => { + response.abort(); + }); + const resultBuffer: string[] = await new Promise(async (resolve, reject) => { const buffer: string[] = []; const timeout = setTimeout(() => { From a620903c105a1424fd20b06167bbaee275bd08d6 Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sat, 19 Jul 2025 20:14:21 -0500 Subject: [PATCH 05/10] added pruning to trie --- src/extension.ts | 23 +++++++++++++++++++---- src/trie.ts | 30 +++++++++++++++++++++++++++++- 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/src/extension.ts b/src/extension.ts index 31e681e..3f3133f 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -2,7 +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 { flattenTrie, trieInsert, trieLookup, TrieNode } from './trie'; +import { flattenTrie, trieInsert, trieLookup, TrieNode, triePrune } from './trie'; const MODEL = 'deepseek-coder:6.7b'; @@ -14,6 +14,7 @@ const SUFFIX_END = ''; const MAX_TOKENS = 50; const GENERATION_TIMEOUT = 200; +const TRIE_PRUNE_TIMEOUT = 10000; const HOST = undefined; @@ -70,9 +71,8 @@ const getSuffix = (document: vscode.TextDocument, position: vscode.Position) => }; let trieRoot: TrieNode = { - isLeaf: false, + isLeaf: true, value: '', - children: {}, }; const trieRootInsert = (text: string) => { @@ -83,6 +83,10 @@ const trieRootLookup = (text: string) => { return trieLookup(trieRoot, text); }; +const trieRootPrune = (text: string) => { + return triePrune(trieRoot, text); +}; + const tokenProvider = async ( document: vscode.TextDocument, position: vscode.Position, @@ -102,6 +106,10 @@ const tokenProvider = async ( return flattenTrie(result); } + if (token.isCancellationRequested) { + return []; + } + const response = await ollama.generate({ model: MODEL, prompt, @@ -114,8 +122,15 @@ const tokenProvider = async ( }, }); + const pruneTimeout = setTimeout(() => { + trieRootPrune(prompt); + }, TRIE_PRUNE_TIMEOUT); + token.onCancellationRequested(() => { - response.abort(); + clearTimeout(pruneTimeout); + try { + response.abort(); + } catch {} }); const resultBuffer: string[] = await new Promise(async (resolve, reject) => { diff --git a/src/trie.ts b/src/trie.ts index 678df44..3ed2ed9 100644 --- a/src/trie.ts +++ b/src/trie.ts @@ -96,6 +96,12 @@ export const trieInsert = (node: TrieNode, text: string): TrieNode => { return node; }; +/** + * Gets a new trie that is the node trie after the text + * @param node + * @param text + * @returns + */ 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 @@ -133,7 +139,6 @@ export const trieLookup = (node: TrieNode, text: string): TrieNode | null => { return child === undefined ? null : trieLookup(child, childText); }; - export const flattenTrie = (node: TrieNode): string[] => { if (node.isLeaf) { return [ @@ -145,3 +150,26 @@ export const flattenTrie = (node: TrieNode): string[] => { return flattenTrie(child).map((value) => node.value + key + value); }); }; + +export const triePrune = (node: TrieNode, text: string): TrieNode => { + const value = trieLookup(node, text); + if (value === null) { + return { + isLeaf: true, + value: text, + }; + } + + if (value.isLeaf) { + return { + isLeaf: true, + value: text + value.value, + }; + } + + return { + isLeaf: false, + value: text + value.value, + children: value.children, + }; +}; From 7fd842dbc0bac62d115c0149a86be71dfd8d41e8 Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sat, 19 Jul 2025 22:00:06 -0500 Subject: [PATCH 06/10] made settings configurable --- src/extension.ts | 130 +++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 108 insertions(+), 22 deletions(-) diff --git a/src/extension.ts b/src/extension.ts index 3f3133f..05632e1 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -16,14 +16,95 @@ const MAX_TOKENS = 50; const GENERATION_TIMEOUT = 200; const TRIE_PRUNE_TIMEOUT = 10000; -const HOST = undefined; +interface ExtensionConfiguration { + ollamaHost: string | undefined + inlineCompletion: { + model: string + prefixStart: string + prefixEnds: string[] + suffixStart: string + suffixEnd: string + maxTokens: number + generationTimeout: number + triePruneTimeout: number + } +}; -// TODO: make these configurable by extension setting -const ollama = new Ollama({ - host: HOST, -}); +interface ExtensionState { + configuration: ExtensionConfiguration + ollama: Ollama +} -const getModelSupportsSuffix = async (model: string) => { +const getExtensionState = (): ExtensionState => { + const extensionConfig = vscode.workspace.getConfiguration('ai-code'); + + const configuration: ExtensionConfiguration = { + ollamaHost: extensionConfig.get('ollamaHost'), + inlineCompletion: { + model: extensionConfig.get('inlineCompletion.prefixStart') ?? MODEL, + prefixStart: extensionConfig.get('inlineCompletion.prefixStart') ?? PREFIX_START, + prefixEnds: extensionConfig.get('inlineCompletion.prefixEnd')?.split(',') ?? PREFIX_ENDS, + suffixStart: extensionConfig.get('inlineCompletion.suffixStart') ?? SUFFIX_START, + suffixEnd: extensionConfig.get('inlineCompletion.suffixEnd') ?? SUFFIX_END, + maxTokens: extensionConfig.get('inlineCompletion.maxTokens') ?? MAX_TOKENS, + generationTimeout: extensionConfig.get('inlineCompletion.generationTimeout') ?? GENERATION_TIMEOUT, + triePruneTimeout: extensionConfig.get('inlineCompletion.triePruneTimeout') ?? TRIE_PRUNE_TIMEOUT, + }, + }; + + const state: ExtensionState = { + ollama: new Ollama({ + host: configuration.ollamaHost, + }), + configuration, + }; + + + vscode.workspace.onDidChangeConfiguration((event) => { + if (event.affectsConfiguration("ai-code.ollamaHost")) { + configuration.ollamaHost = extensionConfig.get('ollamaHost'); + state.ollama = new Ollama({ + host: configuration.ollamaHost, + }); + } + + if (event.affectsConfiguration("inlineCompletion.prefixStart")) { + configuration.inlineCompletion.model = extensionConfig.get('inlineCompletion.prefixStart') ?? MODEL; + } + + if (event.affectsConfiguration("inlineCompletion.prefixStart")) { + configuration.inlineCompletion.prefixStart = extensionConfig.get('inlineCompletion.prefixStart') ?? PREFIX_START; + } + + if (event.affectsConfiguration("inlineCompletion.prefixEnd")) { + configuration.inlineCompletion.prefixEnds = extensionConfig.get('inlineCompletion.prefixEnd')?.split(',') ?? PREFIX_ENDS; + } + + if (event.affectsConfiguration("inlineCompletion.suffixStart")) { + configuration.inlineCompletion.suffixStart = extensionConfig.get('inlineCompletion.suffixStart') ?? SUFFIX_START; + } + + if (event.affectsConfiguration("inlineCompletion.suffixEnd")) { + configuration.inlineCompletion.suffixEnd = extensionConfig.get('inlineCompletion.suffixEnd') ?? SUFFIX_END; + } + + if (event.affectsConfiguration("inlineCompletion.maxTokens")) { + configuration.inlineCompletion.maxTokens = extensionConfig.get('inlineCompletion.maxTokens') ?? MAX_TOKENS; + } + + if (event.affectsConfiguration("inlineCompletion.generationTimeout")) { + configuration.inlineCompletion.generationTimeout = extensionConfig.get('inlineCompletion.generationTimeout') ?? GENERATION_TIMEOUT; + } + + if (event.affectsConfiguration("inlineCompletion.triePruneTimeout")) { + configuration.inlineCompletion.triePruneTimeout = extensionConfig.get('inlineCompletion.triePruneTimeout') ?? TRIE_PRUNE_TIMEOUT; + } + }); + + return state; +}; + +const getModelSupportsSuffix = async (extension: ExtensionState, model: string) => { // TODO: get if model supports suffixes and use that if available // const response = await ollama.show({ @@ -34,13 +115,13 @@ const getModelSupportsSuffix = async (model: string) => { return false; }; -const getPrompt = (document: vscode.TextDocument, position: vscode.Position, prefix: string) => { +const getPrompt = (extension: ExtensionState, 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) .replace("{LANG}", document.languageId); - const message = `File:\n${PREFIX_START}`; + const message = `File:\n${extension.configuration.inlineCompletion.prefixStart}`; const prompt = `${messageHeader}\n${message}\n${prefix}`; @@ -48,11 +129,11 @@ const getPrompt = (document: vscode.TextDocument, position: vscode.Position, pre return prompt; }; -const getPromptWithSuffix = (document: vscode.TextDocument, position: vscode.Position, prefix: string) => { +const getPromptWithSuffix = (extension: ExtensionState, document: vscode.TextDocument, position: vscode.Position, prefix: string) => { 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 messageSuffix = `End of the file:\n${extension.configuration.inlineCompletion.suffixStart}\n${suffix}\n${extension.configuration.inlineCompletion.suffixEnd}\n`; + const messagePrefix = `Start of the file:\n${extension.configuration.inlineCompletion.prefixStart}`; 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") @@ -64,7 +145,7 @@ const getPromptWithSuffix = (document: vscode.TextDocument, position: vscode.Pos return prompt; }; -const getSuffix = (document: vscode.TextDocument, position: vscode.Position) => { +const getSuffix = (extension: ExtensionState, 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; @@ -88,6 +169,7 @@ const trieRootPrune = (text: string) => { }; const tokenProvider = async ( + extension: ExtensionState, document: vscode.TextDocument, position: vscode.Position, _context: vscode.InlineCompletionContext, @@ -95,10 +177,12 @@ const tokenProvider = async ( ) => { const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character)); - const modelSupportsSuffix = await getModelSupportsSuffix(MODEL); - const prompt = modelSupportsSuffix ? getPrompt(document, position, prefix) : getPromptWithSuffix(document, position, prefix); + const model = extension.configuration.inlineCompletion.model; - const suffix = modelSupportsSuffix ? getSuffix(document, position) : undefined; + const modelSupportsSuffix = await getModelSupportsSuffix(extension, model); + const prompt = modelSupportsSuffix ? getPrompt(extension, document, position, prefix) : getPromptWithSuffix(extension, document, position, prefix); + + const suffix = modelSupportsSuffix ? getSuffix(extension, document, position) : undefined; const result = trieRootLookup(prefix); @@ -110,21 +194,21 @@ const tokenProvider = async ( return []; } - const response = await ollama.generate({ - model: MODEL, + const response = await extension.ollama.generate({ + model, prompt, suffix, raw: true, stream: true, options: { - num_predict: MAX_TOKENS, - stop: PREFIX_ENDS, + num_predict: extension.configuration.inlineCompletion.maxTokens, + stop: extension.configuration.inlineCompletion.prefixEnds, }, }); const pruneTimeout = setTimeout(() => { trieRootPrune(prompt); - }, TRIE_PRUNE_TIMEOUT); + }, extension.configuration.inlineCompletion.triePruneTimeout); token.onCancellationRequested(() => { clearTimeout(pruneTimeout); @@ -137,7 +221,7 @@ const tokenProvider = async ( const buffer: string[] = []; const timeout = setTimeout(() => { resolve(buffer); - }, GENERATION_TIMEOUT); + }, extension.configuration.inlineCompletion.generationTimeout); try { for await (const part of response) { @@ -161,10 +245,12 @@ const tokenProvider = async ( export const activate = (context: vscode.ExtensionContext) => { console.log('"ai-code" extensions loaded'); + const extension = getExtensionState(); + const provider: vscode.InlineCompletionItemProvider = { async provideInlineCompletionItems(document, position, context, token) { try { - const completions = await tokenProvider(document, position, context, token); + const completions = await tokenProvider(extension, document, position, context, token); return completions.map((text) => ({ insertText: text, From a747b663618107313ee15a7a233a45930b4a1663 Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sun, 20 Jul 2025 00:31:19 -0500 Subject: [PATCH 07/10] moved config to own file --- src/config.ts | 114 +++++++++++++++++++++++++++++++++++++++++++++++ src/extension.ts | 102 +----------------------------------------- 2 files changed, 115 insertions(+), 101 deletions(-) create mode 100644 src/config.ts diff --git a/src/config.ts b/src/config.ts new file mode 100644 index 0000000..46eaad0 --- /dev/null +++ b/src/config.ts @@ -0,0 +1,114 @@ +import { Ollama } from 'ollama/browser'; +import * as vscode from 'vscode'; + +const CONFIG_NAMESPACE = 'ai-code'; + +const KEY_OLLAMA_HOST = 'ollamaHost'; +const KEY_INLINE_COMPLETION_MODEL = 'inlineCompletion.model'; +const KEY_INLINE_COMPLETION_PREFIX_START = 'inlineCompletion.prefixStart'; +const KEY_INLINE_COMPLETION_PREFIX_END = 'inlineCompletion.prefixEnd'; +const KEY_INLINE_COMPLETION_SUFFIX_START = 'inlineCompletion.suffixStart'; +const KEY_INLINE_COMPLETION_SUFFIX_END = 'inlineCompletion.suffixEnd'; +const KEY_INLINE_COMPLETION_MAX_TOKENS = 'inlineCompletion.maxTokens'; +const KEY_INLINE_COMPLETION_GENERATION_TIMEOUT = 'inlineCompletion.generationTimeout'; +const KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT = 'inlineCompletion.triePruneTimeout'; + +const MODEL = 'deepseek-coder:6.7b'; + +const PREFIX_START = ''; +const PREFIX_ENDS = ['', '', '', '']; + +const SUFFIX_START = ''; +const SUFFIX_END = ''; + +const MAX_TOKENS = 50; +const GENERATION_TIMEOUT = 200; +const TRIE_PRUNE_TIMEOUT = 10000; + +interface ExtensionConfiguration { + ollamaHost: string | undefined + inlineCompletion: { + model: string + prefixStart: string + prefixEnds: string[] + suffixStart: string + suffixEnd: string + maxTokens: number + generationTimeout: number + triePruneTimeout: number + } +}; + +export interface ExtensionState { + configuration: ExtensionConfiguration + ollama: Ollama +} + +export const getExtensionState = (): ExtensionState => { + const extensionConfig = vscode.workspace.getConfiguration(CONFIG_NAMESPACE); + + const configuration: ExtensionConfiguration = { + ollamaHost: extensionConfig.get(KEY_OLLAMA_HOST), + inlineCompletion: { + model: extensionConfig.get(KEY_INLINE_COMPLETION_MODEL) ?? MODEL, + prefixStart: extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_START) ?? PREFIX_START, + prefixEnds: extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_END)?.split(',') ?? PREFIX_ENDS, + suffixStart: extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_START) ?? SUFFIX_START, + suffixEnd: extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_END) ?? SUFFIX_END, + maxTokens: extensionConfig.get(KEY_INLINE_COMPLETION_MAX_TOKENS) ?? MAX_TOKENS, + generationTimeout: extensionConfig.get(KEY_INLINE_COMPLETION_GENERATION_TIMEOUT) ?? GENERATION_TIMEOUT, + triePruneTimeout: extensionConfig.get(KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT) ?? TRIE_PRUNE_TIMEOUT, + }, + }; + + const state: ExtensionState = { + ollama: new Ollama({ + host: configuration.ollamaHost, + }), + configuration, + }; + + + vscode.workspace.onDidChangeConfiguration((event) => { + if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_OLLAMA_HOST}`)) { + configuration.ollamaHost = extensionConfig.get(KEY_OLLAMA_HOST); + state.ollama = new Ollama({ + host: configuration.ollamaHost, + }); + } + + if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_MODEL}`)) { + configuration.inlineCompletion.model = extensionConfig.get(KEY_INLINE_COMPLETION_MODEL) ?? MODEL; + } + + if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_PREFIX_START}`)) { + configuration.inlineCompletion.prefixStart = extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_START) ?? PREFIX_START; + } + + if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_PREFIX_END}`)) { + configuration.inlineCompletion.prefixEnds = extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_END)?.split(',') ?? PREFIX_ENDS; + } + + if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_SUFFIX_START}`)) { + configuration.inlineCompletion.suffixStart = extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_START) ?? SUFFIX_START; + } + + if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_SUFFIX_END}`)) { + configuration.inlineCompletion.suffixEnd = extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_END) ?? SUFFIX_END; + } + + if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_MAX_TOKENS}`)) { + configuration.inlineCompletion.maxTokens = extensionConfig.get(KEY_INLINE_COMPLETION_MAX_TOKENS) ?? MAX_TOKENS; + } + + if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_GENERATION_TIMEOUT}`)) { + configuration.inlineCompletion.generationTimeout = extensionConfig.get(KEY_INLINE_COMPLETION_GENERATION_TIMEOUT) ?? GENERATION_TIMEOUT; + } + + if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT}`)) { + configuration.inlineCompletion.triePruneTimeout = extensionConfig.get(KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT) ?? TRIE_PRUNE_TIMEOUT; + } + }); + + return state; +}; diff --git a/src/extension.ts b/src/extension.ts index 05632e1..5591d49 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -1,108 +1,8 @@ // The module 'vscode' contains the VS Code extensibility API // 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 { flattenTrie, trieInsert, trieLookup, TrieNode, triePrune } from './trie'; - -const MODEL = 'deepseek-coder:6.7b'; - -const PREFIX_START = ''; -const PREFIX_ENDS = ['', '', '', '']; - -const SUFFIX_START = ''; -const SUFFIX_END = ''; - -const MAX_TOKENS = 50; -const GENERATION_TIMEOUT = 200; -const TRIE_PRUNE_TIMEOUT = 10000; - -interface ExtensionConfiguration { - ollamaHost: string | undefined - inlineCompletion: { - model: string - prefixStart: string - prefixEnds: string[] - suffixStart: string - suffixEnd: string - maxTokens: number - generationTimeout: number - triePruneTimeout: number - } -}; - -interface ExtensionState { - configuration: ExtensionConfiguration - ollama: Ollama -} - -const getExtensionState = (): ExtensionState => { - const extensionConfig = vscode.workspace.getConfiguration('ai-code'); - - const configuration: ExtensionConfiguration = { - ollamaHost: extensionConfig.get('ollamaHost'), - inlineCompletion: { - model: extensionConfig.get('inlineCompletion.prefixStart') ?? MODEL, - prefixStart: extensionConfig.get('inlineCompletion.prefixStart') ?? PREFIX_START, - prefixEnds: extensionConfig.get('inlineCompletion.prefixEnd')?.split(',') ?? PREFIX_ENDS, - suffixStart: extensionConfig.get('inlineCompletion.suffixStart') ?? SUFFIX_START, - suffixEnd: extensionConfig.get('inlineCompletion.suffixEnd') ?? SUFFIX_END, - maxTokens: extensionConfig.get('inlineCompletion.maxTokens') ?? MAX_TOKENS, - generationTimeout: extensionConfig.get('inlineCompletion.generationTimeout') ?? GENERATION_TIMEOUT, - triePruneTimeout: extensionConfig.get('inlineCompletion.triePruneTimeout') ?? TRIE_PRUNE_TIMEOUT, - }, - }; - - const state: ExtensionState = { - ollama: new Ollama({ - host: configuration.ollamaHost, - }), - configuration, - }; - - - vscode.workspace.onDidChangeConfiguration((event) => { - if (event.affectsConfiguration("ai-code.ollamaHost")) { - configuration.ollamaHost = extensionConfig.get('ollamaHost'); - state.ollama = new Ollama({ - host: configuration.ollamaHost, - }); - } - - if (event.affectsConfiguration("inlineCompletion.prefixStart")) { - configuration.inlineCompletion.model = extensionConfig.get('inlineCompletion.prefixStart') ?? MODEL; - } - - if (event.affectsConfiguration("inlineCompletion.prefixStart")) { - configuration.inlineCompletion.prefixStart = extensionConfig.get('inlineCompletion.prefixStart') ?? PREFIX_START; - } - - if (event.affectsConfiguration("inlineCompletion.prefixEnd")) { - configuration.inlineCompletion.prefixEnds = extensionConfig.get('inlineCompletion.prefixEnd')?.split(',') ?? PREFIX_ENDS; - } - - if (event.affectsConfiguration("inlineCompletion.suffixStart")) { - configuration.inlineCompletion.suffixStart = extensionConfig.get('inlineCompletion.suffixStart') ?? SUFFIX_START; - } - - if (event.affectsConfiguration("inlineCompletion.suffixEnd")) { - configuration.inlineCompletion.suffixEnd = extensionConfig.get('inlineCompletion.suffixEnd') ?? SUFFIX_END; - } - - if (event.affectsConfiguration("inlineCompletion.maxTokens")) { - configuration.inlineCompletion.maxTokens = extensionConfig.get('inlineCompletion.maxTokens') ?? MAX_TOKENS; - } - - if (event.affectsConfiguration("inlineCompletion.generationTimeout")) { - configuration.inlineCompletion.generationTimeout = extensionConfig.get('inlineCompletion.generationTimeout') ?? GENERATION_TIMEOUT; - } - - if (event.affectsConfiguration("inlineCompletion.triePruneTimeout")) { - configuration.inlineCompletion.triePruneTimeout = extensionConfig.get('inlineCompletion.triePruneTimeout') ?? TRIE_PRUNE_TIMEOUT; - } - }); - - return state; -}; +import { ExtensionState, getExtensionState } from './config'; const getModelSupportsSuffix = async (extension: ExtensionState, model: string) => { // TODO: get if model supports suffixes and use that if available From 6d93d4c45cbc748feace39a1d263defff9964c67 Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sun, 20 Jul 2025 00:32:53 -0500 Subject: [PATCH 08/10] updated const names in config --- src/config.ts | 52 +++++++++++++++++++++++++-------------------------- 1 file changed, 25 insertions(+), 27 deletions(-) diff --git a/src/config.ts b/src/config.ts index 46eaad0..4da9580 100644 --- a/src/config.ts +++ b/src/config.ts @@ -4,6 +4,7 @@ import * as vscode from 'vscode'; const CONFIG_NAMESPACE = 'ai-code'; const KEY_OLLAMA_HOST = 'ollamaHost'; + const KEY_INLINE_COMPLETION_MODEL = 'inlineCompletion.model'; const KEY_INLINE_COMPLETION_PREFIX_START = 'inlineCompletion.prefixStart'; const KEY_INLINE_COMPLETION_PREFIX_END = 'inlineCompletion.prefixEnd'; @@ -13,17 +14,14 @@ const KEY_INLINE_COMPLETION_MAX_TOKENS = 'inlineCompletion.maxTokens'; const KEY_INLINE_COMPLETION_GENERATION_TIMEOUT = 'inlineCompletion.generationTimeout'; const KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT = 'inlineCompletion.triePruneTimeout'; -const MODEL = 'deepseek-coder:6.7b'; - -const PREFIX_START = ''; -const PREFIX_ENDS = ['', '', '', '']; - -const SUFFIX_START = ''; -const SUFFIX_END = ''; - -const MAX_TOKENS = 50; -const GENERATION_TIMEOUT = 200; -const TRIE_PRUNE_TIMEOUT = 10000; +const DEFAULT_INLINE_COMPLETION_MODEL = 'deepseek-coder:6.7b'; +const DEFAULT_INLINE_COMPLETION_PREFIX_START = ''; +const DEFAULT_INLINE_COMPLETION_PREFIX_ENDS = ['', '', '', '']; +const DEFAULT_INLINE_COMPLETION_SUFFIX_START = ''; +const DEFAULT_INLINE_COMPLETION_SUFFIX_END = ''; +const DEFAULT_INLINE_COMPLETION_MAX_TOKENS = 50; +const DEFAULT_INLINE_COMPLETION_GENERATION_TIMEOUT = 200; +const DEFAULT_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT = 10000; interface ExtensionConfiguration { ollamaHost: string | undefined @@ -50,14 +48,14 @@ export const getExtensionState = (): ExtensionState => { const configuration: ExtensionConfiguration = { ollamaHost: extensionConfig.get(KEY_OLLAMA_HOST), inlineCompletion: { - model: extensionConfig.get(KEY_INLINE_COMPLETION_MODEL) ?? MODEL, - prefixStart: extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_START) ?? PREFIX_START, - prefixEnds: extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_END)?.split(',') ?? PREFIX_ENDS, - suffixStart: extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_START) ?? SUFFIX_START, - suffixEnd: extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_END) ?? SUFFIX_END, - maxTokens: extensionConfig.get(KEY_INLINE_COMPLETION_MAX_TOKENS) ?? MAX_TOKENS, - generationTimeout: extensionConfig.get(KEY_INLINE_COMPLETION_GENERATION_TIMEOUT) ?? GENERATION_TIMEOUT, - triePruneTimeout: extensionConfig.get(KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT) ?? TRIE_PRUNE_TIMEOUT, + model: extensionConfig.get(KEY_INLINE_COMPLETION_MODEL) ?? DEFAULT_INLINE_COMPLETION_MODEL, + prefixStart: extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_START) ?? DEFAULT_INLINE_COMPLETION_PREFIX_START, + prefixEnds: extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_END)?.split(',') ?? DEFAULT_INLINE_COMPLETION_PREFIX_ENDS, + suffixStart: extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_START) ?? DEFAULT_INLINE_COMPLETION_SUFFIX_START, + suffixEnd: extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_END) ?? DEFAULT_INLINE_COMPLETION_SUFFIX_END, + maxTokens: extensionConfig.get(KEY_INLINE_COMPLETION_MAX_TOKENS) ?? DEFAULT_INLINE_COMPLETION_MAX_TOKENS, + generationTimeout: extensionConfig.get(KEY_INLINE_COMPLETION_GENERATION_TIMEOUT) ?? DEFAULT_INLINE_COMPLETION_GENERATION_TIMEOUT, + triePruneTimeout: extensionConfig.get(KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT) ?? DEFAULT_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT, }, }; @@ -78,35 +76,35 @@ export const getExtensionState = (): ExtensionState => { } if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_MODEL}`)) { - configuration.inlineCompletion.model = extensionConfig.get(KEY_INLINE_COMPLETION_MODEL) ?? MODEL; + configuration.inlineCompletion.model = extensionConfig.get(KEY_INLINE_COMPLETION_MODEL) ?? DEFAULT_INLINE_COMPLETION_MODEL; } if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_PREFIX_START}`)) { - configuration.inlineCompletion.prefixStart = extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_START) ?? PREFIX_START; + configuration.inlineCompletion.prefixStart = extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_START) ?? DEFAULT_INLINE_COMPLETION_PREFIX_START; } if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_PREFIX_END}`)) { - configuration.inlineCompletion.prefixEnds = extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_END)?.split(',') ?? PREFIX_ENDS; + configuration.inlineCompletion.prefixEnds = extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_END)?.split(',') ?? DEFAULT_INLINE_COMPLETION_PREFIX_ENDS; } if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_SUFFIX_START}`)) { - configuration.inlineCompletion.suffixStart = extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_START) ?? SUFFIX_START; + configuration.inlineCompletion.suffixStart = extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_START) ?? DEFAULT_INLINE_COMPLETION_SUFFIX_START; } if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_SUFFIX_END}`)) { - configuration.inlineCompletion.suffixEnd = extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_END) ?? SUFFIX_END; + configuration.inlineCompletion.suffixEnd = extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_END) ?? DEFAULT_INLINE_COMPLETION_SUFFIX_END; } if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_MAX_TOKENS}`)) { - configuration.inlineCompletion.maxTokens = extensionConfig.get(KEY_INLINE_COMPLETION_MAX_TOKENS) ?? MAX_TOKENS; + configuration.inlineCompletion.maxTokens = extensionConfig.get(KEY_INLINE_COMPLETION_MAX_TOKENS) ?? DEFAULT_INLINE_COMPLETION_MAX_TOKENS; } if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_GENERATION_TIMEOUT}`)) { - configuration.inlineCompletion.generationTimeout = extensionConfig.get(KEY_INLINE_COMPLETION_GENERATION_TIMEOUT) ?? GENERATION_TIMEOUT; + configuration.inlineCompletion.generationTimeout = extensionConfig.get(KEY_INLINE_COMPLETION_GENERATION_TIMEOUT) ?? DEFAULT_INLINE_COMPLETION_GENERATION_TIMEOUT; } if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT}`)) { - configuration.inlineCompletion.triePruneTimeout = extensionConfig.get(KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT) ?? TRIE_PRUNE_TIMEOUT; + configuration.inlineCompletion.triePruneTimeout = extensionConfig.get(KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT) ?? DEFAULT_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT; } }); From b37bac525021f95ce95f18e6512eea6da28e3bc8 Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sun, 20 Jul 2025 01:24:32 -0500 Subject: [PATCH 09/10] moved auto complete to own module --- src/autoComplete.ts | 162 ++++++++++++++++++++++++++++++++++++++++++++ src/extension.ts | 160 ++----------------------------------------- 2 files changed, 166 insertions(+), 156 deletions(-) create mode 100644 src/autoComplete.ts diff --git a/src/autoComplete.ts b/src/autoComplete.ts new file mode 100644 index 0000000..7c77e64 --- /dev/null +++ b/src/autoComplete.ts @@ -0,0 +1,162 @@ +import { ExtensionState } from "./config"; +import * as vscode from 'vscode'; +import { flattenTrie, trieInsert, trieLookup, TrieNode, triePrune } from "./trie"; + +const getModelSupportsSuffix = async (extension: ExtensionState, model: string) => { + // TODO: get if model supports suffixes and use that if available + + // const response = await ollama.show({ + // model: model + // }) + + // model.capabilities.includes('suffix') + return false; +}; + +const getPrompt = (extension: ExtensionState, 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) + .replace("{LANG}", document.languageId); + + const message = `File:\n${extension.configuration.inlineCompletion.prefixStart}`; + + + const prompt = `${messageHeader}\n${message}\n${prefix}`; + + return prompt; +}; + +const getPromptWithSuffix = (extension: ExtensionState, document: vscode.TextDocument, position: vscode.Position, prefix: string) => { + 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${extension.configuration.inlineCompletion.suffixStart}\n${suffix}\n${extension.configuration.inlineCompletion.suffixEnd}\n`; + const messagePrefix = `Start of the file:\n${extension.configuration.inlineCompletion.prefixStart}`; + + 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 = (extension: ExtensionState, 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: true, + value: '', +}; + +const trieRootInsert = (text: string) => { + trieRoot = trieInsert(trieRoot, text); +}; + +const trieRootLookup = (text: string) => { + return trieLookup(trieRoot, text); +}; + +const trieRootPrune = (text: string) => { + return triePrune(trieRoot, text); +}; + +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 model = extension.configuration.inlineCompletion.model; + + const modelSupportsSuffix = await getModelSupportsSuffix(extension, model); + const prompt = modelSupportsSuffix ? getPrompt(extension, document, position, prefix) : getPromptWithSuffix(extension, document, position, prefix); + + 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, + suffix, + raw: true, + stream: true, + options: { + num_predict: extension.configuration.inlineCompletion.maxTokens, + stop: extension.configuration.inlineCompletion.prefixEnds, + }, + }); + + const pruneTimeout = setTimeout(() => { + trieRootPrune(prompt); + }, extension.configuration.inlineCompletion.triePruneTimeout); + + token.onCancellationRequested(() => { + clearTimeout(pruneTimeout); + try { + response.abort(); + } catch { } + }); + + const resultBuffer: string[] = await new Promise(async (resolve, reject) => { + const buffer: string[] = []; + const timeout = setTimeout(() => { + resolve(buffer); + }, extension.configuration.inlineCompletion.generationTimeout); + + try { + for await (const part of response) { + buffer.push(part.response); + trieRootInsert(prefix + buffer.join('')); + } + resolve(buffer); + } catch (err) { + reject(err); + } finally { + response.abort(); + clearTimeout(timeout); + }; + }); + + return [ + resultBuffer.join('') + ]; +}; + +export const getAutoCompleteProvider = (extension: ExtensionState) => { + const provider: vscode.InlineCompletionItemProvider = { + async provideInlineCompletionItems(document, position, context, token) { + try { + const completions = await tokenProvider(extension, document, position, context, token); + + return completions.map((text) => ({ + insertText: text, + range: new vscode.Range(position, position), + })); + } catch (err) { + console.log(err); + } + + return []; + }, + }; + + return provider; +}; \ No newline at end of file diff --git a/src/extension.ts b/src/extension.ts index 5591d49..715ae36 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -1,170 +1,18 @@ // The module 'vscode' contains the VS Code extensibility API // Import the module and reference it with the alias vscode in your code below import * as vscode from 'vscode'; -import { flattenTrie, trieInsert, trieLookup, TrieNode, triePrune } from './trie'; -import { ExtensionState, getExtensionState } from './config'; - -const getModelSupportsSuffix = async (extension: ExtensionState, model: string) => { - // TODO: get if model supports suffixes and use that if available - - // const response = await ollama.show({ - // model: model - // }) - - // model.capabilities.includes('suffix') - return false; -}; - -const getPrompt = (extension: ExtensionState, 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) - .replace("{LANG}", document.languageId); - - const message = `File:\n${extension.configuration.inlineCompletion.prefixStart}`; - - - const prompt = `${messageHeader}\n${message}\n${prefix}`; - - return prompt; -}; - -const getPromptWithSuffix = (extension: ExtensionState, document: vscode.TextDocument, position: vscode.Position, prefix: string) => { - 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${extension.configuration.inlineCompletion.suffixStart}\n${suffix}\n${extension.configuration.inlineCompletion.suffixEnd}\n`; - const messagePrefix = `Start of the file:\n${extension.configuration.inlineCompletion.prefixStart}`; - - 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 = (extension: ExtensionState, 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: true, - value: '', -}; - -const trieRootInsert = (text: string) => { - trieRoot = trieInsert(trieRoot, text); -}; - -const trieRootLookup = (text: string) => { - return trieLookup(trieRoot, text); -}; - -const trieRootPrune = (text: string) => { - return triePrune(trieRoot, text); -}; - -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 model = extension.configuration.inlineCompletion.model; - - const modelSupportsSuffix = await getModelSupportsSuffix(extension, model); - const prompt = modelSupportsSuffix ? getPrompt(extension, document, position, prefix) : getPromptWithSuffix(extension, document, position, prefix); - - 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, - suffix, - raw: true, - stream: true, - options: { - num_predict: extension.configuration.inlineCompletion.maxTokens, - stop: extension.configuration.inlineCompletion.prefixEnds, - }, - }); - - const pruneTimeout = setTimeout(() => { - trieRootPrune(prompt); - }, extension.configuration.inlineCompletion.triePruneTimeout); - - token.onCancellationRequested(() => { - clearTimeout(pruneTimeout); - try { - response.abort(); - } catch {} - }); - - const resultBuffer: string[] = await new Promise(async (resolve, reject) => { - const buffer: string[] = []; - const timeout = setTimeout(() => { - resolve(buffer); - }, extension.configuration.inlineCompletion.generationTimeout); - - try { - for await (const part of response) { - buffer.push(part.response); - trieRootInsert(prefix + buffer.join('')); - } - resolve(buffer); - } catch (err) { - reject(err); - } finally { - response.abort(); - clearTimeout(timeout); - }; - }); - - return [ - resultBuffer.join('') - ]; -}; +import { getExtensionState } from './config'; +import { getAutoCompleteProvider } from './autoComplete'; export const activate = (context: vscode.ExtensionContext) => { console.log('"ai-code" extensions loaded'); const extension = getExtensionState(); - const provider: vscode.InlineCompletionItemProvider = { - async provideInlineCompletionItems(document, position, context, token) { - try { - const completions = await tokenProvider(extension, document, position, context, token); + const autoCompleteProvider = getAutoCompleteProvider(extension); - return completions.map((text) => ({ - insertText: text, - range: new vscode.Range(position, position), - })); - } catch (err) { - console.log(err); - } - return []; - }, - }; - - vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, provider); + vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, autoCompleteProvider); }; // This method is called when your extension is deactivated From bb3e0c59d184b96fed9e1d4e1b3013b341157d2b Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sun, 20 Jul 2025 01:25:20 -0500 Subject: [PATCH 10/10] added TODOs --- src/extension.ts | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/extension.ts b/src/extension.ts index 715ae36..bbba054 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -11,6 +11,13 @@ export const activate = (context: vscode.ExtensionContext) => { const autoCompleteProvider = getAutoCompleteProvider(extension); + // TODO: code suggestion provider + + // TODO: quick fix provider + + // TODO: chat provider + + // TODO: agent mode provider vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, autoCompleteProvider); };