Trim prefixes to increase cache hit rate

This commit is contained in:
mp 2024-11-10 23:52:25 -08:00
parent 652f64cd14
commit aab065139e

View file

@ -1,6 +1,6 @@
import * as vscode from 'vscode';
import { AbortRef, LLMMessage, sendLLMMessage } from '../common/sendLLMMessage';
import { getVoidConfigFromPartial } from '../webviews/common/contextForConfig';
import { getVoidConfigFromPartial, VoidConfig } from '../webviews/common/contextForConfig';
type AutocompletionStatus = 'pending' | 'finished' | 'error';
type Autocompletion = {
@ -14,42 +14,80 @@ type Autocompletion = {
result: string,
}
const DEBOUNCE_TIME = 500
const TIMEOUT_TIME = 60000
const toInlineCompletion = ({ prefix, suffix, autocompletion }: { prefix: string, suffix: string, autocompletion: Autocompletion }): vscode.InlineCompletionItem => {
// trims the end of the prefix to improve cache hit rate
const trimPrefix = (prefix: string) => {
const trimmedPrefix = prefix.trimEnd()
const trailingEnd = prefix.substring(trimmedPrefix.length)
// keep only a single trailing newline
if (trailingEnd.includes('\n')) {
return trimmedPrefix + '\n'
}
// else ignore all spaces and return the trimmed prefix
return trimmedPrefix
}
// finds the text in the autocompletion to display
const toInlineCompletion = ({ prefix, autocompletion }: { prefix: string, autocompletion: Autocompletion }): vscode.InlineCompletionItem => {
const originalPrefix = autocompletion.prefix
const generatedMiddle = autocompletion.result
const trimmedOriginalPrefix = trimPrefix(originalPrefix)
const trimmedCurrentPrefix = trimPrefix(prefix)
const lastMatchupIndex = trimmedCurrentPrefix.length - trimmedOriginalPrefix.length
console.log('generatedMiddle ', generatedMiddle)
console.log('trimmedOriginalPrefix ', trimmedOriginalPrefix)
console.log('trimmedCurrentPrefix ', trimmedCurrentPrefix)
console.log('lastMatchupIndex ', lastMatchupIndex)
if (lastMatchupIndex < 0) {
return new vscode.InlineCompletionItem('')
}
// example:
// originalPrefix = abcd
// generatedMiddle = efgh
// originalSuffix = ijkl
// the user has typed "ef" so prefix = abcdef
// we want to return the rest of the generatedMiddle, which is "gh"
const completionStr = generatedMiddle.substring(lastMatchupIndex)
return new vscode.InlineCompletionItem(completionStr)
}
// returns whether we can use this autocompletion to complete the prefix
const doesPrefixMatchAutocompletion = ({ prefix, autocompletion }: { prefix: string, autocompletion: Autocompletion }): boolean => {
const originalPrefix = autocompletion.prefix
const generatedMiddle = autocompletion.result
const fullPrefix = originalPrefix + generatedMiddle
const trimmedOriginalPrefix = trimPrefix(originalPrefix)
const trimmedCurrentPrefix = trimPrefix(prefix)
// check if the currently generated text matches with the prefix
let remainingText = ''
if (fullPrefix.startsWith(prefix)) {
// example:
// originalPrefix = abcd
// generatedMiddle = efgh
// originalSuffix = ijkl
// the user has typed "ef" so prefix = abcdef
// we want to return the rest of the generatedMiddle, which is "gh"
const index = (prefix.length - originalPrefix.length) - 1
remainingText = generatedMiddle.substring(index + 1)
if (trimmedCurrentPrefix.length < trimmedOriginalPrefix.length) {
return false
}
console.log('generated middle: ', JSON.stringify(generatedMiddle))
console.log('remaining text: ', JSON.stringify(remainingText))
return new vscode.InlineCompletionItem(remainingText)
const isMatch = (trimmedOriginalPrefix + generatedMiddle).startsWith(trimmedCurrentPrefix)
return isMatch
}
export class AutocompleteProvider implements vscode.InlineCompletionItemProvider {
private _extensionContext: vscode.ExtensionContext;
private _autocompletionsOfDocument: { [docUriStr: string]: Autocompletion[] } = {}
private _lastTime = 0
constructor(context: vscode.ExtensionContext) {
this._extensionContext = context
}
@ -65,8 +103,6 @@ export class AutocompleteProvider implements vscode.InlineCompletionItemProvider
const docUriStr = document.uri.toString()
console.log('initial _autocompletionsOfDocument', this._autocompletionsOfDocument[docUriStr])
const fullText = document.getText();
const cursorOffset = document.offsetAt(position);
const prefix = fullText.substring(0, cursorOffset)
@ -78,35 +114,33 @@ export class AutocompleteProvider implements vscode.InlineCompletionItemProvider
const voidConfig = getVoidConfigFromPartial(this._extensionContext.globalState.get('partialVoidConfig') ?? {})
// get autocompletion from cache
let cachedAutocompletion: Autocompletion | undefined = undefined
loop: for (const autocompletion of this._autocompletionsOfDocument[docUriStr]!) {
const originalPrefix = autocompletion.prefix
const generatedMiddle = autocompletion.result
// if the user's change matches up with the generated text
if ((originalPrefix + generatedMiddle).startsWith(prefix)) {
if (doesPrefixMatchAutocompletion({ prefix, autocompletion })) {
cachedAutocompletion = autocompletion
break loop;
}
}
// if there is an autocompletion for this line, return it
// if there is a cached autocompletion, return it
if (cachedAutocompletion) {
if (cachedAutocompletion.status === 'finished') {
console.log('AAA1')
const inlineCompletion = toInlineCompletion({ autocompletion: cachedAutocompletion, prefix, suffix, })
const inlineCompletion = toInlineCompletion({ autocompletion: cachedAutocompletion, prefix, })
return [inlineCompletion]
} else if (cachedAutocompletion.status === 'pending') {
console.log('AAA2')
try {
// await the result; if it hasnt resolved in 10 seconds assume the request is dead
await cachedAutocompletion.promise;
const inlineCompletion = toInlineCompletion({ autocompletion: cachedAutocompletion, prefix, suffix, })
const inlineCompletion = toInlineCompletion({ autocompletion: cachedAutocompletion, prefix, })
return [inlineCompletion]
} catch (e) {
@ -120,9 +154,42 @@ export class AutocompleteProvider implements vscode.InlineCompletionItemProvider
return []
}
// if there is no cached autocompletion, create it and add it to cache
// wait DEBOUNCE_TIME for the user to stop typing
const thisTime = Date.now()
this._lastTime = thisTime
const didTypingHappenDuringDebounce = await new Promise((resolve, reject) =>
setTimeout(() => {
if (this._lastTime === thisTime) {
resolve(false)
} else {
resolve(true)
}
}, DEBOUNCE_TIME)
)
// if more typing happened, then do not go forwards with the request
if (didTypingHappenDuringDebounce) {
return []
}
console.log('BBB')
// if there is no autocomplete for this line, create it and add it to cache
// else if no more typing happens, then go forwards with the request
const newAutocompletion: Autocompletion = {
prefix: prefix,
suffix: suffix,
startTime: Date.now(),
endTime: undefined,
abortRef: { current: () => { } },
status: 'pending',
promise: undefined,
result: '',
}
let messages: LLMMessage[] = []
switch (voidConfig.default.whichApi) {
case 'ollama':
@ -141,19 +208,7 @@ export class AutocompleteProvider implements vscode.InlineCompletionItemProvider
throw new Error(`We do not recommend using autocomplete with your selected provider (${voidConfig.default.whichApi}).`);
}
const newAutocompletion: Autocompletion = {
prefix: prefix,
suffix: suffix,
startTime: Date.now(),
endTime: undefined,
abortRef: { current: () => { } },
status: 'pending',
promise: undefined,
result: '',
}
// set the parameters of `newAutocompletion` appropriately
// set parameters of `newAutocompletion` appropriately
newAutocompletion.promise = new Promise((resolve, reject) => {
sendLLMMessage({
@ -186,20 +241,21 @@ export class AutocompleteProvider implements vscode.InlineCompletionItemProvider
abortRef: newAutocompletion.abortRef,
})
setTimeout(() => { // if the request hasnt resolved in TIMEOUT_TIME seconds, end it
setTimeout(() => { // if the request hasnt resolved in TIMEOUT_TIME seconds, reject it
if (newAutocompletion.status === 'pending') {
reject('Timeout')
}
}, TIMEOUT_TIME)
})
this._autocompletionsOfDocument[docUriStr]!.push(newAutocompletion)
// add autocompletion to cache
this._autocompletionsOfDocument[docUriStr]?.push(newAutocompletion)
// show autocompletion
try {
await newAutocompletion.promise;
const inlineCompletion = toInlineCompletion({ autocompletion: newAutocompletion, prefix, suffix, })
const inlineCompletion = toInlineCompletion({ autocompletion: newAutocompletion, prefix, })
return [inlineCompletion]
} catch (e) {