From cd739f25aace2373222d4ca4bacbcf96fc7cf072 Mon Sep 17 00:00:00 2001 From: Will Lo Date: Fri, 12 Jul 2024 17:49:32 -0700 Subject: [PATCH 1/7] neighbor file implementation, test WIP --- .../util/CodeWhispererConstants.kt | 1 + .../util/CodeWhispererFileContextProvider.kt | 57 +++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt index b7ec26558d..7dd41f4e2d 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt @@ -136,6 +136,7 @@ object CodeWhispererConstants { const val CHUNK_SIZE = 60 const val NUMBER_OF_LINE_IN_CHUNK = 10 const val NUMBER_OF_CHUNK_TO_FETCH = 3 + const val NEIGHBOR_FILES_DISTANCE = 1 } object Utg { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt index 27911fccfd..d9813c6f36 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt @@ -9,6 +9,7 @@ import com.intellij.openapi.components.service import com.intellij.openapi.editor.Editor import com.intellij.openapi.project.Project import com.intellij.openapi.vfs.VirtualFile +import com.intellij.psi.PsiDirectory import com.intellij.psi.PsiFile import com.intellij.util.gist.GistManager import com.intellij.util.io.DataExternalizer @@ -34,6 +35,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextI import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings +import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.CrossFile.NEIGHBOR_FILES_DISTANCE import java.io.DataInput import java.io.DataOutput import java.util.Collections @@ -124,6 +126,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi } } } else { + neighborfiles(psiFile) when (shouldFetchCrossfileContext(language, group)) { true -> extractSupplementalFileContextForSrc(psiFile, targetContext) false -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename) @@ -285,6 +288,60 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi } } + /** + * 1. A: root/util/context/a.ts + * 2. B: root/util/b.ts + * 3. C: root/util/service/c.ts + * 4. D: root/d.ts + * 5. E: root/util/context/e.ts + ** + * neighborfiles(A) = [B, E] + * neighborfiles(B) = [A, C, D] + * neighborfiles(C) = [] + * neighborfiles(D) = [] + * neighborfiles(E) = [] + */ + fun neighborfiles(psiFile: PsiFile): Set = search(psiFile, NEIGHBOR_FILES_DISTANCE) + + private fun search(psiFile: PsiFile, distance: Int): Set { + return search(psiFile.containingDirectory, distance, true) + search(psiFile.containingDirectory, distance, false) + } + + private fun search(psiDir: PsiDirectory, distance: Int, goDown: Boolean): Set { + var d = distance + val res = mutableListOf() + val pendingVisit = mutableListOf(psiDir) + val visisted = mutableMapOf(psiDir to false) + + while (d >= 0 && pendingVisit.isNotEmpty()) { + var toVisit = emptyList() + for (dir in pendingVisit) { + if (visisted[dir] == true) { + continue + } + + val fs = dir.files + res.addAll(fs) + + val dirs = if (goDown) { + dir.subdirectories.toList() + } else { + dir.parentDirectory?.let { + listOf(it) + } ?: emptyList() + } + + toVisit = dirs + visisted[dir] = true + } + + pendingVisit.addAll(toVisit) + d-- + } + + return res.toSet() + } + companion object { private val LOG = getLogger() From c0700518326336692a07f985addf1648cd4b18b4 Mon Sep 17 00:00:00 2001 From: Will Lo Date: Fri, 12 Jul 2024 18:54:34 -0700 Subject: [PATCH 2/7] tst --- .../util/CodeWhispererFileContextProvider.kt | 29 ++++++--- .../CodeWhispererFileContextProviderTest.kt | 60 +++++++++++++++++++ 2 files changed, 80 insertions(+), 9 deletions(-) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt index d9813c6f36..9ccdc6039c 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt @@ -126,7 +126,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi } } } else { - neighborfiles(psiFile) + neighborFiles(psiFile) when (shouldFetchCrossfileContext(language, group)) { true -> extractSupplementalFileContextForSrc(psiFile, targetContext) false -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename) @@ -294,17 +294,28 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi * 3. C: root/util/service/c.ts * 4. D: root/d.ts * 5. E: root/util/context/e.ts - ** + * 6. F: root/util/foo/bar/baz/f.ts + * * neighborfiles(A) = [B, E] - * neighborfiles(B) = [A, C, D] - * neighborfiles(C) = [] - * neighborfiles(D) = [] - * neighborfiles(E) = [] + * neighborfiles(B) = [A, C, D, E] + * neighborfiles(C) = [B,] + * neighborfiles(D) = [B,] + * neighborfiles(E) = [A, B] + * neighborfiles(F) = [] + * + * A B C D E F + * A x 1 2 2 0 4 + * B 1 x 1 1 1 3 + * C 2 1 x 2 2 2 + * D 2 1 2 x 2 2 + * E 0 1 2 2 x 4 + * F 4 3 4 4 4 x */ - fun neighborfiles(psiFile: PsiFile): Set = search(psiFile, NEIGHBOR_FILES_DISTANCE) + fun neighborFiles(psiFile: PsiFile): Set = search(psiFile, NEIGHBOR_FILES_DISTANCE).filterNot { it == psiFile }.toSet() - private fun search(psiFile: PsiFile, distance: Int): Set { - return search(psiFile.containingDirectory, distance, true) + search(psiFile.containingDirectory, distance, false) + private fun search(psiFile: PsiFile, distance: Int): Set = runReadAction { + search(psiFile.containingDirectory, distance, true) + + search(psiFile.containingDirectory, distance, false) } private fun search(psiDir: PsiDirectory, distance: Int, goDown: Boolean): Set { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt index 31573c9434..224c302d1b 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt @@ -26,6 +26,7 @@ import org.mockito.kotlin.spy import org.mockito.kotlin.times import org.mockito.kotlin.verify import org.mockito.kotlin.whenever +import software.aws.toolkits.core.utils.test.aString import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCpp import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCsharp import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererGo @@ -40,6 +41,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants +import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererFileCrawler import software.aws.toolkits.jetbrains.services.codewhisperer.util.DefaultCodeWhispererFileContextProvider import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule @@ -383,6 +385,64 @@ class CodeWhispererFileContextProviderTest { } } + /** + * 1. A: root/util/context/a.ts + * 2. B: root/util/b.ts + * 3. C: root/util/service/c.ts + * 4. D: root/d.ts + * 5. E: root/util/context/e.ts + * 6. F: root/util/foo/bar/baz/f.ts + * + * neighborfiles(A) = [B, E] + * neighborfiles(B) = [A, C, D, E] + * neighborfiles(C) = [B,] + * neighborfiles(D) = [B,] + * neighborfiles(E) = [A, B] + * neighborfiles(F) = [] + * + * A B C D E F + * A x 1 2 2 0 4 + * B 1 x 1 1 1 3 + * C 2 1 x 2 2 2 + * D 2 1 2 x 2 2 + * E 0 1 2 2 x 4 + * F 4 3 4 4 4 x + */ + @Test + fun `neighborFile should return all files except itself with distance less than or equal to 1`() { + val a = fixture.addFileToProject("root/util/context/a.java", aString()) + val b = fixture.addFileToProject("root/util/b.java", aString()) + val c = fixture.addFileToProject("root/util/service/c.java", aString()) + val d = fixture.addFileToProject("root/d.java", aString()) + val e = fixture.addFileToProject("root/util/context/e.java", aString()) + val f = fixture.addFileToProject("root/util/foo/bar/baz/f.java", aString()) + + assertThat(sut.neighborFiles(a)).isEqualTo(setOf(b, e)).also { + assertThat(CodeWhispererFileCrawler.getFileDistance(a.virtualFile, e.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(0) + assertThat(CodeWhispererFileCrawler.getFileDistance(a.virtualFile, b.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) + } + + assertThat(sut.neighborFiles(b)).isEqualTo(setOf(a, c, d, e)).also { + assertThat(CodeWhispererFileCrawler.getFileDistance(b.virtualFile, c.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) + assertThat(CodeWhispererFileCrawler.getFileDistance(b.virtualFile, d.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) + assertThat(CodeWhispererFileCrawler.getFileDistance(b.virtualFile, e.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) + } + + assertThat(sut.neighborFiles(c)).isEqualTo(setOf(b)) + + assertThat(sut.neighborFiles(d)).isEqualTo(setOf(b)) + + assertThat(sut.neighborFiles(e)).isEqualTo(setOf(a, b)) + + assertThat(sut.neighborFiles(f)).isEmpty().also { + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, a.virtualFile)).isGreaterThan(1).isEqualTo(4) + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, b.virtualFile)).isGreaterThan(1).isEqualTo(3) + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, c.virtualFile)).isGreaterThan(1).isEqualTo(4) + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, d.virtualFile)).isGreaterThan(1).isEqualTo(4) + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, e.virtualFile)).isGreaterThan(1).isEqualTo(4) + } + } + private fun setupFixture(fixture: JavaCodeInsightTestFixture): List { val psiFile1 = fixture.addFileToProject("Main.java", JAVA_MAIN) val psiFile2 = fixture.addFileToProject("UtilClass.java", JAVA_UTILCLASS) From 8ce40dbcb4be82560c6ffd81dd540908ccd53fb5 Mon Sep 17 00:00:00 2001 From: Will Lo Date: Fri, 12 Jul 2024 19:15:38 -0700 Subject: [PATCH 3/7] patch --- .../codewhisperer/util/CodeWhispererFileContextProvider.kt | 1 - .../codewhisperer/CodeWhispererFileContextProviderTest.kt | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt index 9ccdc6039c..fd458c3d7d 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt @@ -126,7 +126,6 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi } } } else { - neighborFiles(psiFile) when (shouldFetchCrossfileContext(language, group)) { true -> extractSupplementalFileContextForSrc(psiFile, targetContext) false -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt index 224c302d1b..65e6e3467a 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt @@ -403,8 +403,8 @@ class CodeWhispererFileContextProviderTest { * A B C D E F * A x 1 2 2 0 4 * B 1 x 1 1 1 3 - * C 2 1 x 2 2 2 - * D 2 1 2 x 2 2 + * C 2 1 x 2 2 4 + * D 2 1 2 x 2 4 * E 0 1 2 2 x 4 * F 4 3 4 4 4 x */ From c3b00a317bc3cd1e8096951e46382736da79fc73 Mon Sep 17 00:00:00 2001 From: Will Lo Date: Fri, 12 Jul 2024 19:16:47 -0700 Subject: [PATCH 4/7] patch --- .../codewhisperer/util/CodeWhispererFileContextProvider.kt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt index fd458c3d7d..910bfae48e 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt @@ -305,8 +305,8 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi * A B C D E F * A x 1 2 2 0 4 * B 1 x 1 1 1 3 - * C 2 1 x 2 2 2 - * D 2 1 2 x 2 2 + * C 2 1 x 2 2 4 + * D 2 1 2 x 2 4 * E 0 1 2 2 x 4 * F 4 3 4 4 4 x */ From 3075cb0c5160123e6e209b6c336bff449da33780 Mon Sep 17 00:00:00 2001 From: Will Lo Date: Sat, 3 Aug 2024 03:32:41 -0700 Subject: [PATCH 5/7] use crossfile when UTG is not supported --- .../codewhisperer/language/languages/CodeWhispererC.kt | 2 ++ .../codewhisperer/language/languages/CodeWhispererCpp.kt | 2 ++ .../language/languages/CodeWhispererCsharp.kt | 2 ++ .../codewhisperer/language/languages/CodeWhispererGo.kt | 2 ++ .../language/languages/CodeWhispererKotlin.kt | 2 ++ .../codewhisperer/language/languages/CodeWhispererPhp.kt | 2 ++ .../codewhisperer/language/languages/CodeWhispererRuby.kt | 2 ++ .../codewhisperer/language/languages/CodeWhispererRust.kt | 2 ++ .../codewhisperer/language/languages/CodeWhispererScala.kt | 2 ++ .../codewhisperer/language/languages/CodeWhispererShell.kt | 2 ++ .../service/CodeWhispererUserGroupSettings.kt | 1 + .../codewhisperer/util/CodeWhispererFileContextProvider.kt | 7 +++++-- 12 files changed, 26 insertions(+), 2 deletions(-) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererC.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererC.kt index 6b2bdce130..498a74aca1 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererC.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererC.kt @@ -17,6 +17,8 @@ class CodeWhispererC private constructor() : CodeWhispererProgrammingLanguage() override fun isAutoFileScanSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "c" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCpp.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCpp.kt index d6e57d621e..e2ebfae2a4 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCpp.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCpp.kt @@ -17,6 +17,8 @@ class CodeWhispererCpp private constructor() : CodeWhispererProgrammingLanguage( override fun isAutoFileScanSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "cpp" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCsharp.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCsharp.kt index abea7f95a7..32fe87de77 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCsharp.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCsharp.kt @@ -17,6 +17,8 @@ class CodeWhispererCsharp private constructor() : CodeWhispererProgrammingLangua override fun isAutoFileScanSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "csharp" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererGo.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererGo.kt index 166c3e7334..0628b51c0a 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererGo.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererGo.kt @@ -17,6 +17,8 @@ class CodeWhispererGo private constructor() : CodeWhispererProgrammingLanguage() override fun isAutoFileScanSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "go" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererKotlin.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererKotlin.kt index 0edb7042c1..256588225c 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererKotlin.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererKotlin.kt @@ -13,6 +13,8 @@ class CodeWhispererKotlin private constructor() : CodeWhispererProgrammingLangua override fun isCodeCompletionSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "kotlin" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererPhp.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererPhp.kt index 3f10018865..cca0fa22fc 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererPhp.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererPhp.kt @@ -17,6 +17,8 @@ class CodeWhispererPhp private constructor() : CodeWhispererProgrammingLanguage( override fun isAutoFileScanSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "php" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRuby.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRuby.kt index 3ed9dbf0ef..9be215ef54 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRuby.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRuby.kt @@ -17,6 +17,8 @@ class CodeWhispererRuby private constructor() : CodeWhispererProgrammingLanguage override fun isAutoFileScanSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "ruby" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRust.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRust.kt index cf953ff1da..db92702fd6 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRust.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRust.kt @@ -13,6 +13,8 @@ class CodeWhispererRust private constructor() : CodeWhispererProgrammingLanguage override fun isCodeCompletionSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "rust" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererScala.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererScala.kt index ffd5594706..3d9312d691 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererScala.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererScala.kt @@ -13,6 +13,8 @@ class CodeWhispererScala private constructor() : CodeWhispererProgrammingLanguag override fun isCodeCompletionSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "scala" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererShell.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererShell.kt index bd8fa3450f..96bd561e2f 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererShell.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererShell.kt @@ -13,6 +13,8 @@ class CodeWhispererShell private constructor() : CodeWhispererProgrammingLanguag override fun isCodeCompletionSupported(): Boolean = true + override fun isSupplementalContextSupported() = true + companion object { const val ID = "shell" diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererUserGroupSettings.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererUserGroupSettings.kt index 2dc1bfa69e..7759c9c481 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererUserGroupSettings.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererUserGroupSettings.kt @@ -19,6 +19,7 @@ import kotlin.reflect.KClass /** * Component controlling codewhisperer user group settings */ +@Deprecated("use CodeWhispererFeatureConfigService instead") @Service @State(name = "codewhispererUserGroupSettings", storages = [Storage("aws.xml", roamingType = RoamingType.DISABLED)]) class CodeWhispererUserGroupSettings : PersistentStateComponent { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt index 27911fccfd..71ed1f4cde 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt @@ -114,7 +114,8 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi val language = targetContext.programmingLanguage val group = CodeWhispererUserGroupSettings.getInstance().getUserGroup() - val supplementalContext = if (isTst) { + // if utg is not supported, use crossfile context as fallback + val supplementalContext = if (isTst && language.isUTGSupported()) { when (shouldFetchUtgContext(language, group)) { true -> extractSupplementalFileContextForTst(psiFile, targetContext) false -> SupplementalContextInfo.emptyUtgFileContextInfo(targetContext.filename) @@ -299,6 +300,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi } } + @Suppress("UNUSED_PARAMETER") fun shouldFetchCrossfileContext(language: CodeWhispererProgrammingLanguage, userGroup: CodeWhispererUserGroup): Boolean? { if (!language.isSupplementalContextSupported()) { return null @@ -312,7 +314,8 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi is CodeWhispererJsx, is CodeWhispererTsx -> true - else -> userGroup == CodeWhispererUserGroup.CrossFile + // TODO: languages under A/B, should read feature flag from [CodeWhispererFeatureConfigService] + else -> false } } } From 47d00d8757198c6587dcd65285a1d1c5991d622e Mon Sep 17 00:00:00 2001 From: Will Lo Date: Mon, 5 Aug 2024 12:41:38 -0700 Subject: [PATCH 6/7] todo: service side configuration changes, A/B setup --- .../service/CodeWhispererService.kt | 46 +++--- .../util/CodeWhispererConstants.kt | 8 +- .../util/CodeWhispererFileContextProvider.kt | 91 +++--------- .../util/CodeWhispererFileCrawler.kt | 129 +++++++++++++---- .../util/JavaCodeWhispererFileCrawler.kt | 48 ------- .../JavascriptCodeWhispererFileCrawler.kt | 2 - .../util/PythonCodeWhispererFileCrawler.kt | 2 - .../TypescriptCodeWhispererFileCrawler.kt | 2 - .../CodeWhispererFileContextProviderTest.kt | 60 -------- .../CodeWhispererFileCrawlerTest.kt | 135 ++++++++---------- 10 files changed, 211 insertions(+), 312 deletions(-) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt index 8962abc1cb..3f19adfade 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt @@ -21,7 +21,9 @@ import com.intellij.psi.PsiDocumentManager import com.intellij.psi.PsiFile import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.messages.Topic +import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.TimeoutCancellationException +import kotlinx.coroutines.async import kotlinx.coroutines.delay import kotlinx.coroutines.isActive import kotlinx.coroutines.launch @@ -43,7 +45,6 @@ import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingExce import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info -import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager @@ -91,7 +92,7 @@ import software.aws.toolkits.telemetry.CodewhispererTriggerType import java.util.concurrent.TimeUnit @Service -class CodeWhispererService : Disposable { +class CodeWhispererService(private val coroutineScope: CoroutineScope) : Disposable { private val codeInsightSettingsFacade = CodeInsightsSettingsFacade() private var refreshFailure: Int = 0 @@ -193,7 +194,6 @@ class CodeWhispererService : Disposable { // from CodeWhispererPopupManager.cancelPopup() and CodeWhispererPopupManager.closePopup(). // It's possible and ok that coroutine will keep running until the next time we check it's state. // As long as we don't show to the user extra info we are good. - val coroutineScope = disposableCoroutineScope(popup) var states: InvocationContext? = null var lastRecommendationIndex = -1 @@ -624,27 +624,29 @@ class CodeWhispererService : Disposable { val startFetchingTimestamp = System.currentTimeMillis() val isTstFile = FileContextProvider.getInstance(project).isTestFile(psiFile) val supplementalContext = runBlocking { - try { - withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) { - FileContextProvider.getInstance(project).extractSupplementalFileContext(psiFile, fileContext) - } - } catch (e: Exception) { - if (e is TimeoutCancellationException) { - LOG.debug { - "Supplemental context fetch timed out in ${System.currentTimeMillis() - startFetchingTimestamp}ms" + coroutineScope.async { + try { + withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) { + FileContextProvider.getInstance(project).extractSupplementalFileContext(psiFile, fileContext) + } + } catch (e: Exception) { + if (e is TimeoutCancellationException) { + LOG.debug { + "Supplemental context fetch timed out in ${System.currentTimeMillis() - startFetchingTimestamp}ms" + } + SupplementalContextInfo( + isUtg = isTstFile, + contents = emptyList(), + latency = System.currentTimeMillis() - startFetchingTimestamp, + targetFileName = fileContext.filename, + strategy = if (isTstFile) UtgStrategy.Empty else CrossFileStrategy.Empty + ) + } else { + LOG.debug { "Run into unexpected error when fetching supplemental context, error: ${e.message}" } + null } - SupplementalContextInfo( - isUtg = isTstFile, - contents = emptyList(), - latency = System.currentTimeMillis() - startFetchingTimestamp, - targetFileName = fileContext.filename, - strategy = if (isTstFile) UtgStrategy.Empty else CrossFileStrategy.Empty - ) - } else { - LOG.debug { "Run into unexpected error when fetching supplemental context, error: ${e.message}" } - null } - } + }.await() } // 3. caret position diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt index 0ba5358540..842a3cf697 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt @@ -9,6 +9,7 @@ import com.intellij.ui.JBColor import software.amazon.awssdk.regions.Region import software.amazon.awssdk.services.codewhispererruntime.model.AccessDeniedException import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException +import software.aws.toolkits.jetbrains.isDeveloperMode import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava import software.aws.toolkits.telemetry.CodewhispererGettingStartedTask import java.awt.Font @@ -133,8 +134,11 @@ object CodeWhispererConstants { } } object CrossFile { - const val CHUNK_SIZE = 60 - const val NUMBER_OF_LINE_IN_CHUNK = 10 + val CHUNK_SIZE + get() = if (isDeveloperMode()) 200 else 60 + val NUMBER_OF_LINE_IN_CHUNK + get() = if (isDeveloperMode()) 50 else 10 + // TODO: 3 -> 10 when service side CR is done const val NUMBER_OF_CHUNK_TO_FETCH = 3 const val NEIGHBOR_FILES_DISTANCE = 1 } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt index dac0e9936d..b94229bf85 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt @@ -9,7 +9,6 @@ import com.intellij.openapi.components.service import com.intellij.openapi.editor.Editor import com.intellij.openapi.project.Project import com.intellij.openapi.vfs.VirtualFile -import com.intellij.psi.PsiDirectory import com.intellij.psi.PsiFile import com.intellij.util.gist.GistManager import com.intellij.util.io.DataExternalizer @@ -21,6 +20,7 @@ import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info import software.aws.toolkits.core.utils.warn +import software.aws.toolkits.jetbrains.isDeveloperMode import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava @@ -35,7 +35,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextI import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.CrossFile.NEIGHBOR_FILES_DISTANCE import java.io.DataInput import java.io.DataOutput import java.util.Collections @@ -177,13 +176,19 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi chunks.addAll(file.toCodeChunk(relativePath)) hasUsed.add(file) if (chunks.size > CodeWhispererConstants.CrossFile.CHUNK_SIZE) { - LOG.debug { "finish fetching ${CodeWhispererConstants.CrossFile.CHUNK_SIZE} chunks in ${System.currentTimeMillis() - parseFilesStart} ms" } + LOG.debug { + "finish fetching ${CodeWhispererConstants.CrossFile.CHUNK_SIZE} " + + "chunks in ${System.currentTimeMillis() - parseFilesStart} ms from files ${hasUsed.map { it.name }}" + } return chunks.take(CodeWhispererConstants.CrossFile.CHUNK_SIZE) } } } - LOG.debug { "finish fetching ${CodeWhispererConstants.CrossFile.CHUNK_SIZE} chunks in ${System.currentTimeMillis() - parseFilesStart} ms" } + LOG.debug { + "finish fetching ${CodeWhispererConstants.CrossFile.CHUNK_SIZE} " + + "chunks in ${System.currentTimeMillis() - parseFilesStart} ms from files ${hasUsed.map { it.name }}" + } return chunks.take(CodeWhispererConstants.CrossFile.CHUNK_SIZE) } @@ -200,7 +205,18 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi // step 1: prepare data val first60Chunks: List = try { - runReadAction { codewhispererCodeChunksIndex.getFileData(psiFile) } + runReadAction { + // for dev purpose only, not using cache when it's dev mode to monitor performance + if (isDeveloperMode()) { + runBlocking { + val fileCrawler = psiFile.programmingLanguage().fileCrawler + val fileProducers = listOf List> { psiFile -> fileCrawler.listCrossFileCandidate(psiFile) } + FileContextProvider.getInstance(psiFile.project).extractCodeChunksFromFiles(psiFile, fileProducers) + } + } else { + codewhispererCodeChunksIndex.getFileData(psiFile) + } + } } catch (e: TimeoutCancellationException) { throw e } @@ -288,71 +304,6 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi } } - /** - * 1. A: root/util/context/a.ts - * 2. B: root/util/b.ts - * 3. C: root/util/service/c.ts - * 4. D: root/d.ts - * 5. E: root/util/context/e.ts - * 6. F: root/util/foo/bar/baz/f.ts - * - * neighborfiles(A) = [B, E] - * neighborfiles(B) = [A, C, D, E] - * neighborfiles(C) = [B,] - * neighborfiles(D) = [B,] - * neighborfiles(E) = [A, B] - * neighborfiles(F) = [] - * - * A B C D E F - * A x 1 2 2 0 4 - * B 1 x 1 1 1 3 - * C 2 1 x 2 2 4 - * D 2 1 2 x 2 4 - * E 0 1 2 2 x 4 - * F 4 3 4 4 4 x - */ - fun neighborFiles(psiFile: PsiFile): Set = search(psiFile, NEIGHBOR_FILES_DISTANCE).filterNot { it == psiFile }.toSet() - - private fun search(psiFile: PsiFile, distance: Int): Set = runReadAction { - search(psiFile.containingDirectory, distance, true) + - search(psiFile.containingDirectory, distance, false) - } - - private fun search(psiDir: PsiDirectory, distance: Int, goDown: Boolean): Set { - var d = distance - val res = mutableListOf() - val pendingVisit = mutableListOf(psiDir) - val visisted = mutableMapOf(psiDir to false) - - while (d >= 0 && pendingVisit.isNotEmpty()) { - var toVisit = emptyList() - for (dir in pendingVisit) { - if (visisted[dir] == true) { - continue - } - - val fs = dir.files - res.addAll(fs) - - val dirs = if (goDown) { - dir.subdirectories.toList() - } else { - dir.parentDirectory?.let { - listOf(it) - } ?: emptyList() - } - - toVisit = dirs - visisted[dir] = true - } - - pendingVisit.addAll(toVisit) - d-- - } - - return res.toSet() - } - companion object { private val LOG = getLogger() diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt index 4ce957bb8d..fa9d909e4c 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileCrawler.kt @@ -5,36 +5,27 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.util import com.intellij.openapi.application.runReadAction import com.intellij.openapi.fileEditor.FileEditorManager +import com.intellij.openapi.fileEditor.impl.FileEditorManagerImpl import com.intellij.openapi.project.Project import com.intellij.openapi.project.guessProjectDir import com.intellij.openapi.roots.TestSourcesFilter import com.intellij.openapi.vfs.VfsUtil import com.intellij.openapi.vfs.VirtualFile +import com.intellij.psi.PsiDirectory import com.intellij.psi.PsiFile import com.intellij.psi.PsiManager import software.aws.toolkits.core.utils.tryOrNull +import software.aws.toolkits.jetbrains.isDeveloperMode import software.aws.toolkits.jetbrains.services.codewhisperer.model.ListUtgCandidateResult +import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.CrossFile.NEIGHBOR_FILES_DISTANCE /** * An interface define how do we parse and fetch files provided a psi file or project * since different language has its own way importing other files or its own naming style for test file */ interface FileCrawler { - /** - * parse the import statements provided a file - * @param psiFile of the file we are search with - * @return list of file reference from the import statements - */ - suspend fun listFilesImported(psiFile: PsiFile): List - fun listFilesUnderProjectRoot(project: Project): List - /** - * @param psiFile the file we are searching with, aka target file - * @return Files under the same package as the given file and exclude the given file - */ - fun listFilesWithinSamePackage(psiFile: PsiFile): List - /** * should be invoked at test files e.g. MainTest.java, or test_main.py * @param target psi of the test file we are searching with, e.g. MainTest.java @@ -59,12 +50,9 @@ interface FileCrawler { } class NoOpFileCrawler : FileCrawler { - override suspend fun listFilesImported(psiFile: PsiFile): List = emptyList() - override fun listFilesUnderProjectRoot(project: Project): List = emptyList() - override fun listUtgCandidate(target: PsiFile) = ListUtgCandidateResult(null, UtgStrategy.Empty) - override fun listFilesWithinSamePackage(psiFile: PsiFile): List = emptyList() + override fun listUtgCandidate(target: PsiFile) = ListUtgCandidateResult(null, UtgStrategy.Empty) override fun listCrossFileCandidate(target: PsiFile): List = emptyList() @@ -99,18 +87,24 @@ abstract class CodeWhispererFileCrawler : FileCrawler { } }.orEmpty() - override fun listFilesWithinSamePackage(psiFile: PsiFile): List = runReadAction { - psiFile.containingDirectory?.files?.mapNotNull { - // exclude target file - if (it != psiFile) { - it.virtualFile - } else { - null - } - }.orEmpty() + override fun listCrossFileCandidate(target: PsiFile): List = if (isDeveloperMode()) { + val previousSelected = listPreviousSelectedFile(target) + val neighbors = neighborFiles(target).mapNotNull { it.virtualFile } + val result = previousSelected + neighbors + println("MRU list: ${previousSelected.map { it.name }}") + println("neighbors: ${neighbors.map { it.name }}") + result.distinctBy { it.name } + } else { + listAllOpenedFilesSortedByDist(target) + }.also { + val logStr = it.map { file -> file.name } + println("crossfile candidates: $logStr") } - override fun listCrossFileCandidate(target: PsiFile): List { + /** + * Default strategy will return all opened files sorted with file distance against the target + */ + private fun listAllOpenedFilesSortedByDist(target: PsiFile): List { val targetFile = target.virtualFile val openedFiles = runReadAction { @@ -130,6 +124,22 @@ abstract class CodeWhispererFileCrawler : FileCrawler { return fileToFileDistanceList.sortedBy { it.second }.map { it.first } } + /** + * New strategy will return opened files sorted by timeline the file is used (Most Recently Used), which should be as same as JB's file switcher (ctrl + tab) + * Note: test file is included here unlike the default strategy (thus different predicate inside filter) + */ + private fun listPreviousSelectedFile(target: PsiFile): List { + val targetVFile = target.virtualFile + return runReadAction { + (FileEditorManager.getInstance(target.project) as FileEditorManagerImpl).getSelectionHistory() + .map { it.first } + .filter { + it.name != targetVFile.name && + isSameDialect(it.extension) + } + } + } + override fun listUtgCandidate(target: PsiFile): ListUtgCandidateResult { val byName = findSourceFileByName(target) if (byName != null) { @@ -165,6 +175,71 @@ abstract class CodeWhispererFileCrawler : FileCrawler { dialects.contains(fileExt) } ?: false + /** + * 1. A: root/util/context/a.ts + * 2. B: root/util/b.ts + * 3. C: root/util/service/c.ts + * 4. D: root/d.ts + * 5. E: root/util/context/e.ts + * 6. F: root/util/foo/bar/baz/f.ts + * + * neighborfiles(A) = [B, E] + * neighborfiles(B) = [A, C, D, E] + * neighborfiles(C) = [B,] + * neighborfiles(D) = [B,] + * neighborfiles(E) = [A, B] + * neighborfiles(F) = [] + * + * A B C D E F + * A x 1 2 2 0 4 + * B 1 x 1 1 1 3 + * C 2 1 x 2 2 4 + * D 2 1 2 x 2 4 + * E 0 1 2 2 x 4 + * F 4 3 4 4 4 x + */ + fun neighborFiles(psiFile: PsiFile) = search(psiFile, NEIGHBOR_FILES_DISTANCE).filterNot { it == psiFile }.toSet() + + private fun search(psiFile: PsiFile, distance: Int): Set = runReadAction { + search(psiFile.containingDirectory, distance, true) + + search(psiFile.containingDirectory, distance, false) + } + + private fun search(psiDir: PsiDirectory, distance: Int, goDown: Boolean): Set { + var d = distance + val res = mutableListOf() + val pendingVisit = mutableListOf(psiDir) + val visisted = mutableMapOf(psiDir to false) + + while (d >= 0 && pendingVisit.isNotEmpty()) { + var toVisit = emptyList() + for (dir in pendingVisit) { + if (visisted[dir] == true) { + continue + } + + val fs = dir.files + res.addAll(fs) + + val dirs = if (goDown) { + dir.subdirectories.toList() + } else { + dir.parentDirectory?.let { + listOf(it) + } ?: emptyList() + } + + toVisit = dirs + visisted[dir] = true + } + + pendingVisit.addAll(toVisit) + d-- + } + + return res.toSet() + } + companion object { // TODO: move to CodeWhispererUtils.kt /** diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavaCodeWhispererFileCrawler.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavaCodeWhispererFileCrawler.kt index 5f5eece6a1..93d6b9c7ff 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavaCodeWhispererFileCrawler.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavaCodeWhispererFileCrawler.kt @@ -3,17 +3,11 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.util -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.vfs.VfsUtil import com.intellij.openapi.vfs.VirtualFile import com.intellij.psi.PsiFile -import com.intellij.psi.PsiJavaFile -import com.intellij.psi.PsiPackage -import com.intellij.psi.search.GlobalSearchScope -import kotlinx.coroutines.yield import org.jetbrains.jps.model.java.JavaModuleSourceRootTypes import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.warn @@ -29,48 +23,6 @@ object JavaCodeWhispererFileCrawler : CodeWhispererFileCrawler() { Regex("""^(.+)Tests(\.java)$""") ) - override suspend fun listFilesImported(psiFile: PsiFile): List { - if (psiFile !is PsiJavaFile) return emptyList() - val result = mutableListOf() - val imports = runReadAction { psiFile.importList?.allImportStatements } - val activeFiles = FileEditorManager.getInstance(psiFile.project).openFiles.toSet() - - // only consider imported files which belong users' own package, thus [isInLocalFileSystem && isWritable] - val fileHandleLambda = { virtualFile: VirtualFile -> - if (virtualFile.isInLocalFileSystem && virtualFile.isWritable) { - // prioritize active files on users' editor - if (activeFiles.contains(virtualFile)) { - result.add(0, virtualFile) - } else { - result.add(virtualFile) - } - } - } - - imports?.forEach { - yield() - runReadAction { it.resolve() }?.let { psiElement -> - // case like import javax.swing.*; - if (psiElement is PsiPackage) { - val filesInPackage = psiElement.getFiles(GlobalSearchScope.allScope(psiFile.project)).mapNotNull { it.virtualFile } - filesInPackage.forEach { file -> - fileHandleLambda(file) - } - } else { - // single file import - runReadAction { - psiElement.containingFile.virtualFile?.let { virtualFile -> - // file within users' project - fileHandleLambda(virtualFile) - } - } - } - } - } - - return result - } - // psiFile = "MainTest.java", targetFileName = "Main.java" override fun findSourceFileByName(target: PsiFile): VirtualFile? = guessSourceFileName(target.virtualFile.name)?.let { srcName -> diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavascriptCodeWhispererFileCrawler.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavascriptCodeWhispererFileCrawler.kt index 00217f8894..cf4ddbe650 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavascriptCodeWhispererFileCrawler.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/JavascriptCodeWhispererFileCrawler.kt @@ -14,8 +14,6 @@ object JavascriptCodeWhispererFileCrawler : CodeWhispererFileCrawler() { Regex("""^(.+)\.(?i:s)pec(\.js|\.jsx)$""") ) - override suspend fun listFilesImported(psiFile: PsiFile): List = emptyList() - override fun findSourceFileByName(target: PsiFile): VirtualFile? = null override fun findSourceFileByContent(target: PsiFile): VirtualFile? = null diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/PythonCodeWhispererFileCrawler.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/PythonCodeWhispererFileCrawler.kt index 83216b7640..56130b2dd4 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/PythonCodeWhispererFileCrawler.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/PythonCodeWhispererFileCrawler.kt @@ -17,8 +17,6 @@ object PythonCodeWhispererFileCrawler : CodeWhispererFileCrawler() { Regex("""^(.+)_test(\.py)$""") ) - override suspend fun listFilesImported(psiFile: PsiFile): List = emptyList() - override fun findSourceFileByName(target: PsiFile): VirtualFile? = super.listFilesUnderProjectRoot(target.project).find { !it.isDirectory && it.isWritable && diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/TypescriptCodeWhispererFileCrawler.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/TypescriptCodeWhispererFileCrawler.kt index 462a16501d..b126c6aa31 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/TypescriptCodeWhispererFileCrawler.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/TypescriptCodeWhispererFileCrawler.kt @@ -14,8 +14,6 @@ object TypescriptCodeWhispererFileCrawler : CodeWhispererFileCrawler() { Regex("""^(.+)\.(?i:s)pec(\.ts|\.tsx)$""") ) - override suspend fun listFilesImported(psiFile: PsiFile): List = emptyList() - override fun findSourceFileByName(target: PsiFile): VirtualFile? = null override fun findSourceFileByContent(target: PsiFile): VirtualFile? = null diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt index 65e6e3467a..31573c9434 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt @@ -26,7 +26,6 @@ import org.mockito.kotlin.spy import org.mockito.kotlin.times import org.mockito.kotlin.verify import org.mockito.kotlin.whenever -import software.aws.toolkits.core.utils.test.aString import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCpp import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCsharp import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererGo @@ -41,7 +40,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererFileCrawler import software.aws.toolkits.jetbrains.services.codewhisperer.util.DefaultCodeWhispererFileContextProvider import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule @@ -385,64 +383,6 @@ class CodeWhispererFileContextProviderTest { } } - /** - * 1. A: root/util/context/a.ts - * 2. B: root/util/b.ts - * 3. C: root/util/service/c.ts - * 4. D: root/d.ts - * 5. E: root/util/context/e.ts - * 6. F: root/util/foo/bar/baz/f.ts - * - * neighborfiles(A) = [B, E] - * neighborfiles(B) = [A, C, D, E] - * neighborfiles(C) = [B,] - * neighborfiles(D) = [B,] - * neighborfiles(E) = [A, B] - * neighborfiles(F) = [] - * - * A B C D E F - * A x 1 2 2 0 4 - * B 1 x 1 1 1 3 - * C 2 1 x 2 2 4 - * D 2 1 2 x 2 4 - * E 0 1 2 2 x 4 - * F 4 3 4 4 4 x - */ - @Test - fun `neighborFile should return all files except itself with distance less than or equal to 1`() { - val a = fixture.addFileToProject("root/util/context/a.java", aString()) - val b = fixture.addFileToProject("root/util/b.java", aString()) - val c = fixture.addFileToProject("root/util/service/c.java", aString()) - val d = fixture.addFileToProject("root/d.java", aString()) - val e = fixture.addFileToProject("root/util/context/e.java", aString()) - val f = fixture.addFileToProject("root/util/foo/bar/baz/f.java", aString()) - - assertThat(sut.neighborFiles(a)).isEqualTo(setOf(b, e)).also { - assertThat(CodeWhispererFileCrawler.getFileDistance(a.virtualFile, e.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(0) - assertThat(CodeWhispererFileCrawler.getFileDistance(a.virtualFile, b.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) - } - - assertThat(sut.neighborFiles(b)).isEqualTo(setOf(a, c, d, e)).also { - assertThat(CodeWhispererFileCrawler.getFileDistance(b.virtualFile, c.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) - assertThat(CodeWhispererFileCrawler.getFileDistance(b.virtualFile, d.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) - assertThat(CodeWhispererFileCrawler.getFileDistance(b.virtualFile, e.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) - } - - assertThat(sut.neighborFiles(c)).isEqualTo(setOf(b)) - - assertThat(sut.neighborFiles(d)).isEqualTo(setOf(b)) - - assertThat(sut.neighborFiles(e)).isEqualTo(setOf(a, b)) - - assertThat(sut.neighborFiles(f)).isEmpty().also { - assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, a.virtualFile)).isGreaterThan(1).isEqualTo(4) - assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, b.virtualFile)).isGreaterThan(1).isEqualTo(3) - assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, c.virtualFile)).isGreaterThan(1).isEqualTo(4) - assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, d.virtualFile)).isGreaterThan(1).isEqualTo(4) - assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, e.virtualFile)).isGreaterThan(1).isEqualTo(4) - } - } - private fun setupFixture(fixture: JavaCodeInsightTestFixture): List { val psiFile1 = fixture.addFileToProject("Main.java", JAVA_MAIN) val psiFile2 = fixture.addFileToProject("UtilClass.java", JAVA_UTILCLASS) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileCrawlerTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileCrawlerTest.kt index 6e1dfcdff2..fac3309a2a 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileCrawlerTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileCrawlerTest.kt @@ -6,13 +6,11 @@ package software.aws.toolkits.jetbrains.services.codewhisperer import com.intellij.openapi.application.runReadAction import com.intellij.openapi.editor.EditorFactory import com.intellij.openapi.project.Project -import com.intellij.openapi.vfs.VirtualFile import com.intellij.psi.PsiFile import com.intellij.testFramework.DisposableRule import com.intellij.testFramework.ExtensionTestUtil import com.intellij.testFramework.fixtures.CodeInsightTestFixture import com.intellij.testFramework.runInEdtAndWait -import kotlinx.coroutines.runBlocking import org.assertj.core.api.Assertions.assertThat import org.junit.Before import org.junit.Rule @@ -119,6 +117,64 @@ class CodeWhispererFileCrawlerTest { assertThat(result).isEqualTo(file1.virtualFile) } } + + /** + * 1. A: root/util/context/a.ts + * 2. B: root/util/b.ts + * 3. C: root/util/service/c.ts + * 4. D: root/d.ts + * 5. E: root/util/context/e.ts + * 6. F: root/util/foo/bar/baz/f.ts + * + * neighborfiles(A) = [B, E] + * neighborfiles(B) = [A, C, D, E] + * neighborfiles(C) = [B,] + * neighborfiles(D) = [B,] + * neighborfiles(E) = [A, B] + * neighborfiles(F) = [] + * + * A B C D E F + * A x 1 2 2 0 4 + * B 1 x 1 1 1 3 + * C 2 1 x 2 2 4 + * D 2 1 2 x 2 4 + * E 0 1 2 2 x 4 + * F 4 3 4 4 4 x + */ + @Test + fun `neighborFile should return all files except itself with distance less than or equal to 1`() { + val a = fixture.addFileToProject("root/util/context/a.java", aString()) + val b = fixture.addFileToProject("root/util/b.java", aString()) + val c = fixture.addFileToProject("root/util/service/c.java", aString()) + val d = fixture.addFileToProject("root/d.java", aString()) + val e = fixture.addFileToProject("root/util/context/e.java", aString()) + val f = fixture.addFileToProject("root/util/foo/bar/baz/f.java", aString()) + + assertThat(sut.neighborFiles(a)).isEqualTo(setOf(b, e)).also { + assertThat(CodeWhispererFileCrawler.getFileDistance(a.virtualFile, e.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(0) + assertThat(CodeWhispererFileCrawler.getFileDistance(a.virtualFile, b.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) + } + + assertThat(sut.neighborFiles(b)).isEqualTo(setOf(a, c, d, e)).also { + assertThat(CodeWhispererFileCrawler.getFileDistance(b.virtualFile, c.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) + assertThat(CodeWhispererFileCrawler.getFileDistance(b.virtualFile, d.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) + assertThat(CodeWhispererFileCrawler.getFileDistance(b.virtualFile, e.virtualFile)).isLessThanOrEqualTo(1).isEqualTo(1) + } + + assertThat(sut.neighborFiles(c)).isEqualTo(setOf(b)) + + assertThat(sut.neighborFiles(d)).isEqualTo(setOf(b)) + + assertThat(sut.neighborFiles(e)).isEqualTo(setOf(a, b)) + + assertThat(sut.neighborFiles(f)).isEmpty().also { + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, a.virtualFile)).isGreaterThan(1).isEqualTo(4) + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, b.virtualFile)).isGreaterThan(1).isEqualTo(3) + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, c.virtualFile)).isGreaterThan(1).isEqualTo(4) + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, d.virtualFile)).isGreaterThan(1).isEqualTo(4) + assertThat(CodeWhispererFileCrawler.getFileDistance(f.virtualFile, e.virtualFile)).isGreaterThan(1).isEqualTo(4) + } + } } class JavaCodeWhispererFileCrawlerTest { @@ -262,81 +318,6 @@ class JavaCodeWhispererFileCrawlerTest { } } - @Test - fun `listFilesWithinSamePackage`() { - val targetFile = fixture.addFileToProject("/utils/AnotherClass.java", "") - val file2Package1 = fixture.addFileToProject("/utils/NotImported.java", "") - val file3Package1 = fixture.addFileToProject("/utils/NotImported2.java", "") - fixture.addFileToProject("Main.java", "") - fixture.addFileToProject("service/controllers/MyController.java", "") - - runReadAction { - val actual = sut.listFilesWithinSamePackage(targetFile) - val expected = listOf(file2Package1, file3Package1) - .map { it.virtualFile } - .toSet() - - assertThat(actual).hasSize(expected.size) - actual.forEach { - assertThat(expected.contains(it)).isTrue - } - } - } - - @Test - fun `findFilesImported`() { - val mainClass = fixture.addFileToProject( - "Main.java", - """ - package com.cw.file_crawler_test; - - import java.util.Map; - import java.util.regex.Pattern; - - import com.cw.file_crawler_test.utils.AnotherClass; - import com.cw.file_crawler_test.service.controllers.MyController; - - public class Main { - }; - """.trimIndent() - ) - - val controllerClass = fixture.addFileToProject( - "service/controllers/MyController.java", - """ - package com.cw.file_crawler_test.service.controllers; - - public class MyController {} - """.trimIndent() - ) - - val anotherClass = fixture.addFileToProject( - "/utils/AnotherClass.java", - """ - package com.cw.file_crawler_test.utils; - - public class AnotherClass {} - """.trimIndent() - ) - - fun assertCrawlerFindCorrectFiles(sut: CodeWhispererFileCrawler) { - runReadAction { - val expected = setOf(controllerClass.virtualFile, anotherClass.virtualFile) - val actualFiles = runBlocking { sut.listFilesImported(mainClass) } - - assertThat(actualFiles).hasSize(2) - actualFiles.forEach { - assertThat(expected).contains(it) - } - } - } - - assertCrawlerFindCorrectFiles(JavaCodeWhispererFileCrawler) - // can't make it work right since the temp file created is not in the real file system - // Naive crawler will actually read the file system thun unable to find files - // assertCrawlerFindCorrectFiles(NaiveJavaCodeWhispererFileCrawler(project)) - } - @Test fun `listUtgCandidate by name`() { val mainPsi = fixture.addFileToProject("Main.java", aString()) From 7b4d556ca0549cf505f67b1ef41f1dc30486fa4b Mon Sep 17 00:00:00 2001 From: Will Lo Date: Mon, 12 Aug 2024 19:37:14 -0700 Subject: [PATCH 7/7] ptch test --- .../CodeWhispererFeatureConfigService.kt | 14 +++ .../util/CodeWhispererConstants.kt | 1 + .../util/CodeWhispererFileContextProvider.kt | 23 ++++- .../CodeWhispererFileContextProviderTest.kt | 94 +++++++++---------- 4 files changed, 82 insertions(+), 50 deletions(-) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt index af0bc46e36..abbdb28031 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt @@ -82,6 +82,19 @@ class CodeWhispererFeatureConfigService { fun getCustomizationArnOverride(): String = getFeatureValueForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME).stringValue() + fun getCrissfileConfig(): String { + // TODO: use key instead of random +// val key = getFeatureValueForKey(CROSSFILE_KEY).stringValue() + val randomNum = Math.random() + val group = if (randomNum < 1 / 2.0) { + "control" + } else { + "experiment" + } + + return group + } + // Get the feature value for the given key. // In case of a misconfiguration, it will return a default feature value of Boolean true. private fun getFeatureValueForKey(name: String): FeatureValue = @@ -92,6 +105,7 @@ class CodeWhispererFeatureConfigService { fun getInstance(): CodeWhispererFeatureConfigService = service() private const val TEST_FEATURE_NAME = "testFeature" const val CUSTOMIZATION_ARN_OVERRIDE_NAME = "customizationArnOverride" + const val CROSSFILE_KEY = "crossfile" private val LOG = getLogger() // TODO: add real feature later diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt index 842a3cf697..e7c9d73f5c 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt @@ -138,6 +138,7 @@ object CodeWhispererConstants { get() = if (isDeveloperMode()) 200 else 60 val NUMBER_OF_LINE_IN_CHUNK get() = if (isDeveloperMode()) 50 else 10 + // TODO: 3 -> 10 when service side CR is done const val NUMBER_OF_CHUNK_TO_FETCH = 3 const val NEIGHBOR_FILES_DISTANCE = 1 diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt index b94229bf85..719a0ebb76 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt @@ -23,16 +23,27 @@ import software.aws.toolkits.core.utils.warn import software.aws.toolkits.jetbrains.isDeveloperMode import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererC +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCpp +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCsharp +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererGo import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJavaScript import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJsx +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererKotlin +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPhp import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererRuby +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererRust +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererScala +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererShell import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTsx import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTypeScript import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage import software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo +import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererFeatureConfigService import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings import java.io.DataInput @@ -332,7 +343,17 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi is CodeWhispererJsx, is CodeWhispererTsx -> true - // TODO: languages under A/B, should read feature flag from [CodeWhispererFeatureConfigService] + is CodeWhispererC, + is CodeWhispererGo, + is CodeWhispererPhp, + is CodeWhispererRust, + is CodeWhispererKotlin, + is CodeWhispererCpp, + is CodeWhispererCsharp, + is CodeWhispererRuby, + is CodeWhispererShell, + is CodeWhispererScala -> CodeWhispererFeatureConfigService.getInstance().getCrissfileConfig() == "experiment" + else -> false } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt index 31573c9434..b5e844f7b7 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt @@ -14,6 +14,7 @@ import com.intellij.testFramework.replaceService import com.intellij.testFramework.runInEdtAndGet import com.intellij.testFramework.runInEdtAndWait import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.test.runTest import org.assertj.core.api.Assertions.assertThat import org.junit.Before import org.junit.Ignore @@ -31,12 +32,17 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererGo import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJavaScript +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJson import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJsx import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererKotlin +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPlainText import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererRuby +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTf import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTsx import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTypeScript +import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererYaml +import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererFeatureConfigService import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants @@ -122,72 +128,62 @@ class CodeWhispererFileContextProviderTest { @Test fun `shouldFetchCrossfileContext - fully support`() { - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJava.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJava.INSTANCE, CodeWhispererUserGroup.Control)).isTrue - - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererPython.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererPython.INSTANCE, CodeWhispererUserGroup.Control)).isTrue - - assertThat( - DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext( - CodeWhispererJavaScript.INSTANCE, - CodeWhispererUserGroup.CrossFile - ) - ).isTrue - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJavaScript.INSTANCE, CodeWhispererUserGroup.Control)).isTrue - - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJsx.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJsx.INSTANCE, CodeWhispererUserGroup.Control)).isTrue - - assertThat( - DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext( - CodeWhispererTypeScript.INSTANCE, - CodeWhispererUserGroup.CrossFile - ) - ).isTrue - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererTypeScript.INSTANCE, CodeWhispererUserGroup.Control)).isTrue - - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererTsx.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererTsx.INSTANCE, CodeWhispererUserGroup.Control)).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJava.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererPython.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJavaScript.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJsx.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererTypeScript.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererTsx.INSTANCE, mock())).isTrue } - @Ignore("Reenable this once we have any partial support language") @Test - fun `shouldFetchCrossfileContext - partially support`() { - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererPython.INSTANCE, CodeWhispererUserGroup.Control)).isFalse - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererPython.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue + fun `shouldFetchCrossfileContext - partially support should return true if feature flag is enabled`() { + val mockFeatureConfigService = mock { + on { getCrissfileConfig() } doReturn "experiment" + } + ApplicationManager.getApplication().replaceService(CodeWhispererFeatureConfigService::class.java, mockFeatureConfigService, disposableRule.disposable) + + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererPython.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCsharp.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererKotlin.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererGo.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCpp.INSTANCE, mock())).isTrue + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererRuby.INSTANCE, mock())).isTrue } @Test - fun `shouldFetchCrossfileContext - no support`() { - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCsharp.INSTANCE, CodeWhispererUserGroup.Control)).isNull() - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCsharp.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull() - - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererKotlin.INSTANCE, CodeWhispererUserGroup.Control)).isNull() - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererKotlin.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull() - - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererGo.INSTANCE, CodeWhispererUserGroup.Control)).isNull() - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererGo.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull() + fun `shouldFetchCrossfileContext - partially support should return true if feature flag is disabled`() { + val mockFeatureConfigService = mock { + on { getCrissfileConfig() } doReturn "control" + } + ApplicationManager.getApplication().replaceService(CodeWhispererFeatureConfigService::class.java, mockFeatureConfigService, disposableRule.disposable) - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCpp.INSTANCE, CodeWhispererUserGroup.Control)).isNull() - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCpp.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull() + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCsharp.INSTANCE, mock())).isFalse + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererKotlin.INSTANCE, mock())).isFalse + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererGo.INSTANCE, mock())).isFalse + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCpp.INSTANCE, mock())).isFalse + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererRuby.INSTANCE, mock())).isFalse + } - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererRuby.INSTANCE, CodeWhispererUserGroup.Control)).isNull() - assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererRuby.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull() + @Test + fun `shouldFetchCrossfileContext - no support`() { + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererYaml.INSTANCE, mock())).isNull() + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJson.INSTANCE, mock())).isNull() + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererPlainText.INSTANCE, mock())).isNull() + assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererTf.INSTANCE, mock())).isNull() } @Test - fun `languages not supporting supplemental context will return empty`() { + fun `languages not supported will return empty directly`() { val psiFiles = setupFixture(fixture) val psi = psiFiles[0] - runBlocking { - var context = aFileContextInfo(CodeWhispererCsharp.INSTANCE) - + runTest { + var context = aFileContextInfo(CodeWhispererYaml.INSTANCE) assertThat(sut.extractSupplementalFileContextForSrc(psi, context).contents).isEmpty() assertThat(sut.extractSupplementalFileContextForTst(psi, context).contents).isEmpty() - context = aFileContextInfo(CodeWhispererKotlin.INSTANCE) + context = aFileContextInfo(CodeWhispererTf.INSTANCE) assertThat(sut.extractSupplementalFileContextForSrc(psi, context).contents).isEmpty() assertThat(sut.extractSupplementalFileContextForTst(psi, context).contents).isEmpty() }