a

Author: Will-ShaoHua

packages/core/src/amazonq/lsp/lspController.ts

@@ -64,7 +64,7 @@ export interface Manifest {
         targets: Target[]
     }[]
 }
-const manifestUrl = 'https://ducvaeoffl85c.cloudfront.net/manifest-0.1.28.json'
+const manifestUrl = 'https://aws-toolkit-language-servers.amazonaws.com/q-context/manifest.json'
 // this LSP client in Q extension is only going to work with these LSP server versions
 const supportedLspServerVersions = ['0.1.28']
 
@@ -308,7 +308,7 @@ export class LspController {
         return resp
     }
 
-    async queryInlineProjectContext(query: string, path: string) {
+    async queryInlineProjectContext(query: string, path: string, target: 'bm25' | 'codemap' | 'default') {
         try {
             return await LspClient.instance.queryInlineProjectContext(query, path)
         } catch (e) {

packages/core/src/codewhisperer/models/model.ts

@@ -47,9 +47,9 @@ export const vsCodeState: VsCodeState = {
 
 export type UtgStrategy = 'ByName' | 'ByContent'
 
-export type CrossFileStrategy = 'OpenTabs_BM25'
+export type CrossFileStrategy = 'OpenTabs_BM25' | 'Codemap' | 'Project_BM25' | 'Project'
 
-export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty' | 'LSP'
+export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty'
 
 export interface CodeWhispererSupplementalContext {
     isUtg: boolean

packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts

@@ -53,7 +53,7 @@ interface Chunk {
     score?: number
 }
 
-type SupplementalContextConfig = 'none' | 'v1' | 'v2'
+type SupplementalContextConfig = 'none' | 'openbtabs' | 'codemap' | 'bm25' | 'default'
 
 export async function fetchSupplementalContextForSrc(
     editor: vscode.TextEditor,
@@ -64,43 +64,114 @@ export async function fetchSupplementalContextForSrc(
     if (supplementalContextConfig === 'none') {
         return undefined
     }
-    if (supplementalContextConfig === 'v1') {
-        return fetchSupplementalContextForSrcV1(editor, cancellationToken)
+
+    if (supplementalContextConfig === 'openbtabs') {
+        return {
+            supplementalContextItems: (await fetchOpentabsContext(editor, cancellationToken)) ?? [],
+            strategy: 'OpenTabs_BM25',
+        }
+    }
+
+    if (supplementalContextConfig === 'codemap') {
+        const opentabsContextAndCodemap = await waitUntil(
+            async function () {
+                const result: CodeWhispererSupplementalContextItem[] = []
+                const opentabsContext = await fetchOpentabsContext(editor, cancellationToken)
+                const codemap = await fetchProjectContext(editor, 'codemap')
+
+                if (codemap && codemap.length > 0) {
+                    result.push(...codemap)
+                }
+
+                if (opentabsContext && opentabsContext.length > 0) {
+                    result.push(...opentabsContext)
+                }
+
+                return result
+            },
+            { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
+        )
+
+        return {
+            supplementalContextItems: opentabsContextAndCodemap ?? [],
+            strategy: 'Codemap',
+        }
     }
-    const promiseV1 = waitUntil(
+
+    // fallback to opentabs if projectContext timeout for 'default' | 'bm25'
+    const opentabsContextPromise = waitUntil(
         async function () {
-            return await fetchSupplementalContextForSrcV1(editor, cancellationToken)
+            return await fetchOpentabsContext(editor, cancellationToken)
         },
         { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
     )
-    const promiseV2 = waitUntil(
+
+    if (supplementalContextConfig === 'bm25') {
+        const projectBM25Promise = waitUntil(
+            async function () {
+                return await fetchProjectContext(editor, 'bm25')
+            },
+            { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
+        )
+
+        const [projectContext, opentabsContext] = await Promise.all([projectBM25Promise, opentabsContextPromise])
+        if (projectContext && projectContext.length > 0) {
+            return {
+                supplementalContextItems: projectContext,
+                strategy: 'Project_BM25',
+            }
+        }
+
+        return {
+            supplementalContextItems: opentabsContext ?? [],
+            strategy: 'OpenTabs_BM25',
+        }
+    }
+
+    const projectContextAndCodemapPromise = waitUntil(
         async function () {
-            return await fetchSupplementalContextForSrcV2(editor)
+            return await fetchProjectContext(editor, 'default')
         },
         { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
     )
-    const [resultV1, resultV2] = await Promise.all([promiseV1, promiseV2])
-    return resultV2 ?? resultV1
+
+    const [projectContext, opentabsContext] = await Promise.all([
+        projectContextAndCodemapPromise,
+        opentabsContextPromise,
+    ])
+    if (projectContext && projectContext.length > 0) {
+        return {
+            supplementalContextItems: projectContext,
+            strategy: 'Project',
+        }
+    }
+
+    return {
+        supplementalContextItems: opentabsContext ?? [],
+        strategy: 'OpenTabs_BM25',
+    }
 }
 
-export async function fetchSupplementalContextForSrcV2(
-    editor: vscode.TextEditor
-): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
+export async function fetchProjectContext(
+    editor: vscode.TextEditor,
+    target: 'default' | 'codemap' | 'bm25'
+): Promise<CodeWhispererSupplementalContextItem[] | undefined> {
     const inputChunkContent = getInputChunk(editor)
 
     const inlineProjectContext: { content: string; score: number; filePath: string }[] =
-        await LspController.instance.queryInlineProjectContext(inputChunkContent.content, editor.document.uri.fsPath)
+        await LspController.instance.queryInlineProjectContext(
+            inputChunkContent.content,
+            editor.document.uri.fsPath,
+            target
+        )
 
-    return {
-        supplementalContextItems: [...inlineProjectContext],
-        strategy: 'LSP',
-    }
+    return inlineProjectContext
 }
 
-export async function fetchSupplementalContextForSrcV1(
+export async function fetchOpentabsContext(
     editor: vscode.TextEditor,
     cancellationToken: vscode.CancellationToken
-): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
+): Promise<CodeWhispererSupplementalContextItem[] | undefined> {
     const codeChunksCalculated = crossFileContextConfig.numberOfChunkToFetch
 
     // Step 1: Get relevant cross files to refer
@@ -151,10 +222,7 @@ export async function fetchSupplementalContextForSrcV1(
 
     // DO NOT send code chunk with empty content
     getLogger().debug(`CodeWhisperer finished fetching crossfile context out of ${relevantCrossFilePaths.length} files`)
-    return {
-        supplementalContextItems: supplementalContexts,
-        strategy: 'OpenTabs_BM25',
-    }
+    return supplementalContexts
 }
 
 function findBestKChunkMatches(chunkInput: Chunk, chunkReferences: Chunk[], k: number): Chunk[] {
@@ -201,10 +269,18 @@ function getSupplementalContextConfig(languageId: vscode.TextDocument['languageI
     if (!isCrossFileSupported(languageId)) {
         return 'none'
     }
-    if (FeatureConfigProvider.instance.isNewProjectContextGroup()) {
-        return 'v2'
+
+    const group = FeatureConfigProvider.instance.getProjectContextGroup()
+    switch (group) {
+        case 'control':
+            return 'openbtabs'
+
+        case 't1':
+            return 'codemap'
+
+        case 't2':
+            return 'bm25'
     }
-    return 'v1'
 }
 
 /**

packages/core/src/shared/featureConfig.ts

@@ -64,9 +64,22 @@ export class FeatureConfigProvider {
         return (this.#instance ??= new this())
     }
 
-    isNewProjectContextGroup(): boolean {
-        return true
-        return this.featureConfigs.get(Features.projectContextFeature)?.variation === 'TREATMENT'
+    getProjectContextGroup(): 'control' | 't1' | 't2' {
+        const variation = this.featureConfigs.get(Features.projectContextFeature)?.variation
+
+        switch (variation) {
+            case 'CONTROL':
+                return 'control'
+
+            case 'TREATMENT_1':
+                return 't1'
+
+            case 'TREATMENT_2':
+                return 't2'
+
+            default:
+                return 'control'
+        }
     }
 
     public async listFeatureEvaluations(): Promise<ListFeatureEvaluationsResponse> {