made settings configurable
This commit is contained in:
parent
a620903c10
commit
7fd842dbc0
1 changed files with 108 additions and 22 deletions
130
src/extension.ts
130
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<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;
|
||||
}
|
||||
});
|
||||
|
||||
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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue