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 = '<prefixStart>';
-const PREFIX_ENDS = ['<prefixEnd>', '</prefixEnd>'];
+const PREFIX_ENDS = ['<prefixStart>', '</prefixStart>', '<prefixEnd>', '</prefixEnd>'];
+
+const SUFFIX_START = '<suffixStart>';
+const SUFFIX_END = '</suffixStart>';
 
 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) {