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
|
||||
// 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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue