feat(inline completion): enhance inline completion supplemental context fetching

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt

@@ -25,7 +25,6 @@ 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.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextController
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
@@ -147,12 +146,20 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
                 val latency = System.currentTimeMillis() - startFetchingTimestamp
                 if (it.contents.isNotEmpty()) {
                     val logStr = buildString {
-                        append("Successfully fetched supplemental context with strategy ${it.strategy} with $latency ms")
+                        append(
+                            """Q inline completion supplemental context: 
+                            | Strategy: ${it.strategy},
+                            | Latency: $latency ms,
+                            | Contents: ${it.contents.size} chunks,
+                            | ContentLength: ${it.contentLength} chars,
+                            | TargetFile: ${it.targetFileName},
+                            """.trimMargin()
+                        )
                         it.contents.forEachIndexed { index, chunk ->
                             append(
                                 """
                             |
-                            | Chunk ${index + 1}:
+                            | Chunk $index:
                             |    path = ${chunk.path},
                             |    score = ${chunk.score},
                             |    contentLength = ${chunk.content.length}
@@ -219,23 +226,19 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
         val query = generateQuery(targetContext)
 
         val contexts = withContext(coroutineContext) {
-            val projectContextDeferred1 = if (CodeWhispererFeatureConfigService.getInstance().getInlineCompletion()) {
-                async {
-                    val t0 = System.currentTimeMillis()
-                    val r = fetchProjectContext(query, psiFile, targetContext)
-                    val t1 = System.currentTimeMillis()
-                    LOG.debug {
-                        buildString {
-                            append("time elapse for fetching project context=${t1 - t0}ms; ")
-                            append("numberOfChunks=${r.contents.size}; ")
-                            append("totalLength=${r.contentLength}")
-                        }
+            val projectContextDeferred1 = async {
+                val t0 = System.currentTimeMillis()
+                val r = fetchProjectContext(query, psiFile, targetContext)
+                val t1 = System.currentTimeMillis()
+                LOG.debug {
+                    buildString {
+                        append("time elapse for fetching project context=${t1 - t0}ms; ")
+                        append("numberOfChunks=${r.contents.size}; ")
+                        append("totalLength=${r.contentLength}")
                     }
-
-                    r
                 }
-            } else {
-                null
+
+                r
             }
 
             val openTabsContextDeferred1 = async {
@@ -253,20 +256,65 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
                 r
             }
 
-            if (projectContextDeferred1 != null) {
-                awaitAll(projectContextDeferred1, openTabsContextDeferred1)
-            } else {
-                awaitAll(openTabsContextDeferred1)
-            }
+            awaitAll(projectContextDeferred1, openTabsContextDeferred1)
         }
 
-        val projectContext = contexts.find { it.strategy == CrossFileStrategy.ProjectContext }
+        val projectContext = contexts.find { it.strategy == CrossFileStrategy.Codemap }
         val openTabsContext = contexts.find { it.strategy == CrossFileStrategy.OpenTabsBM25 }
 
-        return if (projectContext != null && projectContext.contents.isNotEmpty()) {
-            projectContext
-        } else {
-            openTabsContext ?: SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
+        return when {
+            projectContext == null && openTabsContext == null -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
+
+            projectContext != null && openTabsContext != null -> {
+                val context1 = projectContext.contents
+                val context2 = openTabsContext.contents
+                val mergedContext = (context1 + context2).filter { it.content.isNotEmpty() }
+
+                val strategy = if (projectContext.contentLength != 0 && openTabsContext.contentLength != 0) {
+                    CrossFileStrategy.Codemap
+                } else if (projectContext.contentLength != 0) {
+                    CrossFileStrategy.Codemap
+                } else if (openTabsContext.contentLength != 0) {
+                    CrossFileStrategy.OpenTabsBM25
+                } else {
+                    CrossFileStrategy.Empty
+                }
+
+                SupplementalContextInfo(
+                    isUtg = false,
+                    contents = mergedContext,
+                    targetFileName = targetContext.filename,
+                    strategy = strategy
+                )
+            }
+
+            projectContext != null -> {
+                return if (projectContext.contentLength == 0) {
+                    SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
+                } else {
+                    SupplementalContextInfo(
+                        isUtg = false,
+                        contents = projectContext.contents,
+                        targetFileName = targetContext.filename,
+                        strategy = CrossFileStrategy.Codemap
+                    )
+                }
+            }
+
+            openTabsContext != null -> {
+                return if (openTabsContext.contentLength == 0) {
+                    SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
+                } else {
+                    SupplementalContextInfo(
+                        isUtg = false,
+                        contents = openTabsContext.contents,
+                        targetFileName = targetContext.filename,
+                        strategy = CrossFileStrategy.OpenTabsBM25
+                    )
+                }
+            }
+
+            else -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
         }
     }
 
@@ -285,7 +333,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
                 )
             },
             targetFileName = targetContext.filename,
-            strategy = CrossFileStrategy.ProjectContext
+            strategy = CrossFileStrategy.Codemap
         )
     }
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/SupplementalContextStrategy.kt

@@ -12,21 +12,23 @@ enum class UtgStrategy : SupplementalContextStrategy {
     ;
 
     override fun toString() = when (this) {
-        ByName -> "ByName"
-        ByContent -> "ByContent"
-        Empty -> "Empty"
+        ByName -> "byname"
+        ByContent -> "bycontent"
+        Empty -> "empty"
     }
 }
 
 enum class CrossFileStrategy : SupplementalContextStrategy {
     OpenTabsBM25,
     Empty,
     ProjectContext,
+    Codemap,
     ;
 
     override fun toString() = when (this) {
-        OpenTabsBM25 -> "OpenTabs_BM25"
-        Empty -> "Empty"
-        ProjectContext -> "ProjectContext"
+        OpenTabsBM25 -> "opentabs"
+        Empty -> "empty"
+        ProjectContext -> "projectcontext"
+        Codemap -> "codemap"
     }
 }

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt

@@ -137,7 +137,7 @@ class CodeWhispererFileContextProviderTest {
 
     @Test
     fun `extractSupplementalFileContext should timeout 50ms`() = runTest {
-        featureConfigService.stub { on { getInlineCompletion() } doReturn false }
+        mockProjectContext.stub { onBlocking { queryInline(any(), any()) }.doReturn(emptyList()) }
         sut = spy(sut)
 
         val files = NaiveSampleCase.setupFixture(fixture)
@@ -159,7 +159,8 @@ class CodeWhispererFileContextProviderTest {
     }
 
     @Test
-    fun `should only call and use openTabsContext if projectContext is disabled`() = runTest {
+    fun `should only use openTabsContext if projectContext is empty`() = runTest {
+        mockProjectContext.stub { onBlocking { queryInline(any(), any()) }.doReturn(emptyList()) }
         featureConfigService.stub { on { getInlineCompletion() } doReturn false }
         sut = spy(sut)
 
@@ -169,7 +170,7 @@ class CodeWhispererFileContextProviderTest {
 
         val result = sut.extractSupplementalFileContextForSrc(queryPsi, mockFileContext)
 
-        verify(sut, times(0)).fetchProjectContext(any(), any(), any())
+        verify(sut, times(1)).fetchProjectContext(any(), any(), any())
         verify(sut, times(1)).fetchOpenTabsContext(any(), any(), any())
 
         assertThat(result.isUtg).isFalse
@@ -208,7 +209,7 @@ class CodeWhispererFileContextProviderTest {
                 assertThat(providerContext.constructed()).hasSize(1)
                 assertThat(serverContext.constructed()).hasSize(1)
 
-                whenever(providerContext.constructed()[0].queryInline(any(), any())).thenThrow(RuntimeException("mock exception"))
+                whenever(providerContext.constructed()[0].queryInline(any(), any(), any())).thenThrow(RuntimeException("mock exception"))
 
                 val result = controller.queryInline("query", "filePath")
                 assertThat(result).isEmpty()
@@ -217,19 +218,16 @@ class CodeWhispererFileContextProviderTest {
     }
 
     @Test
-    fun `should use project context if it is present`() = runTest {
+    fun `should use both project context and open tabs if both are present`() = runTest {
         mockProjectContext.stub {
             runBlocking {
                 doReturn(
                     listOf(
                         InlineBm25Chunk("project_context1", "path1", 0.0),
-                        InlineBm25Chunk("project_context2", "path2", 0.0),
-                        InlineBm25Chunk("project_context3", "path3", 0.0),
                     )
                 ).whenever(it).queryInline(any(), any())
             }
         }
-        featureConfigService.stub { on { getInlineCompletion() } doReturn true }
         sut = spy(sut)
         val files = NaiveSampleCase.setupFixture(fixture)
         val queryPsi = files[0]
@@ -238,8 +236,8 @@ class CodeWhispererFileContextProviderTest {
         val result = sut.extractSupplementalFileContextForSrc(queryPsi, mockFileContext)
 
         assertThat(result.isUtg).isFalse
-        assertThat(result.strategy).isEqualTo(CrossFileStrategy.ProjectContext)
-        assertThat(result.contents).hasSize(3)
+        assertThat(result.strategy).isEqualTo(CrossFileStrategy.Codemap)
+        assertThat(result.contents).hasSize(4)
     }
 
     @Test
@@ -420,6 +418,7 @@ class CodeWhispererFileContextProviderTest {
 
     @Test
     fun `extractSupplementalFileContext from src file should extract src`() = runTest {
+        mockProjectContext.stub { onBlocking { queryInline(any(), any()) }.doReturn(emptyList()) }
         val files = NaiveSampleCase.setupFixture(fixture)
         val queryPsi = files[0]
 

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/LspMessage.kt

@@ -52,6 +52,7 @@
 data class QueryInlineCompletionRequest(
     val query: String,
     val filePath: String,
+    val target: String,
 ) : LspRequest
 
 data class LspResponse(
@@ -70,6 +71,14 @@
     DEFAULT("default"),
 }
 
+enum class InlineContextTarget(private val v: String) {
+    DEFAULT("default"),
+    CODEMAP("codemap"),
+    BM25("bm25"), ;
+
+    override fun toString(): String = this.v
+}
+
 // TODO: unify with [software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk]
 data class InlineBm25Chunk(
     val content: String,

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextController.kt

@@ -58,7 +58,7 @@ class ProjectContextController(private val project: Project, private val cs: Cor
 
     suspend fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> =
         try {
-            projectContextProvider.queryInline(query, filePath)
+            projectContextProvider.queryInline(query, filePath, InlineContextTarget.CODEMAP)
         } catch (e: Exception) {
             var logStr = "error while querying inline for project context $e.message"
             if (e is TimeoutCancellationException || e is TimeoutException) {

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt

@@ -171,9 +171,9 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         }.await()
     }
 
-    suspend fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> = withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
+    suspend fun queryInline(query: String, filePath: String, target: InlineContextTarget): List<InlineBm25Chunk> = withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
         cs.async {
-            val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath))
+            val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath, target.toString()))
             val r = sendMsgToLsp(LspMessage.QueryInlineCompletion, encrypted)
             return@async mapper.readValue<List<InlineBm25Chunk>>(r.responseBody)
         }.await()