moved auto complete to own module
This commit is contained in:
parent
6d93d4c45c
commit
b37bac5250
2 changed files with 166 additions and 156 deletions
162
src/autoComplete.ts
Normal file
162
src/autoComplete.ts
Normal file
|
|
@ -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;
|
||||||
|
};
|
||||||
160
src/extension.ts
160
src/extension.ts
|
|
@ -1,170 +1,18 @@
|
||||||
// The module 'vscode' contains the VS Code extensibility API
|
// The module 'vscode' contains the VS Code extensibility API
|
||||||
// Import the module and reference it with the alias vscode in your code below
|
// Import the module and reference it with the alias vscode in your code below
|
||||||
import * as vscode from 'vscode';
|
import * as vscode from 'vscode';
|
||||||
import { flattenTrie, trieInsert, trieLookup, TrieNode, triePrune } from './trie';
|
import { getExtensionState } from './config';
|
||||||
import { ExtensionState, getExtensionState } from './config';
|
import { getAutoCompleteProvider } from './autoComplete';
|
||||||
|
|
||||||
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 activate = (context: vscode.ExtensionContext) => {
|
export const activate = (context: vscode.ExtensionContext) => {
|
||||||
console.log('"ai-code" extensions loaded');
|
console.log('"ai-code" extensions loaded');
|
||||||
|
|
||||||
const extension = getExtensionState();
|
const extension = getExtensionState();
|
||||||
|
|
||||||
const provider: vscode.InlineCompletionItemProvider = {
|
const autoCompleteProvider = getAutoCompleteProvider(extension);
|
||||||
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 [];
|
vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, autoCompleteProvider);
|
||||||
},
|
|
||||||
};
|
|
||||||
|
|
||||||
vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, provider);
|
|
||||||
};
|
};
|
||||||
|
|
||||||
// This method is called when your extension is deactivated
|
// This method is called when your extension is deactivated
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue