From 7fd842dbc0bac62d115c0149a86be71dfd8d41e8 Mon Sep 17 00:00:00 2001 From: Leyla Becker Date: Sat, 19 Jul 2025 22:00:06 -0500 Subject: [PATCH] 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,