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