patch

Author: Will-ShaoHua

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

@@ -51,7 +51,7 @@ class ProjectContextProviderTest {
     fun `test query payload is encrypted`() = runTest {
         whenever(encoderServer.port).thenReturn(3000)
         try {
-            sut.query("what does this project do")
+            sut.queryChat("what does this project do")
         } catch (e: ConnectException) {
             // no-op
         }

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

@@ -53,7 +53,7 @@ class ProjectContextController(private val project: Project, private val cs: Cor
 
     fun query(prompt: String): List<RelevantDocument> {
         try {
-            return projectContextProvider.query(prompt)
+            return projectContextProvider.queryChat(prompt)
         } catch (e: Exception) {
             logger.warn { "error while querying for project context $e.message" }
             return emptyList()
@@ -62,7 +62,7 @@ class ProjectContextController(private val project: Project, private val cs: Cor
 
     fun queryBM25(prompt: String, filePath: String): Any {
         try {
-            return projectContextProvider.queryBM25(prompt, filePath)
+            return projectContextProvider.queryInline(prompt, filePath)
         } catch (e: Exception) {
             logger.warn { "error while querying for project context $e.message" }
             return emptyList<Any>()

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

@@ -181,36 +181,38 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         }
     }
 
-    fun query(prompt: String): List<RelevantDocument> {
+    fun queryChat(prompt: String): List<RelevantDocument> {
         logger.info { "project context: querying ${project.name} on port ${encoderServer.port}" }
         val url = URL("http://localhost:${encoderServer.port}/query")
         val payload = QueryRequestPayload(prompt)
         val payloadJson = mapper.writeValueAsString(payload)
         val encrypted = encoderServer.encrypt(payloadJson)
 
-        val connection = url.openConnection() as HttpURLConnection
-        setConnectionProperties(connection)
-        setConnectionTimeout(connection)
-        setConnectionRequest(connection, encrypted)
+        val response = with(url.openConnection() as HttpURLConnection) {
+            setConnectionProperties(this)
+            setConnectionTimeout(this)
+            setConnectionRequest(this, encrypted)
+            val responseCode = this.responseCode
 
-        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 {
-            ""
+            logger.info { "project context query response code: $responseCode for $prompt" }
+
+            if (responseCode == 200) {
+                this.inputStream.bufferedReader().use { reader -> reader.readText() }
+            } else {
+                ""
+            }
         }
-        connection.disconnect()
-        try {
-            val parsedResponse = mapper.readValue<List<Chunk>>(responseBody)
-            return queryResultToRelevantDocuments(parsedResponse)
+
+        return try {
+            val parsedResponse = mapper.readValue<List<Chunk>>(response)
+            queryResultToRelevantDocuments(parsedResponse)
         } catch (e: Exception) {
             logger.warn { "error parsing query response ${e.message}" }
-            return emptyList()
+            emptyList()
         }
     }
 
-    fun queryBM25(prompt: String, filePath: String): List<BM25Chunk> {
+    fun queryInline(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)
@@ -233,8 +235,8 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
             responseBody
         }
 
-        try {
-            return mapper.readValue<List<BM25Chunk>>(response)
+        return try {
+            mapper.readValue<List<BM25Chunk>>(response)
         } catch (e: Exception) {
             logger.warn { "error parsing query response ${e.message}" }
             throw e