From 9e55aeaad7f9f8f87eca71195f40fe143018c7d3 Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Wed, 16 Oct 2024 15:17:03 +0300
Subject: [PATCH] fix: OpenAI edit code request building (fixes #737)

---
 .../codegpt/actions/editor/EditCodeAction.kt           |  7 ++++++-
 .../codegpt/completions/CompletionRequestFactory.kt    |  2 +-
 .../completions/factory/OpenAIRequestFactory.kt        |  5 +++--
 src/main/resources/META-INF/plugin.xml                 | 10 ++++++----
 4 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeAction.kt b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeAction.kt
index acbfa3305..2d60a819a 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeAction.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/actions/editor/EditCodeAction.kt
@@ -5,11 +5,16 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import ee.carlrobert.codegpt.Icons
 import ee.carlrobert.codegpt.ui.EditCodePopover
+import javax.swing.Icon
 
-class EditCodeAction : BaseEditorAction(Icons.Sparkle) {
+open class EditCodeAction(icon: Icon) : BaseEditorAction(icon) {
     override fun actionPerformed(project: Project, editor: Editor, selectedText: String) {
         runInEdt {
             EditCodePopover(editor).show()
         }
     }
 }
+
+class EditCodeFloatingMenuAction : EditCodeAction(Icons.DefaultSmall)
+
+class EditCodeContextMenuAction : EditCodeAction(Icons.Sparkle)
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt
index 50fde9853..b7b012866 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt
@@ -31,7 +31,7 @@ interface CompletionRequestFactory {
 
 abstract class BaseRequestFactory : CompletionRequestFactory {
     override fun createEditCodeRequest(params: EditCodeRequestParameters): CompletionRequest {
-        val prompt = "${params.prompt}\n\n${params.selectedText}"
+        val prompt = "Code to modify:\n${params.selectedText}\n\nInstructions: ${params.prompt}"
         return createBasicCompletionRequest(EDIT_CODE_SYSTEM_PROMPT, prompt, 8192, true)
     }
 
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt
index d83d71a2b..3035c6837 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt
@@ -50,10 +50,11 @@ class OpenAIRequestFactory : CompletionRequestFactory {
 
     override fun createEditCodeRequest(params: EditCodeRequestParameters): OpenAIChatCompletionRequest {
         val model = service<OpenAISettings>().state.model
+        val prompt = "Code to modify:\n${params.selectedText}\n\nInstructions: ${params.prompt}"
         if (model == "o1-mini" || model == "o1-preview") {
-            return buildBasicO1Request(model, params.prompt, EDIT_CODE_SYSTEM_PROMPT)
+            return buildBasicO1Request(model, prompt, EDIT_CODE_SYSTEM_PROMPT)
         }
-        return createBasicCompletionRequest(EDIT_CODE_SYSTEM_PROMPT, params.prompt, model, true)
+        return createBasicCompletionRequest(EDIT_CODE_SYSTEM_PROMPT, prompt, model, true)
     }
 
     override fun createCommitMessageRequest(params: CommitMessageRequestParameters): OpenAIChatCompletionRequest {
diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml
index 697fe4260..71b03a3fe 100644
--- a/src/main/resources/META-INF/plugin.xml
+++ b/src/main/resources/META-INF/plugin.xml
@@ -105,9 +105,10 @@
                     <keyboard-shortcut keymap="$default" first-keystroke="ctrl shift I" replace-all="true"/>
                 </action>
                 <action
-                  id="CodeGPT.EditCode"
+                  id="CodeGPT.ContextMenuEditCodeAction"
                   text="Edit Code"
-                  class="ee.carlrobert.codegpt.actions.editor.EditCodeAction">
+                  description="Edit code in natural language"
+                  class="ee.carlrobert.codegpt.actions.editor.EditCodeContextMenuAction">
                     <keyboard-shortcut keymap="$default" first-keystroke="ctrl shift K" replace-all="true"/>
                 </action>
                 <action
@@ -156,9 +157,10 @@
 
         <group id="CodeGPT.FloatingCodeToolbarMenuRootGroup">
             <action
-              id="CodeGPT.ModifySelectionAction"
+              id="CodeGPT.FloatingMenuEditCodeAction"
               text="CodeGPT: Edit Code"
-              class="ee.carlrobert.codegpt.actions.editor.EditCodeAction">
+              description="Edit code in natural language"
+              class="ee.carlrobert.codegpt.actions.editor.EditCodeFloatingMenuAction">
                 <keyboard-shortcut keymap="$default" first-keystroke="ctrl shift K" replace-all="true"/>
             </action>
             <action