From 75c82fa7c610f730dcf3b54548d28d6157741635 Mon Sep 17 00:00:00 2001 From: Zhiming Ma Date: Thu, 11 May 2023 22:55:34 +0800 Subject: [PATCH] VSCode extensions: add completion cache. (#128) --- clients/vscode/package.json | 1 + clients/vscode/src/CompletionCache.ts | 75 +++++++++++++++++++ clients/vscode/src/TabbyCompletionProvider.ts | 53 ++++++++----- clients/vscode/yarn.lock | 5 ++ 4 files changed, 117 insertions(+), 17 deletions(-) create mode 100644 clients/vscode/src/CompletionCache.ts diff --git a/clients/vscode/package.json b/clients/vscode/package.json index 3d610d6..6239395 100644 --- a/clients/vscode/package.json +++ b/clients/vscode/package.json @@ -107,6 +107,7 @@ "axios": "^1.3.4", "events": "^3.3.0", "form-data": "^4.0.0", + "linked-list-typescript": "^1.0.15", "process": "^0.11.10" } } diff --git a/clients/vscode/src/CompletionCache.ts b/clients/vscode/src/CompletionCache.ts new file mode 100644 index 0000000..027875f --- /dev/null +++ b/clients/vscode/src/CompletionCache.ts @@ -0,0 +1,75 @@ +import { LinkedList } from "linked-list-typescript"; +import { CompletionResponse, Choice } from "./generated"; + +type Range = { + start: number; + end: number; +}; + +export type CompletionCacheEntry = { + documentId: any; + promptRange: Range; + prompt: string; + completion: CompletionResponse; +}; + +export class CompletionCache { + public static cacheSize = 10; + private cache = new LinkedList(); + + constructor() {} + + private evict() { + while (this.cache.length > CompletionCache.cacheSize) { + this.cache.removeTail(); + } + } + + private pop(entry: CompletionCacheEntry) { + this.cache.remove(entry); + this.cache.prepend(entry); + } + + public add(entry: CompletionCacheEntry) { + this.evict(); + this.cache.prepend(entry); + } + + public findCompatible(documentId: any, text: string, cursor: number): CompletionResponse | null { + let hit: { entry: CompletionCacheEntry; compatibleChoices: Choice[] } | null = null; + for (const entry of this.cache) { + if (entry.documentId !== documentId) { + continue; + } + // Check if text in prompt range has not changed + if (text.slice(entry.promptRange.start, entry.promptRange.end) !== entry.prompt) { + continue; + } + // Filter choices that start with inputed text after prompt + const compatibleChoices = entry.completion.choices + .filter((choice) => choice.text.startsWith(text.slice(entry.promptRange.end, cursor))) + .map((choice) => { + return { + index: choice.index, + text: choice.text.substring(cursor - entry.promptRange.end), + }; + }); + if (compatibleChoices.length > 0) { + hit = { + entry, + compatibleChoices, + }; + break; + } + } + if (hit) { + this.pop(hit.entry); + return { + id: hit.entry.completion.id, + created: hit.entry.completion.created, + choices: hit.compatibleChoices, + }; + } + return null; + } +} diff --git a/clients/vscode/src/TabbyCompletionProvider.ts b/clients/vscode/src/TabbyCompletionProvider.ts index 2846a79..ab07560 100644 --- a/clients/vscode/src/TabbyCompletionProvider.ts +++ b/clients/vscode/src/TabbyCompletionProvider.ts @@ -12,6 +12,7 @@ import { } from "vscode"; import { CompletionResponse, EventType, ChoiceEvent, ApiError, CancelablePromise, CancelError } from "./generated"; import { TabbyClient } from "./TabbyClient"; +import { CompletionCache } from "./CompletionCache"; import { sleep } from "./utils"; export class TabbyCompletionProvider implements InlineCompletionItemProvider { @@ -20,6 +21,7 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider { private pendingCompletion: CancelablePromise | null = null; private tabbyClient = TabbyClient.getInstance(); + private completionCache = new CompletionCache(); // User Settings private enabled: boolean = true; private suggestionDelay: number = 150; @@ -42,7 +44,8 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider { return emptyResponse; } - const prompt = this.getPrompt(document, position); + const promptRange = this.calculatePromptRange(position); + const prompt = document.getText(promptRange); if (this.isNil(prompt)) { console.debug("Prompt is empty, skipping"); return emptyResponse; @@ -56,6 +59,15 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider { return emptyResponse; } + const replaceRange = this.calculateReplaceRange(document, position); + + const compatibleCache = this.completionCache.findCompatible(document.uri, document.getText(), document.offsetAt(position)); + if (compatibleCache) { + const completions = this.toInlineCompletions(compatibleCache, replaceRange); + console.debug("Use cached completions: ", compatibleCache); + return Promise.resolve(completions); + } + console.debug( "Requesting: ", { @@ -86,15 +98,14 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider { }); this.pendingCompletion = null; - const hasSuffixParen = this.hasSuffixParen(document, position); - const replaceRange = hasSuffixParen - ? new Range( - position.line, - position.character, - position.line, - position.character + 1 - ) - : new Range(position, position); + if (completion) { + this.completionCache.add({ + documentId: document.uri, + promptRange: { start: document.offsetAt(promptRange.start), end: document.offsetAt(promptRange.end) }, + prompt, + completion, + }); + } const completions = this.toInlineCompletions(completion, replaceRange); console.debug("Result completions: ", completions); return Promise.resolve(completions); @@ -106,13 +117,6 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider { this.suggestionDelay = configuration.get("suggestionDelay", 150); } - private getPrompt(document: TextDocument, position: Position): string | undefined { - const maxLines = 20; - const firstLine = Math.max(position.line - maxLines, 0); - - return document.getText(new Range(firstLine, 0, position.line, position.character)); - } - private isNil(value: string | undefined | null): boolean { return value === undefined || value === null || value.length === 0; } @@ -140,4 +144,19 @@ export class TabbyCompletionProvider implements InlineCompletionItemProvider { ); return ")]}".indexOf(suffix) > -1; } + + private calculatePromptRange(position: Position): Range { + const maxLines = 20; + const firstLine = Math.max(position.line - maxLines, 0); + return new Range(firstLine, 0, position.line, position.character); + } + + private calculateReplaceRange(document: TextDocument, position: Position): Range { + const hasSuffixParen = this.hasSuffixParen(document, position); + if (hasSuffixParen) { + return new Range(position.line, position.character, position.line, position.character + 1); + } else { + return new Range(position, position); + } + } } diff --git a/clients/vscode/yarn.lock b/clients/vscode/yarn.lock index f7c48fa..979e5b5 100644 --- a/clients/vscode/yarn.lock +++ b/clients/vscode/yarn.lock @@ -1740,6 +1740,11 @@ lie@~3.3.0: dependencies: immediate "~3.0.5" +linked-list-typescript@^1.0.15: + version "1.0.15" + resolved "https://registry.yarnpkg.com/linked-list-typescript/-/linked-list-typescript-1.0.15.tgz#faeed93cf9203f102e2158c29edcddda320abe82" + integrity sha512-RIyUu9lnJIyIaMe63O7/aFv/T2v3KsMFuXMBbUQCHX+cgtGro86ETDj5ed0a8gQL2+DFjzYYsgVG4I36/cUwgw== + linkify-it@^3.0.1: version "3.0.3" resolved "https://registry.npmmirror.com/linkify-it/-/linkify-it-3.0.3.tgz#a98baf44ce45a550efb4d49c769d07524cc2fa2e"