config(amazonq): Add project context to inline completion

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/workspace/context/ProjectContextProviderTest.kt

@@ -3,37 +3,218 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.workspace.context
 
+import com.github.tomakehurst.wiremock.client.WireMock.aResponse
+import com.github.tomakehurst.wiremock.client.WireMock.any
+import com.github.tomakehurst.wiremock.client.WireMock.stubFor
+import com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo
+import com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig
+import com.github.tomakehurst.wiremock.http.Body
+import com.github.tomakehurst.wiremock.junit.WireMockRule
 import com.intellij.openapi.project.Project
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
+import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
 import org.junit.Rule
+import org.junit.jupiter.api.assertThrows
 import org.mockito.kotlin.any
-import org.mockito.kotlin.mock
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.stub
 import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
 import software.aws.toolkits.jetbrains.services.amazonq.project.EncoderServer
+import software.aws.toolkits.jetbrains.services.amazonq.project.InlineBm25Chunk
+import software.aws.toolkits.jetbrains.services.amazonq.project.LspMessage
 import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextProvider
+import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument
 import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
 import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
 import java.net.ConnectException
+import java.util.concurrent.TimeoutException
 import kotlin.test.Test
 
 class ProjectContextProviderTest {
     @Rule
     @JvmField
     val projectRule: CodeInsightTestFixtureRule = JavaCodeInsightTestFixtureRule()
 
+    @Rule
+    @JvmField
+    val wireMock: WireMockRule = createMockServer()
+
     private val project: Project
         get() = projectRule.project
 
-    private val encoderServer: EncoderServer = mock()
+    private lateinit var encoderServer: EncoderServer
     private lateinit var sut: ProjectContextProvider
 
     @Before
     fun setup() {
+        encoderServer = spy(EncoderServer(project))
+        encoderServer.stub { on { port } doReturn wireMock.port() }
+
         sut = ProjectContextProvider(project, encoderServer, TestScope())
+
+        // initialization
+        stubFor(any(urlPathEqualTo("/initialize")).willReturn(aResponse().withStatus(200).withResponseBody(Body("initialize response"))))
+
+        // build index
+        stubFor(any(urlPathEqualTo("/indexFiles")).willReturn(aResponse().withStatus(200).withResponseBody(Body("initialize response"))))
+        stubFor(any(urlPathEqualTo("/buildIndex")).willReturn(aResponse().withStatus(200).withResponseBody(Body("initialize response"))))
+
+        // update index
+        stubFor(any(urlPathEqualTo("/updateIndex")).willReturn(aResponse().withStatus(200).withResponseBody(Body("initialize response"))))
+        stubFor(any(urlPathEqualTo("/updateIndexV2")).willReturn(aResponse().withStatus(200).withResponseBody(Body("initialize response"))))
+
+        // query
+        stubFor(
+            any(urlPathEqualTo("/query")).willReturn(
+                aResponse()
+                    .withStatus(200)
+                    .withResponseBody(Body(validQueryChatResponse))
+            )
+        )
+        stubFor(
+            any(urlPathEqualTo("/queryInlineProjectContext")).willReturn(
+                aResponse()
+                    .withStatus(200)
+                    .withResponseBody(
+                        Body(validQueryInlineResponse)
+                    )
+            )
+        )
+
+        stubFor(
+            any(urlPathEqualTo("/getUsage"))
+                .willReturn(
+                    aResponse()
+                        .withStatus(200)
+                        .withResponseBody(Body(validGetUsageResponse))
+                )
+        )
+    }
+
+    @Test
+    fun `Lsp endpoint correctness`() {
+        assertThat(LspMessage.Initialize.endpoint).isEqualTo("initialize")
+        assertThat(LspMessage.Index.endpoint).isEqualTo("buildIndex")
+        assertThat(LspMessage.UpdateIndex.endpoint).isEqualTo("updateIndexV2")
+        assertThat(LspMessage.QueryChat.endpoint).isEqualTo("query")
+        assertThat(LspMessage.QueryInlineCompletion.endpoint).isEqualTo("queryInlineProjectContext")
+        assertThat(LspMessage.GetUsageMetrics.endpoint).isEqualTo("getUsage")
+    }
+
+    @Test
+    fun `query chat should return empty if resultset non deserializable`() = runTest {
+        stubFor(
+            any(urlPathEqualTo("/query")).willReturn(
+                aResponse().withStatus(200).withResponseBody(
+                    Body(
+                        """
+                            [
+                                "foo", "bar"
+                            ]
+                        """.trimIndent()
+                    )
+                )
+            )
+        )
+
+        assertThrows<Exception> {
+            sut.query("foo")
+        }
+    }
+
+    @Test
+    fun `query chat should return deserialized relevantDocument`() = runTest {
+        val r = sut.query("foo")
+        assertThat(r).hasSize(2)
+        assertThat(r[0]).isEqualTo(
+            RelevantDocument(
+                "relativeFilePath1",
+                "context1"
+            )
+        )
+        assertThat(r[1]).isEqualTo(
+            RelevantDocument(
+                "relativeFilePath2",
+                "context2"
+            )
+        )
+    }
+
+    @Test
+    fun `query inline should throw if resultset not deserializable`() {
+        stubFor(
+            any(urlPathEqualTo("/queryInlineProjectContext")).willReturn(
+                aResponse().withStatus(200).withResponseBody(
+                    Body(
+                        """
+                            [
+                                "foo", "bar"
+                            ]
+                        """.trimIndent()
+                    )
+                )
+            )
+        )
+
+        assertThrows<Exception> {
+            sut.queryInline("foo", "filepath")
+        }
+    }
+
+    @Test
+    fun `query inline should return deserialized bm25 chunks`() = runTest {
+        val r = sut.queryInline("foo", "filepath")
+        assertThat(r).hasSize(3)
+        assertThat(r[0]).isEqualTo(
+            InlineBm25Chunk(
+                "content1",
+                "file1",
+                0.1
+            )
+        )
+        assertThat(r[1]).isEqualTo(
+            InlineBm25Chunk(
+                "content2",
+                "file2",
+                0.2
+            )
+        )
+        assertThat(r[2]).isEqualTo(
+            InlineBm25Chunk(
+                "content3",
+                "file3",
+                0.3
+            )
+        )
+    }
+
+    @Test
+    fun `get usage should return memory, cpu usage`() = runTest {
+        val r = sut.getUsage()
+        assertThat(r).isEqualTo(ProjectContextProvider.Usage(123, 456))
+    }
+
+    @Test
+    fun `should return empty if timeout with 50ms`() {
+        stubFor(
+            any(urlPathEqualTo("/queryInlineProjectContext")).willReturn(
+                aResponse()
+                    .withStatus(200)
+                    .withResponseBody(
+                        Body(validQueryInlineResponse)
+                    )
+                    .withFixedDelay(51) // 10 sec
+            )
+        )
+
+        assertThrows<TimeoutException> {
+            sut.queryInline("foo", "bar")
+        }
     }
 
     @Test
@@ -57,4 +238,70 @@ class ProjectContextProviderTest {
         }
         verify(encoderServer, times(1)).encrypt(any())
     }
