fix(codewhisperer): refactor BM25 algorithm impl, prioritize opening files when generating supplemental context (#3583)

* fix cross file not prioritize opened files in editors
* trimEnd() on every code chunk of cross file
* fix BM25 algorithm counting logic throwing NaN

Author: Will-ShaoHua

src/codewhisperer/models/constants.ts

@@ -267,3 +267,13 @@ export enum UserGroup {
 export const isClassifierEnabledKey = 'CODEWHISPERER_CLASSIFIER_TRIGGER_ENABLED'
 
 export const supplemetalContextFetchingTimeoutMsg = 'codewhisperer supplemental context fetching timeout'
+
+export const crossFileContextConfig = {
+    numberOfChunkToFetch: 60,
+    topK: 3,
+    numberOfLinesEachChunk: 10,
+}
+
+export const utgConfig = {
+    maxSegmentSize: 10200,
+}

src/codewhisperer/service/recommendationHandler.ts

@@ -53,7 +53,7 @@ export class RecommendationHandler {
     public isGenerateRecommendationInProgress: boolean
     private _onDidReceiveRecommendation: vscode.EventEmitter<void> = new vscode.EventEmitter<void>()
     public readonly onDidReceiveRecommendation: vscode.Event<void> = this._onDidReceiveRecommendation.event
-    private supplementalContextMetadata: Omit<CodeWhispererSupplementalContext, 'contents'> | undefined
+    private supplementalContextMetadata: Omit<CodeWhispererSupplementalContext, 'supplementalContextItems'> | undefined
 
     constructor() {
         this.requestId = ''
@@ -149,7 +149,7 @@ export class RecommendationHandler {
         let nextToken = ''
         let errorCode = ''
         let req: codewhispererClient.ListRecommendationsRequest | codewhispererClient.GenerateRecommendationsRequest
-        let supplementalContextMetadata: Omit<CodeWhispererSupplementalContext, 'contents'> | undefined
+        let supplementalContextMetadata: Omit<CodeWhispererSupplementalContext, 'supplementalContextItems'> | undefined
         let shouldRecordServiceInvocation = false
 
         if (pagination) {

src/codewhisperer/util/editorContext.ts

@@ -19,6 +19,7 @@ import { supplementalContextTimeoutInMs } from '../models/constants'
 import { CodeWhispererUserGroupSettings } from './userGroupUtil'
 import { isTestFile } from './supplementalContext/codeParsingUtil'
 import { DependencyGraphFactory } from './dependencyGraph/dependencyGraphFactory'
+import { selectFrom } from '../../shared/utilities/tsUtils'
 
 let tabSize: number = getTabSizeSetting()
 
@@ -80,7 +81,7 @@ export async function buildListRecommendationRequest(
     allowCodeWithReference: boolean | undefined = undefined
 ): Promise<{
     request: codewhispererClient.ListRecommendationsRequest
-    supplementalMetadata: Omit<CodeWhispererSupplementalContext, 'contents'> | undefined
+    supplementalMetadata: Omit<CodeWhispererSupplementalContext, 'supplementalContextItems'> | undefined
 }> {
     const fileContext = extractContextForCodeWhisperer(editor)
 
@@ -98,23 +99,30 @@ export async function buildListRecommendationRequest(
             ? await fetchSupplementalContext(editor, tokenSource.token)
             : undefined
 
-    const suppelmetalMetadata: Omit<CodeWhispererSupplementalContext, 'contents'> | undefined = supplementalContexts
-        ? {
-              isUtg: supplementalContexts.isUtg,
-              isProcessTimeout: supplementalContexts.isProcessTimeout,
-              contentsLength: supplementalContexts.contentsLength,
-              latency: supplementalContexts.latency,
-          }
-        : undefined
+    const suppelmetalMetadata: Omit<CodeWhispererSupplementalContext, 'supplementalContextItems'> | undefined =
+        supplementalContexts
+            ? {
+                  isUtg: supplementalContexts.isUtg,
+                  isProcessTimeout: supplementalContexts.isProcessTimeout,
+                  contentsLength: supplementalContexts.contentsLength,
+                  latency: supplementalContexts.latency,
+              }
+            : undefined
 
     logSupplementalContext(supplementalContexts)
 
+    const supplementalContext: codewhispererClient.SupplementalContext[] = supplementalContexts
+        ? supplementalContexts.supplementalContextItems.map(v => {
+              return selectFrom(v, 'content', 'filePath')
+          })
+        : []
+
     if (allowCodeWithReference === undefined) {
         return {
             request: {
                 fileContext: fileContext,
                 nextToken: nextToken,
-                supplementalContexts: supplementalContexts ? supplementalContexts.contents : [],
+                supplementalContexts: supplementalContext,
             },
             supplementalMetadata: suppelmetalMetadata,
         }
@@ -127,15 +135,15 @@ export async function buildListRecommendationRequest(
             referenceTrackerConfiguration: {
                 recommendationsWithReferences: allowCodeWithReference ? 'ALLOW' : 'BLOCK',
             },
-            supplementalContexts: supplementalContexts ? supplementalContexts.contents : [],
+            supplementalContexts: supplementalContext,
         },
         supplementalMetadata: suppelmetalMetadata,
     }
 }
 
 export async function buildGenerateRecommendationRequest(editor: vscode.TextEditor): Promise<{
     request: codewhispererClient.GenerateRecommendationsRequest
-    supplementalMetadata: Omit<CodeWhispererSupplementalContext, 'contents'> | undefined
+    supplementalMetadata: Omit<CodeWhispererSupplementalContext, 'supplementalContextItems'> | undefined
 }> {
     const fileContext = extractContextForCodeWhisperer(editor)
 
@@ -144,7 +152,7 @@ export async function buildGenerateRecommendationRequest(editor: vscode.TextEdit
         tokenSource.cancel()
     }, supplementalContextTimeoutInMs)
     const supplementalContexts = await fetchSupplementalContext(editor, tokenSource.token)
-    let supplemetalMetadata: Omit<CodeWhispererSupplementalContext, 'contents'> | undefined
+    let supplemetalMetadata: Omit<CodeWhispererSupplementalContext, 'supplementalContextItems'> | undefined
 
     if (supplementalContexts) {
         supplemetalMetadata = {
@@ -161,7 +169,7 @@ export async function buildGenerateRecommendationRequest(editor: vscode.TextEdit
         request: {
             fileContext: fileContext,
             maxResults: CodeWhispererConstants.maxRecommendations,
-            supplementalContexts: supplementalContexts?.contents ?? [],
+            supplementalContexts: supplementalContexts?.supplementalContextItems ?? [],
         },
         supplementalMetadata: supplemetalMetadata,
     }
@@ -228,10 +236,12 @@ function logSupplementalContext(supplementalContext: CodeWhispererSupplementalCo
             latency: ${supplementalContext.latency},
         `)
 
-    supplementalContext.contents.forEach((context, index) => {
+    supplementalContext.supplementalContextItems.forEach((context, index) => {
         getLogger().verbose(`
                 -----------------------------------------------
-                Chunk ${index}:${context.content}
+                Path: ${context.filePath}
+                Score: ${context.score}
+                Chunk: ${index}:${context.content}
                 -----------------------------------------------
             `)
     })

src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts

@@ -4,30 +4,28 @@
  */
 
 import * as vscode from 'vscode'
-import * as codewhispererClient from '../../client/codewhisperer'
 import * as fs from 'fs-extra'
 import { DependencyGraph } from '../dependencyGraph/dependencyGraph'
-import { BMDocument, performBM25Scoring } from './rankBm25'
-import { getRelevantFilesFromEditor, isRelevant } from './editorFilesUtil'
+import { BM25Document, BM25Okapi } from './rankBm25'
+import { isRelevant } from './editorFilesUtil'
 import { ToolkitError } from '../../../shared/errors'
-import { supplemetalContextFetchingTimeoutMsg } from '../../models/constants'
+import { crossFileContextConfig, supplemetalContextFetchingTimeoutMsg } from '../../models/constants'
 import { CancellationError } from '../../../shared/utilities/timeoutUtils'
+import { CodeWhispererSupplementalContextItem } from './supplementalContextUtil'
 
 const crossFileLanguageConfigs = ['java']
 interface Chunk {
     fileName: string
     content: string
     nextContent: string
+    score?: number
 }
-const chunkSize = 10
-const chunkCount = 60
-const topK = 3
 
 export async function fetchSupplementalContextForSrc(
     editor: vscode.TextEditor,
     dependencyGraph: DependencyGraph,
     cancellationToken: vscode.CancellationToken
-) {
+): Promise<CodeWhispererSupplementalContextItem[] | undefined> {
     if (crossFileLanguageConfigs.includes(editor.document.languageId) === false) {
         return undefined
     }
@@ -38,59 +36,66 @@ export async function fetchSupplementalContextForSrc(
     // Step 2: Split files to chunks with upper bound on chunkCount
     // We restrict the total number of chunks to improve on latency.
     // Chunk linking is required as we want to pass the next chunk value for matched chunk.
-    const chunkList: Chunk[] = []
+    let chunkList: Chunk[] = []
     for (const relevantFile of relevantCrossFilePaths) {
         throwIfCancelled(cancellationToken)
-
-        const chunks: Chunk[] = splitFileToChunks(relevantFile, chunkSize)
+        const chunks: Chunk[] = splitFileToChunks(relevantFile, crossFileContextConfig.numberOfLinesEachChunk)
         const linkedChunks = linkChunks(chunks)
         chunkList.push(...linkedChunks)
-        if (chunkList.length >= chunkCount) {
+        if (chunkList.length >= crossFileContextConfig.numberOfChunkToFetch) {
             break
         }
     }
 
+    // it's required since chunkList.push(...) is likely giving us a list of size > 60
+    chunkList = chunkList.slice(0, crossFileContextConfig.numberOfChunkToFetch)
+
     // Step 3: Generate Input chunk (10 lines left of cursor position)
     // and Find Best K chunks w.r.t input chunk using BM25
-    const inputChunk: Chunk = getInputChunk(editor, chunkSize)
-    const bestChunks: Chunk[] = findBestKChunkMatches(inputChunk, chunkList, topK)
+    const inputChunk: Chunk = getInputChunk(editor, crossFileContextConfig.numberOfLinesEachChunk)
+    const bestChunks: Chunk[] = findBestKChunkMatches(inputChunk, chunkList, crossFileContextConfig.topK)
     throwIfCancelled(cancellationToken)
 
     // Step 4: Transform best chunks to supplemental contexts
-    const supplementalContexts: codewhispererClient.SupplementalContext[] = []
+    const supplementalContexts: CodeWhispererSupplementalContextItem[] = []
     for (const chunk of bestChunks) {
         throwIfCancelled(cancellationToken)
 
-        const context = {
+        supplementalContexts.push({
             filePath: chunk.fileName,
             content: chunk.nextContent,
-        } as codewhispererClient.SupplementalContext
-        supplementalContexts.push(context)
+            score: chunk.score,
+        })
     }
 
     return supplementalContexts
 }
 
-function findBestKChunkMatches(chunkInput: Chunk, chunkReferences: Chunk[], k: number) {
+function findBestKChunkMatches(chunkInput: Chunk, chunkReferences: Chunk[], k: number): Chunk[] {
     const chunkContentList = chunkReferences.map(chunk => chunk.content)
+
     //performBM25Scoring returns the output in a sorted order (descending of scores)
-    const output: BMDocument[] = performBM25Scoring(chunkContentList, chunkInput.content) as BMDocument[]
-    const bestChunks: Chunk[] = []
-    //pick Top 3
-    for (let i = 0; i < Math.min(k, output.length); i++) {
-        const chunkIndex = output[i].index
+    const top3: BM25Document[] = new BM25Okapi(chunkContentList).topN(chunkInput.content, crossFileContextConfig.topK)
+
+    return top3.map(doc => {
+        // reference to the original metadata since BM25.top3 will sort the result
+        const chunkIndex = doc.index
         const chunkReference = chunkReferences[chunkIndex]
-        bestChunks.push(chunkReference)
-    }
-    return bestChunks
+        return {
+            content: chunkReference.content,
+            fileName: chunkReference.fileName,
+            nextContent: chunkReference.nextContent,
+            score: doc.score,
+        }
+    })
 }
 
 /* This extract 10 lines to the left of the cursor from trigger file.
  * This will be the inputquery to bm25 matching against list of cross-file chunks
  */
 function getInputChunk(editor: vscode.TextEditor, chunkSize: number) {
     const cursorPosition = editor.selection.active
-    const startLine = Math.max(cursorPosition.line - 10, 0)
+    const startLine = Math.max(cursorPosition.line - chunkSize, 0)
     const endLine = Math.max(cursorPosition.line - 1, 0)
     const inputChunkContent = editor.document.getText(
         new vscode.Range(startLine, 0, endLine, editor.document.lineAt(endLine).text.length)
@@ -109,7 +114,7 @@ function linkChunks(chunks: Chunk[]) {
 
     // This additional chunk is needed to create a next pointer to chunk 0.
     const firstChunk = chunks[0]
-    const firstChunkSubContent = firstChunk.content.split('\n').slice(0, 3).join('\n')
+    const firstChunkSubContent = firstChunk.content.split('\n').slice(0, 3).join('\n').trimEnd()
     const newFirstChunk = {
         fileName: firstChunk.fileName,
         content: firstChunkSubContent,
@@ -132,12 +137,12 @@ function linkChunks(chunks: Chunk[]) {
 function splitFileToChunks(filePath: string, chunkSize: number): Chunk[] {
     const chunks: Chunk[] = []
 
-    const fileContent = fs.readFileSync(filePath, 'utf-8')
+    const fileContent = fs.readFileSync(filePath, 'utf-8').trimEnd()
     const lines = fileContent.split('\n')
 
     for (let i = 0; i < lines.length; i += chunkSize) {
         const chunkContent = lines.slice(i, Math.min(i + chunkSize, lines.length)).join('\n')
-        const chunk = { fileName: filePath, content: chunkContent, nextContent: '' }
+        const chunk = { fileName: filePath, content: chunkContent.trimEnd(), nextContent: '' }
         chunks.push(chunk)
     }
     return chunks
@@ -148,7 +153,10 @@ function splitFileToChunks(filePath: string, chunkSize: number): Chunk[] {
  * by referencing open files, imported files and same package files.
  */
 async function getRelevantCrossFiles(editor: vscode.TextEditor, dependencyGraph: DependencyGraph): Promise<string[]> {
-    const srcDependencies = await dependencyGraph.getSourceDependencies(editor.document.uri, editor.document.getText())
+    const openedFilesInEditor = new Set(getOpenFilesInWindow())
+
+    let srcDependencies = await dependencyGraph.getSourceDependencies(editor.document.uri, editor.document.getText())
+    srcDependencies = moveToFront(srcDependencies, openedFilesInEditor)
 
     const samePackageFiles = await dependencyGraph.getSamePackageFiles(
         editor.document.uri,
@@ -158,21 +166,31 @@ async function getRelevantCrossFiles(editor: vscode.TextEditor, dependencyGraph:
         return isRelevant(editor.document.fileName, file, editor.document.languageId)
     })
 
-    const relevantOpenFiles: vscode.Uri[] = await getRelevantFilesFromEditor(
-        editor.document.fileName,
-        editor.document.languageId
-    )
+    const mergedCrossFileList = [...new Set([...srcDependencies, ...samePackageRelevantFiles])]
 
-    // We refer to only those open files which are in srcDependencies
-    const filteredRelevantOpenFiles = relevantOpenFiles
-        .filter(file => srcDependencies.includes(file.fsPath))
-        .map(file => file.fsPath)
+    return mergedCrossFileList
+}
 
-    const mergedCrossFileList = [
-        ...new Set([...filteredRelevantOpenFiles, ...srcDependencies, ...samePackageRelevantFiles]),
-    ]
+// Util to move selected files to the front of the input array if it exists
+function moveToFront<T>(arr: T[], picked: Set<T>) {
+    return [...arr].sort((a, b) => (picked.has(b) ? 1 : 0) - (picked.has(a) ? 1 : 0))
+}
 
-    return mergedCrossFileList
+function getOpenFilesInWindow(): string[] {
+    const filesOpenedInEditor: string[] = []
+
+    try {
+        const tabArrays = vscode.window.tabGroups.all
+        tabArrays.forEach(tabArray => {
+            tabArray.tabs.forEach(tab => {
+                filesOpenedInEditor.push((tab.input as any).uri.path)
+            })
+        })
+    } catch (e) {
+        // Older versions of VSC do not have the tab API
+    }
+
+    return filesOpenedInEditor
 }
 
 function throwIfCancelled(token: vscode.CancellationToken): void | never {