Merge pull request #106 from devoxx/issue-99

Fix #99 : Chat memory correctly cleaned when removing a msg + Notific…

build.gradle.kts

@@ -5,7 +5,7 @@ plugins {
 }
 
 group = "com.devoxx.genie"
-version = "0.1.17"
+version = "0.1.18"
 
 repositories {
     mavenCentral()

core/build/tmp/jar/MANIFEST.MF

@@ -0,0 +1,2 @@
+Manifest-Version: 1.0
+

src/main/java/com/devoxx/genie/service/ChatMemoryService.java

@@ -1,5 +1,6 @@
 package com.devoxx.genie.service;
 
+import com.devoxx.genie.model.request.ChatMessageContext;
 import com.devoxx.genie.ui.listener.ChatMemorySizeListener;
 import com.devoxx.genie.ui.settings.DevoxxGenieStateService;
 import com.devoxx.genie.ui.topic.AppTopics;
@@ -26,12 +27,19 @@ public void init() {
         createChangeListener();
     }
 
+    /**
+     * Create the change listener.
+     */
     private void createChangeListener() {
         ApplicationManager.getApplication().getMessageBus()
             .connect()
             .subscribe(AppTopics.CHAT_MEMORY_SIZE_TOPIC, this);
     }
 
+    /**
+     * Get the chat memory service instance.
+     * @return the chat memory service instance
+     */
     @NotNull
     public static ChatMemoryService getInstance() {
         return ApplicationManager.getApplication().getService(ChatMemoryService.class);
@@ -41,23 +49,68 @@ public void clear() {
         chatMemory.clear();
     }
 
+    /**
+     * Add the chat message to the chat memory.
+     * @param chatMessage the chat message
+     */
     public void add(ChatMessage chatMessage) {
         chatMemory.add(chatMessage);
     }
 
+    /**
+     * Remove the chat message from the chat memory.
+     * @param chatMessageContext the chat message context
+     */
+    public void remove(@NotNull ChatMessageContext chatMessageContext) {
+        List<ChatMessage> messages = chatMemory.messages();
+        messages.remove(chatMessageContext.getAiMessage());
+        messages.remove(chatMessageContext.getUserMessage());
+        chatMemory.clear();
+        messages.forEach(this::add);
+    }
+
+    /**
+     * Remove the last message from the chat memory.
+     * This is used when an exception occurs and the last message is not valid.
+     */
+    public void removeLast() {
+        List<ChatMessage> messages = chatMemory.messages();
+        if (!messages.isEmpty()) {
+            messages.remove(messages.size() - 1);
+            chatMemory.clear();
+            messages.forEach(this::add);
+        }
+    }
+
+    /**
+     * Get the messages from the chat memory.
+     * @return the list of chat messages
+     */
     public List<ChatMessage> messages() {
         return chatMemory.messages();
     }
 
+    /**
+     * Check if the chat memory is empty.
+     * @return true if the chat memory is empty
+     */
     public boolean isEmpty() {
         return chatMemory.messages().isEmpty();
     }
 
+    /**
+     * On chat memory size changed.
+     * @param chatMemorySize the chat memory size
+     */
     @Override
     public void onChatMemorySizeChanged(int chatMemorySize) {
         createChatMemory(chatMemorySize);
     }
 
+    /**
+     * Create the chat memory.
+     * @param chatMemorySize the chat memory size
+     */
     private void createChatMemory(int chatMemorySize) {
         chatMemory = MessageWindowChatMemory.builder()
             .id("devoxxgenie")

src/main/java/com/devoxx/genie/service/ChatPromptExecutor.java

@@ -1,6 +1,7 @@
 package com.devoxx.genie.service;
 
 import com.devoxx.genie.model.request.ChatMessageContext;
+import com.devoxx.genie.service.exception.ProviderUnavailableException;
 import com.devoxx.genie.ui.panel.PromptOutputPanel;
 import com.devoxx.genie.ui.settings.DevoxxGenieStateService;
 import com.devoxx.genie.ui.util.NotificationUtil;
@@ -11,6 +12,8 @@
 import dev.langchain4j.model.chat.StreamingChatLanguageModel;
 import org.jetbrains.annotations.NotNull;
 
+import javax.swing.*;
+import java.net.ConnectException;
 import java.util.Optional;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.TimeoutException;
@@ -163,10 +166,16 @@ private void runPrompt(@NotNull ChatMessageContext chatMessageContext,
                     return null;
                 }
                 if (e.getCause() instanceof TimeoutException) {
-                    promptOutputPanel.addWarningText(chatMessageContext, "Timeout occurred. Please try again.");
+                    NotificationUtil.sendNotification(chatMessageContext.getProject(),
+                        "Timeout occurred. Please increase the timeout setting.");
+                    return null;
+                } else if (e.getCause() instanceof ProviderUnavailableException) {
+                    NotificationUtil.sendNotification(chatMessageContext.getProject(),
+                        "LLM provider not available. Please select another provider or make sure it's running.");
                     return null;
                 }
-                promptOutputPanel.addWarningText(chatMessageContext, e.getMessage());
+                String message = e.getMessage() + ". Maybe create an issue on GitHub?";
+                NotificationUtil.sendNotification(chatMessageContext.getProject(), "Error occurred: " + message);
                 return null;
             });
 

src/main/java/com/devoxx/genie/service/MessageCreationService.java

@@ -42,7 +42,7 @@ public static MessageCreationService getInstance() {
         if ((selectedText != null && !selectedText.isEmpty()) || (context != null && !context.isEmpty())) {
             userMessage = constructUserMessage(chatMessageContext, context);
         } else {
-            userMessage = new UserMessage(QUESTION + " " + chatMessageContext.getUserPrompt());
+            userMessage = new UserMessage(chatMessageContext.getUserPrompt());
         }
         return userMessage;
     }

src/main/java/com/devoxx/genie/service/PromptExecutionService.java

@@ -2,6 +2,7 @@
 
 import com.devoxx.genie.model.Constant;
 import com.devoxx.genie.model.request.ChatMessageContext;
+import com.devoxx.genie.service.exception.ProviderUnavailableException;
 import com.devoxx.genie.ui.settings.DevoxxGenieStateService;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.diagnostic.Logger;
@@ -96,7 +97,8 @@ private boolean isCanceled() {
             ChatMemoryService.getInstance().add(response.content());
             return Optional.of(response.content());
         } catch (Exception e) {
-            throw new CompletionException(e);
+            ChatMemoryService.getInstance().removeLast();
+            throw new ProviderUnavailableException(e.getMessage());
         }
     }
 }

