patch

Author: Will-ShaoHua

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/LspApi.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
+
+sealed interface LspApi {
+    val command: String
+
+
+}

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

@@ -181,27 +181,29 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         }
     }
 
-    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)
+    fun updateIndex(filePaths: List<String>, mode: String) {
+        if (!isIndexComplete.get()) return
 
-        val response = with(url.openConnection() as HttpURLConnection) {
+        val url = URL("http://localhost:${encoderServer.port}/updateIndexV2")
+        val payload = UpdateIndexRequestV2(filePaths, mode)
+        val encrypted = encoderServer.encrypt(mapper.writeValueAsString(payload))
+
+        with(url.openConnection() as HttpURLConnection) {
             setConnectionProperties(this)
             setConnectionTimeout(this)
             setConnectionRequest(this, encrypted)
-            val responseCode = this.responseCode
+            val responseCode = responseCode
+            logger.debug { "project context update index response code: $responseCode" }
+            return
+        }
+    }
 
-            logger.info { "project context query response code: $responseCode for $prompt" }
+    fun queryChat(prompt: String): List<RelevantDocument> {
+        val payload = QueryRequestPayload(prompt)
+        val payloadJson = mapper.writeValueAsString(payload)
+        val encrypted = encoderServer.encrypt(payloadJson)
 
-            if (responseCode == 200) {
-                this.inputStream.bufferedReader().use { reader -> reader.readText() }
-            } else {
-                ""
-            }
-        }
+        val response = sendMsgToLsp("query", encrypted)
 
         return try {
             val parsedResponse = mapper.readValue<List<Chunk>>(response)
@@ -213,27 +215,11 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
     }
 
     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)
         val payloadJson = mapper.writeValueAsString(payload)
         val encrypted = encoderServer.encrypt(payloadJson)
 
-        val response = with(url.openConnection() as HttpURLConnection) {
-            setConnectionProperties(this)
-            setConnectionTimeout(this)
-            setConnectionRequest(this, encrypted)
-            val responseCode = this.responseCode
-            logger.info { "project context query response code: $responseCode for $prompt" }
-
-            val responseBody = if (responseCode == 200) {
-                this.inputStream.bufferedReader().use { reader -> reader.readText() }
-            } else {
-                ""
-            }
-
-            responseBody
-        }
+        val response = sendMsgToLsp("queryInlineProjectContext", encrypted)
 
         return try {
             mapper.readValue<List<BM25Chunk>>(response)
@@ -287,20 +273,24 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         }
     }
 
-    fun updateIndex(filePaths: List<String>, mode: String) {
-        if (!isIndexComplete.get()) return
-
-        val url = URL("http://localhost:${encoderServer.port}/updateIndexV2")
-        val payload = UpdateIndexRequestV2(filePaths, mode)
-        val encrypted = encoderServer.encrypt(mapper.writeValueAsString(payload))
+    fun sendMsgToLsp(msgType: String, request: String): String {
+        logger.info { "sending message: ${msgType} to lsp on port ${encoderServer.port}" }
+        val url = URL("http://localhost:${encoderServer.port}/$msgType")
 
-        with(url.openConnection() as HttpURLConnection) {
+        return with(url.openConnection() as HttpURLConnection) {
             setConnectionProperties(this)
             setConnectionTimeout(this)
-            setConnectionRequest(this, encrypted)
-            val responseCode = responseCode
-            logger.debug { "project context update index response code: $responseCode" }
-            return
+            setConnectionRequest(this, request)
+            val responseCode = this.responseCode
+            logger.info { "receiving response for $msgType with responseCode $responseCode" }
+
+            val responseBody = if (responseCode == 200) {
+                this.inputStream.bufferedReader().use { reader -> reader.readText() }
+            } else {
+                ""
+            }
+
+            responseBody
         }
     }