fix: undo action after completion

This commit is contained in:
Carl-Robert Linnupuu 2024-06-30 16:45:15 +03:00
parent 0e8d6a187d
commit d41bb317e5
5 changed files with 58 additions and 56 deletions

View file

@ -71,14 +71,20 @@ import org.jetbrains.annotations.Nullable;
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 class CompletionRequestProvider {
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 class CompletionRequestProvider {
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 class CompletionRequestProvider {
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();
}

View file

@ -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)
}
replacedLength += message.length
editor.caretModel.moveToOffset(startOffset + replacedLength)
updateHighlighter(editor)
}
private fun cleanupAndFormat() {
val project = editor.project ?: return
val document = editor.document
val psiDocumentManager = project.service<PsiDocumentManager>()
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<CodeStyleManager>().reformatText(
psiFile,
listOf(TextRange(startOffset, newEndOffset))
)
}
editor.caretModel.moveToOffset(newEndOffset)
psiDocumentManager.doPostponedOperationsAndUnblockDocument(document)
cleanupHighlighter()
}
private fun cleanupHighlighter() {
currentHighlighter?.let { editor.markupModel.removeHighlighter(it) }
currentHighlighter = null
}
private fun cleanupAndFormat() {
val project = editor.project ?: return
runWriteCommandAction(project) {
val document = editor.document
val psiDocumentManager = project.service<PsiDocumentManager>()
val psiFile = psiDocumentManager.getPsiFile(document)
?: return@runWriteCommandAction
val startOffset = selectionTextRange.startOffset
val endOffset = selectionTextRange.endOffset
val newEndOffset = startOffset + replacedLength
if (newEndOffset < endOffset) {
document.deleteString(newEndOffset, endOffset)
}
psiDocumentManager.commitDocument(document)
project.service<CodeStyleManager>().reformatText(
psiFile,
listOf(TextRange(startOffset, newEndOffset))
)
editor.caretModel.moveToOffset(newEndOffset)
psiDocumentManager.doPostponedOperationsAndUnblockDocument(document)
cleanupHighlighter()
}
}
}

View file

@ -41,7 +41,7 @@ class EditCodeSubmissionHandler(
"$userPrompt\n\n$selectedText",
service<CodeGPTServiceSettings>().state.chatCompletionSettings.model
),
EditCodeCompletionHandler(editor, observableProperties, selectionTextRange)
EditCodeCompletionListener(editor, observableProperties, selectionTextRange)
)
} finally {
observableProperties.loading.set(false)