Skip to content

Commit 00d3a5f

Browse files
Will-ShaoHuaWill-ShaoHua
committed
update crossfile context to use entire project but not limited to open tabs
1 parent 3f80141 commit 00d3a5f

File tree

6 files changed

+179
-23
lines changed

6 files changed

+179
-23
lines changed

‎packages/core/src/amazonq/lsp/lspClient.ts‎

Lines changed: 84 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,19 @@ import * as jose from 'jose'
1717
import { Disposable, ExtensionContext } from 'vscode'
1818

1919
import { LanguageClient, LanguageClientOptions, ServerOptions, TransportKind } from 'vscode-languageclient'
20-
import { GetUsageRequestType, IndexRequestType, QueryRequestType, UpdateIndexRequestType, Usage } from './types'
20+
import {
21+
BuildIndexRequestPayload,
22+
BuildIndexRequestType,
23+
GetUsageRequestType,
24+
IndexConfig,
25+
IndexRequestType,
26+
QueryInlineProjectContextRequestType,
27+
QueryRequestType,
28+
UpdateIndexRequestType,
29+
UpdateIndexV2RequestPayload,
30+
UpdateIndexV2RequestType,
31+
Usage,
32+
} from './types'
2133
import { Writable } from 'stream'
2234
import { CodeWhispererSettings } from '../../codewhisperer/util/codewhispererSettings'
2335
import { fs, getLogger } from '../../shared'
@@ -78,6 +90,24 @@ export class LspClient {
7890
}
7991
}
8092

