From fbf88bc75aa7c5b466e5cab64072ae6719228944 Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Sun, 30 Jun 2024 16:45:15 +0300 Subject: [PATCH] fix: undo action after completion --- .../CompletionRequestProvider.java | 32 +++++----- ...ndler.kt => EditCodeCompletionListener.kt} | 60 +++++++++---------- .../editor/EditCodeSubmissionHandler.kt | 2 +- ...stem-prompt.txt => default-completion.txt} | 0 ...prompt.txt => generate-commit-message.txt} | 0 5 files changed, 48 insertions(+), 46 deletions(-) rename src/main/kotlin/ee/carlrobert/codegpt/actions/editor/{EditCodeCompletionHandler.kt => EditCodeCompletionListener.kt} (74%) rename src/main/resources/prompts/{default-completion-system-prompt.txt => default-completion.txt} (100%) rename src/main/resources/prompts/{generate-commit-message-system-prompt.txt => generate-commit-message.txt} (100%) diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java index f670e689..1af6f0e3 100644 --- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java +++ b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java @@ -71,14 +71,20 @@ public class CompletionRequestProvider { - public static final String COMPLETION_SYSTEM_PROMPT = getResourceContent( - "/prompts/default-completion-system-prompt.txt"); + public static final String COMPLETION_SYSTEM_PROMPT = + getResourceContent("/prompts/default-completion.txt"); - public static final String GENERATE_COMMIT_MESSAGE_SYSTEM_PROMPT = getResourceContent( - "/prompts/generate-commit-message-system-prompt.txt"); + public static final String GENERATE_COMMIT_MESSAGE_SYSTEM_PROMPT = + getResourceContent("/prompts/generate-commit-message.txt"); - public static final String FIX_COMPILE_ERRORS_SYSTEM_PROMPT = getResourceContent( - "/prompts/fix-compile-errors.txt"); + public static final String FIX_COMPILE_ERRORS_SYSTEM_PROMPT = + getResourceContent("/prompts/fix-compile-errors.txt"); + + public static final String GENERATE_METHOD_NAMES_SYSTEM_PROMPT = + getResourceContent("/prompts/method-name-generator.txt"); + + public static final String EDIT_CODE_SYSTEM_PROMPT = + getResourceContent("/prompts/edit-code.txt"); private final EncodingManager encodingManager = EncodingManager.getInstance(); private final Conversation conversation; @@ -113,9 +119,7 @@ public static OpenAIChatCompletionRequest buildOpenAILookupCompletionRequest( String model) { return new OpenAIChatCompletionRequest.Builder( List.of( - new OpenAIChatCompletionStandardMessage( - "system", - getResourceContent("/prompts/method-name-generator.txt")), + new OpenAIChatCompletionStandardMessage("system", GENERATE_METHOD_NAMES_SYSTEM_PROMPT), new OpenAIChatCompletionStandardMessage("user", context))) .setModel(model) .setStream(false) @@ -127,9 +131,7 @@ public static OpenAIChatCompletionRequest buildEditCodeRequest( String model) { return new OpenAIChatCompletionRequest.Builder( List.of( - new OpenAIChatCompletionStandardMessage( - "system", - getResourceContent("/prompts/edit-code.txt")), + new OpenAIChatCompletionStandardMessage("system", EDIT_CODE_SYSTEM_PROMPT), new OpenAIChatCompletionStandardMessage("user", context))) .setModel(model) .setStream(true) @@ -169,14 +171,14 @@ public static Request buildCustomOpenAILookupCompletionRequest(String context) { List.of( new OpenAIChatCompletionStandardMessage( "system", - getResourceContent("/prompts/method-name-generator.txt")), + GENERATE_COMMIT_MESSAGE_SYSTEM_PROMPT), new OpenAIChatCompletionStandardMessage("user", context)), false); } public static LlamaCompletionRequest buildLlamaLookupCompletionRequest(String context) { - return new LlamaCompletionRequest.Builder(PromptTemplate.LLAMA - .buildPrompt(getResourceContent("/prompts/method-name-generator.txt"), context, List.of())) + return new LlamaCompletionRequest.Builder( + PromptTemplate.LLAMA.buildPrompt(GENERATE_COMMIT_MESSAGE_SYSTEM_PROMPT, context, List.of())) .setStream(false) .build(); } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt similarity index 74% rename from src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionHandler.kt rename to src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt index 50b12c34..5ba4f19d 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionHandler.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeCompletionListener.kt @@ -1,6 +1,8 @@ package ee.carlrobert.codegpt.actions.editor import com.intellij.notification.NotificationType +import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.application.runUndoTransparentWriteAction import com.intellij.openapi.command.WriteCommandAction.runWriteCommandAction import com.intellij.openapi.components.service import com.intellij.openapi.editor.Editor @@ -15,7 +17,7 @@ import ee.carlrobert.llm.client.openai.completion.ErrorDetails import ee.carlrobert.llm.completion.CompletionEventListener import okhttp3.sse.EventSource -class EditCodeCompletionHandler( +class EditCodeCompletionListener( private val editor: Editor, private val observableProperties: ObservableProperties, private val selectionTextRange: TextRange @@ -25,11 +27,11 @@ class EditCodeCompletionHandler( private var currentHighlighter: RangeHighlighter? = null override fun onMessage(message: String, eventSource: EventSource) { - handleDiff(message) + runInEdt { handleDiff(message) } } override fun onComplete(messageBuilder: StringBuilder) { - cleanupAndFormat() + runInEdt { cleanupAndFormat() } observableProperties.loading.set(false) } @@ -40,8 +42,8 @@ class EditCodeCompletionHandler( ) } - private fun highlightCurrentRow(editor: Editor) { - currentHighlighter?.let { editor.markupModel.removeHighlighter(it) } + private fun updateHighlighter(editor: Editor) { + cleanupHighlighter() val document = editor.document val lineNumber = document.getLineNumber(editor.caretModel.offset) @@ -59,13 +61,12 @@ class EditCodeCompletionHandler( ) } - private fun handleDiff(message: String) { - runWriteCommandAction(editor.project) { - val document = editor.document - val startOffset = selectionTextRange.startOffset - val endOffset = selectionTextRange.endOffset + val document = editor.document + val startOffset = selectionTextRange.startOffset + val endOffset = selectionTextRange.endOffset + runUndoTransparentWriteAction { val remainingOriginalLength = endOffset - (startOffset + replacedLength) if (remainingOriginalLength > 0) { document.replaceString( @@ -79,41 +80,40 @@ class EditCodeCompletionHandler( } else { document.insertString(startOffset + replacedLength, message) } - - replacedLength += message.length - editor.caretModel.moveToOffset(startOffset + replacedLength) - highlightCurrentRow(editor) } - } - private fun cleanupHighlighter() { - currentHighlighter?.let { editor.markupModel.removeHighlighter(it) } - currentHighlighter = null + replacedLength += message.length + editor.caretModel.moveToOffset(startOffset + replacedLength) + updateHighlighter(editor) } private fun cleanupAndFormat() { val project = editor.project ?: return - runWriteCommandAction(project) { - val document = editor.document - val psiDocumentManager = project.service() - val psiFile = psiDocumentManager.getPsiFile(document) - ?: return@runWriteCommandAction - val startOffset = selectionTextRange.startOffset - val endOffset = selectionTextRange.endOffset - val newEndOffset = startOffset + replacedLength + val document = editor.document + val psiDocumentManager = project.service() + val psiFile = psiDocumentManager.getPsiFile(document) ?: return + val startOffset = selectionTextRange.startOffset + val endOffset = selectionTextRange.endOffset + val newEndOffset = startOffset + replacedLength + runWriteCommandAction(project) { if (newEndOffset < endOffset) { document.deleteString(newEndOffset, endOffset) } - psiDocumentManager.commitDocument(document) project.service().reformatText( psiFile, listOf(TextRange(startOffset, newEndOffset)) ) - editor.caretModel.moveToOffset(newEndOffset) - psiDocumentManager.doPostponedOperationsAndUnblockDocument(document) - cleanupHighlighter() } + + editor.caretModel.moveToOffset(newEndOffset) + psiDocumentManager.doPostponedOperationsAndUnblockDocument(document) + cleanupHighlighter() + } + + private fun cleanupHighlighter() { + currentHighlighter?.let { editor.markupModel.removeHighlighter(it) } + currentHighlighter = null } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeSubmissionHandler.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeSubmissionHandler.kt index 1eaf4316..dda10942 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeSubmissionHandler.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeSubmissionHandler.kt @@ -41,7 +41,7 @@ class EditCodeSubmissionHandler( "$userPrompt\n\n$selectedText", service().state.chatCompletionSettings.model ), - EditCodeCompletionHandler(editor, observableProperties, selectionTextRange) + EditCodeCompletionListener(editor, observableProperties, selectionTextRange) ) } finally { observableProperties.loading.set(false) diff --git a/src/main/resources/prompts/default-completion-system-prompt.txt b/src/main/resources/prompts/default-completion.txt similarity index 100% rename from src/main/resources/prompts/default-completion-system-prompt.txt rename to src/main/resources/prompts/default-completion.txt diff --git a/src/main/resources/prompts/generate-commit-message-system-prompt.txt b/src/main/resources/prompts/generate-commit-message.txt similarity index 100% rename from src/main/resources/prompts/generate-commit-message-system-prompt.txt rename to src/main/resources/prompts/generate-commit-message.txt