feat(chat): enhance chat handling with history and response tracking

phodal · phodal · commit 803e184d8388 · 2024-12-28T15:08:06.000+08:00
#251 - Modify `MessageView` to expose `message` property. - Update `CustomAgentChatProcessor` to return LLM response and handle chat history. - Add `getHistoryMessages` to `ChatCodingPanel` for retrieving chat history. - Enhance `ChatCodingService` to manage chat history and agent state transitions.
diff --git a/core/src/main/kotlin/cc/unitmesh/devti/agent/CustomAgentChatProcessor.kt b/core/src/main/kotlin/cc/unitmesh/devti/agent/CustomAgentChatProcessor.kt
@@ -2,30 +2,29 @@ package cc.unitmesh.devti.agent
 
 import cc.unitmesh.devti.AutoDevBundle
 import cc.unitmesh.devti.agent.model.CustomAgentConfig
-import cc.unitmesh.devti.agent.model.CustomAgentState
 import cc.unitmesh.devti.agent.model.CustomAgentResponseAction
+import cc.unitmesh.devti.agent.model.CustomAgentState
 import cc.unitmesh.devti.gui.chat.ChatCodingPanel
 import cc.unitmesh.devti.gui.chat.ChatRole
 import cc.unitmesh.devti.llms.LLMProvider
 import cc.unitmesh.devti.provider.ContextPrompter
-import cc.unitmesh.devti.provider.devins.LanguagePromptProcessor
 import cc.unitmesh.devti.provider.devins.CustomAgentContext