+
+    private fun createMockServer() = WireMockRule(wireMockConfig().dynamicPort())
 }
+
+val validQueryInlineResponse = """
+                            [
+                                {
+                                    "content": "content1",
+                                    "filePath": "file1",
+                                    "score": 0.1
+                                },
+                                {
+                                    "content": "content2",
+                                    "filePath": "file2",
+                                    "score": 0.2
+                                },
+                                {
+                                    "content": "content3",
+                                    "filePath": "file3",
+                                    "score": 0.3
+                                }
+                            ]    
+""".trimIndent()
+
+val validQueryChatResponse = """
+                            [
+                                {
+                                    "filePath": "file1",
+                                    "content": "content1",
+                                    "id": "id1",
+                                    "index": "index1",
+                                    "vec": [
+                                        "vec_1-1",
+                                        "vec_1-2",
+                                        "vec_1-3"
+                                    ],
+                                    "context": "context1",
+                                    "prev": "prev1",
+                                    "next": "next1",
+                                    "relativePath": "relativeFilePath1",
+                                    "programmingLanguage": "language1"
+                                },
+                                {
+                                    "filePath": "file2",
+                                    "content": "content2",
+                                    "id": "id2",
+                                    "index": "index2",
+                                    "vec": [
+                                        "vec_2-1",
+                                        "vec_2-2",
+                                        "vec_2-3"
+                                    ],
+                                    "context": "context2",
+                                    "prev": "prev2",
+                                    "next": "next2",
+                                    "relativePath": "relativeFilePath2",
+                                    "programmingLanguage": "language2"
+                                }
+                            ]
+""".trimIndent()
+
+val validGetUsageResponse = """
+                                {
+                                "memoryUsage":123,
+                                "cpuUsage":456
+                                } 
+""".trimIndent()

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

