ai-code/src/extension.ts

// 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 { 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 provider: vscode.InlineCompletionItemProvider = {
		async provideInlineCompletionItems(document, position, context, token) {
			try {
				const response = await tokenProvider(document, position, context, token);

				const resultBuffer: string[] = await new Promise(async (resolve, reject) => {
					const buffer: string[] = [];
					const timeout = setTimeout(() => {
						resolve(buffer);
					}, GENERATION_TIMEOUT);

					try {
						for await (const part of response) {
							console.log(part.response);
							buffer.push(part.response);
						}
						resolve(buffer);
					} catch (err) {
						reject(err);
					} finally {
						clearTimeout(timeout);
					};
				});

				const text = resultBuffer.join('');

				return [
					{
						insertText: text,
						range: new vscode.Range(position, position),
					}
				];
			} catch (err) {
				console.log(err);
			}

			return [];
		},
	};

	vscode.languages.registerInlineCompletionItemProvider({ pattern: '**' }, provider);
};

// This method is called when your extension is deactivated
export function deactivate() {}