fix(chat): handle custom agent state and add support for custom varia…

…ble compilation #51
unit-mesh · Mar 6, 2024 · 93a5e52 · 93a5e52
1 parent 71cf4f2
commit 93a5e52
Show file tree

Hide file tree

Showing 6 changed files with 58 additions and 9 deletions.
diff --git a/src/main/kotlin/cc/unitmesh/devti/counit/view/WebViewWindow.kt b/src/main/kotlin/cc/unitmesh/devti/counit/view/WebViewWindow.kt
@@ -14,7 +14,7 @@ import java.awt.Component
  * for custom webview can refs: https://github.com/mucharafal/jcef_example
  */
 class WebViewWindow {
-    // offical doc: https://plugins.jetbrains.com/docs/intellij/jcef.html#executing-javascript
+    // official doc: https://plugins.jetbrains.com/docs/intellij/jcef.html#executing-javascript
     private val browser: JBCefBrowser
 
     init {

diff --git a/src/main/kotlin/cc/unitmesh/devti/custom/compile/CustomVariable.kt b/src/main/kotlin/cc/unitmesh/devti/custom/compile/CustomVariable.kt
@@ -17,7 +17,7 @@ enum class CustomVariable(val variable: String, val description: String) {
         fun all(): List<CustomVariable> = values().toList()
 
         fun hasVariable(content: String): Boolean {
-            return all().any { content.contains("\${${it.variable}}") }
+            return all().any { content.contains("\$${it.variable}") }
         }
 
         fun compile(content: String, compiler: VariableTemplateCompiler): String {

diff --git a/src/main/kotlin/cc/unitmesh/devti/gui/chat/AutoDevInputSection.kt b/src/main/kotlin/cc/unitmesh/devti/gui/chat/AutoDevInputSection.kt
@@ -40,6 +40,7 @@ import com.intellij.util.ui.JBEmptyBorder
 import com.intellij.util.ui.JBUI
 import com.intellij.util.ui.UIUtil
 import com.intellij.util.ui.components.BorderLayoutPanel
+import kotlinx.serialization.decodeFromString
 import kotlinx.serialization.json.Json
 import java.awt.Color
 import java.awt.Component
@@ -199,6 +200,7 @@ class AutoDevInputSection(private val project: Project, val disposable: Disposab
                             text += "${selectedItem.customVariable.variable} "
                         }
                         this@AutoDevInputSection.popup?.cancel()
+                        this@AutoDevInputSection.requestFocus()
                     }
 
                     KeyEvent.VK_DOWN -> {
@@ -219,6 +221,11 @@ class AutoDevInputSection(private val project: Project, val disposable: Disposab
                             list.setSelectedIndex(list.getItemsCount() - 1)
                         }
                     }
+
+                    // Esc
+                    KeyEvent.VK_ESCAPE -> {
+                        this@AutoDevInputSection.requestFocus()
+                    }
                 }
             }
         })

diff --git a/src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingService.kt b/src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingService.kt
@@ -6,6 +6,8 @@ import cc.unitmesh.devti.util.LLMCoroutineScope
 import cc.unitmesh.devti.counit.CustomAgentChatProcessor
 import cc.unitmesh.devti.counit.configurable.customAgentSetting
 import cc.unitmesh.devti.counit.model.CustomAgentState
+import cc.unitmesh.devti.custom.compile.CustomVariable
+import cc.unitmesh.devti.custom.compile.VariableTemplateCompiler
 import cc.unitmesh.devti.llms.LlmFactory
 import cc.unitmesh.devti.util.parser.PostCodeProcessor
 import cc.unitmesh.devti.provider.ContextPrompter
