patch

Author: Will-ShaoHua

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

@@ -22,6 +22,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.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
@@ -214,6 +215,8 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
 
         // 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 bm25 = ProjectContextController.getInstance(project).queryBM25(query, targetContext.filename)
+        println("---------------------------------${bm25} --------------------------------------------")
 
         // step 1: prepare data
         val first60Chunks: List<Chunk> = try {

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

@@ -0,0 +1,10 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.project
+
+data class BM25Chunk(
+    val content: String,
+    val filePath: String,
+    val score: Double
+)

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

@@ -37,6 +37,15 @@ class ProjectContextController(private val project: Project, private val cs: Cor
         }
     }
 
+    fun queryBM25(prompt: String, filePath: String): Any {
+        try {
+            return projectContextProvider.queryBM25(prompt, filePath)
+        } catch (e: Exception) {
+            logger.warn { "error while querying for project context $e.message" }
+            return emptyList<Any>()
+        }
+    }
+
     fun updateIndex(filePath: String) {
         try {
             return projectContextProvider.updateIndex(filePath)

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

@@ -20,6 +20,7 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.yield
 import software.aws.toolkits.core.utils.debug
+import software.aws.toolkits.core.utils.error
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.warn
@@ -59,6 +60,13 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         val refresh: Boolean,
     )
 
+    data class IndexRequestV2(
+        val filePaths: List<String>,
+        val projectRoot: String,
+        val config: String,
+        val language: String = ""
+    )
+
     data class FileCollectionResult(
         val files: List<String>,
         val fileSize: Int,
@@ -68,10 +76,20 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         val query: String,
     )
 
+    data class QueryRequestV2(
+        val query: String,
+        val filePath: String,
+    )
+
     data class UpdateIndexRequestPayload(
         val filePath: String,
     )
 
+    data class UpdateIndexRequestV2(
+        val filePath: String,
+        val content: String,
+    )
+
     data class Usage(
         @JsonIgnoreProperties(ignoreUnknown = true)
         @JsonProperty("memoryUsage")
@@ -118,6 +136,7 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
                         return@launch
                     }
                 } catch (e: Exception) {
+                    logger.error { "project context index error: message=${e.message} stack=${e.stackTraceToString()}"}
                     if (e.stackTraceToString().contains("Connection refused")) {
                         retryCount.incrementAndGet()
                         delay(10000)
@@ -143,13 +162,14 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
     fun index(): Boolean {
         logger.info { "project context: indexing ${project.name} on port ${encoderServer.port}" }
         val indexStartTime = System.currentTimeMillis()
-        val url = URL("http://localhost:${encoderServer.port}/indexFiles")
+        val url = URL("http://localhost:${encoderServer.port}/buildIndex")
         val filesResult = collectFiles()
         var duration = (System.currentTimeMillis() - indexStartTime).toDouble()
         logger.debug { "project context file collection time: ${duration}ms" }
         logger.debug { "list of files collected: ${filesResult.files.joinToString("\n")}" }
         val projectRoot = project.guessProjectDir()?.path ?: return false
-        val payload = IndexRequestPayload(filesResult.files, projectRoot, false)
+//        val payload = IndexRequestPayload(filesResult.files, projectRoot, false)
+        val payload = IndexRequestV2(filesResult.files, projectRoot, "all", "")
         val payloadJson = mapper.writeValueAsString(payload)
         val encrypted = encoderServer.encrypt(payloadJson)
 
@@ -200,6 +220,34 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         }
     }
 
+    fun queryBM25(prompt: String, filePath: String): List<BM25Chunk> {
+        logger.info { "project context: querying ${project.name} on port ${encoderServer.port}" }
+        val url = URL("http://localhost:${encoderServer.port}/queryInlineProjectContext")
+        val payload = QueryRequestV2(prompt, filePath)
+        val payloadJson = mapper.writeValueAsString(payload)
+        val encrypted = encoderServer.encrypt(payloadJson)
+
+        val connection = url.openConnection() as HttpURLConnection
+        setConnectionProperties(connection)
+        setConnectionTimeout(connection)
+        setConnectionRequest(connection, encrypted)
+
+        val responseCode = connection.responseCode
+        logger.info { "project context query response code: $responseCode for $prompt" }
+        val responseBody = if (responseCode == 200) {
+            connection.inputStream.bufferedReader().use { reader -> reader.readText() }
+        } else {
+            ""
+        }
+        connection.disconnect()
+        try {
+            return mapper.readValue<List<BM25Chunk>>(responseBody)
+        } catch (e: Exception) {
+            logger.warn { "error parsing query response ${e.message}" }
+            throw e
+        }
+    }
+
     private fun recordIndexWorkspace(
         duration: Double,
         fileCount: Int = 0,