feat(custom_agent): add state management for custom agent flow #51

The custom agent chat flow now includes state management, allowing for a more robust and interactive experience. The `CustomAgentState` enum has been added to track the state of the custom agent, which can be `START`, `HANDLING`, or `FINISHED`.

Author: phodal

example/custom_agent/server.py

@@ -22,7 +22,9 @@ class Messages(BaseModel):
 
 @app.post("/api/agent/api-market")
 def mock_market(messages: Messages):
-    return """GET /wp/v2/posts
+    return """
+```markdown
+GET /wp/v2/posts
 GET /wp/v2/posts/{id}
 POST /wp/v2/posts
 PUT /wp/v2/posts/{id}
@@ -32,12 +34,15 @@ def mock_market(messages: Messages):
 GET /wp/v2/pages/{id}
 POST /wp/v2/pages
 PUT /wp/v2/pages/{id}
-DELETE /wp/v2/pages/{id}"""
+DELETE /wp/v2/pages/{id}
+```
+"""
 
 
 @app.post("/api/agent/frontend-gen")
 def mock_frontend(messages: Messages):
-    return """Button:可快速创建不同样式的按钮。
+    return """```markdown
+Button:可快速创建不同样式的按钮。
 Checkbox:提供多选框组件，通常用于某选项的打开或关闭。
 CheckboxGroup:多选框群组，用于控制多选框全选或者不全选状态。
 CustomDialog:自定义弹窗（CustomDialog）可用于广告、中奖、警告、软件更新等与用户交互响应操作。
@@ -50,6 +55,7 @@ def mock_frontend(messages: Messages):
 TextInput:单行文本输入框组件。
 Radio:提供相应的用户交互选择项。
 Toggle:Toggle为开关组件，常用于在应用中进行开关操作。
+```
 """
 
 

src/main/kotlin/cc/unitmesh/devti/counit/CustomAgentChatProcessor.kt

@@ -1,24 +1,18 @@
 package cc.unitmesh.devti.counit
 
 import cc.unitmesh.devti.counit.model.CustomAgentConfig
+import cc.unitmesh.devti.counit.model.CustomAgentState
 import cc.unitmesh.devti.counit.model.ResponseAction
 import cc.unitmesh.devti.gui.chat.ChatCodingPanel
 import cc.unitmesh.devti.gui.chat.ChatContext
-import cc.unitmesh.devti.gui.chat.ChatRole
-import cc.unitmesh.devti.llms.LlmFactory
 import cc.unitmesh.devti.provider.ContextPrompter
-import cc.unitmesh.devti.util.LLMCoroutineScope
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.project.Project
-import kotlinx.coroutines.launch
 
 @Service(Service.Level.PROJECT)
 class CustomAgentChatProcessor(val project: Project) {
-    private val llmFactory = LlmFactory()
-
     private val customAgentExecutor = CustomAgentExecutor(project)
-    private val llmProvider = llmFactory.create(project)
 
     fun handleChat(prompter: ContextPrompter, ui: ChatCodingPanel, context: ChatContext?) {
         val originPrompt = prompter.requestPrompt()
@@ -27,33 +21,29 @@ class CustomAgentChatProcessor(val project: Project) {
         val request = originPrompt.trim()
         val selectedAgent: CustomAgentConfig = ui.getSelectedCustomAgent()
 
+        selectedAgent.state = CustomAgentState.HANDLING
+
         val response = customAgentExecutor.execute(request, selectedAgent)
         if (response == null) {
             logger.error("error for custom agent: $selectedAgent with request: $request")
             return
         }
 
+        selectedAgent.state = CustomAgentState.FINISHED
         when (selectedAgent.responseAction) {
             ResponseAction.Direct -> {
-                ui.addMessage(response, true, response)
+                ui.addMessage(response, false, response)
+                ui.hiddenProgressBar()
             }
 
             ResponseAction.TextChunk -> {
                 ui.setInput(response)
+                ui.hiddenProgressBar()
             }
 
             ResponseAction.Flow -> {
-                ui.addMessage(response, true, response)
-
-                // loading
-                LLMCoroutineScope.scope(project).launch {
-                    llmProvider.appendLocalMessage(response, ChatRole.User)
-
-                    val intentionFlow = llmProvider.stream(response, "")
-                    val result = ui.updateMessage(intentionFlow)
-
-                    llmProvider.appendLocalMessage(result, ChatRole.Assistant)
-                }
+                ui.addMessage(response, false, response)
+                ui.hiddenProgressBar()
             }
 
             ResponseAction.WebView -> TODO()

src/main/kotlin/cc/unitmesh/devti/counit/model/CustomAgentConfig.kt

@@ -78,7 +78,16 @@ data class CustomAgentConfig(
     val transition: List<CustomFlowTransition> = emptyList(),
     val interactive: InteractionType = InteractionType.ChatPanel,
     val auth: CustomAgentAuth? = null
-)
+) {
+    var state: CustomAgentState = CustomAgentState.START
+}
+
+@Serializable
+enum class CustomAgentState {
+    START,
+    HANDLING,
+    FINISHED
+}
 
 @Serializable
 data class CustomAgentAuth(

src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingPanel.kt

@@ -244,4 +244,8 @@ class ChatCodingPanel(private val chatCodingService: ChatCodingService, val disp
     fun getSelectedCustomAgent(): CustomAgentConfig {
         return inputSection.getSelectedAgent()
     }
+
+    fun hiddenProgressBar() {
+        progressBar.isVisible = false
+    }
 }

src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingService.kt

@@ -5,6 +5,7 @@ import cc.unitmesh.devti.AutoDevBundle
 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.llms.LlmFactory
 import cc.unitmesh.devti.util.parser.PostCodeProcessor
 import cc.unitmesh.devti.provider.ContextPrompter
@@ -35,8 +36,10 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
         val displayPrompt = prompter.displayPrompt()
 
         if (project.customAgentSetting.enableCustomRag && ui.hasSelectedCustomAgent()) {
-            counitProcessor.handleChat(prompter, ui, context)
-            return
+            if (ui.getSelectedCustomAgent().state === CustomAgentState.START) {
+                counitProcessor.handleChat(prompter, ui, context)
+                return
+            }
         }
 
         ui.addMessage(requestPrompt, true, displayPrompt)