config(amazonq): inline project context AB #5994

Will-ShaoHua · web-flow · commit 127a7ff8f54d · 2024-11-14T09:17:53.000-08:00
## Problem
update inline suggestion project context ab test config

#### CONTROL
- use bm25 opentabs context 

#### TREATMENT_1
- use repomap + bm25 open tabs context 

#### TREATMENT_2
- use bm25 global project context
diff --git a/packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts
@@ -19,6 +19,7 @@ import {
 } from 'aws-core-vscode/test'
 import { areEqual, normalize } from 'aws-core-vscode/shared'
 import * as path from 'path'
+import { LspController } from 'aws-core-vscode/amazonq'
 
 let tempFolder: string
 
@@ -39,8 +40,8 @@ describe('crossFileContextUtil', function () {
             tempFolder = (await createTestWorkspaceFolder()).uri.fsPath
         })
 
-        it('opentabs context should fetch 3 chunks and each chunk should contains 50 lines', async function () {
-            sinon.stub(FeatureConfigProvider.instance, 'isNewProjectContextGroup').alwaysReturned(false)
+        it('for control group, should return opentabs context where there will be 3 chunks and each chunk should contains 50 lines', async function () {
+            sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').alwaysReturned('control')
             await toTextEditor(aStringWithLineCount(200), 'CrossFile.java', tempFolder, { preview: false })
             const myCurrentEditor = await toTextEditor('', 'TargetFile.java', tempFolder, {
                 preview: false,
@@ -53,6 +54,110 @@ describe('crossFileContextUtil', function () {
             assert.strictEqual(actual.supplementalContextItems[1].content.split('\n').length, 50)
             assert.strictEqual(actual.supplementalContextItems[2].content.split('\n').length, 50)
         })
+
+        it('for t1 group, should return repomap + opentabs context', async function () {
+            await toTextEditor(aStringWithLineCount(200), 'CrossFile.java', tempFolder, { preview: false })
+            const myCurrentEditor = await toTextEditor('', 'TargetFile.java', tempFolder, {
+                preview: false,
+            })
+            sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').returns('t1')
+            sinon
+                .stub(LspController.instance, 'queryInlineProjectContext')
+                .withArgs(sinon.match.any, sinon.match.any, 'codemap')
+                .resolves([
+                    {
+                        content: 'foo',
+                        score: 0,
+                        filePath: 'q-inline',
+                    },
+                ])
+
+            const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor, fakeCancellationToken)
+            assert.ok(actual)
+            assert.ok(actual.supplementalContextItems.length === 4)
+            assert.strictEqual(actual?.strategy, 'codemap')
+            assert.deepEqual(actual?.supplementalContextItems[0], {
+                content: 'foo',
+                score: 0,
+                filePath: 'q-inline',
+            })
+
+            assert.strictEqual(actual.supplementalContextItems[1].content.split('\n').length, 50)
+            assert.strictEqual(actual.supplementalContextItems[2].content.split('\n').length, 50)
+            assert.strictEqual(actual.supplementalContextItems[3].content.split('\n').length, 50)
+        })
+
+        it('for t2 group, should return global bm25 context and no repomap', async function () {
+            await toTextEditor(aStringWithLineCount(200), 'CrossFile.java', tempFolder, { preview: false })
+            const myCurrentEditor = await toTextEditor('', 'TargetFile.java', tempFolder, {
+                preview: false,
+            })
+            sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').returns('t2')
+            sinon
+                .stub(LspController.instance, 'queryInlineProjectContext')
+                .withArgs(sinon.match.any, sinon.match.any, 'bm25')
+                .resolves([
+                    {
+                        content: 'foo',
+                        score: 5,
+                        filePath: 'foo.java',
+                    },
+                    {
+                        content: 'bar',
+                        score: 4,
+                        filePath: 'bar.java',
+                    },
+                    {
+                        content: 'baz',
+                        score: 3,
+                        filePath: 'baz.java',
+                    },
+                    {
+                        content: 'qux',
+                        score: 2,
+                        filePath: 'qux.java',
+                    },
+                    {
+                        content: 'quux',
+                        score: 1,
+                        filePath: 'quux.java',
+                    },
+                ])
+
+            const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor, fakeCancellationToken)
+            assert.ok(actual)
+            assert.ok(actual.supplementalContextItems.length === 5)
+            assert.strictEqual(actual?.strategy, 'bm25')
+
+            assert.deepEqual(actual?.supplementalContextItems[0], {
+                content: 'foo',
+                score: 5,
+                filePath: 'foo.java',
+            })
+
+            assert.deepEqual(actual?.supplementalContextItems[1], {
+                content: 'bar',
+                score: 4,
+                filePath: 'bar.java',
+            })
+            assert.deepEqual(actual?.supplementalContextItems[2], {
+                content: 'baz',
+                score: 3,
+                filePath: 'baz.java',
+            })
+
+            assert.deepEqual(actual?.supplementalContextItems[3], {
+                content: 'qux',
+                score: 2,
+                filePath: 'qux.java',
+            })
+
+            assert.deepEqual(actual?.supplementalContextItems[4], {
+                content: 'quux',
+                score: 1,
+                filePath: 'quux.java',
+            })
+        })
     })
 
     describe('non supported language should return undefined', function () {
@@ -212,7 +317,7 @@ describe('crossFileContextUtil', function () {
 
         fileExtLists.forEach((fileExt) => {
             it('should be non empty', async function () {
-                sinon.stub(FeatureConfigProvider.instance, 'isNewProjectContextGroup').alwaysReturned(false)
+                sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').alwaysReturned('control')
                 const editor = await toTextEditor('content-1', `file-1.${fileExt}`, tempFolder)
                 await toTextEditor('content-2', `file-2.${fileExt}`, tempFolder, { preview: false })
                 await toTextEditor('content-3', `file-3.${fileExt}`, tempFolder, { preview: false })
diff --git a/packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts
@@ -21,7 +21,7 @@ describe('supplementalContextUtil', function () {
 
     beforeEach(async function () {
         testFolder = await TestFolder.create()
-        sinon.stub(FeatureConfigProvider.instance, 'isNewProjectContextGroup').alwaysReturned(false)
+        sinon.stub(FeatureConfigProvider.instance, 'getProjectContextGroup').alwaysReturned('control')
     })
 
     afterEach(function () {
diff --git a/packages/core/src/amazonq/lsp/lspClient.ts b/packages/core/src/amazonq/lsp/lspClient.ts
@@ -103,11 +103,12 @@ export class LspClient {
         }
     }
 
-    async queryInlineProjectContext(query: string, path: string) {
+    async queryInlineProjectContext(query: string, path: string, target: 'default' | 'codemap' | 'bm25') {
         try {
             const request = JSON.stringify({
                 query: query,
                 filePath: path,
+                target,
             })
             const encrypted = await this.encrypt(request)
             const resp: any = await this.client?.sendRequest(QueryInlineProjectContextRequestType, encrypted)
diff --git a/packages/core/src/amazonq/lsp/lspController.ts b/packages/core/src/amazonq/lsp/lspController.ts
@@ -66,7 +66,7 @@ export interface Manifest {
 }
 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.27']
+const supportedLspServerVersions = ['0.1.28']
 
 const nodeBinName = process.platform === 'win32' ? 'node.exe' : 'node'
 
@@ -308,10 +308,13 @@ 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)
+            return await LspClient.instance.queryInlineProjectContext(query, path, target)
         } catch (e) {
+            if (e instanceof Error) {
+                getLogger().error(`unexpected error while querying inline project context, e=${e.message}`)
+            }
             return []
         }
     }
diff --git a/packages/core/src/amazonq/lsp/types.ts b/packages/core/src/amazonq/lsp/types.ts
@@ -51,6 +51,7 @@ export type QueryInlineProjectContextRequest = string
 export type QueryInlineProjectContextRequestPayload = {
     query: string
     filePath: string
+    target: string
 }
 export const QueryInlineProjectContextRequestType: RequestType<QueryInlineProjectContextRequest, any, any> =
     new RequestType('lsp/queryInlineProjectContext')
diff --git a/packages/core/src/codewhisperer/models/model.ts b/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' | 'codemap' | 'bm25' | 'default'
 
-export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty' | 'LSP'
+export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty'
 
 export interface CodeWhispererSupplementalContext {
     isUtg: boolean
diff --git a/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts b/packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts
@@ -53,54 +53,137 @@ interface Chunk {
     score?: number
 }
 
-type SupplementalContextConfig = 'none' | 'v1' | 'v2'
+/**
+ * `none`: supplementalContext is not supported
+ * `opentabs`: opentabs_BM25
+ * `codemap`: repomap + opentabs BM25
+ * `bm25`: global_BM25
+ * `default`: repomap + global_BM25
+ */
+type SupplementalContextConfig = 'none' | 'opentabs' | 'codemap' | 'bm25' | 'default'
 
 export async function fetchSupplementalContextForSrc(
     editor: vscode.TextEditor,
     cancellationToken: vscode.CancellationToken
 ): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
     const supplementalContextConfig = getSupplementalContextConfig(editor.document.languageId)
 
+    // not supported case
     if (supplementalContextConfig === 'none') {
         return undefined
     }
-    if (supplementalContextConfig === 'v1') {
-        return fetchSupplementalContextForSrcV1(editor, cancellationToken)
+
+    // opentabs context will use bm25 and users' open tabs to fetch supplemental context
+    if (supplementalContextConfig === 'opentabs') {
+        return {
+            supplementalContextItems: (await fetchOpentabsContext(editor, cancellationToken)) ?? [],
+            strategy: 'opentabs',
+        }
+    }
+
+    // codemap will use opentabs context plus repomap if it's present
+    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(
+
+    // global bm25 without repomap
+    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: 'bm25',
+            }
+        }
+
+        return {
+            supplementalContextItems: opentabsContext ?? [],
+            strategy: 'opentabs',
+        }
+    }
+
+    // global bm25 with repomap
+    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: 'default',
+        }
+    }
+
+    return {
+        supplementalContextItems: opentabsContext ?? [],
+        strategy: 'opentabs',
+    }
 }
 
-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[]> {
     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 +234,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 +281,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 'opentabs'
+
+        case 't1':
+            return 'codemap'
+
+        case 't2':
+            return 'bm25'
     }
-    return 'v1'
 }
 
 /**
diff --git a/packages/core/src/shared/featureConfig.ts b/packages/core/src/shared/featureConfig.ts
diff --git a/packages/core/src/test/shared/featureConfig.test.ts b/packages/core/src/test/shared/featureConfig.test.ts

Original file line number	Diff line number	Diff line change
`@@ -103,11 +103,12 @@ export class LspClient {`
`103`	`103`	`}`
`104`	`104`	`}`
`105`	`105`
`106`		`- async queryInlineProjectContext(query: string, path: string) {`
	`106`	`+ async queryInlineProjectContext(query: string, path: string, target: 'default' \| 'codemap' \| 'bm25') {`
`107`	`107`	`try {`
`108`	`108`	`const request = JSON.stringify({`
`109`	`109`	`query: query,`
`110`	`110`	`filePath: path,`
	`111`	`+ target,`
`111`	`112`	`})`
`112`	`113`	`const encrypted = await this.encrypt(request)`
`113`	`114`	`const resp: any = await this.client?.sendRequest(QueryInlineProjectContextRequestType, encrypted)`
Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ export interface Manifest {`
`66`	`66`	`}`
`67`	`67`	`const manifestUrl = 'https://aws-toolkit-language-servers.amazonaws.com/q-context/manifest.json'`
`68`	`68`	`// this LSP client in Q extension is only going to work with these LSP server versions`
`69`		`-const supportedLspServerVersions = ['0.1.27']`
	`69`	`+const supportedLspServerVersions = ['0.1.28']`
`70`	`70`
`71`	`71`	`const nodeBinName = process.platform === 'win32' ? 'node.exe' : 'node'`
`72`	`72`
`@@ -308,10 +308,13 @@ export class LspController {`
`308`	`308`	`return resp`
`309`	`309`	`}`
`310`	`310`
`311`		`- async queryInlineProjectContext(query: string, path: string) {`
	`311`	`+ async queryInlineProjectContext(query: string, path: string, target: 'bm25' \| 'codemap' \| 'default') {`
`312`	`312`	`try {`
`313`		`- return await LspClient.instance.queryInlineProjectContext(query, path)`
	`313`	`+ return await LspClient.instance.queryInlineProjectContext(query, path, target)`
`314`	`314`	`} catch (e) {`
	`315`	`+ if (e instanceof Error) {`
	`316`	+ getLogger().error(`unexpected error while querying inline project context, e=${e.message}`)
	`317`	`+ }`
`315`	`318`	`return []`
`316`	`319`	`}`
`317`	`320`	`}`
Original file line number	Diff line number	Diff line change
`@@ -51,6 +51,7 @@ export type QueryInlineProjectContextRequest = string`
`51`	`51`	`export type QueryInlineProjectContextRequestPayload = {`
`52`	`52`	`query: string`
`53`	`53`	`filePath: string`
	`54`	`+ target: string`
`54`	`55`	`}`
`55`	`56`	`export const QueryInlineProjectContextRequestType: RequestType<QueryInlineProjectContextRequest, any, any> =`
`56`	`57`	`new RequestType('lsp/queryInlineProjectContext')`