Autocomplete: Fix codebase context inference for embeddings

vscode/CHANGELOG.md

@@ -10,6 +10,8 @@ Starting from `0.2.0`, Cody is using `major.EVEN_NUMBER.patch` for release versi
 
 ### Fixed
 
+- Bug: Fixes an issue where the codebase context was not correctly inferred to load embeddings context for autocomplete. [pull/525](https://github.com/sourcegraph/cody/pull/525)
+
 ### Changed
 
 ## [0.6.4]

vscode/src/chat/ContextProvider.ts

@@ -89,6 +89,7 @@ export class ContextProvider implements vscode.Disposable {
     }
 
     public async init(): Promise<void> {
+        await this.updateCodebaseContext()
         await this.publishContextStatus()
     }
 

vscode/src/completions/context-embeddings.ts

@@ -4,7 +4,7 @@ import { differenceInMinutes } from 'date-fns'
 import { LRUCache } from 'lru-cache'
 import * as vscode from 'vscode'
 
-import { CodebaseContext } from '@sourcegraph/cody-shared/src/codebase-context'
+import type { CodebaseContext } from '@sourcegraph/cody-shared/src/codebase-context'
 
 import type { ReferenceSnippet } from './context'
 import { logCompletionEvent } from './logger'
@@ -13,7 +13,7 @@ interface Options {
     document: vscode.TextDocument
     prefix: string
     suffix: string
-    codebaseContext: CodebaseContext
+    getCodebaseContext: () => CodebaseContext
 }
 
 interface EmbeddingsForFile {
@@ -26,23 +26,25 @@ const embeddingsPerFile = new LRUCache<string, EmbeddingsForFile>({
 })
 
 export function getContextFromEmbeddings(options: Options): ReferenceSnippet[] {
-    const { document, codebaseContext, prefix, suffix } = options
+    const { document, getCodebaseContext, prefix, suffix } = options
 
     const currentFilePath = path.normalize(document.fileName)
     const embeddingsForCurrentFile = embeddingsPerFile.get(currentFilePath)
 
     /**
-     * Fetch embeddings if we don't have any or if the last fetch was more than 5 minutes ago.
-     * Ideally, we should fetch embeddings in the background if file significantly changed.
-     * We can use the `onDidChangeTextDocument` event with some diffing logic for that in the improved version.
+     * Fetch embeddings if we don't have any or if the last fetch was more than
+     * 5 minutes ago. Ideally, we should fetch embeddings in the background if
+     * file significantly changed. We can use the `onDidChangeTextDocument`
+     * event with some diffing logic for that in the improved version.
      */
     if (!embeddingsForCurrentFile || differenceInMinutes(embeddingsForCurrentFile.lastChange, new Date()) > 5) {
         fetchAndSaveEmbeddings({
-            codebaseContext,
+            getCodebaseContext,
             currentFilePath,
-            // Use preifx + suffix to limit number of lines we send to the server.
-            // We can use the fullText here via `currentEditor.document.getText()` but
-            // it can negatively affect the embeddings quality and price.
+            // Use prefix + suffix to limit number of lines we send to the
+            // server. We can use the fullText here via
+            // `currentEditor.document.getText()` but it can negatively affect
+            // the embeddings quality and price.
             text: prefix + suffix,
         }).catch(console.error)
     }
@@ -54,7 +56,7 @@ export function getContextFromEmbeddings(options: Options): ReferenceSnippet[] {
 interface FetchEmbeddingsOptions {
     currentFilePath: string
     text: string
-    codebaseContext: CodebaseContext
+    getCodebaseContext: () => CodebaseContext
 }
 
 const NUM_CODE_RESULTS = 2
@@ -63,7 +65,8 @@ const NUM_CODE_RESULTS_EXTRA = 5
 const NUM_TEXT_RESULTS = 1
 
 async function fetchAndSaveEmbeddings(options: FetchEmbeddingsOptions): Promise<void> {
-    const { currentFilePath, text, codebaseContext } = options
+    const { currentFilePath, text, getCodebaseContext } = options
+    const codebaseContext = getCodebaseContext()
 
     if (!codebaseContext.checkEmbeddingsConnection()) {
         return

vscode/src/completions/context.ts

@@ -21,7 +21,7 @@ export interface GetContextOptions {
     suffix: string
     jaccardDistanceWindowSize: number
     maxChars: number
-    codebaseContext: CodebaseContext
+    getCodebaseContext: () => CodebaseContext
     isEmbeddingsContextEnabled?: boolean
 }
 
@@ -36,7 +36,7 @@ export interface GetContextResult {
 
 export async function getContext(options: GetContextOptions): Promise<GetContextResult> {
     const { maxChars, isEmbeddingsContextEnabled } = options
-    const start = Date.now()
+    const start = performance.now()
 
     /**
      * The embeddings context is sync to retrieve to keep the completions latency minimal. If it's
@@ -53,6 +53,8 @@ export async function getContext(options: GetContextOptions): Promise<GetContext
     const context: ReferenceSnippet[] = []
     let totalChars = 0
     function addMatch(match: ReferenceSnippet): boolean {
+        // TODO(@philipp-spiess): We should de-dupe on the snippet range and not
+        // the file name to allow for more than one snippet of the same file
         if (usedFilenames.has(match.fileName)) {
             return false
         }
@@ -66,6 +68,8 @@ export async function getContext(options: GetContextOptions): Promise<GetContext
         return true
     }
 
+    // TODO(@philipp-spiess): Rank embedding results against local context
+    // snippets instead of always appending embedding results at the top.
     let includedEmbeddingsMatches = 0
     for (const match of embeddingsMatches) {
         if (addMatch(match)) {
@@ -84,7 +88,7 @@ export async function getContext(options: GetContextOptions): Promise<GetContext
         logSummary: {
             ...(includedEmbeddingsMatches ? { embeddings: includedEmbeddingsMatches } : {}),
             ...(includedLocalMatches ? { local: includedLocalMatches } : {}),
-            duration: Date.now() - start,
+            duration: performance.now() - start,
         },
     }
 }

vscode/src/completions/getInlineCompletions.ts

@@ -39,7 +39,7 @@ export interface InlineCompletionsParams {
 
     // Injected
     contextFetcher?: (options: GetContextOptions) => Promise<GetContextResult>
-    codebaseContext?: CodebaseContext
+    getCodebaseContext?: () => CodebaseContext
     documentHistory?: DocumentHistory
 
     // Shared
@@ -144,7 +144,7 @@ async function doGetInlineCompletions({
     isEmbeddingsContextEnabled,
     toWorkspaceRelativePath,
     contextFetcher,
-    codebaseContext,
+    getCodebaseContext,
     documentHistory,
     requestManager,
     lastCandidate: lastCandidate,
@@ -211,7 +211,7 @@ async function doGetInlineCompletions({
         promptChars,
         isEmbeddingsContextEnabled,
         contextFetcher,
-        codebaseContext,
+        getCodebaseContext,
         documentHistory,
         docContext,
     })
@@ -393,7 +393,7 @@ interface GetCompletionContextParams
         | 'promptChars'
         | 'isEmbeddingsContextEnabled'
         | 'contextFetcher'
-        | 'codebaseContext'
+        | 'getCodebaseContext'
         | 'documentHistory'
     > {
     docContext: DocumentContext
@@ -404,15 +404,15 @@ async function getCompletionContext({
     promptChars,
     isEmbeddingsContextEnabled,
     contextFetcher,
-    codebaseContext,
+    getCodebaseContext,
     documentHistory,
     docContext: { prefix, suffix },
 }: GetCompletionContextParams): Promise<GetContextResult | null> {
     if (!contextFetcher) {
         return null
     }
-    if (!codebaseContext) {
-        throw new Error('codebaseContext is required if contextFetcher is provided')
+    if (!getCodebaseContext) {
+        throw new Error('getCodebaseContext is required if contextFetcher is provided')
     }
     if (!documentHistory) {
         throw new Error('documentHistory is required if contextFetcher is provided')
@@ -425,7 +425,7 @@ async function getCompletionContext({
         history: documentHistory,
         jaccardDistanceWindowSize: SNIPPET_WINDOW_SIZE,
         maxChars: promptChars,
-        codebaseContext,
+        getCodebaseContext,
         isEmbeddingsContextEnabled,
     })
 }

vscode/src/completions/vscodeInlineCompletionItemProvider.test.ts

@@ -43,7 +43,7 @@ class MockableInlineCompletionItemProvider extends InlineCompletionItemProvider
             // we can just make them `null`.
             //
             // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-explicit-any
-            codebaseContext: null as any,
+            getCodebaseContext: null as any,
             // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-explicit-any
             history: null as any,
             // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-explicit-any

vscode/src/completions/vscodeInlineCompletionItemProvider.ts

@@ -22,7 +22,7 @@ interface CodyCompletionItemProviderConfig {
     providerConfig: ProviderConfig
     history: DocumentHistory
     statusBar: CodyStatusBar
-    codebaseContext: CodebaseContext
+    getCodebaseContext: () => CodebaseContext
     responsePercentage?: number
     prefixPercentage?: number
     suffixPercentage?: number
@@ -138,7 +138,7 @@ export class InlineCompletionItemProvider implements vscode.InlineCompletionItem
             isEmbeddingsContextEnabled: this.config.isEmbeddingsContextEnabled,
             toWorkspaceRelativePath: uri => vscode.workspace.asRelativePath(uri),
             contextFetcher: this.config.contextFetcher,
-            codebaseContext: this.config.codebaseContext,
+            getCodebaseContext: this.config.getCodebaseContext,
             documentHistory: this.config.history,
             requestManager: this.requestManager,
             lastCandidate: this.lastCandidate,

vscode/src/main.ts

@@ -1,7 +1,6 @@
 import * as vscode from 'vscode'
 
 import { RecipeID } from '@sourcegraph/cody-shared/src/chat/recipes/recipe'
-import { CodebaseContext } from '@sourcegraph/cody-shared/src/codebase-context'
 import { Configuration, ConfigurationWithAccessToken } from '@sourcegraph/cody-shared/src/configuration'
 import { SourcegraphCompletionsClient } from '@sourcegraph/cody-shared/src/sourcegraph-api/completions/client'
 
@@ -113,7 +112,7 @@ const register = async (
 
     const {
         intentDetector,
-        codebaseContext,
+        codebaseContext: initialCodebaseContext,
         chatClient,
         completionsClient,
         guardrails,
@@ -126,7 +125,7 @@ const register = async (
     const contextProvider = new ContextProvider(
         initialConfig,
         chatClient,
-        codebaseContext,
+        initialCodebaseContext,
         editor,
         secretStorage,
         localStorage,
@@ -136,6 +135,7 @@ const register = async (
         platform
     )
     disposables.push(contextProvider)
+    await contextProvider.init()
 
     // Shared configuration that is required for chat views to send and receive messages
     const messageProviderOptions: MessageProviderOptions = {
@@ -339,7 +339,7 @@ const register = async (
             webviewErrorMessenger,
             completionsClient,
             statusBar,
-            codebaseContext
+            contextProvider
         )
     }
 
@@ -371,7 +371,7 @@ const register = async (
                 webviewErrorMessenger,
                 completionsClient,
                 statusBar,
-                codebaseContext
+                contextProvider
             )
         }
     })
@@ -411,7 +411,7 @@ function createCompletionsProvider(
     webviewErrorMessenger: (error: string) => Promise<void>,
     completionsClient: SourcegraphCompletionsClient,
     statusBar: CodyStatusBar,
-    codebaseContext: CodebaseContext
+    contextProvider: ContextProvider
 ): vscode.Disposable {
     const disposables: vscode.Disposable[] = []
 
@@ -421,7 +421,7 @@ function createCompletionsProvider(
         providerConfig,
         history,
         statusBar,
-        codebaseContext,
+        getCodebaseContext: () => contextProvider.context,
         isEmbeddingsContextEnabled: config.autocompleteAdvancedEmbeddings,
         completeSuggestWidgetSelection: config.autocompleteExperimentalCompleteSuggestWidgetSelection,
     })

vscode/test/completions/run-code-completions-on-dataset.ts

@@ -61,7 +61,7 @@ async function initCompletionsProvider(context: GetContextResult): Promise<Inlin
             dispose: () => {},
         },
         history,
-        codebaseContext,
+        getCodebaseContext: () => codebaseContext,
         isEmbeddingsContextEnabled: true,
         contextFetcher: () => Promise.resolve(context),
     })