refactor(sketch): simplify message handling and improve markdown preview

- Replace mutable `text` with immutable `context` in `MarkdownPreviewHighlightSketch`. - Simplify message retrieval in `SketchToolWindow` by using `getAllMessages` from `ChatCodingService`. - Add `getAllMessages` method to `LLMProvider` and `CustomLLMProvider` for consistent message handling. - Recreate `llmProvider` in `ChatCodingService` to ensure history is cleared properly.
unit-mesh · Mar 10, 2025 · 39779b5 · 39779b5
1 parent 12ac068
commit 39779b5
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 11 deletions.
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
@@ -11,6 +11,7 @@ import cc.unitmesh.devti.gui.chat.message.ChatActionType
 import cc.unitmesh.devti.gui.chat.message.ChatContext
 import cc.unitmesh.devti.gui.chat.message.ChatRole
 import cc.unitmesh.devti.llms.LlmFactory
+import cc.unitmesh.devti.llms.custom.Message
 import cc.unitmesh.devti.provider.ContextPrompter
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runInEdt
@@ -21,7 +22,7 @@ import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.launch
 
 class ChatCodingService(var actionType: ChatActionType, val project: Project) {
-    private val llmProvider = LlmFactory.create(project)
+    private var llmProvider = LlmFactory.create(project)
     private val counitProcessor = project.service<CustomAgentChatProcessor>()
     private var currentJob: Job? = null
 
@@ -131,11 +132,16 @@ class ChatCodingService(var actionType: ChatActionType, val project: Project) {
     }
 
     fun clearSession() {
-        llmProvider.clearMessage()
+        // recreate session to make sure the history is cleared
+        llmProvider = LlmFactory.create(project)
     }
 
     fun request(systemPrompt: String, userPrompt: String, isFromSketch: Boolean = true): Flow<String> {
         /// is from sketch the first model should use Plan then use Act
         return llmProvider.stream(userPrompt, systemPrompt, keepHistory = true, isFromSketch)
     }
+
+    fun getAllMessages(): List<Message> {
+        return llmProvider.getAllMessages()
+    }
 }
diff --git a/core/src/main/kotlin/cc/unitmesh/devti/llms/LLMProvider.kt b/core/src/main/kotlin/cc/unitmesh/devti/llms/LLMProvider.kt
@@ -1,6 +1,7 @@
 package cc.unitmesh.devti.llms
 
 import cc.unitmesh.devti.gui.chat.message.ChatRole
+import cc.unitmesh.devti.llms.custom.Message
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 
@@ -15,5 +16,7 @@ interface LLMProvider {
      */
     fun clearMessage() {}
 
+    fun getAllMessages() = emptyList<Message>()
+
     fun appendLocalMessage(msg: String, role: ChatRole) {}
 }
diff --git a/core/src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomLLMProvider.kt b/core/src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomLLMProvider.kt
@@ -34,6 +34,10 @@ class CustomLLMProvider(val project: Project, var llmConfig: LlmConfig = LlmConf
 
     override fun clearMessage() = messages.clear()
 
+    override fun getAllMessages(): List<Message> {
+        return messages
+    }
+
     override fun appendLocalMessage(msg: String, role: ChatRole) {
         if (msg.isEmpty()) return
         messages += Message(role.roleName(), msg)

diff --git a/core/src/main/kotlin/cc/unitmesh/devti/sketch/SketchToolWindow.kt b/core/src/main/kotlin/cc/unitmesh/devti/sketch/SketchToolWindow.kt
@@ -94,10 +94,7 @@ open class SketchToolWindow(
 
         addMouseListener(object : MouseAdapter() {
             override fun mouseClicked(e: MouseEvent?) {
-                var allText =
-                    historyPanel.components.filterIsInstance<LangSketch>().joinToString("\n") { it.getViewText() }
-                allText += myList.components.filterIsInstance<LangSketch>().joinToString("\n") { it.getViewText() }
-
+                var allText = chatCodingService.getAllMessages().joinToString("\n") { it.content }
                 val selection = StringSelection(allText)
                 val clipboard = Toolkit.getDefaultToolkit().systemClipboard
                 clipboard.setContents(selection, null)
@@ -308,8 +305,7 @@ open class SketchToolWindow(
                             ?.create(project, codeFence.text)
                     }
 
-                    val isCanHtml =
-                        codeFence.language == PlainTextLanguage.INSTANCE || codeFence.language.displayName.lowercase() == "markdown"
+                    val isCanHtml = codeFence.language.displayName.lowercase() == "markdown"
                     if (isCanHtml && codeFence.isComplete && blockViews[index] !is ExtensionLangSketch) {
                         langSketch = MarkdownPreviewHighlightSketch(project, codeFence.text)
                     }

diff --git a/core/src/main/kotlin/cc/unitmesh/devti/sketch/ui/MarkdownPreviewHighlightSketch.kt b/core/src/main/kotlin/cc/unitmesh/devti/sketch/ui/MarkdownPreviewHighlightSketch.kt
@@ -12,9 +12,11 @@ import java.awt.event.MouseEvent
 import javax.swing.JComponent
 import javax.swing.event.HyperlinkEvent
 
-class MarkdownPreviewHighlightSketch(val project: Project, var text: String) : ExtensionLangSketch {
+class MarkdownPreviewHighlightSketch(val project: Project, val text: String) : ExtensionLangSketch {
     override fun getExtensionName(): String = "Markdown Highlight"
-    override fun getViewText(): String = text
+    override fun getViewText(): String = context
+
+    private var context = text
 
     private val editorPane = MarkdownViewer.createBaseComponent().apply {
         addHyperlinkListener {
@@ -42,7 +44,7 @@ class MarkdownPreviewHighlightSketch(val project: Project, var text: String) : E
         editorPane.text = convertMarkdownToHtml(text)
         editorPane.invalidate()
         editorPane.repaint()
-        this.text = text
+        this.context = text
     }
 
     override fun getComponent(): JComponent {