getInlineCompletions, separate provider, update tests

vscode/src/completions/logger.ts

@@ -83,7 +83,7 @@ export function networkRequestStarted(
         embeddings?: number
         local?: number
         duration: number
-    }
+    } | null
 ): void {
     const event = displayedCompletions.get(id)
     if (event) {

vscode/src/completions/processInlineCompletions.test.ts

@@ -0,0 +1,48 @@
+import { describe, expect, test } from 'vitest'
+
+import { range } from '../testutils/textDocument'
+
+import { adjustRangeToOverwriteOverlappingCharacters } from './processInlineCompletions'
+import { documentAndPosition } from './testHelpers'
+import { InlineCompletionItem } from './types'
+
+describe('adjustRangeToOverwriteOverlappingCharacters', () => {
+    test('no adjustment at end of line', () => {
+        const item: InlineCompletionItem = { insertText: 'array) {' }
+        const { position } = documentAndPosition('function sort(█')
+        expect(
+            adjustRangeToOverwriteOverlappingCharacters(item, {
+                position,
+                docContext: { currentLineSuffix: '' },
+            })
+        ).toEqual<InlineCompletionItem>(item)
+    })
+
+    test('handles non-empty currentLineSuffix', () => {
+        const item: InlineCompletionItem = { insertText: 'array) {' }
+        const { position } = documentAndPosition('function sort(█)')
+        expect(
+            adjustRangeToOverwriteOverlappingCharacters(item, {
+                position,
+                docContext: { currentLineSuffix: ')' },
+            })
+        ).toEqual<InlineCompletionItem>({
+            ...item,
+            range: range(0, 14, 0, 15),
+        })
+    })
+
+    test('handles whitespace in currentLineSuffix', () => {
+        const item: InlineCompletionItem = { insertText: 'array) {' }
+        const { position } = documentAndPosition('function sort(█)')
+        expect(
+            adjustRangeToOverwriteOverlappingCharacters(item, {
+                position,
+                docContext: { currentLineSuffix: ') ' },
+            })
+        ).toEqual<InlineCompletionItem>({
+            ...item,
+            range: range(0, 14, 0, 16),
+        })
+    })
+})

vscode/src/completions/processInlineCompletions.ts

@@ -0,0 +1,99 @@
+import type { Position, TextDocument } from 'vscode'
+
+import { DocumentContext } from './document'
+import { truncateMultilineCompletion } from './multiline'
+import { collapseDuplicativeWhitespace, removeTrailingWhitespace, trimUntilSuffix } from './text-processing'
+import { InlineCompletionItem } from './types'
+
+export interface ProcessInlineCompletionsParams {
+    document: Pick<TextDocument, 'languageId'>
+    position: Position
+    multiline: boolean
+    docContext: DocumentContext
+}
+
+/**
+ * This function implements post-processing logic that is applied regardless of
+ * which provider is chosen.
+ */
+export function processInlineCompletions(
+    items: InlineCompletionItem[],
+    { document, position, multiline, docContext }: ProcessInlineCompletionsParams
+): InlineCompletionItem[] {
+    // Shared post-processing logic
+    const processedCompletions = items.map(item => processItem(item, { document, position, multiline, docContext }))
+
+    // Filter results
+    const visibleResults = filterCompletions(processedCompletions)
+
+    // Remove duplicate results
+    const uniqueResults = [...new Map(visibleResults.map(item => [item.insertText, item])).values()]
+
+    // Rank results
+    const rankedResults = rankCompletions(uniqueResults)
+
+    return rankedResults
+}
+
+function processItem(
+    item: InlineCompletionItem,
+    {
+        document,
+        position,
+        multiline,
+        docContext: { prefix, suffix, currentLineSuffix },
+    }: Pick<ProcessInlineCompletionsParams, 'document' | 'position' | 'multiline' | 'docContext'>
+): InlineCompletionItem {
+    // Make a copy to avoid unexpected behavior.
+    item = { ...item }
+
+    if (typeof item.insertText !== 'string') {
+        throw new TypeError('SnippetText not supported')
+    }
+
+    item = adjustRangeToOverwriteOverlappingCharacters(item, { position, docContext: { currentLineSuffix } })
+    if (multiline) {
+        item.insertText = truncateMultilineCompletion(item.insertText, prefix, suffix, document.languageId)
+        item.insertText = removeTrailingWhitespace(item.insertText)
+    }
+    item.insertText = trimUntilSuffix(item.insertText, prefix, suffix, document.languageId)
+    item.insertText = collapseDuplicativeWhitespace(prefix, item.insertText)
+
+    return item
+}
+
+/**
+ * Return a copy of item with an adjusted range to overwrite duplicative characters after the
+ * completion on the first line.
+ *
+ * For example, with position `function sort(█)` and completion `array) {`, the range should be
+ * adjusted to span the `)` so it is overwritten by the `insertText` (so that we don't end up with
+ * the invalid `function sort(array) {)`).
+ */
+export function adjustRangeToOverwriteOverlappingCharacters(
+    item: InlineCompletionItem,
+    {
+        position,
+        docContext: { currentLineSuffix },
+    }: Pick<ProcessInlineCompletionsParams, 'position'> & {
+        docContext: Pick<DocumentContext, 'currentLineSuffix'>
+    }
+): InlineCompletionItem {
+    // TODO(sqs): This is a very naive implementation that will not work for many cases. It always
+    // just clobbers the rest of the line.
+
+    if (!item.range && currentLineSuffix !== '') {
+        return { ...item, range: { start: position, end: position.translate(undefined, currentLineSuffix.length) } }
+    }
+
+    return item
+}
+
+function rankCompletions(completions: InlineCompletionItem[]): InlineCompletionItem[] {
+    // TODO(philipp-spiess): Improve ranking to something more complex then just length
+    return completions.sort((a, b) => b.insertText.split('\n').length - a.insertText.split('\n').length)
+}
+
+function filterCompletions(completions: InlineCompletionItem[]): InlineCompletionItem[] {
+    return completions.filter(c => c.insertText.trim() !== '')
+}

vscode/src/completions/providers/anthropic.ts

@@ -145,7 +145,7 @@ export class AnthropicProvider extends Provider {
         // Create prompt
         const { messages: prompt } = this.createPrompt(snippets)
         if (prompt.length > this.promptChars) {
-            throw new Error('prompt length exceeded maximum alloted chars')
+            throw new Error(`prompt length (${prompt.length}) exceeded maximum character length (${this.promptChars})`)
         }
 
         const args: CompletionParameters = this.options.multiline

vscode/src/completions/providers/unstable-azure-openai.ts

@@ -1,7 +1,7 @@
 import { logger } from '../../log'
 import { ReferenceSnippet } from '../context'
 import { getHeadAndTail } from '../text-processing'
-import { Completion } from '../vscodeInlineCompletionItemProvider'
+import { Completion } from '../types'
 
 import { Provider, ProviderConfig, ProviderOptions } from './provider'
 

vscode/src/completions/request-manager.test.ts

@@ -3,12 +3,11 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 import { vsCodeMocks } from '../testutils/mocks'
 
 import { Provider } from './providers/provider'
-import { RequestManager } from './request-manager'
+import { RequestManager, RequestManagerResult } from './request-manager'
 import { Completion } from './types'
 
 vi.mock('vscode', () => vsCodeMocks)
 
-const DOCUMENT_URI = 'file:///path/to/file.ts'
 const LOG_ID = 'some-log-id'
 
 class MockProvider extends Provider {
@@ -48,12 +47,12 @@ function createProvider(prefix: string) {
 }
 
 describe('RequestManager', () => {
-    let createRequest: (prefix: string, provider: Provider) => Promise<Completion[]>
+    let createRequest: (prefix: string, provider: Provider) => Promise<RequestManagerResult>
     beforeEach(() => {
         const requestManager = new RequestManager()
 
         createRequest = (prefix: string, provider: Provider) =>
-            requestManager.request(DOCUMENT_URI, LOG_ID, prefix, [provider], [], new AbortController().signal)
+            requestManager.request({ prefix }, [provider], [], new AbortController().signal)
     })
 
     it('resolves a single request', async () => {
@@ -63,12 +62,15 @@ describe('RequestManager', () => {
         setTimeout(() => provider.resolveRequest(["'hello')"]), 0)
 
         await expect(createRequest(prefix, provider)).resolves.toMatchInlineSnapshot(`
-          [
-            {
-              "content": "'hello')",
-              "prefix": "console.log(",
-            },
-          ]
+          {
+            "cacheHit": false,
+            "completions": [
+              {
+                "content": "'hello')",
+                "prefix": "console.log(",
+              },
+            ],
+          }
         `)
     })
 
@@ -86,14 +88,14 @@ describe('RequestManager', () => {
 
         provider2.resolveRequest(["'hello')"])
 
-        expect((await promise2)[0].content).toBe("'hello')")
+        expect((await promise2).completions[0].content).toBe("'hello')")
 
         expect(provider1.didFinishNetworkRequest).toBe(false)
         expect(provider2.didFinishNetworkRequest).toBe(true)
 
         provider1.resolveRequest(['log();'])
 
-        expect((await promise1)[0].content).toBe('log();')
+        expect((await promise1).completions[0].content).toBe('log();')
         expect(provider1.didFinishNetworkRequest).toBe(true)
     })
 })

vscode/src/completions/request-manager.ts

@@ -1,23 +1,46 @@
+import { LRUCache } from 'lru-cache'
+
+import { debug } from '../log'
+
 import { ReferenceSnippet } from './context'
 import { CompletionProviderTracer, Provider } from './providers/provider'
 import { Completion } from './types'
 
+export interface RequestParams {
+    // TODO(sqs): This is not a unique enough cache key. We should cache based on the params wrapped
+    // into generateCompletions instead of requiring callers to separately pass cache-key-able
+    // params to RequestManager.
+    prefix: string
+}
+
+export interface RequestManagerResult {
+    completions: Completion[]
+    cacheHit: boolean
+}
+
 /**
  * This class can handle concurrent requests for code completions. The idea is
  * that requests are not cancelled even when the user continues typing in the
  * document. This allows us to cache the results of expensive completions and
  * return them when the user triggers a completion again.
  */
 export class RequestManager {
+    private cache = new RequestCache()
+
     public async request(
-        documentUri: string,
-        logId: string,
-        prefix: string,
+        params: RequestParams,
         providers: Provider[],
         context: ReferenceSnippet[],
-        signal: AbortSignal,
+        signal?: AbortSignal,
         tracer?: CompletionProviderTracer
-    ): Promise<Completion[]> {
+    ): Promise<RequestManagerResult> {
+        const existing = this.cache.get({ params })
+        if (existing) {
+            debug('RequestManager', 'cache hit', { verbose: { params, existing } })
+            return { ...existing.result, cacheHit: true }
+        }
+        debug('RequestManager', 'cache miss', { verbose: { params } })
+
         // We forward a different abort controller to the network request so we
         // can cancel the network request independently of the user cancelling
         // the completion.
@@ -28,10 +51,38 @@ export class RequestManager {
         )
             .then(res => res.flat())
             .then(completions => {
-                if (signal.aborted) {
+                // Cache even if the request was aborted.
+                this.cache.set({ params }, { result: { completions } })
+
+                if (signal?.aborted) {
                     throw new Error('aborted')
                 }
-                return completions
+
+                return { completions, cacheHit: false }
             })
     }
 }
+
+interface RequestCacheKey {
+    params: RequestParams
+}
+
+interface RequestCacheEntry {
+    result: Omit<RequestManagerResult, 'cacheHit'>
+}
+
+class RequestCache {
+    private cache = new LRUCache<string, RequestCacheEntry>({ max: 50 })
+
+    private toCacheKey(key: RequestCacheKey): string {
+        return key.params.prefix
+    }
+
+    public get(key: RequestCacheKey): RequestCacheEntry | undefined {
+        return this.cache.get(this.toCacheKey(key))
+    }
+
+    public set(key: RequestCacheKey, entry: RequestCacheEntry): void {
+        this.cache.set(this.toCacheKey(key), entry)
+    }
+}

vscode/src/completions/testHelpers.ts

@@ -36,7 +36,8 @@ const CURSOR_MARKER = '█'
 
 export function documentAndPosition(
     textWithCursor: string,
-    languageId = 'typescript'
+    languageId = 'typescript',
+    uriString = 'file:///test.ts'
 ): { document: VSCodeTextDocument; position: VSCodePosition } {
     const cursorIndex = textWithCursor.indexOf(CURSOR_MARKER)
     if (cursorIndex === -1) {
@@ -45,7 +46,7 @@ export function documentAndPosition(
     const prefix = textWithCursor.slice(0, cursorIndex)
     const suffix = textWithCursor.slice(cursorIndex + CURSOR_MARKER.length)
     const codeWithoutCursor = prefix + suffix
-    const document = wrapVSCodeTextDocument(TextDocument.create('file:///test.ts', languageId, 0, codeWithoutCursor))
+    const document = wrapVSCodeTextDocument(TextDocument.create(uriString, languageId, 0, codeWithoutCursor))
     const position = document.positionAt(cursorIndex)
     return { document, position }
 }