config(amazonq): Add project context to inline completion

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/editor/context/project/InlineBm25Chunk.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.cwc.editor.context.project
+
+data class InlineBm25Chunk(
+    val content: String,
+    val filePath: String,
+    val score: Double,
+)

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/editor/context/project/LspMessage.kt

@@ -0,0 +1,60 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.cwc.editor.context.project
+
+sealed interface LspMessage {
+    val endpoint: String
+
+    data object Initialize : LspMessage {
+        override val endpoint: String = "initialize"
+    }
+
+    data object Index : LspMessage {
+        override val endpoint: String = "buildIndex"
+    }
+
+    data object UpdateIndex : LspMessage {
+        override val endpoint: String = "updateIndexV2"
+    }
+
+    data object QueryChat : LspMessage {
+        override val endpoint: String = "query"
+    }
+
+    data object QueryInlineCompletion : LspMessage {
+        override val endpoint: String = "queryInlineProjectContext"
+    }
+
+    data object GetUsageMetrics : LspMessage {
+        override val endpoint: String = "getUsage"
+    }
+}
+
+interface LspRequest
+
+data class IndexRequest(
+    val filePaths: List<String>,
+    val projectRoot: String,
+    val config: String,
+    val language: String = "",
+) : LspRequest
+
+data class UpdateIndexRequest(
+    val filePaths: List<String>,
+    val mode: String,
+) : LspRequest
+
+data class QueryChatRequest(
+    val query: String,
+) : LspRequest
+
+data class QueryInlineCompletionRequest(
+    val query: String,
+    val filePath: String,
+) : LspRequest
+
+data class LspResponse(
+    val responseCode: Int,
+    val responseBody: String,
+)

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/editor/context/project/ProjectContextController.kt

@@ -8,6 +8,11 @@
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.vfs.VirtualFileManager
+import com.intellij.openapi.vfs.newvfs.BulkFileListener
+import com.intellij.openapi.vfs.newvfs.events.VFileCreateEvent
+import com.intellij.openapi.vfs.newvfs.events.VFileDeleteEvent
+import com.intellij.openapi.vfs.newvfs.events.VFileEvent
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
 import software.aws.toolkits.core.utils.getLogger
@@ -18,12 +23,26 @@
 class ProjectContextController(private val project: Project, private val cs: CoroutineScope) : Disposable {
     private val encoderServer: EncoderServer = EncoderServer(project)
     private val projectContextProvider: ProjectContextProvider = ProjectContextProvider(project, encoderServer, cs)
+
     init {
         cs.launch {
             if (CodeWhispererSettings.getInstance().isProjectContextEnabled()) {
                 encoderServer.downloadArtifactsAndStartServer()
             }
         }
+
+        project.messageBus.connect(this).subscribe(
+            VirtualFileManager.VFS_CHANGES,
+            object : BulkFileListener {
+                override fun after(events: MutableList<out VFileEvent>) {
+                    val createdFiles = events.filterIsInstance<VFileCreateEvent>().mapNotNull { it.file?.path }
+                    val deletedFiles = events.filterIsInstance<VFileDeleteEvent>().map { it.file.path }
+
+                    updateIndex(createdFiles, ProjectContextProvider.IndexUpdateMode.ADD)
+                    updateIndex(deletedFiles, ProjectContextProvider.IndexUpdateMode.REMOVE)
+                }
+            }
+        )
     }
 
     fun getProjectContextIndexComplete() = projectContextProvider.isIndexComplete.get()
@@ -37,9 +56,22 @@
         }
     }
 
