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

Compare commits

base: jan-leila:bb3e0c59d184b96fed9e1d4e1b3013b341157d2b
Branches Tags
jan-leila:main
..
compare: jan-leila:6743ea6ece103f17f0a0207aee104987c6cbffff
Branches Tags
jan-leila:main

No commits in common. "bb3e0c59d184b96fed9e1d4e1b3013b341157d2b" and "6743ea6ece103f17f0a0207aee104987c6cbffff" have entirely different histories.
bb3e0c59d1 ... 6743ea6ece

6 changed files with 152 additions and 459 deletions
Download patch file Download diff file Expand all files Collapse all files

7
package.json
View file

@ -15,8 +15,15 @@
"main": "./dist/extension.js",
"contributes": {
"commands": [
{
"command": "ai-code.helloWorld",
"title": "Hello World"
}
]
},
"enabledApiProposals": [
"inlineCompletionsAdditions"
],
"scripts": {
"vscode:prepublish": "npm run package",
"compile": "webpack",

162
src/autoComplete.ts
View file

@ -1,162 +0,0 @@
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;
};

112
src/config.ts
View file

@ -1,112 +0,0 @@
import { Ollama } from 'ollama/browser';
import * as vscode from 'vscode';
const CONFIG_NAMESPACE = 'ai-code';
const KEY_OLLAMA_HOST = 'ollamaHost';
const KEY_INLINE_COMPLETION_MODEL = 'inlineCompletion.model';
const KEY_INLINE_COMPLETION_PREFIX_START = 'inlineCompletion.prefixStart';
const KEY_INLINE_COMPLETION_PREFIX_END = 'inlineCompletion.prefixEnd';
const KEY_INLINE_COMPLETION_SUFFIX_START = 'inlineCompletion.suffixStart';
const KEY_INLINE_COMPLETION_SUFFIX_END = 'inlineCompletion.suffixEnd';
const KEY_INLINE_COMPLETION_MAX_TOKENS = 'inlineCompletion.maxTokens';
const KEY_INLINE_COMPLETION_GENERATION_TIMEOUT = 'inlineCompletion.generationTimeout';
const KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT = 'inlineCompletion.triePruneTimeout';
const DEFAULT_INLINE_COMPLETION_MODEL = 'deepseek-coder:6.7b';
const DEFAULT_INLINE_COMPLETION_PREFIX_START = '<prefixStart>';
const DEFAULT_INLINE_COMPLETION_PREFIX_ENDS = ['<prefixStart>', '</prefixStart>', '<prefixEnd>', '</prefixEnd>'];
const DEFAULT_INLINE_COMPLETION_SUFFIX_START = '<suffixStart>';
const DEFAULT_INLINE_COMPLETION_SUFFIX_END = '</suffixStart>';
const DEFAULT_INLINE_COMPLETION_MAX_TOKENS = 50;
const DEFAULT_INLINE_COMPLETION_GENERATION_TIMEOUT = 200;
const DEFAULT_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT = 10000;
interface ExtensionConfiguration {
ollamaHost: string | undefined
inlineCompletion: {
model: string
prefixStart: string
prefixEnds: string[]
suffixStart: string
suffixEnd: string
maxTokens: number
generationTimeout: number
triePruneTimeout: number
}
};
export interface ExtensionState {
configuration: ExtensionConfiguration
ollama: Ollama
}
export const getExtensionState = (): ExtensionState => {
const extensionConfig = vscode.workspace.getConfiguration(CONFIG_NAMESPACE);
const configuration: ExtensionConfiguration = {
ollamaHost: extensionConfig.get(KEY_OLLAMA_HOST),
inlineCompletion: {
model: extensionConfig.get(KEY_INLINE_COMPLETION_MODEL) ?? DEFAULT_INLINE_COMPLETION_MODEL,
prefixStart: extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_START) ?? DEFAULT_INLINE_COMPLETION_PREFIX_START,
prefixEnds: extensionConfig.get<string>(KEY_INLINE_COMPLETION_PREFIX_END)?.split(',') ?? DEFAULT_INLINE_COMPLETION_PREFIX_ENDS,
suffixStart: extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_START) ?? DEFAULT_INLINE_COMPLETION_SUFFIX_START,
suffixEnd: extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_END) ?? DEFAULT_INLINE_COMPLETION_SUFFIX_END,
maxTokens: extensionConfig.get(KEY_INLINE_COMPLETION_MAX_TOKENS) ?? DEFAULT_INLINE_COMPLETION_MAX_TOKENS,
generationTimeout: extensionConfig.get(KEY_INLINE_COMPLETION_GENERATION_TIMEOUT) ?? DEFAULT_INLINE_COMPLETION_GENERATION_TIMEOUT,
triePruneTimeout: extensionConfig.get(KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT) ?? DEFAULT_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT,
},
};
const state: ExtensionState = {
ollama: new Ollama({
host: configuration.ollamaHost,
}),
configuration,
};
vscode.workspace.onDidChangeConfiguration((event) => {
if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_OLLAMA_HOST}`)) {
configuration.ollamaHost = extensionConfig.get(KEY_OLLAMA_HOST);
state.ollama = new Ollama({
host: configuration.ollamaHost,
});
}
if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_MODEL}`)) {
configuration.inlineCompletion.model = extensionConfig.get(KEY_INLINE_COMPLETION_MODEL) ?? DEFAULT_INLINE_COMPLETION_MODEL;
}
if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_PREFIX_START}`)) {
configuration.inlineCompletion.prefixStart = extensionConfig.get(KEY_INLINE_COMPLETION_PREFIX_START) ?? DEFAULT_INLINE_COMPLETION_PREFIX_START;
}
if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_PREFIX_END}`)) {
configuration.inlineCompletion.prefixEnds = extensionConfig.get<string>(KEY_INLINE_COMPLETION_PREFIX_END)?.split(',') ?? DEFAULT_INLINE_COMPLETION_PREFIX_ENDS;
}
if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_SUFFIX_START}`)) {
configuration.inlineCompletion.suffixStart = extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_START) ?? DEFAULT_INLINE_COMPLETION_SUFFIX_START;
}
if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_SUFFIX_END}`)) {
configuration.inlineCompletion.suffixEnd = extensionConfig.get(KEY_INLINE_COMPLETION_SUFFIX_END) ?? DEFAULT_INLINE_COMPLETION_SUFFIX_END;
}
if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_MAX_TOKENS}`)) {
configuration.inlineCompletion.maxTokens = extensionConfig.get(KEY_INLINE_COMPLETION_MAX_TOKENS) ?? DEFAULT_INLINE_COMPLETION_MAX_TOKENS;
}
if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_GENERATION_TIMEOUT}`)) {
configuration.inlineCompletion.generationTimeout = extensionConfig.get(KEY_INLINE_COMPLETION_GENERATION_TIMEOUT) ?? DEFAULT_INLINE_COMPLETION_GENERATION_TIMEOUT;
}
if (event.affectsConfiguration(`${CONFIG_NAMESPACE}.${KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT}`)) {
configuration.inlineCompletion.triePruneTimeout = extensionConfig.get(KEY_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT) ?? DEFAULT_INLINE_COMPLETION_TRIE_PRUNE_TIMEOUT;
}
});
return state;
};

