diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java index 9c71c613..0a45984b 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatMessageResponseBody.java @@ -123,11 +123,12 @@ public class ChatMessageResponseBody extends JPanel { return; } - if (handleThinking(partialMessage)) { + var processedPartialMessage = processThinkingOutput(partialMessage); + if (processedPartialMessage.isEmpty()) { return; } - for (var item : streamParser.parse(partialMessage)) { + for (var item : streamParser.parse(processedPartialMessage)) { processResponse(item.response(), CODE.equals(item.type()), true); } } @@ -240,9 +241,8 @@ public class ChatMessageResponseBody extends JPanel { revalidate(); } - private boolean handleThinking(String partialMessage) { - thinkingOutputParser.processChunk(partialMessage); - + private String processThinkingOutput(String partialMessage) { + var processedChunk = thinkingOutputParser.processChunk(partialMessage); var thoughtProcessPanel = getExistingThoughtProcessPanel(); if (thinkingOutputParser.isThinking()) { @@ -254,14 +254,13 @@ public class ChatMessageResponseBody extends JPanel { } else { thoughtProcessPanel.updateText(thinkingOutputParser.getThoughtProcess()); } - return true; } - if (thoughtProcessPanel != null && !thoughtProcessPanel.getFinished()) { + if (thoughtProcessPanel != null && thinkingOutputParser.isFinished()) { thoughtProcessPanel.setFinished(); } - return false; + return processedChunk; } private ThoughtProcessPanel getExistingThoughtProcessPanel() { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParser.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParser.kt index e79e28bb..b01dca3b 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParser.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParser.kt @@ -1,39 +1,54 @@ package ee.carlrobert.codegpt.toolwindow.chat -import java.util.regex.Pattern - class ThinkingOutputParser { - private val buffer = StringBuilder() - - var isThinking: Boolean = false - private set - - var isFinished: Boolean = false - private set + companion object { + private const val OPEN_TAG = "" + private const val CLOSE_TAG = "" + } var thoughtProcess: String = "" private set + var isThinking: Boolean = false + private set + var isFinished: Boolean = false + private set - fun processChunk(chunk: String) { + private val buffer = StringBuilder() + + fun processChunk(chunk: String): String { if (isFinished) { - return + return chunk + } + + if (buffer.isEmpty() && chunk.isNotEmpty() && !OPEN_TAG.contains(chunk.take(OPEN_TAG.length))) { + isFinished = true + return chunk } buffer.append(chunk) + val current = buffer.toString() - val thinkPattern = Pattern.compile("(.*?)", Pattern.DOTALL) - val matcher = thinkPattern.matcher(buffer.toString()) - if (matcher.find()) { + + val indexOpen = current.indexOf(OPEN_TAG) + if (indexOpen == -1) { + return "" + } + + isThinking = true + val startContent = indexOpen + OPEN_TAG.length + + val indexClose = current.indexOf(CLOSE_TAG, startContent) + if (indexClose != -1) { + thoughtProcess = current.substring(startContent, indexClose).trim() isFinished = true isThinking = false - thoughtProcess = matcher.group(1).trim { it <= ' ' } - } else if (buffer.isNotBlank() && "".contains(buffer)) { - thoughtProcess = "" - isThinking = true - } else if (buffer.toString().startsWith("")) { - thoughtProcess = buffer.toString().replaceFirst("".toRegex(), "") - isThinking = true + + val responseStart = indexClose + CLOSE_TAG.length + return current.substring(responseStart) + } else { + thoughtProcess = current.substring(startContent) + return "" } } -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ResponseBodyProgressPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ResponseBodyProgressPanel.kt index 3fc7979e..ad5992da 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ResponseBodyProgressPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ResponseBodyProgressPanel.kt @@ -23,11 +23,14 @@ class ResponseBodyProgressPanel : JPanel() { init { layout = BoxLayout(this, BoxLayout.Y_AXIS) - border = JBUI.Borders.empty(4, 0, 8, 0) + border = JBUI.Borders.empty(4, 0) + isVisible = false } fun updateProgressContainer(text: String, icon: Icon?) { runInEdt { + isVisible = true + removeAll() val wrapper = if (icon != null) { JBLabel( diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ResponseMessagePanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ResponseMessagePanel.kt index 69a969b9..f2a8a683 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ResponseMessagePanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/ResponseMessagePanel.kt @@ -16,5 +16,8 @@ open class ResponseMessagePanel : BaseMessagePanel() { ) .setAllowAutoWrapping(true) .withFont(JBFont.label().asBold()) + .apply { + iconTextGap = 6 + } } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/UserMessagePanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/UserMessagePanel.kt index 4c971ffb..daf579f1 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/UserMessagePanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/ui/UserMessagePanel.kt @@ -51,7 +51,7 @@ class UserMessagePanel( Icons.User } else { val originalIcon = ImageIcon(Base64.getDecoder().decode(avatarBase64)) - val resizedImage = originalIcon.image.getScaledInstance(20, 20, Image.SCALE_SMOOTH) + val resizedImage = originalIcon.image.getScaledInstance(24, 24, Image.SCALE_SMOOTH) RoundedIcon(resizedImage, 1.0) } } catch (ex: Exception) { @@ -67,6 +67,9 @@ class UserMessagePanel( ) .setAllowAutoWrapping(true) .withFont(JBFont.label().asBold()) + .apply { + iconTextGap = 6 + } } fun addReloadAction(onReload: Runnable) { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/ThoughtProcessPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/ThoughtProcessPanel.kt index 562f3172..838ca566 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/ThoughtProcessPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/ThoughtProcessPanel.kt @@ -11,10 +11,10 @@ import javax.swing.* class ThoughtProcessPanel : JPanel(BorderLayout()) { - var finished: Boolean = false - private set - - private val responseBodyContent = UIUtil.createTextPane("", false) + private var finished: Boolean = false + private val responseBodyContent = UIUtil.createTextPane("", false).apply { + foreground = JBUI.CurrentTheme.Label.disabledForeground() + } private val contentPanel = createContentPanel() private val toggleButton: JToggleButton = createToggleButton() @@ -26,6 +26,8 @@ class ThoughtProcessPanel : JPanel(BorderLayout()) { } fun setFinished() { + if (finished) return + toggleButton.text = "Thought Process" toggleButton.isSelected = false finished = true @@ -58,21 +60,18 @@ class ThoughtProcessPanel : JPanel(BorderLayout()) { return panel } - private fun createToggleButton(): JToggleButton { - return JToggleButton("Thinking...", AllIcons.General.ArrowUp) - .apply { - isFocusPainted = false - isContentAreaFilled = false - background = background - selectedIcon = AllIcons.General.ArrowDown - border = null - isSelected = true - horizontalAlignment = SwingConstants.LEFT - horizontalTextPosition = SwingConstants.RIGHT - iconTextGap = 4 - addItemListener { e: ItemEvent -> - contentPanel.isVisible = e.stateChange == ItemEvent.SELECTED - } + private fun createToggleButton() = + JToggleButton("Thinking...", AllIcons.General.ArrowUp, true).apply { + isFocusPainted = false + isContentAreaFilled = false + background = background + selectedIcon = AllIcons.General.ArrowDown + border = null + horizontalAlignment = SwingConstants.LEFT + horizontalTextPosition = SwingConstants.RIGHT + iconTextGap = 4 + addItemListener { e: ItemEvent -> + contentPanel.isVisible = e.stateChange == ItemEvent.SELECTED } - } + } } \ No newline at end of file diff --git a/src/main/resources/icons/codegpt.svg b/src/main/resources/icons/codegpt.svg index 93f662a6..e8b05d08 100644 --- a/src/main/resources/icons/codegpt.svg +++ b/src/main/resources/icons/codegpt.svg @@ -1 +1 @@ - + diff --git a/src/main/resources/icons/codegpt_dark.svg b/src/main/resources/icons/codegpt_dark.svg index 804bf273..35194334 100644 --- a/src/main/resources/icons/codegpt_dark.svg +++ b/src/main/resources/icons/codegpt_dark.svg @@ -1 +1 @@ - + diff --git a/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParserTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParserTest.kt new file mode 100644 index 00000000..3c553a40 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParserTest.kt @@ -0,0 +1,54 @@ +package ee.carlrobert.codegpt.toolwindow.chat + +import org.assertj.core.api.Assertions.assertThat +import org.junit.Test + +class ThinkingOutputParserTest { + + @Test + fun `when not processing then return empty string`() { + val parser = ThinkingOutputParser() + assertThat(parser.processChunk("Some text")).isEmpty() + assertThat(parser.thoughtProcess).isEmpty() + } + + @Test + fun `when processing chunk but thinking not finished then return empty string`() { + val parser = ThinkingOutputParser() + assertThat(parser.processChunk("starting")).isEmpty() + assertThat(parser.thoughtProcess).isEqualTo("starting") + + val finalOutput = parser.processChunk(" some processing...") + + assertThat(finalOutput).isEmpty() + assertThat(parser.thoughtProcess).isEqualTo("starting some processing...") + } + + @Test + fun `when thinking finished then return everything after the last closing think tag`() { + val parser = ThinkingOutputParser() + parser.processChunk("the internal thought") + assertThat(parser.thoughtProcess).isEqualTo("the internal thought") + + val finalOutput = parser.processChunk("Here is the user response.") + + assertThat(finalOutput).isEqualTo("Here is the user response.") + assertThat(parser.thoughtProcess).isEqualTo("the internal thought") + } + + @Test + fun `accumulate chunks and return response only after final chunk with closing tag`() { + val parser = ThinkingOutputParser() + assertThat(parser.processChunk("")).isEmpty() + assertThat(parser.thoughtProcess).isEqualTo("") + assertThat(parser.processChunk("some internal processing")).isEmpty() + assertThat(parser.thoughtProcess).isEqualTo("some internal processing") + assertThat(parser.processChunk(" with even more details... ")).isEmpty() + assertThat(parser.thoughtProcess).isEqualTo("some internal processing with even more details... ") + + val finalOutput = parser.processChunk("The final answer.") + + assertThat(finalOutput).isEqualTo("The final answer.") + assertThat(parser.thoughtProcess).isEqualTo("some internal processing with even more details...") + } +} \ No newline at end of file