+    fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> {
+        try {
+            return projectContextProvider.queryInline(query, filePath)
+        } catch (e: Exception) {
+            logger.warn { "error while querying inline for project context $e.message" }
+            return emptyList()
+        }
+    }
+
     fun updateIndex(filePath: String) {
+        updateIndex(listOf(filePath), ProjectContextProvider.IndexUpdateMode.UPDATE)
+    }
+
+    fun updateIndex(filePaths: List<String>, mode: ProjectContextProvider.IndexUpdateMode) {
         try {
-            return projectContextProvider.updateIndex(filePath)
+            return projectContextProvider.updateIndex(filePaths, mode)
         } catch (e: Exception) {
             logger.warn { "error while updating index for project context $e.message" }
         }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/editor/context/project/ProjectContextEditorListener.kt

@@ -2,15 +2,19 @@
 // SPDX-License-Identifier: Apache-2.0
 
 package software.aws.toolkits.jetbrains.services.cwc.editor.context.project
-import com.intellij.openapi.fileEditor.FileEditorManager
+import com.intellij.openapi.fileEditor.FileDocumentManager
+import com.intellij.openapi.fileEditor.FileEditorManagerEvent
 import com.intellij.openapi.fileEditor.FileEditorManagerListener
-import com.intellij.openapi.vfs.VirtualFile
-import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
 
 class ProjectContextEditorListener : FileEditorManagerListener {
-    override fun fileClosed(source: FileEditorManager, file: VirtualFile) {
-        if (CodeWhispererSettings.getInstance().isProjectContextEnabled()) {
-            ProjectContextController.getInstance(source.project).updateIndex(file.path)
+    override fun selectionChanged(event: FileEditorManagerEvent) {
+        val oldFile = event.oldFile ?: return
+
+        // TODO: should respect isIdeAutosave config
+        with(FileDocumentManager.getInstance()) {
+            getDocument(oldFile)?.let {
+                saveDocument(it)
+            }
         }
     }
 }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/editor/context/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
@@ -53,25 +54,18 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
             }
         }
     }
-    data class IndexRequestPayload(
-        val filePaths: List<String>,
-        val projectRoot: String,
-        val refresh: Boolean,
-    )
+
+    enum class IndexUpdateMode(val value: String) {
+        UPDATE("update"),
+        REMOVE("remove"),
+        ADD("add"),
+    }
 
     data class FileCollectionResult(
         val files: List<String>,
         val fileSize: Int,
     )
 
