projectContextController & abtest

Author: Will-ShaoHua

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.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,14 +23,29 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhisp
 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()
 
     fun query(prompt: String): List<RelevantDocument> {
@@ -37,9 +57,22 @@ class ProjectContextController(private val project: Project, private val cs: Cor
         }
     }
 
+    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,22 @@
 // SPDX-License-Identifier: Apache-2.0
 
 package software.aws.toolkits.jetbrains.services.cwc.editor.context.project
+import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.fileEditor.FileEditorManager
+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

@@ -54,6 +54,12 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         }
     }
 
+    enum class IndexUpdateMode(val value: String) {
+        UPDATE("update"),
+        REMOVE("remove"),
+        ADD("add"),
+    }
+
     data class FileCollectionResult(
         val files: List<String>,
         val fileSize: Int,
@@ -185,9 +191,9 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         }
     }
 
-    fun updateIndex(filePaths: List<String>, mode: String) {
+    fun updateIndex(filePaths: List<String>, mode: IndexUpdateMode) {
         if (!isIndexComplete.get()) return
-        val encrypted = encryptRequest(UpdateIndexRequest(filePaths, mode))
+        val encrypted = encryptRequest(UpdateIndexRequest(filePaths, mode.value))
         sendMsgToLsp(LspMessage.UpdateIndex, encrypted)
     }
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt

@@ -10,6 +10,7 @@ import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
 import software.amazon.awssdk.services.codewhispererruntime.model.FeatureValue
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
+import software.aws.toolkits.jetbrains.isDeveloperMode
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfBIDConnection
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfIamIdentityCenterConnection
@@ -94,6 +95,9 @@ class CodeWhispererFeatureConfigService {
 
     fun getNewAutoTriggerUX(): Boolean = getFeatureValueForKey(NEW_AUTO_TRIGGER_UX).boolValue()
 
+    // TODO: remove dev mode flag
+    fun getInlineCompletion(): Boolean = if(isDeveloperMode()) true else getFeatureValueForKey(INLINE_COMPLETION).boolValue()
+
     // Get the feature value for the given key.
     // In case of a misconfiguration, it will return a default feature value of Boolean false.
     private fun getFeatureValueForKey(name: String): FeatureValue =
@@ -103,6 +107,8 @@ class CodeWhispererFeatureConfigService {
     companion object {
         fun getInstance(): CodeWhispererFeatureConfigService = service()
         private const val TEST_FEATURE_NAME = "testFeature"
+        // TODO: update
+        private const val INLINE_COMPLETION = "inlineCompletion"
         private const val DATA_COLLECTION_FEATURE = "IDEProjectContextDataCollection"
         const val CUSTOMIZATION_ARN_OVERRIDE_NAME = "customizationArnOverride"
         private const val NEW_AUTO_TRIGGER_UX = "newAutoTriggerUX"