From 93a5e527df18b03e754ceb3b9b246c3671c2db52 Mon Sep 17 00:00:00 2001 From: Phodal Huang Date: Wed, 6 Mar 2024 21:53:01 +0800 Subject: [PATCH] fix(chat): handle custom agent state and add support for custom variable compilation #51 --- .../devti/counit/view/WebViewWindow.kt | 2 +- .../devti/custom/compile/CustomVariable.kt | 2 +- .../devti/gui/chat/AutoDevInputSection.kt | 7 +++++ .../devti/gui/chat/ChatCodingService.kt | 26 +++++++++++++++---- .../devti/provider/ContextPrompter.kt | 26 +++++++++++++++++++ .../custom/variable/CustomVariableTest.kt | 4 +-- 6 files changed, 58 insertions(+), 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 index 81c0640052..74ec406f3e 100644 --- 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 index 7c286f74d3..e7a2a41890 100644 --- 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 = 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 index 9fdd79d88d..89d85bf072 100644 --- 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 index fe45a55517..837a3594c6 100644 --- 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 index 6942aa7512..d09365258e 100644 --- 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() { 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 = 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 index 02fd3b0587..1ff3c16d4d 100644 --- 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 } } \ No newline at end of file