-    data class QueryRequestPayload(
-        val query: String,
-    )
-
-    data class UpdateIndexRequestPayload(
-        val filePath: String,
-    )
-
     data class Usage(
         @JsonIgnoreProperties(ignoreUnknown = true)
         @JsonProperty("memoryUsage")
@@ -130,37 +124,27 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
     }
 
     private fun initEncryption(): Boolean {
-        logger.info { "project context: init key for ${project.guessProjectDir()} on port ${encoderServer.port}" }
-        val url = URL("http://localhost:${encoderServer.port}/initialize")
-        val payload = encoderServer.getEncryptionRequest()
-        val connection = url.openConnection() as HttpURLConnection
-        setConnectionProperties(connection)
-        setConnectionRequest(connection, payload)
-        logger.info { "project context initialize response code: ${connection.responseCode} for ${project.name}" }
-        return connection.responseCode == 200
+        val request = encoderServer.getEncryptionRequest()
+        val response = sendMsgToLsp(LspMessage.Initialize, request)
+        return response.responseCode == 200
     }
 
     fun index(): Boolean {
-        logger.info { "project context: indexing ${project.name} on port ${encoderServer.port}" }
+        val projectRoot = project.guessProjectDir()?.path ?: return false
+
         val indexStartTime = System.currentTimeMillis()
-        val url = URL("http://localhost:${encoderServer.port}/indexFiles")
         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 payloadJson = mapper.writeValueAsString(payload)
-        val encrypted = encoderServer.encrypt(payloadJson)
-
-        val connection = url.openConnection() as HttpURLConnection
-        setConnectionProperties(connection)
-        setConnectionRequest(connection, encrypted)
-        logger.info { "project context index response code: ${connection.responseCode} for ${project.name}" }
+        logger.debug { "time elapsed to collect project context files: ${duration}ms, collected ${filesResult.files.size} files" }
+
+        val encrypted = encryptRequest(IndexRequest(filesResult.files, projectRoot, "all", ""))
+        val response = sendMsgToLsp(LspMessage.Index, encrypted)
+
         duration = (System.currentTimeMillis() - indexStartTime).toDouble()
-        val startUrl = getStartUrl(project)
         logger.debug { "project context index time: ${duration}ms" }
-        if (connection.responseCode == 200) {
+
+        val startUrl = getStartUrl(project)
+        if (response.responseCode == 200) {
             val usage = getUsage()
             TelemetryHelper.recordIndexWorkspace(duration, filesResult.files.size, filesResult.fileSize, true, usage?.memoryUsage, usage?.cpuUsage, startUrl)
             logger.debug { "project context index finished for ${project.name}" }
@@ -171,73 +155,47 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         }
     }
 
+    // TODO: rename queryChat
     fun query(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 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 {
-            val parsedResponse = mapper.readValue<List<Chunk>>(responseBody)
-            return queryResultToRelevantDocuments(parsedResponse)
+        val encrypted = encryptRequest(QueryChatRequest(prompt))
+        val response = sendMsgToLsp(LspMessage.QueryChat, encrypted)
+
+        return try {
+            val parsedResponse = mapper.readValue<List<Chunk>>(response.responseBody)
+            queryResultToRelevantDocuments(parsedResponse)
         } catch (e: Exception) {
-            logger.warn { "error parsing query response ${e.message}" }
-            return emptyList()
+            logger.error { "error parsing query response ${e.message}" }
+            throw e
         }
     }
 
-    private fun getUsage(): Usage? {
-        logger.info { "project context: getting usage for ${project.name} on port ${encoderServer.port}" }
-        val url = URL("http://localhost:${encoderServer.port}/getUsage")
-        val connection = url.openConnection() as HttpURLConnection
-        setConnectionProperties(connection)
-        val responseCode = connection.responseCode
+    fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> {
+        val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath))
+        val response = sendMsgToLsp(LspMessage.QueryInlineCompletion, encrypted)
 
-        logger.info { "project context getUsage response code: $responseCode for ${project.name} " }
-        val responseBody = if (responseCode == 200) {
-            connection.inputStream.bufferedReader().use { reader -> reader.readText() }
-        } else {
-            ""
+        return try {
+            mapper.readValue<List<InlineBm25Chunk>>(response.responseBody)
+        } catch (e: Exception) {
+            logger.error { "error parsing query response ${e.message}" }
+            throw e
         }
-        connection.disconnect()
-        try {
-            val parsedResponse = mapper.readValue<Usage>(responseBody)
-            return parsedResponse
+    }
+
+    fun getUsage(): Usage? {
+        val response = sendMsgToLsp(LspMessage.GetUsageMetrics, request = null)
+        return try {
+            val parsedResponse = mapper.readValue<Usage>(response.responseBody)
+            parsedResponse
         } catch (e: Exception) {
             logger.warn { "error parsing query response ${e.message}" }
-            return null
+            null
         }
     }
 
-    fun updateIndex(filePath: String) {
+    fun updateIndex(filePaths: List<String>, mode: IndexUpdateMode) {
         if (!isIndexComplete.get()) return
-        logger.info { "project context: updating index for $filePath on port ${encoderServer.port}" }
-        val url = URL("http://localhost:${encoderServer.port}/updateIndex")
-        val payload = UpdateIndexRequestPayload(filePath)
-        val payloadJson = mapper.writeValueAsString(payload)
-        val encrypted = encoderServer.encrypt(payloadJson)
-        with(url.openConnection() as HttpURLConnection) {
-            setConnectionProperties(this)
-            setConnectionTimeout(this)
-            setConnectionRequest(this, encrypted)
-            val responseCode = responseCode
-            logger.debug { "project context update index response code: $responseCode for $filePath" }
-            return
-        }
+        val encrypted = encryptRequest(UpdateIndexRequest(filePaths, mode.value))
+        sendMsgToLsp(LspMessage.UpdateIndex, encrypted)
     }
 
     private fun setConnectionTimeout(connection: HttpURLConnection) {
@@ -325,6 +283,34 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         return documents
     }
 
+    private fun encryptRequest(r: LspRequest): String {
+        val payloadJson = mapper.writeValueAsString(r)
+        return encoderServer.encrypt(payloadJson)
+    }
+
+    private fun sendMsgToLsp(msgType: LspMessage, request: String?): LspResponse {
+        logger.info { "sending message: ${msgType.endpoint} to lsp on port ${encoderServer.port}" }
+        val url = URL("http://localhost:${encoderServer.port}/${msgType.endpoint}")
+
+        return with(url.openConnection() as HttpURLConnection) {
+            setConnectionProperties(this)
+            setConnectionTimeout(this)
+            request?.let { r ->
+                setConnectionRequest(this, r)
+            }
+            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 {
+                ""
+            }
+
+            LspResponse(responseCode, responseBody)
+        }
+    }
+
     override fun dispose() {
         retryCount.set(0)
     }