@@ -22,6 +22,8 @@
 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
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
@@ -206,15 +208,49 @@
 
     override fun isTestFile(psiFile: PsiFile) = psiFile.programmingLanguage().fileCrawler.isTestFile(psiFile.virtualFile, psiFile.project)
 
-    @VisibleForTesting
     suspend fun extractSupplementalFileContextForSrc(psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo {
         if (!targetContext.programmingLanguage.isSupplementalContextSupported()) {
             return SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
         }
 
-        // takeLast(11) will extract 10 lines (exclusing current line) of left context as the query parameter
-        val query = targetContext.caretContext.leftFileContext.split("\n").takeLast(11).joinToString("\n")
+        val query = generateQuery(targetContext)
+
+        val projectContext = if (CodeWhispererFeatureConfigService.getInstance().getInlineCompletion()) {
+            fetchProjectContext(query, psiFile, targetContext)
+        } else {
+            null
+        }
+
+        val openTabsContext = fetchOpentabsContext(query, psiFile, targetContext)
+
+        return if (projectContext == null || projectContext.contents.isEmpty()) {
+            openTabsContext
+        } else {
+            projectContext
+        }
+    }
+
+    @VisibleForTesting
+    suspend fun fetchProjectContext(query: String, psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo {
+        val response = ProjectContextController.getInstance(project).queryInline(query, psiFile.virtualFile?.path ?: "")
 
+        return SupplementalContextInfo(
+            isUtg = false,
+            contents = response.map {
+                Chunk(
+                    content = it.content,
+                    path = it.filePath,
+                    nextChunk = it.content,
+                    score = it.score
+                )
+            },
+            targetFileName = targetContext.filename,
+            strategy = CrossFileStrategy.ProjectContext
+        )
+    }
+
+    @VisibleForTesting
+    suspend fun fetchOpentabsContext(query: String, psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo {
         // step 1: prepare data
         val first60Chunks: List<Chunk> = try {
             runReadAction { codewhispererCodeChunksIndex.getFileData(psiFile) }
@@ -305,6 +341,9 @@
         }
     }
 
+    // takeLast(11) will extract 10 lines (exclusing current line) of left context as the query parameter
+    fun generateQuery(fileContext: FileContextInfo) = fileContext.caretContext.leftFileContext.split("\n").takeLast(11).joinToString("\n")
+
     companion object {
         private val LOG = getLogger<DefaultCodeWhispererFileContextProvider>()
 

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

@@ -84,11 +84,11 @@ abstract class CodeWhispererFileCrawler : FileCrawler {
     }.orEmpty()
 
     override fun listCrossFileCandidate(target: PsiFile): List<VirtualFile> {
-        val targetFile = target.virtualFile
+        val targetFile = target.viewProvider.virtualFile
 
         val openedFiles = runReadAction {
             FileEditorManager.getInstance(target.project).openFiles.toList().filter {
-                it.name != target.virtualFile.name &&
+                it.name != targetFile.name &&
                     isSameDialect(it.extension) &&
                     !isTestFile(it, target.project)
             }

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

@@ -21,10 +21,12 @@ enum class UtgStrategy : SupplementalContextStrategy {
 enum class CrossFileStrategy : SupplementalContextStrategy {
     OpenTabsBM25,
     Empty,
+    ProjectContext,
     ;
 
     override fun toString() = when (this) {
         OpenTabsBM25 -> "OpenTabs_BM25"
         Empty -> "Empty"
+        ProjectContext -> "ProjectContext"
     }
 }