fix(chat): handle custom agent state and add support for custom variable compilation #51

phodal · phodal · commit 93a5e527df18 · 2024-03-06T21:53:01.000+08:00
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
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@ enum class CustomVariable(val variable: String, val description: String) {`
`17`	`17`	`fun all(): List<CustomVariable> = values().toList()`
`18`	`18`
`19`	`19`	`fun hasVariable(content: String): Boolean {`
`20`		`- return all().any { content.contains("\${${it.variable}}") }`
	`20`	`+ return all().any { content.contains("\$${it.variable}") }`
`21`	`21`	`}`
`22`	`22`
`23`	`23`	`fun compile(content: String, compiler: VariableTemplateCompiler): String {`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@ import com.intellij.util.ui.JBEmptyBorder`
`40`	`40`	`import com.intellij.util.ui.JBUI`
`41`	`41`	`import com.intellij.util.ui.UIUtil`
`42`	`42`	`import com.intellij.util.ui.components.BorderLayoutPanel`
	`43`	`+import kotlinx.serialization.decodeFromString`
`43`	`44`	`import kotlinx.serialization.json.Json`
`44`	`45`	`import java.awt.Color`
`45`	`46`	`import java.awt.Component`
`@@ -199,6 +200,7 @@ class AutoDevInputSection(private val project: Project, val disposable: Disposab`
`199`	`200`	`text += "${selectedItem.customVariable.variable} "`
`200`	`201`	`}`
`201`	`202`	`this@AutoDevInputSection.popup?.cancel()`
	`203`	`+ this@AutoDevInputSection.requestFocus()`
`202`	`204`	`}`
`203`	`205`
`204`	`206`	`KeyEvent.VK_DOWN -> {`
`@@ -219,6 +221,11 @@ class AutoDevInputSection(private val project: Project, val disposable: Disposab`
`219`	`221`	`list.setSelectedIndex(list.getItemsCount() - 1)`
`220`	`222`	`}`
`221`	`223`	`}`
	`224`	`+`
	`225`	`+ // Esc`
	`226`	`+ KeyEvent.VK_ESCAPE -> {`
	`227`	`+ this@AutoDevInputSection.requestFocus()`
	`228`	`+ }`
`222`	`229`	`}`
`223`	`230`	`}`
`224`	`231`	`})`
Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@ import org.junit.Test`
`8`	`8`	`class CustomVariableTest {`
`9`	`9`	`@Test`
`10`	`10`	`fun should_parse_variable_from_content() {`
`11`		`- CustomVariable.hasVariable("解释一下代码：\${selection}") shouldBe true`
`12`		`- CustomVariable.hasVariable("解释一下代码：\${selection") shouldBe false`
	`11`	`+ CustomVariable.hasVariable("解释一下代码：\$selection") shouldBe true`
	`12`	`+ CustomVariable.hasVariable("解释一下代码：\$selectio") shouldBe false`
`13`	`13`	`}`
`14`	`14`	`}`