From 668c8ed27ef389bc6b910f180fe157007385ce79 Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Mon, 3 Feb 2025 16:47:22 +0000 Subject: [PATCH] feat: support code editing and commit message generation for Deepseek R1 model (closes #835) --- .../codegpt/actions/BaseCommitWorkflowAction.kt | 16 +++++++++++----- .../actions/editor/EditCodeCompletionListener.kt | 7 ++++++- .../toolwindow/chat/ThinkingOutputParser.kt | 3 +-- .../toolwindow/chat/ThinkingOutputParserTest.kt | 6 +++--- 4 files changed, 21 insertions(+), 11 deletions(-) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/BaseCommitWorkflowAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/BaseCommitWorkflowAction.kt index a8ce0e30..a3ca7e98 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/BaseCommitWorkflowAction.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/BaseCommitWorkflowAction.kt @@ -19,6 +19,7 @@ import com.intellij.openapi.vcs.changes.Change import com.intellij.vcs.commit.CommitWorkflowUi import ee.carlrobert.codegpt.EncodingManager import ee.carlrobert.codegpt.completions.CompletionRequestService +import ee.carlrobert.codegpt.toolwindow.chat.ThinkingOutputParser import ee.carlrobert.codegpt.ui.OverlayUtil import ee.carlrobert.codegpt.util.CommitWorkflowChanges import ee.carlrobert.codegpt.util.GitUtil.getProjectRepository @@ -123,12 +124,17 @@ abstract class BaseCommitWorkflowAction : DumbAwareAction() { class CommitMessageEventListener( private val project: Project, private val commitWorkflowUi: CommitWorkflowUi -) : CompletionEventListener { - private val messageBuilder = StringBuilder() +) : CompletionEventListener { - override fun onMessage(message: String?, eventSource: EventSource) { - messageBuilder.append(message) - updateCommitMessage(messageBuilder.toString()) + private val messageBuilder = StringBuilder() + private val thinkingOutputParser = ThinkingOutputParser() + + override fun onMessage(message: String, eventSource: EventSource) { + val processedChunk = thinkingOutputParser.processChunk(message) + if (processedChunk.isNotEmpty() && thinkingOutputParser.isFinished) { + messageBuilder.append(message) + updateCommitMessage(messageBuilder.toString()) + } } override fun onComplete(result: StringBuilder) { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt index a1f99582..94268e77 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt @@ -14,6 +14,7 @@ import com.intellij.openapi.util.text.StringUtil import com.intellij.psi.PsiDocumentManager import com.intellij.psi.codeStyle.CodeStyleManager import com.intellij.ui.JBColor +import ee.carlrobert.codegpt.toolwindow.chat.ThinkingOutputParser import ee.carlrobert.codegpt.ui.ObservableProperties import ee.carlrobert.codegpt.ui.OverlayUtil import ee.carlrobert.llm.client.openai.completion.ErrorDetails @@ -28,9 +29,13 @@ class EditCodeCompletionListener( private var replacedLength = 0 private var currentHighlighter: RangeHighlighter? = null + private val thinkingOutputParser = ThinkingOutputParser() override fun onMessage(message: String, eventSource: EventSource) { - runInEdt { handleDiff(message) } + val processedChunk = thinkingOutputParser.processChunk(message) + if (processedChunk.isNotEmpty() && thinkingOutputParser.isFinished) { + runInEdt { handleDiff(processedChunk) } + } } override fun onComplete(messageBuilder: StringBuilder) { 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 b01dca3b..d2fd0fd5 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParser.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParser.kt @@ -4,7 +4,7 @@ class ThinkingOutputParser { companion object { private const val OPEN_TAG = "" - private const val CLOSE_TAG = "" + private const val CLOSE_TAG = "\n\n" } var thoughtProcess: String = "" @@ -29,7 +29,6 @@ class ThinkingOutputParser { buffer.append(chunk) val current = buffer.toString() - val indexOpen = current.indexOf(OPEN_TAG) if (indexOpen == -1) { return "" diff --git a/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParserTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParserTest.kt index 3c553a40..82d97aff 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParserTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ThinkingOutputParserTest.kt @@ -8,7 +8,7 @@ class ThinkingOutputParserTest { @Test fun `when not processing then return empty string`() { val parser = ThinkingOutputParser() - assertThat(parser.processChunk("Some text")).isEmpty() + assertThat(parser.processChunk("Some text")).isEqualTo("Some text") assertThat(parser.thoughtProcess).isEmpty() } @@ -27,7 +27,7 @@ class ThinkingOutputParserTest { @Test fun `when thinking finished then return everything after the last closing think tag`() { val parser = ThinkingOutputParser() - parser.processChunk("the internal thought") + parser.processChunk("the internal thought\n\n") assertThat(parser.thoughtProcess).isEqualTo("the internal thought") val finalOutput = parser.processChunk("Here is the user response.") @@ -46,7 +46,7 @@ class ThinkingOutputParserTest { 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.") + val finalOutput = parser.processChunk("\n\nThe final answer.") assertThat(finalOutput).isEqualTo("The final answer.") assertThat(parser.thoughtProcess).isEqualTo("some internal processing with even more details...")