made settings configurable

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
+interface ExtensionState {
-const ollama = new Ollama({
+	configuration: ExtensionConfiguration
-	host: HOST,
+	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<string>('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<string>('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;
+		}
 	});
 
-const getModelSupportsSuffix = async (model: string) => {
+	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 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${PREFIX_START}`;
+	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 model = extension.configuration.inlineCompletion.model;
-	const prompt = modelSupportsSuffix ? getPrompt(document, position, prefix) : getPromptWithSuffix(document, position, prefix);
 
-	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({
+	const response = await extension.ollama.generate({
-		model: MODEL,
+		model,
 		prompt,
 		suffix,
 		raw: true,
 		stream: true,
 		options: {
-			num_predict: MAX_TOKENS,
+			num_predict: extension.configuration.inlineCompletion.maxTokens,
-			stop: PREFIX_ENDS,
+			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,