Merge pull request #208 from smallcloudai/fix/codelens-support-multip…

…le-messages

Feature: Support for multiple messages from CodeLens

src/main/kotlin/com/smallcloud/refactai/code_lens/CodeLensAction.kt

@@ -12,22 +12,54 @@ import com.smallcloud.refactai.Resources
 import com.smallcloud.refactai.panes.RefactAIToolboxPaneFactory
 import com.smallcloud.refactai.statistic.UsageStatistic
 import com.smallcloud.refactai.statistic.UsageStats
+import com.smallcloud.refactai.struct.ChatMessage
 import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.io.path.relativeTo
 
 class CodeLensAction(
     private val editor: Editor,
     private val line1: Int,
     private val line2: Int,
-    private val contentMsg: String,
+    private val messages: Array<ChatMessage>,
     private val sendImmediately: Boolean,
     private val openNewTab: Boolean
 ) : DumbAwareAction(Resources.Icons.LOGO_RED_16x16) {
     override fun actionPerformed(p0: AnActionEvent) {
         actionPerformed()
     }
 
-    private fun formatMessage(): String {
+    private fun replaceVariablesInText(
+        text: String,
+        relativePath: String,
+        cursor: Int?,
+        codeSelection: String
+    ): String {
+        return text
+            .replace("%CURRENT_FILE%", relativePath)
+            .replace("%CURSOR_LINE%", cursor?.plus(1)?.toString() ?: "")
+            .replace("%CODE_SELECTION%", codeSelection)
+            .replace("%PROMPT_EXPLORATION_TOOLS%", "")
+    }
+
+    private fun formatMultipleMessagesForCodeLens(
+        messages: Array<ChatMessage>,
+        relativePath: String,
+        cursor: Int?,
+        text: String
+    ): Array<ChatMessage> {
+        val formattedMessages = messages.map { message ->
+            if (message.role == "user") {
+                message.copy(
+                    content = replaceVariablesInText(message.content, relativePath, cursor, text)
+                )
+            } else {
+                message
+            }
+        }.toTypedArray()
+        return formattedMessages
+    }
+
+    private fun formatMessages(): Array<ChatMessage> {
         val pos1 = LogicalPosition(line1, 0)
         val text = editor.document.text.slice(
             editor.logicalPositionToOffset(pos1) until editor.document.getLineEndOffset(line2)
@@ -39,10 +71,9 @@ class CodeLensAction(
             }.minBy { it.toString().length }
         }
 
-        return contentMsg
-            .replace("%CURRENT_FILE%", relativePath?.toString() ?: filePath.toString())
-            .replace("%CURSOR_LINE%", line1.toString())
-            .replace("%CODE_SELECTION%", text)
+        val formattedMessages = formatMultipleMessagesForCodeLens(messages, relativePath?.toString() ?: filePath.toString(), line1, text);
+
+        return formattedMessages
     }
 
     private val isActionRunning = AtomicBoolean(false)
@@ -52,12 +83,12 @@ class CodeLensAction(
 
         chat?.activate {
             RefactAIToolboxPaneFactory.chat?.requestFocus()
-            RefactAIToolboxPaneFactory.chat?.executeCodeLensCommand(formatMessage(), sendImmediately, openNewTab)
+            RefactAIToolboxPaneFactory.chat?.executeCodeLensCommand(formatMessages(), sendImmediately, openNewTab)
             editor.project?.service<UsageStats>()?.addChatStatistic(true, UsageStatistic("openChatByCodelens"), "")
         }
 
         // If content is empty, then it's "Open Chat" instruction, selecting range of code in active tab
-        if (contentMsg.isEmpty() && isActionRunning.compareAndSet(false, true)) {
+        if (messages.isEmpty() && isActionRunning.compareAndSet(false, true)) {
             ApplicationManager.getApplication().invokeLater {
                 try {
                     val pos1 = LogicalPosition(line1, 0)

src/main/kotlin/com/smallcloud/refactai/code_lens/RefactCodeVisionProvider.kt

@@ -70,16 +70,19 @@ class RefactCodeVisionProvider(
                     val value = allCodeLenses.get(commandKey).asJsonObject
                     val msgs = value.asJsonObject.get("messages").asJsonArray.map {
                         gson.fromJson(it.asJsonObject, ChatMessage::class.java)
-                    }.toList()
-                    val msg = msgs.find { it.role == "user" }
+                    }.toTypedArray()
+                    val userMsg = msgs.find { it.role == "user" }
+
                     val sendImmediately = value.asJsonObject.get("auto_submit").asBoolean
                     val openNewTab = value.asJsonObject.get("new_tab")?.asBoolean ?: true
-                    if (msg != null || msgs.isEmpty()) {
+
+                    val isValidCodeLen = msgs.isEmpty() || userMsg != null
+                    if (isValidCodeLen) {
                         resCodeLenses.add(
                             CodeLen(
                                 range,
                                 value.asJsonObject.get("label").asString,
-                                CodeLensAction(editor, line1, line2, msg?.content ?: "", sendImmediately, openNewTab)
+                                CodeLensAction(editor, line1, line2, msgs, sendImmediately, openNewTab)
                             )
                         )
                     }

src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/ChatPanes.kt

@@ -3,6 +3,7 @@ package com.smallcloud.refactai.panes.sharedchat
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.invokeLater
 import com.intellij.openapi.project.Project
+import com.smallcloud.refactai.struct.ChatMessage
 import java.awt.BorderLayout
 import javax.swing.JComponent
 import javax.swing.JPanel
@@ -31,8 +32,8 @@ class ChatPanes(val project: Project) : Disposable {
         return holder
     }
 
-    fun executeCodeLensCommand(command: String, sendImmediately: Boolean, openNewTab: Boolean) {
-        pane?.executeCodeLensCommand(command, sendImmediately, openNewTab)
+    fun executeCodeLensCommand(messages: Array<ChatMessage>, sendImmediately: Boolean, openNewTab: Boolean) {
+        pane?.executeCodeLensCommand(messages, sendImmediately, openNewTab)
     }
 
     fun requestFocus() {

src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/Events.kt

@@ -7,6 +7,7 @@ import com.google.gson.JsonElement
 import com.google.gson.annotations.SerializedName
 import com.smallcloud.refactai.settings.Host
 import com.smallcloud.refactai.settings.HostDeserializer
+import com.smallcloud.refactai.struct.ChatMessage
 import java.io.Serializable
 import java.lang.reflect.Type
 
@@ -366,7 +367,21 @@ class Events {
     data class CodeLensCommandPayload(
         val value: String = "",
         @SerializedName("send_immediately") val sendImmediately: Boolean = false,
-    ) : Payload()
+        val messages: Array<ChatMessage>
+    ) : Payload() {
+        override fun equals(other: Any?): Boolean {
+            if (this === other) return true
+            if (javaClass != other?.javaClass) return false
+
+            other as CodeLensCommandPayload
+
+            return messages.contentEquals(other.messages)
+        }
+
+        override fun hashCode(): Int {
+            return messages.contentHashCode()
+        }
+    }
 
     class CodeLensCommand(payload: CodeLensCommandPayload) : ToChat<Payload>(EventNames.ToChat.CODE_LENS_EXEC, payload)
 

src/main/kotlin/com/smallcloud/refactai/panes/sharedchat/SharedChatPane.kt

@@ -40,6 +40,7 @@ import com.smallcloud.refactai.panes.sharedchat.Events.Editor
 import com.smallcloud.refactai.panes.sharedchat.browser.ChatWebView
 import com.smallcloud.refactai.settings.AppSettingsConfigurable
 import com.smallcloud.refactai.settings.Host
+import com.smallcloud.refactai.struct.ChatMessage
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
@@ -55,7 +56,7 @@ class SharedChatPane(val project: Project) : JPanel(), Disposable {
     private val editor = Editor(project)
     private var currentPage: String = ""
     private var isChatStreaming: Boolean = false
-    var id: String? = null;
+    var id: String? = null
     private val animatedFiles = mutableSetOf<String>()
     private val scheduler = AppExecutorUtil.createBoundedScheduledExecutorService("SMCRainbowScheduler", 2)
 
@@ -98,18 +99,18 @@ class SharedChatPane(val project: Project) : JPanel(), Disposable {
         }
     }
 
-    fun executeCodeLensCommand(command: String, sendImmediately: Boolean, openNewTab: Boolean) {
+    fun executeCodeLensCommand(messages: Array<ChatMessage>, sendImmediately: Boolean, openNewTab: Boolean) {
         if (isChatStreaming) return
         if (openNewTab || this.currentPage != "chat") {
             newChat()
         }
-        if (command.isEmpty()) {
+        if (messages.isEmpty()) {
             // Just opening a new chat, no codelens execution
             newChat()
             return
         }
         isChatStreaming = true
-        this.postMessage(Events.CodeLensCommand(Events.CodeLensCommandPayload(command, sendImmediately)))
+        this.postMessage(Events.CodeLensCommand(Events.CodeLensCommandPayload("", sendImmediately, messages)))
     }
 
     private fun sendUserConfig() {
@@ -152,7 +153,7 @@ class SharedChatPane(val project: Project) : JPanel(), Disposable {
                 val out = process.getResultStdoutStr().getOrNull()
                 if (out == null) {
                     println("Save btok file output is null")
-                    return;
+                    return
                 }
                 val fileName = out.lines().last()
 

src/main/kotlin/com/smallcloud/refactai/struct/ChatMessage.kt

@@ -4,5 +4,5 @@ import com.google.gson.annotations.SerializedName
 
 data class ChatMessage(val role: String,
                        val content: String,
-                       @SerializedName("tool_call_id") val toolCallId: String,
+                       @SerializedName("tool_call_id") val toolCallId: String?,
                        val usage: String?)