From 63e2ef2a2d7efa5c6c9c8e3ff15be69cb8497bd2 Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Thu, 12 Sep 2024 02:16:12 +0300
Subject: [PATCH] fix: highlighted text referencing in prompts

---
 .../toolwindow/chat/ChatToolWindowTabPanel.java   | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
index 1d07a37e..28a37e72 100644
--- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
+++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java
@@ -300,8 +300,19 @@ public class ChatToolWindowTabPanel implements Disposable {
     }
     promptBuilder.append(remainingText);
 
-    message.setUserMessage(promptBuilder.toString());
-    message.setPrompt(promptBuilder.toString());
+    String highlightedTextMd = "";
+    if (editor != null) {
+      var selectionModel = editor.getSelectionModel();
+      var selectedText = selectionModel.getSelectedText();
+      if (selectedText != null && !selectedText.isEmpty()) {
+        var fileExtension = FileUtil.getFileExtension(editor.getVirtualFile().getName());
+        highlightedTextMd = format("\n```%s\n%s\n```\n", fileExtension, selectedText);
+        selectionModel.removeSelection();
+      }
+    }
+
+    message.setUserMessage(highlightedTextMd + promptBuilder);
+    message.setPrompt(highlightedTextMd + promptBuilder);
 
     sendMessage(message, ConversationType.DEFAULT, processEditorSelection(editor, message));
     return Unit.INSTANCE;