-import cc.unitmesh.devti.util.LLMCoroutineScope
+import cc.unitmesh.devti.provider.devins.LanguagePromptProcessor
 import cc.unitmesh.devti.util.parser.CodeFence
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.project.Project
-import kotlinx.coroutines.flow.*
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.runBlocking
+import java.util.concurrent.CompletableFuture
 
 @Service(Service.Level.PROJECT)
 class CustomAgentChatProcessor(val project: Project) {
     private val customAgentExecutor = project.service<CustomAgentExecutor>()
     private val logger = logger<CustomAgentChatProcessor>()
 
-    fun handleChat(prompter: ContextPrompter, ui: ChatCodingPanel, llmProvider: LLMProvider) {
+    fun handleChat(prompter: ContextPrompter, ui: ChatCodingPanel, llmProvider: LLMProvider): String? {
         val originPrompt = prompter.requestPrompt()
         ui.addMessage(originPrompt, true, originPrompt)
 
@@ -37,9 +36,10 @@ class CustomAgentChatProcessor(val project: Project) {
         val response: Flow<String>? = customAgentExecutor.execute(request, selectedAgent)
         if (response == null) {
             logger.error("error for custom agent: $selectedAgent with request: $request")
-            return
+            return null
         }
 
+        var llmResponse = ""
         selectedAgent.state = CustomAgentState.FINISHED
 
         var devInCode: String? = ""
@@ -63,24 +63,32 @@ class CustomAgentChatProcessor(val project: Project) {
 
                 ui.hiddenProgressBar()
                 ui.updateUI()
+
+                llmResponse = content
             }
 
             CustomAgentResponseAction.Stream -> {
                 ui.addMessage(AutoDevBundle.message("autodev.loading"))
-                var msg = ""
-                LLMCoroutineScope.scope(project).launch {
-                    msg = ui.updateMessage(response)
+                val future: CompletableFuture<String> = CompletableFuture()
+                val sb = StringBuilder()
+                runBlocking {
+                    val result = ui.updateMessage(response)
+                    sb.append(result)
                 }
 
-                llmProvider.appendLocalMessage(msg, ChatRole.Assistant)
+                val content = sb.toString()
+                llmProvider.appendLocalMessage(content, ChatRole.Assistant)
                 ui.hiddenProgressBar()
                 ui.updateUI()
                 ui.moveCursorToStart()
 
-                val code = CodeFence.parse(msg)
+                val code = CodeFence.parse(content)
                 if (code.language.displayName == "DevIn") {
                     devInCode = code.text
                 }
+
+                future.complete(content)
+                llmResponse = future.get()
             }
 
             CustomAgentResponseAction.TextChunk -> {
@@ -97,6 +105,8 @@ class CustomAgentChatProcessor(val project: Project) {
                 ui.moveCursorToStart()
                 ui.setInput(content)
                 ui.hiddenProgressBar()
+
+                llmResponse = content
             }
 
             CustomAgentResponseAction.Flow -> {
@@ -116,6 +126,8 @@ class CustomAgentChatProcessor(val project: Project) {
 
                 ui.appendWebView(content, project)
                 ui.hiddenProgressBar()
+
+                llmResponse = content
             }
 
             CustomAgentResponseAction.DevIns -> {
@@ -129,13 +141,17 @@ class CustomAgentChatProcessor(val project: Project) {
                 ui.updateUI()
 
                 devInCode = msg
+
+                llmResponse = msg
             }
         }
 
         if (!devInCode.isNullOrEmpty()) {
             LanguagePromptProcessor.instance("DevIn").forEach {
-                it.execute(project, CustomAgentContext(selectedAgent, devInCode))
+                it.execute(project, CustomAgentContext(selectedAgent, devInCode!!))
             }
         }
+
+        return llmResponse
     }
 }
diff --git a/core/src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingPanel.kt b/core/src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingPanel.kt
@@ -1,5 +1,6 @@
 package cc.unitmesh.devti.gui.chat
 
+import cc.unitmesh.cf.core.llms.LlmMsg
 import cc.unitmesh.devti.*
 import cc.unitmesh.devti.agent.model.CustomAgentConfig
 import cc.unitmesh.devti.agent.view.WebBlock
@@ -11,7 +12,6 @@ import cc.unitmesh.devti.settings.AutoDevSettingsState
 import cc.unitmesh.devti.settings.LanguageChangedCallback.componentStateChanged
 import com.intellij.lang.html.HTMLLanguage
 import com.intellij.openapi.Disposable
-import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.DialogPanel
 import com.intellij.openapi.ui.NullableComponent
@@ -169,6 +169,18 @@ class ChatCodingPanel(private val chatCodingService: ChatCodingService, val disp
         }
     }
 
+    fun getHistoryMessages(): List<LlmMsg.ChatMessage> {
+        val messages = mutableListOf<LlmMsg.ChatMessage>()
+        for (i in 0 until myList.componentCount) {
+            val component = myList.getComponent(i)
+            if (component is MessageView) {
+                val role = LlmMsg.ChatRole.valueOf(component.role.name)
+                messages.add(LlmMsg.ChatMessage(role, component.message, null))
+            }
+        }
+        return messages
+    }
+
     suspend fun updateMessage(content: Flow<String>): String {
         if (myList.componentCount > 0) {
             myList.remove(myList.componentCount - 1)
@@ -215,7 +227,7 @@ class ChatCodingPanel(private val chatCodingService: ChatCodingService, val disp
         postAction(text)
     }
 
-    private suspend fun updateMessageInUi(content: Flow<String>): String {
+    suspend fun updateMessageInUi(content: Flow<String>): String {
         val messageView = MessageView("", ChatRole.Assistant, "")
         myList.add(messageView)
         val startTime = System.currentTimeMillis()
diff --git a/core/src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingService.kt b/core/src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingService.kt
@@ -10,6 +10,7 @@ import cc.unitmesh.devti.custom.compile.CustomVariable
 import cc.unitmesh.devti.llms.LlmFactory
 import cc.unitmesh.devti.provider.ContextPrompter
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import kotlinx.coroutines.Job
@@ -27,12 +28,15 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
         currentJob?.cancel()
     }
 
+    private var isLastAgent: Boolean = false
+
     fun handlePromptAndResponse(
         ui: ChatCodingPanel,
         prompter: ContextPrompter,
         context: ChatContext? = null,
         keepHistory: Boolean
     ) {
+        var chatHistory: List<LlmMsg.ChatMessage> = emptyList()
         currentJob?.cancel()
         var requestPrompt = prompter.requestPrompt()
         var displayPrompt = prompter.displayPrompt()
@@ -41,9 +45,18 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
             val selectedCustomAgent = ui.getSelectedCustomAgent()
             when {
                 selectedCustomAgent.state === CustomAgentState.START -> {
-                    counitProcessor.handleChat(prompter, ui, llmProvider)
-                    ui.resetAgent()
-                    ui.moveCursorToStart()
+                    isLastAgent = true
+                    val response = ApplicationManager.getApplication().executeOnPooledThread<String?> {
+                        val output = counitProcessor.handleChat(prompter, ui, llmProvider)
+
+                        runInEdt {
+                            ui.resetAgent()
+                            ui.moveCursorToStart()
+                        }
+
+                        output
+                    }
+
                     return
                 }
 
@@ -59,13 +72,16 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
             }
         }
 
+        if (isLastAgent) {
+            isLastAgent = false
+            chatHistory = ui.getHistoryMessages()
+        }
 
         ui.addMessage(requestPrompt, true, displayPrompt)
         ui.addMessage(AutoDevBundle.message("autodev.loading"))
 
-
         ApplicationManager.getApplication().executeOnPooledThread {
-            val response = this.makeChatBotRequest(requestPrompt, keepHistory)
+            val response = this.makeChatBotRequest(requestPrompt, keepHistory, chatHistory)
             currentJob = LLMCoroutineScope.scope(project).launch {
                 when {
                     actionType === ChatActionType.REFACTOR -> ui.updateReplaceableContent(response) {
@@ -97,7 +113,17 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
         }
     }
 
-    private fun makeChatBotRequest(requestPrompt: String, newChatContext: Boolean): Flow<String> {
+    private fun makeChatBotRequest(
+        requestPrompt: String,
+        newChatContext: Boolean,
+        chatHistory: List<LlmMsg.ChatMessage>
+    ): Flow<String> {
+        if (chatHistory.isNotEmpty()) {
+            chatHistory.forEach {
+                llmProvider.appendLocalMessage(it.content, ChatRole.valueOf(it.role.name))
+            }
+        }
+
         return llmProvider.stream(requestPrompt, "", keepHistory = !newChatContext)
     }
 
diff --git a/core/src/main/kotlin/cc/unitmesh/devti/gui/chat/MessageView.kt b/core/src/main/kotlin/cc/unitmesh/devti/gui/chat/MessageView.kt
@@ -20,7 +20,7 @@ import java.awt.*
 import javax.swing.*
 import kotlin.jvm.internal.Ref
 
-class MessageView(private val message: String, val role: ChatRole, private val displayText: String) :
+class MessageView(val message: String, val role: ChatRole, private val displayText: String) :
     JBPanel<MessageView>(), DataProvider {
     private val myNameLabel: Component
     private val component: DisplayComponent = DisplayComponent(message)
