Automatically start Embeddings indexing (#4091)

sourcegraph · May 20, 2024 · b94620e · b94620e
1 parent 03805ce
commit b94620e
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 6 deletions.
diff --git a/lib/shared/src/experimentation/FeatureFlagProvider.ts b/lib/shared/src/experimentation/FeatureFlagProvider.ts
@@ -53,6 +53,9 @@ export enum FeatureFlag {
 
     /** Interactive tutorial, primarily for onboarding */
     CodyInteractiveTutorial = 'cody-interactive-tutorial',
+
+    /** Automatically start indexing using embeddings. */
+    CodyEmbeddingsAutoIndexing = 'cody-embeddings-auto-indexing',
 }
 
 const ONE_HOUR = 60 * 60 * 1000

diff --git a/vscode/playwright.config.ts b/vscode/playwright.config.ts
@@ -9,6 +9,6 @@ export default defineConfig({
     testDir: 'test/e2e',
     timeout: 30000,
     expect: {
-        timeout: isWin ? 5000 : 3000,
+        timeout: isWin ? 5000 : 5000,
     },
 })
diff --git a/vscode/src/local-context/local-embeddings.ts b/vscode/src/local-context/local-embeddings.ts
@@ -36,7 +36,10 @@ export async function createLocalEmbeddingsController(
         ((await featureFlagProvider.evaluateFeatureFlag(FeatureFlag.CodyUseSourcegraphEmbeddings))
             ? sourcegraphModelConfig
             : openaiModelConfig)
-    return new LocalEmbeddingsController(context, config, modelConfig)
+    const autoIndexingEnabled = await featureFlagProvider.evaluateFeatureFlag(
+        FeatureFlag.CodyEmbeddingsAutoIndexing
+    )
+    return new LocalEmbeddingsController(context, config, modelConfig, autoIndexingEnabled)
 }
 
 export type LocalEmbeddingsConfig = Pick<
@@ -129,7 +132,8 @@ export class LocalEmbeddingsController
     constructor(
         private readonly context: vscode.ExtensionContext,
         config: LocalEmbeddingsConfig,
-        private readonly modelConfig: EmbeddingsModelConfig
+        private readonly modelConfig: EmbeddingsModelConfig,
+        private readonly autoIndexingEnabled: boolean
     ) {
         logDebug('LocalEmbeddingsController', 'constructor')
         this.disposables.push(this.changeEmitter, this.statusEmitter)
@@ -161,7 +165,13 @@ export class LocalEmbeddingsController
         await this.getService()
         const repoUri = vscode.workspace.workspaceFolders?.[0]?.uri
         if (repoUri && isFileURI(repoUri)) {
-            await this.eagerlyLoad(repoUri)
+            const loadedOk = await this.eagerlyLoad(repoUri)
+            if (!loadedOk) {
+                // failed to load the index, let's see if we should start indexing
+                if (this.canAutoIndex()) {
+                    this.index()
+                }
+            }
         }
     }
 
@@ -495,9 +505,20 @@ export class LocalEmbeddingsController
         }
         this.lastHealth = health
         const hasIssue = health.numItemsNeedEmbedding > 0
-        await vscode.commands.executeCommand('setContext', 'cody.embeddings.hasIssue', hasIssue)
         if (hasIssue) {
-            this.updateIssueStatusBar()
+            const canRetry = this.canAutoIndex() && !this.lastError
+            let retrySucceeded = true
+            if (canRetry) {
+                try {
+                    await this.indexRetry()
+                } catch {
+                    retrySucceeded = false
+                }
+            }
+            if (!canRetry || !retrySucceeded) {
+                await vscode.commands.executeCommand('setContext', 'cody.embeddings.hasIssue', hasIssue)
+                this.updateIssueStatusBar()
+            }
         }
     }
 
@@ -514,6 +535,11 @@ export class LocalEmbeddingsController
         } is only ${percentDone.toFixed(0)}% complete.${options?.suffix || ''}`
     }
 
+    // Check if auto-indexing is enabled and if we're using the Sourcegraph provider.
+    private canAutoIndex(): boolean {
+        return this.autoIndexingEnabled && this.modelConfig.provider === 'sourcegraph'
+    }
+
     private updateIssueStatusBar(): void {
         this.statusBar?.dispose()
         this.statusBar = vscode.window.createStatusBarItem(

diff --git a/vscode/src/main.ts b/vscode/src/main.ts
@@ -313,6 +313,11 @@ const register = async (
         searchViewProvider.initialize()
     }
 
+    if (localEmbeddings) {
+        // kick-off embeddings initialization
+        localEmbeddings.start()
+    }
+
     if (config.experimentalSupercompletions) {
         disposables.push(new SupercompletionProvider({ statusBar, chat: chatClient }))
     }