Skip to content

Commit

Permalink
config(amazonq): inline project context AB aws#5994
Browse files Browse the repository at this point in the history
## 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
  • Loading branch information
Will-ShaoHua authored Nov 14, 2024
1 parent 4531668 commit 127a7ff
Show file tree
Hide file tree
Showing 9 changed files with 290 additions and 38 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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,
Expand All @@ -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 () {
Expand Down Expand Up @@ -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 })
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 () {
Expand Down
3 changes: 2 additions & 1 deletion packages/core/src/amazonq/lsp/lspClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
9 changes: 6 additions & 3 deletions packages/core/src/amazonq/lsp/lspController.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'

Expand Down Expand Up @@ -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 []
}
}
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/amazonq/lsp/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down
4 changes: 2 additions & 2 deletions packages/core/src/codewhisperer/models/model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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[] {
Expand Down Expand Up @@ -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'
}

/**
Expand Down
Loading

0 comments on commit 127a7ff

Please sign in to comment.