feat(custom-agent): refactor to use LlmProvider for chat processing #51

The chat service has been refactored to use a new LlmProvider class for creating and interacting with the language model. This change simplifies the code by removing the need for direct instantiation of the LlmFactory and instead relying on the provider to create and manage the language model instance. The provider also allows for easier integration of custom logic or configuration specific to the project. Additionally, the chat processor now uses a StringBuilder to accumulate the chat response when using the stream response action, and the response is appended to the local message history of the LlmProvider. This ensures that the chat history is correctly maintained and can be accessed for future reference.

Author: phodal

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

@@ -6,20 +6,21 @@ 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.LLMProvider
 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.flow.collect
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 
 @Service(Service.Level.PROJECT)
 class CustomAgentChatProcessor(val project: Project) {
     private val customAgentExecutor = CustomAgentExecutor(project)
 
-    fun handleChat(prompter: ContextPrompter, ui: ChatCodingPanel, context: ChatContext?) {
+    fun handleChat(prompter: ContextPrompter, ui: ChatCodingPanel, context: ChatContext?, llmProvider: LLMProvider) {
         val originPrompt = prompter.requestPrompt()
         ui.addMessage(originPrompt, true, originPrompt)
 
@@ -38,9 +39,13 @@ class CustomAgentChatProcessor(val project: Project) {
         when (selectedAgent.responseAction) {
             ResponseAction.Direct -> {
                 val message = ui.addMessage("loading", false, "")
+                val sb = StringBuilder()
                 runBlocking {
-                    ui.updateMessage(response)
+                    val result = ui.updateMessage(response)
+                    sb.append(result)
                 }
+
+                llmProvider.appendLocalMessage(sb.toString(), ChatRole.Assistant)
                 message.reRenderAssistantOutput()
                 ui.hiddenProgressBar()
                 ui.updateUI()
@@ -49,7 +54,7 @@ class CustomAgentChatProcessor(val project: Project) {
             ResponseAction.Stream -> {
                 ui.addMessage(AutoDevBundle.message("autodev.loading"))
                 LLMCoroutineScope.scope(project).launch {
-//                    ui.updateMessage(response)
+                    ui.updateMessage(response)
                 }
             }
 
@@ -60,8 +65,11 @@ class CustomAgentChatProcessor(val project: Project) {
                         sb.append(it)
                     }
                 }
+
+                val content = sb.toString()
+                llmProvider.appendLocalMessage(content, ChatRole.Assistant)
                 ui.removeLastMessage()
-                ui.setInput(sb.toString())
+                ui.setInput(content)
                 ui.hiddenProgressBar()
             }
 
@@ -78,6 +86,7 @@ class CustomAgentChatProcessor(val project: Project) {
                 }
                 // TODO: add decode support
                 val content = sb.toString()
+                llmProvider.appendLocalMessage(content, ChatRole.Assistant)
 
                 ui.appendWebView(content, project)
                 ui.hiddenProgressBar()

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

@@ -16,15 +16,12 @@ import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.launch
 
 class ChatCodingService(var actionType: ChatActionType, val project: Project) {
-    private val llmFactory = LlmFactory()
+    private val llmProvider = LlmFactory().create(project)
     private val counitProcessor = project.service<CustomAgentChatProcessor>()
 
     val action = actionType.instruction(project = project)
 
-    fun getLabel(): String {
-        val capitalizedAction = actionType
-        return "$capitalizedAction Code"
-    }
+    fun getLabel(): String = "$actionType Code"
 
     fun handlePromptAndResponse(
         ui: ChatCodingPanel,
@@ -37,7 +34,7 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
 
         if (project.customAgentSetting.enableCustomRag && ui.hasSelectedCustomAgent()) {
             if (ui.getSelectedCustomAgent().state === CustomAgentState.START) {
-                counitProcessor.handleChat(prompter, ui, context)
+                counitProcessor.handleChat(prompter, ui, context, llmProvider)
                 return
             }
         }
@@ -74,7 +71,7 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
         ui.addMessage(AutoDevBundle.message("autodev.loading"))
 
         ApplicationManager.getApplication().executeOnPooledThread {
-            val response = llmFactory.create(project).stream(requestPrompt, systemPrompt)
+            val response = llmProvider.stream(requestPrompt, systemPrompt)
 
             LLMCoroutineScope.scope(project).launch {
                 ui.updateMessage(response)
@@ -83,7 +80,7 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
     }
 
     private fun makeChatBotRequest(requestPrompt: String, newChatContext: Boolean): Flow<String> {
-        return llmFactory.create(project).stream(requestPrompt, "", keepHistory = !newChatContext)
+        return llmProvider.stream(requestPrompt, "", keepHistory = !newChatContext)
     }
 
     private fun getCodeSection(content: String, prefixText: String, suffixText: String): String {
@@ -96,6 +93,6 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
     }
 
     fun clearSession() {
-        llmFactory.create(project).clearMessage()
+        llmProvider.clearMessage()
     }
 }