@@ -29,16 +31,30 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
         context: ChatContext? = null,
         newChatContext: Boolean,
     ) {
-        val requestPrompt = prompter.requestPrompt()
-        val displayPrompt = prompter.displayPrompt()
+        var requestPrompt = prompter.requestPrompt()
+        var displayPrompt = prompter.displayPrompt()
 
         if (project.customAgentSetting.enableCustomRag && ui.hasSelectedCustomAgent()) {
-            if (ui.getSelectedCustomAgent().state === CustomAgentState.START) {
-                counitProcessor.handleChat(prompter, ui, context, llmProvider)
-                return
+            val selectedCustomAgent = ui.getSelectedCustomAgent()
+            when {
+                selectedCustomAgent.state === CustomAgentState.START -> {
+                    counitProcessor.handleChat(prompter, ui, context, llmProvider)
+                    return
+                }
+
+                selectedCustomAgent.state === CustomAgentState.FINISHED -> {
+                    if (CustomVariable.hasVariable(requestPrompt)) {
+                        val compiler = prompter.toTemplateCompiler()
+                        compiler?.also {
+                            requestPrompt = CustomVariable.compile(requestPrompt, it)
+                            displayPrompt = CustomVariable.compile(displayPrompt, it)
+                        }
+                    }
+                }
             }
         }
 
+
         ui.addMessage(requestPrompt, true, displayPrompt)
         ui.addMessage(AutoDevBundle.message("autodev.loading"))
 

diff --git a/src/main/kotlin/cc/unitmesh/devti/provider/ContextPrompter.kt b/src/main/kotlin/cc/unitmesh/devti/provider/ContextPrompter.kt
@@ -1,13 +1,19 @@
 package cc.unitmesh.devti.provider
 
+import cc.unitmesh.devti.custom.compile.VariableTemplateCompiler
 import cc.unitmesh.devti.gui.chat.ChatActionType
 import cc.unitmesh.devti.provider.builtin.DefaultContextPrompter
 import cc.unitmesh.devti.provider.context.ChatContextProvider
 import cc.unitmesh.devti.provider.context.ChatCreationContext
 import cc.unitmesh.devti.settings.coder.coderSetting
+import com.intellij.lang.html.HTMLLanguage
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.project.ProjectManager
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.serviceContainer.LazyExtensionInstance
@@ -80,6 +86,26 @@ abstract class ContextPrompter : LazyExtensionInstance<ContextPrompter>() {
     open fun displayPrompt(): String = ""
     open fun requestPrompt(): String = ""
 
+    fun toTemplateCompiler(): VariableTemplateCompiler? {
+        val project = project ?: ProjectManager.getInstance().openProjects.firstOrNull() ?: return null
+        val editor = FileEditorManager.getInstance(project).selectedTextEditor ?: return null
+
+        val file: PsiFile = file ?: PsiDocumentManager.getInstance(project).getPsiFile(editor.document) ?: return null
+
+        val selectedText = selectedText.ifEmpty {
+            editor.selectionModel.selectedText ?: ""
+        }
+
+        return VariableTemplateCompiler(
+            language = file.language,
+            file = file,
+            element = element,
+            editor = editor,
+            selectedText = selectedText
+        )
+
+    }
+
     companion object {
         private val EP_NAME: ExtensionPointName<ContextPrompter> =
             ExtensionPointName.create("cc.unitmesh.contextPrompter")

diff --git a/src/test/kotlin/cc/unitmesh/devti/custom/variable/CustomVariableTest.kt b/src/test/kotlin/cc/unitmesh/devti/custom/variable/CustomVariableTest.kt
@@ -8,7 +8,7 @@ import org.junit.Test
 class CustomVariableTest {
     @Test
     fun should_parse_variable_from_content() {
-        CustomVariable.hasVariable("解释一下代码：\${selection}") shouldBe true
-        CustomVariable.hasVariable("解释一下代码：\${selection") shouldBe false
+        CustomVariable.hasVariable("解释一下代码：\$selection") shouldBe true
+        CustomVariable.hasVariable("解释一下代码：\$selectio") shouldBe false
     }
 }