jan-leila/ai-code
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

added back suffix support

Browse source
This commit is contained in:
Leyla Becker 2025-07-19 19:45:24 -05:00
parent 55c72ff98c
commit 02e65c5700
1 changed files with 21 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

25
src/extension.ts
View file

@ -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) {