138
src/extension.ts
View file

@ -1,25 +1,145 @@
// 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 { getExtensionState } from './config';
import { getAutoCompleteProvider } from './autoComplete';
import { Ollama } from 'ollama/browser';
const MODEL = 'deepseek-coder:6.7b';
const PREFIX_START = '<prefixStart>';
const PREFIX_END = '<prefixEnd>';
const SUFFIX_START = '<suffixStart>';
const SUFFIX_END = '<suffixEnd>';
const MAX_TOKENS = 50;
const GENERATION_TIMEOUT = 200;
const HOST = undefined;
// TODO: make these configurable by extension setting
const ollama = new Ollama({
host: HOST,
});
const getModelSupportsSuffix = async (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 = (document: vscode.TextDocument, position: vscode.Position) => {
const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character));
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 prompt = `${messageHeader}\n${message}\n${prefix}`;
return prompt;
};
const getPromptWithSuffix = (document: vscode.TextDocument, position: vscode.Position) => {
const prefix = document.getText(new vscode.Range(0, 0, position.line, position.character));
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));
return suffix;
};
const tokenProvider = async (
document: vscode.TextDocument,
position: vscode.Position,
context: vscode.InlineCompletionContext,
_token: vscode.CancellationToken,
) => {
const modelSupportsSuffix = await getModelSupportsSuffix(MODEL);
const prompt = modelSupportsSuffix ? getPrompt(document, position) : getPromptWithSuffix(document, position);
const suffix = modelSupportsSuffix ? getSuffix(document, position) : undefined;
const response = await ollama.generate({
model: MODEL,
prompt,
suffix,
raw: true,
stream: true,
options: {
num_predict: MAX_TOKENS,
stop: [PREFIX_END]
},
});
return response;
};
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 response = await tokenProvider(document, position, context, token);
const autoCompleteProvider = getAutoCompleteProvider(extension);
const resultBuffer: string[] = await new Promise(async (resolve, reject) => {
const buffer: string[] = [];
const timeout = setTimeout(() => {
resolve(buffer);
}, GENERATION_TIMEOUT);
// TODO: code suggestion provider
try {
for await (const part of response) {
console.log(part.response);
buffer.push(part.response);
}
resolve(buffer);
} catch (err) {
reject(err);
} finally {
clearTimeout(timeout);
};
});
// TODO: quick fix provider
const text = resultBuffer.join('');
// TODO: chat provider
return [
{
insertText: text,
range: new vscode.Range(position, position),
}
];
} catch (err) {
console.log(err);
}
// TODO: agent mode provider
return [];
},
};
vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, autoCompleteProvider);
vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, provider);
};
// This method is called when your extension is deactivated

