change file fetch strategy to use open files only (#3734)

Author: Will-ShaoHua

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt

@@ -39,10 +39,7 @@ private val codewhispererCodeChunksIndex = GistManager.getInstance()
     .newPsiFileGist("psi to code chunk index", 0, CodeWhispererCodeChunkExternalizer) { psiFile ->
         runBlocking {
             val fileCrawler = getFileCrawlerForLanguage(psiFile.programmingLanguage())
-            val fileProducers = listOf<suspend (PsiFile) -> List<VirtualFile>>(
-                { psiFile -> fileCrawler.listFilesImported(psiFile) },
-                { psiFile -> fileCrawler.listFilesWithinSamePackage(psiFile) }
-            )
+            val fileProducers = listOf<suspend (PsiFile) -> List<VirtualFile>> { psiFile -> fileCrawler.listRelevantFilesInEditors(psiFile) }
             FileContextProvider.getInstance(psiFile.project).extractCodeChunksFromFiles(psiFile, fileProducers)
         }
     }
@@ -127,9 +124,12 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
                 it.forEachIndexed { index, chunk ->
                     LOG.info {
                         """
-                            | Chunk ${index + 1}:
+                            |---------------------------------------------------------------
+                            | Chunk $index:
+                            |    path = ${chunk.path},
+                            |    score = ${chunk.score},
                             |    content = ${chunk.content}
-                            |    path = ${chunk.path}
+                            |----------------------------------------------------------------
                         """.trimMargin()
                     }
                 }
@@ -210,7 +210,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
         // step 2: bm25 calculation
         val timeBeforeBm25 = System.currentTimeMillis()
         val top3Chunks: List<BM25Result> = BM250kapi(first60Chunks.map { it.content }).topN(query)
-        LOG.info { "Time ellapsed for BM25 algorithm: ${System.currentTimeMillis() - timeBeforeBm25} ms; \nResult: $top3Chunks" }
+        LOG.info { "Time ellapsed for BM25 algorithm: ${System.currentTimeMillis() - timeBeforeBm25} ms" }
 
         yield()
 

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt

@@ -37,6 +37,8 @@ interface FileCrawler {
      * @return its source file e.g. Main.java, main.py or most relevant file if any
      */
     fun findFocalFileForTest(psiFile: PsiFile): VirtualFile?
+
+    fun listRelevantFilesInEditors(psiFile: PsiFile): List<VirtualFile>
 }
 
 class NoOpFileCrawler : FileCrawler {
@@ -46,6 +48,8 @@ class NoOpFileCrawler : FileCrawler {
     override fun findFocalFileForTest(psiFile: PsiFile): VirtualFile? = null
 
     override fun listFilesWithinSamePackage(psiFile: PsiFile): List<VirtualFile> = emptyList()
+
+    override fun listRelevantFilesInEditors(psiFile: PsiFile): List<VirtualFile> = emptyList()
 }
 
 abstract class CodeWhispererFileCrawler : FileCrawler {

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavaCodeWhispererFileCrawler.kt

@@ -7,6 +7,7 @@ import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.module.ModuleUtilCore
 import com.intellij.openapi.project.rootManager
+import com.intellij.openapi.roots.TestSourcesFilter
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiClassOwner
@@ -24,6 +25,12 @@ object JavaCodeWhispererFileCrawler : CodeWhispererFileCrawler() {
 
     override fun guessSourceFileName(tstFileName: String): String = tstFileName.substring(0, tstFileName.length - "Test.java".length) + ".java"
 
+    override fun listRelevantFilesInEditors(psiFile: PsiFile): List<VirtualFile> = FileEditorManager.getInstance(psiFile.project).openFiles.toList().filter {
+        it.name != psiFile.virtualFile.name &&
+            it.extension == psiFile.virtualFile.extension &&
+            !TestSourcesFilter.isTestSources(it, psiFile.project)
+    }
+
     override suspend fun listFilesImported(psiFile: PsiFile): List<VirtualFile> {
         if (psiFile !is PsiJavaFile) return emptyList()
         val result = mutableListOf<VirtualFile>()

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/PythonCodeWhispererFileCrawler.kt

@@ -16,6 +16,9 @@ object PythonCodeWhispererFileCrawler : CodeWhispererFileCrawler() {
     override val testFilenamePattern: Regex = Regex("""(?:test_([^/\\]+)\.py|([^/\\]+)_test\.py)${'$'}""")
     override suspend fun listFilesImported(psiFile: PsiFile): List<VirtualFile> = emptyList()
 
+    // TODO: Add implementation once we enable python cross file
+    override fun listRelevantFilesInEditors(psiFile: PsiFile): List<VirtualFile> = emptyList()
+
     override fun guessSourceFileName(tstFileName: String): String {
         assert(testFilenamePattern.matches(tstFileName))
         return tstFileName.substring(5)