93+
// v2
94+
async indexFilesV2(paths: string[], rootPath: string, config: IndexConfig) {
95+
const payload: BuildIndexRequestPayload = {
96+
filePaths: paths,
97+
projectRoot: rootPath,
98+
config: config,
99+
language: '',
100+
}
101+
try {
102+
const encryptedRequest = await this.encrypt(JSON.stringify(payload))
103+
const resp = await this.client?.sendRequest(BuildIndexRequestType, encryptedRequest)
104+
return resp
105+
} catch (e) {
106+
getLogger().error(`LspClient: indexFilesV2 error: ${e}`)
107+
return undefined
108+
}
109+
}
110+
81111
async query(request: string) {
82112
try {
83113
const encryptedRequest = await this.encrypt(
@@ -93,6 +123,23 @@ export class LspClient {
93123
}
94124
}
95125

126+
async queryBM25(query: string, path: string) {
127+
try {
128+
const request = JSON.stringify({
129+
query: query,
130+
filePath: path,
131+
})
132+
133+
const encrpted = await this.encrypt(request)
134+
135+
let resp: any = await this.client?.sendRequest(QueryInlineProjectContextRequestType, encrpted)
136+
return resp
137+
} catch (e) {
138+
getLogger().error(`LspClient: query error: ${e}`)
139+
return []
140+
}
141+
}
142+
96143
async getLspServerUsage(): Promise<Usage | undefined> {
97144
if (this.client) {
98145
return (await this.client.sendRequest(GetUsageRequestType, '')) as Usage
@@ -113,6 +160,23 @@ export class LspClient {
113160
return undefined
114161
}
115162
}
163+
164+
// not yet account for file move
165+
// v2
166+
async updateIndexV2(filePath: string[], mode: 'update' | 'remove' | 'add') {
167+
const payload: UpdateIndexV2RequestPayload = {
168+
filePaths: filePath,
169+
updateMode: mode,
170+
}
171+
try {
172+
const encryptedRequest = await this.encrypt(JSON.stringify(payload))
173+
const resp = await this.client?.sendRequest(UpdateIndexV2RequestType, encryptedRequest)
174+
return resp
175+
} catch (e) {
176+
getLogger().error(`LspClient: updateIndexV2 error: ${e}`)
177+
return undefined
178+
}
179+
}
116180
}
117181
/**
118182
* Activates the language server, this will start LSP server running over IPC protocol.
@@ -197,15 +261,30 @@ export async function activate(extensionContext: ExtensionContext) {
197261
return
198262
}
199263
savedDocument = document.uri
200-
})
201-
)
202-
toDispose.push(
264+
void LspClient.instance.updateIndexV2([document.uri.fsPath], 'update')
265+
}),
203266
vscode.window.onDidChangeActiveTextEditor((editor) => {
204267
if (savedDocument && editor && editor.document.uri.fsPath !== savedDocument.fsPath) {
205-
void LspClient.instance.updateIndex(savedDocument.fsPath)
268+
// void LspClient.instance.updateIndexV2([editor.document.uri.fsPath], 'update')
206269
}
270+
}),
271+
vscode.workspace.onDidCreateFiles((e) => {
272+
void LspClient.instance.updateIndexV2(
273+
e.files.map((f) => f.fsPath),
274+
'add'
275+
)
276+
}),
277+
vscode.workspace.onDidDeleteFiles((e) => {
278+
void LspClient.instance.updateIndexV2(
279+
e.files.map((f) => f.fsPath),
280+
'remove'
281+
)
282+
}),
283+
vscode.workspace.onDidRenameFiles((e) => {
284+
// void LspClient.instance.updateIndexV2(e.files.map((f) => f.newUri.fsPath), 'rename')
207285
})
208286
)
287+
209288
return LspClient.instance.client.onReady().then(() => {
210289
const disposableFunc = { dispose: () => rangeFormatting?.dispose() as void }
211290
toDispose.push(disposableFunc)

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

Lines changed: 7 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -304,7 +304,7 @@ export class LspController {
304304
}
305305

306306
async buildIndex() {
307-
getLogger().info(`LspController: Starting to build vector index of project`)
307+
getLogger().info(`LspController: Starting LSP`)
308308
const start = performance.now()
309309
const projPaths = getProjectPaths()
310310
projPaths.sort()
@@ -325,13 +325,11 @@ export class LspController {
325325
0
326326
)
327327
getLogger().info(`LspController: Found ${files.length} files in current project ${getProjectPaths()}`)
328-
const resp = await LspClient.instance.indexFiles(
329-
files.map((f) => f.fileUri.fsPath),
330-
projRoot,
331-
false
332-
)
328+
const config = CodeWhispererSettings.instance.isLocalIndexEnabled() ? 'all' : 'default'
329+
const r = files.map((f) => f.fileUri.fsPath)
330+
const resp = await LspClient.instance.indexFilesV2(r, projRoot, config)
333331
if (resp) {
334-
getLogger().debug(`LspController: Finish building vector index of project`)
332+
getLogger().debug(`LspController: Finish building index of project`)
335333
const usage = await LspClient.instance.getLspServerUsage()
336334
telemetry.amazonq_indexWorkspace.emit({
337335
duration: performance.now() - start,
@@ -343,7 +341,7 @@ export class LspController {
343341
credentialStartUrl: AuthUtil.instance.startUrl,
344342
})
345343
} else {
346-
getLogger().error(`LspController: Failed to build vector index of project`)
344+
getLogger().error(`LspController: Failed to build index of project`)
347345
telemetry.amazonq_indexWorkspace.emit({
348346
duration: performance.now() - start,
349347
result: 'Failed',
@@ -352,7 +350,7 @@ export class LspController {
352350
})
353351
}
354352
} catch (e) {
355-
getLogger().error(`LspController: Failed to build vector index of project`)
353+
getLogger().error(`LspController: Failed to build index of project ${e}`)
356354
telemetry.amazonq_indexWorkspace.emit({
357355
duration: performance.now() - start,
358356
result: 'Failed',
@@ -371,12 +369,6 @@ export class LspController {
371369
return
372370
}
373371
setImmediate(async () => {
374-
if (!CodeWhispererSettings.instance.isLocalIndexEnabled()) {
375-
// only download LSP for users who did not turn on this feature
376-
// do not start LSP server
377-
await LspController.instance.tryInstallLsp(context)
378-
return
379-
}
380372
const ok = await LspController.instance.tryInstallLsp(context)
381373
if (!ok) {
382374
return

‎packages/core/src/amazonq/lsp/types.ts‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,3 +35,40 @@ export interface Usage {
3535
memoryUsage: number
3636
cpuUsage: number
3737
}
38+
39+
export type BuildIndexRequestPayload = {
40+
filePaths: string[]
41+
projectRoot: string
42+
config: string
43+
language: string
44+
}
45+
46+
export type BuildIndexRequest = string
47+
48+
export const BuildIndexRequestType: RequestType<BuildIndexRequest, any, any> = new RequestType('lsp/buildIndex')
49+
50+
export type UpdateIndexV2Request = string
51+
52+
export type UpdateIndexV2RequestPayload = { filePaths: string[]; updateMode: string }
53+
54+
export const UpdateIndexV2RequestType: RequestType<UpdateIndexV2Request, any, any> = new RequestType(
55+
'lsp/updateIndexV2'
56+
)
57+
58+
export type QueryInlineProjectContextRequest = string
59+
export type QueryInlineProjectContextRequestPayload = {
60+
query: string
61+
filePath: string
62+
}
63+
export const QueryInlineProjectContextRequestType: RequestType<QueryInlineProjectContextRequest, any, any> =
64+
new RequestType('lsp/queryInlineProjectContext')
65+
66+
export type QueryVectorIndexRequestPayload = { query: string }
67+
68+
export type QueryVectorIndexRequest = string
69+
70+
export const QueryVectorIndexRequestType: RequestType<QueryVectorIndexRequest, any, any> = new RequestType(
71+
'lsp/queryVectorIndex'
72+
)
73+
74+
export type IndexConfig = 'all' | 'default'

‎packages/core/src/codewhisperer/models/constants.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -690,6 +690,7 @@ export const crossFileContextConfig = {
690690
numberOfChunkToFetch: 60,
691691
topK: 3,
692692
numberOfLinesEachChunk: 10,
693+
codemapMark: 'q-inline',
693694
}
694695

695696
export const utgConfig = {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ export type UtgStrategy = 'ByName' | 'ByContent'
4848

4949
export type CrossFileStrategy = 'OpenTabs_BM25'
5050

51-
export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty'
51+
export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty' | 'LSP'
5252

5353
export interface CodeWhispererSupplementalContext {
5454
isUtg: boolean

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

Lines changed: 49 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import { getFileDistance } from '../../../shared/filesystemUtilities'
1616
import { getOpenFilesInWindow } from '../../../shared/utilities/editorUtilities'
1717
import { getLogger } from '../../../shared/logger/logger'
1818
import { CodeWhispererSupplementalContext, CodeWhispererSupplementalContextItem } from '../../models/model'
19+
import { LspClient } from '../../../amazonq'
1920

2021
type CrossFileSupportedLanguage =
2122
| 'java'
@@ -66,6 +67,51 @@ export async function fetchSupplementalContextForSrc(
6667
}
6768
}
6869

70+
// TODO:
71+
if (false) {
72+
return fetchSupplementalContextForSrcV1(editor, cancellationToken)
73+
} else {
74+
return fetchSupplementalContextForSrcV2(editor)
75+
}
76+
}
77+
78+
export async function fetchSupplementalContextForSrcV2(
79+
editor: vscode.TextEditor
80+
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
81+
const inputChunkContent = getInputChunk(editor)
82+
83+
const bm25Response: { content: string; score: number; filePath: string }[] = await LspClient.instance.queryBM25(
84+
inputChunkContent.content,
85+
editor.document.uri.fsPath
86+
)
87+
getLogger().info(JSON.stringify(bm25Response))
88+
console.log(bm25Response)
89+
90+
const supContextItems: CodeWhispererSupplementalContextItem[] = bm25Response
91+
return {
92+
supplementalContextItems: [...supContextItems],
93+
strategy: 'LSP',
94+
}
95+
}
96+
97+
export async function fetchSupplementalContextForSrcV1(
98+
editor: vscode.TextEditor,
99+
cancellationToken: vscode.CancellationToken
100+
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
101+
const shouldProceed = shouldFetchCrossFileContext(
102+
editor.document.languageId,
103+
CodeWhispererUserGroupSettings.instance.userGroup
104+
)
105+
106+
if (!shouldProceed) {
107+
return shouldProceed === undefined
108+
? undefined
109+
: {
110+
supplementalContextItems: [],
111+
strategy: 'Empty',
112+
}
113+
}
114+
69115
const codeChunksCalculated = crossFileContextConfig.numberOfChunkToFetch
70116

71117
// Step 1: Get relevant cross files to refer
@@ -91,7 +137,7 @@ export async function fetchSupplementalContextForSrc(
91137

92138
// Step 3: Generate Input chunk (10 lines left of cursor position)
93139
// and Find Best K chunks w.r.t input chunk using BM25
94-
const inputChunk: Chunk = getInputChunk(editor, crossFileContextConfig.numberOfLinesEachChunk)
140+
const inputChunk: Chunk = getInputChunk(editor)
95141
const bestChunks: Chunk[] = findBestKChunkMatches(inputChunk, chunkList, crossFileContextConfig.topK)
96142
throwIfCancelled(cancellationToken)
97143

@@ -137,7 +183,8 @@ function findBestKChunkMatches(chunkInput: Chunk, chunkReferences: Chunk[], k: n
137183
/* This extract 10 lines to the left of the cursor from trigger file.
138184
* This will be the inputquery to bm25 matching against list of cross-file chunks
139185
*/
140-
function getInputChunk(editor: vscode.TextEditor, chunkSize: number) {
186+
function getInputChunk(editor: vscode.TextEditor) {
187+
const chunkSize = crossFileContextConfig.numberOfLinesEachChunk
141188
const cursorPosition = editor.selection.active
142189
const startLine = Math.max(cursorPosition.line - chunkSize, 0)
143190
const endLine = Math.max(cursorPosition.line - 1, 0)

0 commit comments

Comments
 (0)