15
src/test/extension.test.ts Normal file
View file

@ -0,0 +1,15 @@
import * as assert from 'assert';
// You can import and use all API from the 'vscode' module
// as well as import your extension to test it
import * as vscode from 'vscode';
// import * as myExtension from '../../extension';
suite('Extension Test Suite', () => {
vscode.window.showInformationMessage('Start all tests.');
test('Sample test', () => {
assert.strictEqual(-1, [1, 2, 3].indexOf(5));
assert.strictEqual(-1, [1, 2, 3].indexOf(0));
});
});

175
src/trie.ts
View file

@ -1,175 +0,0 @@
interface TrieLeaf {
isLeaf: true
value: string
children?: never
}
interface TrieBranch {
isLeaf: false
value: string
children: { [key: string]: TrieNode }
}
export type TrieNode = TrieLeaf | TrieBranch
/**
* Creates a new TrieNode based on node that has text added to it
* @param node node that we are basing the update on
* @param text text that is being added to the node
* @returns a new node with text added to it
*/
export const trieInsert = (node: TrieNode, text: string): TrieNode => {
// TODO: mutate node to add text to it
for (let index = 0; index < text.length; index++) {
// If the inserted text is longer then the nodes text update the node with the new text
if (index >= node.value.length) {
// If the current node is a leaf we can just replace it with a larger leaf
if (node.isLeaf) {
const newLeaf: TrieLeaf = {
isLeaf: true,
value: text,
};
return newLeaf;
}
// If the current node is a branch then we need add the remaining text to one of its children
const childKey = text[index];
const childText = text.substring(index + 1);
const child = node.children[childKey];
const newBranch: TrieBranch = {
isLeaf: false,
value: node.value,
children: {
...node.children,
[childKey]: child === undefined ? {
isLeaf: true,
value: childText,
} : trieInsert(child, childText)
},
};
return newBranch;
}
// If our inserted text does not match the node then we need to split the node
if (node.value[index] !== text[index]) {
// If the node is a leaf we need to split it into a branch
if (node.isLeaf) {
const newBranch: TrieBranch = {
isLeaf: false,
value: text.substring(0, index),
children: {
[text[index]]: {
isLeaf: true,
value: text.substring(index + 1),
},
[node.value[index]]: {
isLeaf: true,
value: node.value.substring(index + 1),
},
},
};
return newBranch;
}
// If the node is a branch then we need to create a new leaf on it
const newBranch: TrieBranch = {
isLeaf: false,
value: text.substring(0, index),
children: {
[text[index]]: {
isLeaf: true,
value: text.substring(index + 1),
},
[node.value[index]]: {
isLeaf: false,
value: node.value.substring(index + 1),
children: node.children,
},
},
};
return newBranch;
}
}
return node;
};
/**
* Gets a new trie that is the node trie after the text
* @param node
* @param text
* @returns
*/
export const trieLookup = (node: TrieNode, text: string): TrieNode | null => {
for (let index = 0; index < node.value.length; index++) {
// If our node still has text left but we have no more search query return a node starting at where we ran out of characters
if (index >= text.length) {
if (node.isLeaf) {
const newNode: TrieLeaf = {
isLeaf: true,
value: node.value.substring(index),
};
return newNode;
}
const newNode: TrieBranch = {
isLeaf: false,
value: node.value.substring(index),
children: node.children,
};
return newNode;
}
// If we have a difference then there is no match
if (node.value[index] !== text[index]) {
return null;
}
}
// If we get past the end of the node and it is a leaf then there is no match
if (node.isLeaf) {
return null;
}
// Continue matching on the child node
const childKey = text[node.value.length];
const childText = text.substring(node.value.length + 1);
const child = node.children[childKey];
return child === undefined ? null : trieLookup(child, childText);
};
export const flattenTrie = (node: TrieNode): string[] => {
if (node.isLeaf) {
return [
node.value
];
}
return Object.entries(node.children).flatMap(([key, child]) => {
return flattenTrie(child).map((value) => node.value + key + value);
});
};
export const triePrune = (node: TrieNode, text: string): TrieNode => {
const value = trieLookup(node, text);
if (value === null) {
return {
isLeaf: true,
value: text,
};
}
if (value.isLeaf) {
return {
isLeaf: true,
value: text + value.value,
};
}
return {
isLeaf: false,
value: text + value.value,
children: value.children,
};
};