internal data collection for chat workspace (#5330)

## Problem

The relevancy of workspace context retrieved from the LSP has unknown quality because of lack of data.

## Solution

As required, by default enable workspace index for a fraction of amzn users and always send relevantDocument to the service. There is no customer facing changes. This change is purely for data collection and will be removed soon.

To only apply this to a fraction of amzn users, we rely on the listFeatureEvaluations API in the backend.

Tested with builder id login, amzn login, with and without @workspace.

Author: leigaol

packages/core/src/codewhisperer/service/featureConfigProvider.ts

@@ -8,6 +8,8 @@ import { codeWhispererClient as client } from '../client/codewhisperer'
 import { AuthUtil } from '../util/authUtil'
 import { getLogger } from '../../shared/logger'
 import { isBuilderIdConnection, isIdcSsoConnection } from '../../auth/connection'
+import { CodeWhispererSettings } from '../util/codewhispererSettings'
+import globals from '../../shared/extensionGlobals'
 
 export class FeatureContext {
     constructor(
@@ -20,6 +22,7 @@ export class FeatureContext {
 const testFeatureName = 'testFeature'
 const customizationArnOverrideName = 'customizationArnOverride'
 const featureConfigPollIntervalInMs = 30 * 60 * 1000 // 30 mins
+const dataCollectionFeatureName = 'IDEProjectContextDataCollection'
 
 // TODO: add real feature later
 export const featureDefinitions = new Map([
@@ -35,6 +38,8 @@ export class FeatureConfigProvider {
 
     static #instance: FeatureConfigProvider
 
+    private _isDataCollectionGroup = false
+
     constructor() {
         this.fetchFeatureConfigs().catch((e) => {
             getLogger().error('fetchFeatureConfigs failed: %s', (e as Error).message)
@@ -47,6 +52,10 @@ export class FeatureConfigProvider {
         return (this.#instance ??= new this())
     }
 
+    isAmznDataCollectionGroup(): boolean {
+        return this._isDataCollectionGroup
+    }
+
     async fetchFeatureConfigs(): Promise<void> {
         if (AuthUtil.instance.isConnectionExpired()) {
             return
@@ -63,7 +72,6 @@ export class FeatureConfigProvider {
                     new FeatureContext(evaluation.feature, evaluation.variation, evaluation.value)
                 )
             })
-
             const customizationArnOverride = this.featureConfigs.get(customizationArnOverrideName)?.value?.stringValue
             if (customizationArnOverride !== undefined) {
                 // Double check if server-side wrongly returns a customizationArn to BID users
@@ -96,6 +104,17 @@ export class FeatureConfigProvider {
                     }
                 }
             }
+
+            const dataCollectionValue = this.featureConfigs.get(dataCollectionFeatureName)?.value.stringValue
+            if (dataCollectionValue === 'data-collection') {
+                this._isDataCollectionGroup = true
+                // Enable local workspace index by default, for Amzn users.
+                const isSet = globals.globalState.get<boolean>('aws.amazonq.workspaceIndexToggleOn') || false
+                if (!isSet) {
+                    await CodeWhispererSettings.instance.enableLocalIndex()
+                    globals.globalState.tryUpdate('aws.amazonq.workspaceIndexToggleOn', true)
+                }
+            }
         } catch (e) {
             getLogger().error(`CodeWhisperer: Error when fetching feature configs ${e}`, e)
         }

packages/core/src/codewhisperer/util/codewhispererSettings.ts

@@ -46,6 +46,10 @@ export class CodeWhispererSettings extends fromExtensionManifest('amazonQ', desc
         return this.get('workspaceIndex', false)
     }
 
+    public async enableLocalIndex() {
+        await this.update('workspaceIndex', true)
+    }
+
     public isLocalIndexGPUEnabled(): boolean {
         return this.get('workspaceIndexUseGPU', false)
     }

packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts

@@ -95,6 +95,7 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): Gen
     const relevantDocuments: RelevantTextDocument[] = triggerPayload.relevantTextDocuments
         ? triggerPayload.relevantTextDocuments
         : []
+    const useRelevantDocuments = triggerPayload.useRelevantDocuments
     // service will throw validation exception if string is empty
     const customizationArn: string | undefined = undefinedIfEmpty(triggerPayload.customization.arn)
 
@@ -110,6 +111,7 @@ export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): Gen
                             document,
                             cursorState,
                             relevantDocuments,
+                            useRelevantDocuments,
                         },
                     },
                     userIntent: triggerPayload.userIntent,

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -47,6 +47,7 @@ import { randomUUID } from '../../../shared/crypto'
 import { LspController } from '../../../amazonq/lsp/lspController'
 import { CodeWhispererSettings } from '../../../codewhisperer/util/codewhispererSettings'
 import { getSelectedCustomization } from '../../../codewhisperer/util/customizationUtil'
+import { FeatureConfigProvider } from '../../../codewhisperer/service/featureConfigProvider'
 import { getHttpStatusCode, getReasonFromSyntaxError } from '../../../shared/errors'
 
 export interface ChatControllerMessagePublishers {
@@ -557,9 +558,10 @@ export class ChatController {
             await this.messenger.sendAuthNeededExceptionMessage(credentialsState, tabID, triggerID)
             return
         }
+        triggerPayload.useRelevantDocuments = false
         if (triggerPayload.message) {
-            const userIntentEnableProjectContext = triggerPayload.message.includes(`@workspace`)
-            if (userIntentEnableProjectContext) {
+            triggerPayload.useRelevantDocuments = triggerPayload.message.includes(`@workspace`)
+            if (triggerPayload.useRelevantDocuments) {
                 triggerPayload.message = triggerPayload.message.replace(/@workspace/g, '')
                 if (CodeWhispererSettings.instance.isLocalIndexEnabled()) {
                     const start = performance.now()
@@ -575,6 +577,19 @@ export class ChatController {
                     return
                 }
             }
+            // if user does not have @workspace in the prompt, but user is in the data collection group
+            // If the user is in the data collection group but turned off local index to opt-out, do not collect data.
+            // TODO: Remove this entire block of code in one month as requested
+            else if (
+                FeatureConfigProvider.instance.isAmznDataCollectionGroup() &&
+                !LspController.instance.isIndexingInProgress() &&
+                CodeWhispererSettings.instance.isLocalIndexEnabled()
+            ) {
+                getLogger().info(`amazonq: User is in data collection group`)
+                const start = performance.now()
+                triggerPayload.relevantTextDocuments = await LspController.instance.query(triggerPayload.message)
+                triggerPayload.projectContextQueryLatencyMs = performance.now() - start
+            }
         }
 
         const request = triggerPayloadToChatRequest(triggerPayload)

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -128,7 +128,11 @@ export class Messenger {
             )
         }
         this.telemetryHelper.setResponseStreamStartTime(tabID)
-        if (triggerPayload.relevantTextDocuments && triggerPayload.relevantTextDocuments.length > 0) {
+        if (
+            triggerPayload.relevantTextDocuments &&
+            triggerPayload.relevantTextDocuments.length > 0 &&
+            triggerPayload.useRelevantDocuments === true
+        ) {
             this.telemetryHelper.setResponseFromProjectContext(messageID)
         }
 

packages/core/src/codewhispererChat/controllers/chat/model.ts

@@ -142,6 +142,7 @@ export interface TriggerPayload {
     readonly userIntent: UserIntent | undefined
     readonly customization: Customization
     relevantTextDocuments?: RelevantTextDocument[]
+    useRelevantDocuments?: boolean
 }
 
 export interface InsertedCode {

packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts

@@ -324,7 +324,7 @@ export class CWCTelemetryHelper {
             cwsprChatProgrammingLanguage: triggerPayload.fileLanguage,
             credentialStartUrl: AuthUtil.instance.startUrl,
             cwsprChatHasProjectContext: triggerPayload.relevantTextDocuments
-                ? triggerPayload.relevantTextDocuments.length > 0
+                ? triggerPayload.relevantTextDocuments.length > 0 && triggerPayload.useRelevantDocuments === true
                 : false,
             cwsprChatProjectContextQueryMs: triggerPayload.projectContextQueryLatencyMs,
         })
@@ -357,7 +357,7 @@ export class CWCTelemetryHelper {
             credentialStartUrl: AuthUtil.instance.startUrl,
             codewhispererCustomizationArn: triggerPayload.customization.arn,
             cwsprChatHasProjectContext: triggerPayload.relevantTextDocuments
-                ? triggerPayload.relevantTextDocuments.length > 0
+                ? triggerPayload.relevantTextDocuments.length > 0 && triggerPayload.useRelevantDocuments === true
                 : false,
         }
 

packages/core/src/shared/globalState.ts

@@ -37,6 +37,7 @@ type globalKey =
     | 'aws.redshift.connections'
     | 'aws.toolkit.amazonq.dismissed'
     | 'aws.toolkit.amazonqInstall.dismissed'
+    | 'aws.amazonq.workspaceIndexToggleOn'
     | 'aws.toolkit.separationPromptCommand'
     | 'aws.toolkit.separationPromptDismissed'
     // Deprecated/legacy names. New keys should start with "aws.".