Skip to content

Commit f71de47

Browse files
Will-ShaoHuaWill-ShaoHua
committed
refactor
1 parent 043296d commit f71de47

File tree

2 files changed

+128
-91
lines changed

2 files changed

+128
-91
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ export const vsCodeState: VsCodeState = {
4747

4848
export type UtgStrategy = 'ByName' | 'ByContent'
4949

50-
export type CrossFileStrategy = 'OpenTabs_BM25' | 'Codemap' | 'Project_BM25' | 'Project'
50+
export type CrossFileStrategy = 'opentabs' | 'codemap' | 'bm25' | 'default'
5151

5252
export type SupplementalContextStrategy = CrossFileStrategy | UtgStrategy | 'Empty'
5353

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

Lines changed: 127 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -68,99 +68,22 @@ export async function fetchSupplementalContextForSrc(
6868
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
6969
const supplementalContextConfig = getSupplementalContextConfig(editor.document.languageId)
7070

71-
// not supported case
72-
if (supplementalContextConfig === 'none') {
73-
return undefined
74-
}
75-
76-
// opentabs context will use bm25 and users' open tabs to fetch supplemental context
77-
if (supplementalContextConfig === 'opentabs') {
78-
return {
79-
supplementalContextItems: (await fetchOpentabsContext(editor, cancellationToken)) ?? [],
80-
strategy: 'OpenTabs_BM25',
81-
}
82-
}
83-
84-
// codemap will use opentabs context plus repomap if it's present
85-
if (supplementalContextConfig === 'codemap') {
86-
const opentabsContextAndCodemap = await waitUntil(
87-
async function () {
88-
const result: CodeWhispererSupplementalContextItem[] = []
89-
const opentabsContext = await fetchOpentabsContext(editor, cancellationToken)
90-
const codemap = await fetchProjectContext(editor, 'codemap')
91-
92-
if (codemap && codemap.length > 0) {
93-
result.push(...codemap)
94-
}
95-
96-
if (opentabsContext && opentabsContext.length > 0) {
97-
result.push(...opentabsContext)
98-
}
99-
100-
return result
101-
},
102-
{ timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
103-
)
104-
105-
return {
106-
supplementalContextItems: opentabsContextAndCodemap ?? [],
107-
strategy: 'Codemap',
108-
}
109-
}
110-
111-
// fallback to opentabs if projectContext timeout for 'default' | 'bm25'
112-
const opentabsContextPromise = waitUntil(
113-
async function () {
114-
return await fetchOpentabsContext(editor, cancellationToken)
115-
},
116-
{ timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
117-
)
118-
119-
// global bm25 without repomap
120-
if (supplementalContextConfig === 'bm25') {
121-
const projectBM25Promise = waitUntil(
122-
async function () {
123-
return await fetchProjectContext(editor, 'bm25')
124-
},
125-
{ timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
126-
)
71+
switch (supplementalContextConfig) {
72+
// not supported case
73+
case 'none':
74+
return undefined
12775

128-
const [projectContext, opentabsContext] = await Promise.all([projectBM25Promise, opentabsContextPromise])
129-
if (projectContext && projectContext.length > 0) {
130-
return {
131-
supplementalContextItems: projectContext,
132-
strategy: 'Project_BM25',
133-
}
134-
}
135-
136-
return {
137-
supplementalContextItems: opentabsContext ?? [],
138-
strategy: 'OpenTabs_BM25',
139-
}
140-
}
76+
case 'opentabs':
77+
return await new OpenTabsBm25().fetchProjectContext(editor, cancellationToken)
14178

142-
// global bm25 with repomap
143-
const projectContextAndCodemapPromise = waitUntil(
144-
async function () {
145-
return await fetchProjectContext(editor, 'default')
146-
},
147-
{ timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
148-
)
79+
case 'codemap':
80+
return await new OpenTabsBm25Codemap().fetchProjectContext(editor, cancellationToken)
14981

150-
const [projectContext, opentabsContext] = await Promise.all([
151-
projectContextAndCodemapPromise,
152-
opentabsContextPromise,
153-
])
154-
if (projectContext && projectContext.length > 0) {
155-
return {
156-
supplementalContextItems: projectContext,
157-
strategy: 'Project',
158-
}
159-
}
82+
case 'bm25':
83+
return await new GlobalBm25().fetchProjectContext(editor, cancellationToken)
16084

161-
return {
162-
supplementalContextItems: opentabsContext ?? [],
163-
strategy: 'OpenTabs_BM25',
85+
case 'default':
86+
return await new GlobalBm25Codemap().fetchProjectContext(editor, cancellationToken)
16487
}
16588
}
16689

@@ -183,7 +106,7 @@ export async function fetchProjectContext(
183106
export async function fetchOpentabsContext(
184107
editor: vscode.TextEditor,
185108
cancellationToken: vscode.CancellationToken
186-
): Promise<CodeWhispererSupplementalContextItem[] | undefined> {
109+
): Promise<CodeWhispererSupplementalContextItem[]> {
187110
const codeChunksCalculated = crossFileContextConfig.numberOfChunkToFetch
188111

189112
// Step 1: Get relevant cross files to refer
@@ -383,3 +306,117 @@ function throwIfCancelled(token: vscode.CancellationToken): void | never {
383306
throw new ToolkitError(supplemetalContextFetchingTimeoutMsg, { cause: new CancellationError('timeout') })
384307
}
385308
}
309+
310+
const timeoutConfig = { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
311+
312+
interface ProjectContextVariation {
313+
fetchProjectContext(
314+
editor: vscode.TextEditor,
315+
cancellationToken: vscode.CancellationToken
316+
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'>>
317+
}
318+
319+
export class OpenTabsBm25 implements ProjectContextVariation {
320+
async fetchProjectContext(
321+
editor: vscode.TextEditor,
322+
cancellationToken: vscode.CancellationToken
323+
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'>> {
324+
const contextItems =
325+
(await waitUntil(async function () {
326+
return await fetchOpentabsContext(editor, cancellationToken)
327+
}, timeoutConfig)) ?? []
328+
return {
329+
supplementalContextItems: contextItems,
330+
strategy: 'opentabs',
331+
}
332+
}
333+
}
334+
335+
export class OpenTabsBm25Codemap implements ProjectContextVariation {
336+
async fetchProjectContext(
337+
editor: vscode.TextEditor,
338+
cancellationToken: vscode.CancellationToken
339+
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'>> {
340+
const contextItems =
341+
(await waitUntil(async function () {
342+
const result: CodeWhispererSupplementalContextItem[] = []
343+
const opentabsContext = await fetchOpentabsContext(editor, cancellationToken)
344+
const codemap = await fetchProjectContext(editor, 'codemap')
345+
346+
if (codemap && codemap.length > 0) {
347+
result.push(...codemap)
348+
}
349+
350+
if (opentabsContext && opentabsContext.length > 0) {
351+
result.push(...opentabsContext)
352+
}
353+
354+
return result
355+
}, timeoutConfig)) ?? []
356+
return {
357+
supplementalContextItems: contextItems,
358+
strategy: 'codemap',
359+
}
360+
}
361+
}
362+
363+
export class GlobalBm25 implements ProjectContextVariation {
364+
async fetchProjectContext(
365+
editor: vscode.TextEditor,
366+
cancellationToken: vscode.CancellationToken
367+
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'>> {
368+
// fallback to opentabs if projectContext timeout
369+
const opentabsContextPromise = waitUntil(async function () {
370+
return await fetchOpentabsContext(editor, cancellationToken)
371+
}, timeoutConfig)
372+
373+
const projectBM25Promise = waitUntil(async function () {
374+
return await fetchProjectContext(editor, 'bm25')
375+
}, timeoutConfig)
376+
377+
const [projectContext, opentabsContext] = await Promise.all([projectBM25Promise, opentabsContextPromise])
378+
if (projectContext && projectContext.length > 0) {
379+
return {
380+
supplementalContextItems: projectContext,
381+
strategy: 'bm25',
382+
}
383+
}
384+
385+
return {
386+
supplementalContextItems: opentabsContext ?? [],
387+
strategy: 'opentabs',
388+
}
389+
}
390+
}
391+
392+
export class GlobalBm25Codemap implements ProjectContextVariation {
393+
async fetchProjectContext(
394+
editor: vscode.TextEditor,
395+
cancellationToken: vscode.CancellationToken
396+
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'>> {
397+
// fallback to opentabs if projectContext timeout
398+
const opentabsContextPromise = waitUntil(async function () {
399+
return await fetchOpentabsContext(editor, cancellationToken)
400+
}, timeoutConfig)
401+
402+
const projectContextAndCodemapPromise = waitUntil(async function () {
403+
return await fetchProjectContext(editor, 'default')
404+
}, timeoutConfig)
405+
406+
const [projectContext, opentabsContext] = await Promise.all([
407+
projectContextAndCodemapPromise,
408+
opentabsContextPromise,
409+
])
410+
if (projectContext && projectContext.length > 0) {
411+
return {
412+
supplementalContextItems: projectContext,
413+
strategy: 'default',
414+
}
415+
}
416+
417+
return {
418+
supplementalContextItems: opentabsContext ?? [],
419+
strategy: 'opentabs',
420+
}
421+
}
422+
}

0 commit comments

Comments
 (0)