src/main/java/com/devoxx/genie/service/StreamingResponseHandler.java

@@ -1,14 +1,18 @@
 package com.devoxx.genie.service;
 
 import com.devoxx.genie.model.request.ChatMessageContext;
+import com.devoxx.genie.service.exception.ProviderUnavailableException;
 import com.devoxx.genie.ui.component.ExpandablePanel;
 import com.devoxx.genie.ui.panel.ChatStreamingResponsePanel;
 import com.devoxx.genie.ui.panel.PromptOutputPanel;
+import com.devoxx.genie.ui.util.NotificationUtil;
 import dev.langchain4j.data.message.AiMessage;
 import dev.langchain4j.model.output.Response;
 import org.jetbrains.annotations.NotNull;
 
 import javax.swing.*;
+import java.net.ConnectException;
+import java.util.concurrent.TimeoutException;
 
 public class StreamingResponseHandler implements dev.langchain4j.model.StreamingResponseHandler<AiMessage> {
 
@@ -34,8 +38,8 @@ public void onNext(String token) {
 
     @Override
     public void onComplete(@NotNull Response<AiMessage> response) {
-        AiMessage content = response.content();
-        ChatMemoryService.getInstance().add(content);
+        chatMessageContext.setAiMessage(response.content());
+        ChatMemoryService.getInstance().add(response.content());
         enableButtons.run();
         if (chatMessageContext.hasFiles()) {
             SwingUtilities.invokeLater(() -> {
@@ -48,8 +52,17 @@ public void onComplete(@NotNull Response<AiMessage> response) {
 
     @Override
     public void onError(Throwable error) {
-        // TODO Show error message
         enableButtons.run();
+        if (error.getCause() instanceof TimeoutException) {
+            NotificationUtil.sendNotification(chatMessageContext.getProject(),
+                "Timeout occurred. Please increase the timeout setting.");
+        } else if (error.getCause() instanceof ConnectException) {
+            NotificationUtil.sendNotification(chatMessageContext.getProject(),
+                "LLM provider not available. Please select another provider or make sure it's running.");
+        } else {
+            NotificationUtil.sendNotification(chatMessageContext.getProject(),
+                "An error occurred. Please try again.");
+        }
     }
 }
 

src/main/java/com/devoxx/genie/service/exception/ProviderUnavailableException.java

@@ -0,0 +1,8 @@
+package com.devoxx.genie.service.exception;
+
+public class ProviderUnavailableException extends RuntimeException {
+
+    public ProviderUnavailableException(String message) {
+        super(message);
+    }
+}

src/main/java/com/devoxx/genie/ui/component/FileEntryComponent.java

@@ -1,7 +1,6 @@
 package com.devoxx.genie.ui.component;
 
 import com.devoxx.genie.ui.listener.FileRemoveListener;
-import com.devoxx.genie.ui.util.FileTypeIconUtil;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.fileEditor.FileEditorManager;
 import com.intellij.openapi.project.Project;

src/main/java/com/devoxx/genie/ui/panel/UserPromptPanel.java

@@ -1,22 +1,19 @@
 package com.devoxx.genie.ui.panel;
 
 import com.devoxx.genie.model.request.ChatMessageContext;
+import com.devoxx.genie.service.ChatMemoryService;
 import com.devoxx.genie.ui.component.JEditorPaneUtils;
 import com.devoxx.genie.ui.component.JHoverButton;
 import com.devoxx.genie.ui.component.StyleSheetsFactory;
-import com.devoxx.genie.ui.listener.ChatMessageManagementService;
-import com.devoxx.genie.ui.topic.AppTopics;
-import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.ui.components.JBLabel;
-import com.intellij.util.messages.MessageBus;
 import org.jetbrains.annotations.NotNull;
 
 import javax.swing.*;
 import java.awt.*;
-import java.util.Arrays;
 
 import static com.devoxx.genie.ui.util.DevoxxGenieIcons.DevoxxIcon;
 import static com.devoxx.genie.ui.util.DevoxxGenieIcons.TrashIcon;
+import static java.util.Arrays.*;
 
 public class UserPromptPanel extends BackgroundPanel {
 
@@ -36,7 +33,7 @@ public UserPromptPanel(JPanel container,
 
         JPanel headerPanel = new JPanel(new BorderLayout());
         headerPanel.setOpaque(false);
-        headerPanel.add(createHeaderLabel(chatMessageContext), BorderLayout.WEST);
+        headerPanel.add(createHeaderLabel(), BorderLayout.WEST);
         headerPanel.add(createDeleteButton(chatMessageContext), BorderLayout.EAST);
 
         // User prompt setup
@@ -49,10 +46,8 @@ public UserPromptPanel(JPanel container,
 
     /**
      * Create the header label.
-     *
-     * @param chatMessageContext the chat message context
      */
-    private @NotNull JBLabel createHeaderLabel(@NotNull ChatMessageContext chatMessageContext) {
+    private @NotNull JBLabel createHeaderLabel() {
         JBLabel createdOnLabel = new JBLabel("DevoxxGenie", DevoxxIcon, SwingConstants.LEFT);
         createdOnLabel.setFont(createdOnLabel.getFont().deriveFont(12f));
         createdOnLabel.setBorder(BorderFactory.createEmptyBorder(0, 5, 0, 10));
@@ -80,28 +75,16 @@ public UserPromptPanel(JPanel container,
     private void removeChat(ChatMessageContext chatMessageContext) {
 
         // Get all container components and delete by name
-        Arrays.stream(container.getComponents())
+        stream(container.getComponents())
             .filter(c -> c.getName() != null && c.getName().equals(chatMessageContext.getName()))
             .forEach(container::remove);
 
         // Repaint the container
         container.revalidate();
         container.repaint();
 
-        // Broadcast that the chat message has been removed, this way the chat memory can be updated
-        notifyChatMessageRemoval(chatMessageContext);
-    }
-
-    /**
-     * Notify the chat message removal.
-     *
-     * @param chatMessageContext the chat message context
-     */
-    private void notifyChatMessageRemoval(ChatMessageContext chatMessageContext) {
-        // Trigger the chat message change listener
-        MessageBus bus = ApplicationManager.getApplication().getMessageBus();
-        ChatMessageManagementService chatChangeListener = bus.syncPublisher(AppTopics.CHAT_MESSAGES_CHANGED_TOPIC);
-        chatChangeListener.removeMessagePair(chatMessageContext);
+        // Remove the chat from memory
+        ChatMemoryService.getInstance().remove(chatMessageContext);
     }
 }
 

src/main/java/com/devoxx/genie/ui/topic/AppTopics.java

@@ -1,7 +1,6 @@
 package com.devoxx.genie.ui.topic;
 
 import com.devoxx.genie.ui.listener.ChatMemorySizeListener;
-import com.devoxx.genie.ui.listener.ChatMessageManagementService;
 import com.devoxx.genie.ui.listener.SettingsChangeListener;
 import com.intellij.util.messages.Topic;
 
@@ -10,9 +9,6 @@ public class AppTopics {
     public static final Topic<SettingsChangeListener> SETTINGS_CHANGED_TOPIC =
         Topic.create("SettingsChanged", SettingsChangeListener.class);
 
-    public static final Topic<ChatMessageManagementService> CHAT_MESSAGES_CHANGED_TOPIC =
-        Topic.create("chatChanged", ChatMessageManagementService.class);
-
     public static final Topic<ChatMemorySizeListener> CHAT_MEMORY_SIZE_TOPIC =
         new Topic<>("CHAT_MEMORY_SIZE_TOPIC", ChatMemorySizeListener.class);
 }