From 318dd4286a1c947f945661ebb97294813550afa0 Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Wed, 15 Nov 2023 00:44:13 +0200 Subject: [PATCH] Fix minor issues related to total tokens calculation --- .../DeleteAllConversationsAction.java | 2 +- .../completions/CompletionRequestHandler.java | 35 ++- .../CompletionRequestProvider.java | 16 +- ...a => CompletionResponseEventListener.java} | 4 +- .../codegpt/conversations/Conversation.java | 4 + .../conversations/ConversationService.java | 4 +- .../settings/SettingsConfigurable.java | 2 +- .../chat/BaseChatToolWindowTabPanel.java | 295 +++--------------- .../chat/ChatToolWindowScrollablePanel.java | 3 +- .../chat/ChatToolWindowTabPanel.java | 11 +- ...WindowCompletionResponseEventListener.java | 152 +++++++++ .../chat/components/TotalTokensPanel.java | 46 ++- .../chat/components/UserPromptTextArea.java | 35 ++- .../components/UserPromptTextAreaHeader.java | 4 +- .../ContextualChatToolWindowPanel.java | 20 +- .../ContextualChatToolWindowTabPanel.java | 21 +- .../StandardChatToolWindowContentManager.java | 24 +- .../standard/StandardChatToolWindowPanel.java | 27 +- .../StandardChatToolWindowTabPanel.java | 21 +- .../StandardChatToolWindowTabbedPane.java | 32 +- .../conversations/ConversationPanel.java | 11 +- .../DefaultCompletionRequestHandlerTest.java | 6 +- .../StandardChatToolWindowTabPanelTest.java | 6 +- .../StandardChatToolWindowTabbedPaneTest.java | 40 +-- 24 files changed, 424 insertions(+), 397 deletions(-) rename src/main/java/ee/carlrobert/codegpt/completions/{ToolWindowCompletionEventListener.java => CompletionResponseEventListener.java} (89%) create mode 100644 src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java diff --git a/src/main/java/ee/carlrobert/codegpt/actions/toolwindow/DeleteAllConversationsAction.java b/src/main/java/ee/carlrobert/codegpt/actions/toolwindow/DeleteAllConversationsAction.java index 14f0726a..50625b95 100644 --- a/src/main/java/ee/carlrobert/codegpt/actions/toolwindow/DeleteAllConversationsAction.java +++ b/src/main/java/ee/carlrobert/codegpt/actions/toolwindow/DeleteAllConversationsAction.java @@ -43,7 +43,7 @@ public class DeleteAllConversationsAction extends AnAction { if (project != null) { try { ConversationService.getInstance().clearAll(); - project.getService(StandardChatToolWindowContentManager.class).resetActiveTab(); + project.getService(StandardChatToolWindowContentManager.class).resetAll(); } finally { TelemetryAction.IDE_ACTION.createActionMessage() .property("action", ActionType.DELETE_ALL_CONVERSATIONS.name()) diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestHandler.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestHandler.java index 33ab9351..e3d173db 100644 --- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestHandler.java +++ b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestHandler.java @@ -20,19 +20,19 @@ public class CompletionRequestHandler { private final StringBuilder messageBuilder = new StringBuilder(); private final boolean useContextualSearch; - private final ToolWindowCompletionEventListener toolWindowCompletionEventListener; + private final CompletionResponseEventListener completionResponseEventListener; private SwingWorker swingWorker; private EventSource eventSource; public CompletionRequestHandler( boolean useContextualSearch, - ToolWindowCompletionEventListener toolWindowCompletionEventListener) { + CompletionResponseEventListener completionResponseEventListener) { this.useContextualSearch = useContextualSearch; - this.toolWindowCompletionEventListener = toolWindowCompletionEventListener; + this.completionResponseEventListener = completionResponseEventListener; } - public void call(Conversation conversation, Message message, boolean isRetry) { - swingWorker = new CompletionRequestWorker(conversation, message, isRetry); + public void call(Conversation conversation, Message message, boolean retry) { + swingWorker = new CompletionRequestWorker(conversation, message, retry); swingWorker.execute(); } @@ -56,7 +56,7 @@ public class CompletionRequestHandler { if (ex instanceof TotalUsageExceededException) { errorMessage = "The length of the context exceeds the maximum limit that the model can handle. Try reducing the input message or maximum completion token size."; } - toolWindowCompletionEventListener.handleError(new ErrorDetails(errorMessage), ex); + completionResponseEventListener.handleError(new ErrorDetails(errorMessage), ex); throw ex; } } @@ -65,12 +65,12 @@ public class CompletionRequestHandler { private final Conversation conversation; private final Message message; - private final boolean isRetry; + private final boolean retry; - public CompletionRequestWorker(Conversation conversation, Message message, boolean isRetry) { + public CompletionRequestWorker(Conversation conversation, Message message, boolean retry) { this.conversation = conversation; this.message = message; - this.isRetry = isRetry; + this.retry = retry; } protected Void doInBackground() { @@ -79,10 +79,10 @@ public class CompletionRequestHandler { eventSource = startCall( conversation, message, - isRetry, + retry, new YouRequestCompletionEventListener()); } catch (TotalUsageExceededException e) { - toolWindowCompletionEventListener.handleTokensExceeded(conversation, message); + completionResponseEventListener.handleTokensExceeded(conversation, message); } finally { sendInfo(settings); } @@ -93,7 +93,7 @@ public class CompletionRequestHandler { message.setResponse(messageBuilder.toString()); for (String text : chunks) { messageBuilder.append(text); - toolWindowCompletionEventListener.handleMessage(text); + completionResponseEventListener.handleMessage(text); } } @@ -101,7 +101,7 @@ public class CompletionRequestHandler { @Override public void onSerpResults(List results) { - toolWindowCompletionEventListener.handleSerpResults(results, message); + completionResponseEventListener.handleSerpResults(results, message); } @Override @@ -111,14 +111,17 @@ public class CompletionRequestHandler { @Override public void onComplete(StringBuilder messageBuilder) { - toolWindowCompletionEventListener.handleCompleted(messageBuilder.toString(), message, - conversation, isRetry); + completionResponseEventListener.handleCompleted( + messageBuilder.toString(), + message, + conversation, + retry); } @Override public void onError(ErrorDetails error, Throwable ex) { try { - toolWindowCompletionEventListener.handleError(error, ex); + completionResponseEventListener.handleError(error, ex); } finally { sendError(error, ex); } diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java index 6c2586fa..a4aee70d 100644 --- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java +++ b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java @@ -100,18 +100,18 @@ public class CompletionRequestProvider { public OpenAIChatCompletionRequest buildOpenAIChatCompletionRequest( String model, Message message, - boolean isRetry) { - return buildOpenAIChatCompletionRequest(model, message, isRetry, false, null); + boolean retry) { + return buildOpenAIChatCompletionRequest(model, message, retry, false, null); } public OpenAIChatCompletionRequest buildOpenAIChatCompletionRequest( @Nullable String model, Message message, - boolean isRetry, + boolean retry, boolean useContextualSearch, @Nullable String overriddenPath) { var builder = new OpenAIChatCompletionRequest.Builder( - buildMessages(model, message, isRetry, useContextualSearch)) + buildMessages(model, message, retry, useContextualSearch)) .setModel(model) .setMaxTokens(ConfigurationState.getInstance().getMaxTokens()) .setTemperature(ConfigurationState.getInstance().getTemperature()); @@ -125,7 +125,7 @@ public class CompletionRequestProvider { public List buildMessages( Message message, - boolean isRetry, + boolean retry, boolean useContextualSearch) { var messages = new ArrayList(); if (useContextualSearch) { @@ -138,7 +138,7 @@ public class CompletionRequestProvider { systemPrompt.isEmpty() ? COMPLETION_SYSTEM_PROMPT : systemPrompt)); for (var prevMessage : conversation.getMessages()) { - if (isRetry && prevMessage.getId().equals(message.getId())) { + if (retry && prevMessage.getId().equals(message.getId())) { break; } messages.add(new OpenAIChatCompletionMessage("user", prevMessage.getPrompt())); @@ -152,9 +152,9 @@ public class CompletionRequestProvider { private List buildMessages( @Nullable String model, Message message, - boolean isRetry, + boolean retry, boolean useContextualSearch) { - var messages = buildMessages(message, isRetry, useContextualSearch); + var messages = buildMessages(message, retry, useContextualSearch); if (model == null || SettingsState.getInstance().getSelectedService() == ServiceType.YOU) { return messages; diff --git a/src/main/java/ee/carlrobert/codegpt/completions/ToolWindowCompletionEventListener.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionResponseEventListener.java similarity index 89% rename from src/main/java/ee/carlrobert/codegpt/completions/ToolWindowCompletionEventListener.java rename to src/main/java/ee/carlrobert/codegpt/completions/CompletionResponseEventListener.java index 94d669af..a86a1ab2 100644 --- a/src/main/java/ee/carlrobert/codegpt/completions/ToolWindowCompletionEventListener.java +++ b/src/main/java/ee/carlrobert/codegpt/completions/CompletionResponseEventListener.java @@ -6,7 +6,7 @@ import ee.carlrobert.llm.client.openai.completion.ErrorDetails; import ee.carlrobert.llm.client.you.completion.YouSerpResult; import java.util.List; -public interface ToolWindowCompletionEventListener { +public interface CompletionResponseEventListener { default void handleMessage(String message) {} @@ -18,7 +18,7 @@ public interface ToolWindowCompletionEventListener { String fullMessage, Message message, Conversation conversation, - boolean isRetry) {} + boolean retry) {} default void handleSerpResults(List results, Message message) {} } diff --git a/src/main/java/ee/carlrobert/codegpt/conversations/Conversation.java b/src/main/java/ee/carlrobert/codegpt/conversations/Conversation.java index 0aba2cbc..3a2d444e 100644 --- a/src/main/java/ee/carlrobert/codegpt/conversations/Conversation.java +++ b/src/main/java/ee/carlrobert/codegpt/conversations/Conversation.java @@ -83,4 +83,8 @@ public class Conversation { .filter(message -> !message.getId().equals(messageId)) .collect(toList())); } + + public void removeMessages() { + messages.clear(); + } } diff --git a/src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java b/src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java index 13756c20..e1d77bb9 100644 --- a/src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java +++ b/src/main/java/ee/carlrobert/codegpt/conversations/ConversationService.java @@ -64,9 +64,9 @@ public final class ConversationService { String response, Message message, Conversation conversation, - boolean isRetry) { + boolean retry) { var conversationMessages = conversation.getMessages(); - if (isRetry && !conversationMessages.isEmpty()) { + if (retry && !conversationMessages.isEmpty()) { var messageToBeSaved = conversationMessages.stream() .filter(item -> item.getId().equals(message.getId())) .findFirst().orElseThrow(); diff --git a/src/main/java/ee/carlrobert/codegpt/settings/SettingsConfigurable.java b/src/main/java/ee/carlrobert/codegpt/settings/SettingsConfigurable.java index d52de8dd..c2118821 100644 --- a/src/main/java/ee/carlrobert/codegpt/settings/SettingsConfigurable.java +++ b/src/main/java/ee/carlrobert/codegpt/settings/SettingsConfigurable.java @@ -161,6 +161,6 @@ public class SettingsConfigurable implements Configurable { throw new RuntimeException("Could not find current project."); } - project.getService(StandardChatToolWindowContentManager.class).resetActiveTab(); + project.getService(StandardChatToolWindowContentManager.class).resetAll(); } } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/BaseChatToolWindowTabPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/BaseChatToolWindowTabPanel.java index 01aa29e5..20ed80ec 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/BaseChatToolWindowTabPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/BaseChatToolWindowTabPanel.java @@ -1,35 +1,22 @@ package ee.carlrobert.codegpt.toolwindow.chat; -import static com.intellij.openapi.ui.Messages.OK; import static ee.carlrobert.codegpt.util.SwingUtils.createScrollPaneWithSmartScroller; import static ee.carlrobert.codegpt.util.ThemeUtils.getPanelBackgroundColor; import static java.lang.String.format; -import com.intellij.openapi.application.ApplicationManager; import com.intellij.openapi.diagnostic.Logger; -import com.intellij.openapi.editor.EditorFactory; -import com.intellij.openapi.editor.event.EditorFactoryEvent; -import com.intellij.openapi.editor.event.EditorFactoryListener; -import com.intellij.openapi.editor.event.SelectionEvent; -import com.intellij.openapi.editor.event.SelectionListener; import com.intellij.openapi.editor.impl.EditorImpl; import com.intellij.openapi.project.Project; -import com.intellij.ui.DocumentAdapter; import com.intellij.ui.JBColor; -import com.intellij.ui.components.JBCheckBox; import com.intellij.util.ui.JBUI; import ee.carlrobert.codegpt.EncodingManager; import ee.carlrobert.codegpt.actions.ActionType; import ee.carlrobert.codegpt.completions.CompletionRequestHandler; import ee.carlrobert.codegpt.completions.CompletionRequestService; -import ee.carlrobert.codegpt.completions.ToolWindowCompletionEventListener; import ee.carlrobert.codegpt.conversations.Conversation; import ee.carlrobert.codegpt.conversations.ConversationService; import ee.carlrobert.codegpt.conversations.message.Message; -import ee.carlrobert.codegpt.settings.service.ServiceType; -import ee.carlrobert.codegpt.settings.state.OpenAISettingsState; import ee.carlrobert.codegpt.settings.state.SettingsState; -import ee.carlrobert.codegpt.settings.state.YouSettingsState; import ee.carlrobert.codegpt.telemetry.TelemetryAction; import ee.carlrobert.codegpt.toolwindow.chat.components.ChatMessageResponseBody; import ee.carlrobert.codegpt.toolwindow.chat.components.ResponsePanel; @@ -37,101 +24,68 @@ import ee.carlrobert.codegpt.toolwindow.chat.components.TotalTokensPanel; import ee.carlrobert.codegpt.toolwindow.chat.components.UserMessagePanel; import ee.carlrobert.codegpt.toolwindow.chat.components.UserPromptTextArea; import ee.carlrobert.codegpt.toolwindow.chat.components.UserPromptTextAreaHeader; -import ee.carlrobert.codegpt.toolwindow.chat.components.YouProCheckbox; import ee.carlrobert.codegpt.util.EditorUtils; -import ee.carlrobert.codegpt.util.OverlayUtils; import ee.carlrobert.codegpt.util.file.FileUtils; -import ee.carlrobert.llm.client.openai.completion.ErrorDetails; -import ee.carlrobert.llm.client.you.completion.YouSerpResult; import java.awt.BorderLayout; import java.awt.GridBagConstraints; import java.awt.GridBagLayout; -import java.util.HashMap; -import java.util.List; -import java.util.Map; import java.util.UUID; import javax.swing.JComponent; import javax.swing.JPanel; -import javax.swing.event.DocumentEvent; -import javax.swing.text.BadLocationException; import org.jetbrains.annotations.NotNull; -import org.jetbrains.annotations.Nullable; public abstract class BaseChatToolWindowTabPanel implements ChatToolWindowTabPanel { private static final Logger LOG = Logger.getInstance(BaseChatToolWindowTabPanel.class); - private final SettingsState settings; private final boolean useContextualSearch; private final JPanel rootPanel; - private final Map> serpResultsMapping = new HashMap<>(); - private final JBCheckBox gpt4CheckBox; - protected final TotalTokensPanel totalTokensPanel; - protected final Project project; - protected final UserPromptTextArea userPromptTextArea; - protected final ConversationService conversationService; - protected final ChatToolWindowScrollablePanel toolWindowScrollablePanel; - private final EncodingManager encodingManager; + private final Conversation conversation; + private final UserPromptTextArea userPromptTextArea; + private final ConversationService conversationService; - private boolean streaming; - protected @Nullable Conversation conversation; + protected final Project project; + protected final TotalTokensPanel totalTokensPanel; + protected final ChatToolWindowScrollablePanel toolWindowScrollablePanel; protected abstract JComponent getLandingView(); - public BaseChatToolWindowTabPanel(@NotNull Project project, boolean useContextualSearch) { + public BaseChatToolWindowTabPanel( + @NotNull Project project, + @NotNull Conversation conversation, + boolean useContextualSearch) { this.project = project; + this.conversation = conversation; this.useContextualSearch = useContextualSearch; conversationService = ConversationService.getInstance(); - encodingManager = EncodingManager.getInstance(); - settings = SettingsState.getInstance(); + var settings = SettingsState.getInstance(); toolWindowScrollablePanel = new ChatToolWindowScrollablePanel(settings); - gpt4CheckBox = new YouProCheckbox(project); - userPromptTextArea = new UserPromptTextArea(this::handleSubmit, getUserPromptDocumentAdapter()); totalTokensPanel = new TotalTokensPanel( - null, - userPromptTextArea.getText(), - EditorUtils.getSelectedEditorSelectedText(project)); - rootPanel = createRootPanel(); - - addSelectionListeners(); - + conversation, + EditorUtils.getSelectedEditorSelectedText(project), + this); + userPromptTextArea = new UserPromptTextArea(this::handleSubmit, totalTokensPanel); + rootPanel = createRootPanel(settings); userPromptTextArea.requestFocusInWindow(); userPromptTextArea.requestFocus(); } + @Override + public void dispose() { + } + @Override public JPanel getContent() { return rootPanel; } @Override - public @Nullable Conversation getConversation() { + public Conversation getConversation() { return conversation; } - @Override - public void setConversation(@Nullable Conversation conversation) { - this.conversation = conversation; - } - - @Override - public void displayLandingView() { - toolWindowScrollablePanel.displayLandingView(getLandingView()); - } - - @Override - public void startNewConversation(Message message) { - conversation = conversationService.startConversation(); - sendMessage(message); - } - @Override public void sendMessage(Message message) { - streaming = true; - if (conversation == null) { - conversation = conversationService.startConversation(); - } - var messagePanel = toolWindowScrollablePanel.addMessage(message.getId()); messagePanel.add(new UserMessagePanel(project, message, this)); var responsePanel = new ResponsePanel() @@ -140,9 +94,11 @@ public abstract class BaseChatToolWindowTabPanel implements ChatToolWindowTabPan .addContent(new ChatMessageResponseBody(project, true, this)); messagePanel.add(responsePanel); - totalTokensPanel.updateUserPromptTokens(message.getPrompt()); + var userPromptTokens = EncodingManager.getInstance().countTokens(message.getPrompt()); + var conversationTokens = EncodingManager.getInstance().countConversationTokens(conversation); + totalTokensPanel.updateConversationTokens(conversationTokens + userPromptTokens); - call(conversation, message, responsePanel, false); + call(message, responsePanel, false); } @Override @@ -151,21 +107,16 @@ public abstract class BaseChatToolWindowTabPanel implements ChatToolWindowTabPan } @Override - public void dispose() { - } - public void requestFocusForTextArea() { userPromptTextArea.focus(); } - public void updateConversationTokens() { + @Override + public void displayLandingView() { + toolWindowScrollablePanel.displayLandingView(getLandingView()); totalTokensPanel.updateConversationTokens(conversation); } - public boolean isStreaming() { - return streaming; - } - protected void reloadMessage(Message message, Conversation conversation) { ResponsePanel responsePanel = null; try { @@ -180,7 +131,7 @@ public abstract class BaseChatToolWindowTabPanel implements ChatToolWindowTabPan if (responsePanel != null) { message.setResponse(""); conversationService.saveMessage(conversation, message); - call(conversation, message, responsePanel, true); + call(message, responsePanel, true); } totalTokensPanel.updateConversationTokens(conversation); @@ -195,10 +146,9 @@ public abstract class BaseChatToolWindowTabPanel implements ChatToolWindowTabPan toolWindowScrollablePanel.removeMessage(messageId); conversation.removeMessage(messageId); conversationService.saveConversation(conversation); + totalTokensPanel.updateConversationTokens(conversation); if (conversation.getMessages().isEmpty()) { - conversationService.deleteConversation(conversation); - setConversation(null); displayLandingView(); } } @@ -206,14 +156,9 @@ public abstract class BaseChatToolWindowTabPanel implements ChatToolWindowTabPan protected void clearWindow() { toolWindowScrollablePanel.clearAll(); totalTokensPanel.updateConversationTokens(conversation); - updateConversationTokens(); } - private void call( - Conversation conversation, - Message message, - ResponsePanel responsePanel, - boolean isRetry) { + private void call(Message message, ResponsePanel responsePanel, boolean retry) { ChatMessageResponseBody responseContainer = (ChatMessageResponseBody) responsePanel.getContent(); if (!CompletionRequestService.getInstance().isRequestAllowed()) { @@ -223,10 +168,19 @@ public abstract class BaseChatToolWindowTabPanel implements ChatToolWindowTabPan var requestHandler = new CompletionRequestHandler( useContextualSearch, - new ChatToolWindowCompletionEventListener(responsePanel)); + new ToolWindowCompletionResponseEventListener( + conversationService, + responsePanel, + totalTokensPanel, + userPromptTextArea) { + @Override + public void handleTokensExceededPolicyAccepted() { + call(message, responsePanel, true); + } + }); userPromptTextArea.setRequestHandler(requestHandler); userPromptTextArea.setSubmitEnabled(false); - requestHandler.call(conversation, message, isRetry); + requestHandler.call(conversation, message, retry); } private void handleSubmit(String text) { @@ -244,27 +198,23 @@ public abstract class BaseChatToolWindowTabPanel implements ChatToolWindowTabPan } } - if (conversation == null) { - startNewConversation(message); - } else { - sendMessage(message); - } + sendMessage(message); } - private JPanel createUserPromptPanel() { + private JPanel createUserPromptPanel(SettingsState settings) { var panel = new JPanel(new BorderLayout()); panel.setBorder(JBUI.Borders.compound( JBUI.Borders.customLine(JBColor.border(), 1, 0, 0, 0), JBUI.Borders.empty(8))); panel.setBackground(getPanelBackgroundColor()); panel.add( - new UserPromptTextAreaHeader(project, settings, totalTokensPanel, gpt4CheckBox), + new UserPromptTextAreaHeader(project, settings, totalTokensPanel), BorderLayout.NORTH); panel.add(userPromptTextArea, BorderLayout.SOUTH); return panel; } - private JPanel createRootPanel() { + private JPanel createRootPanel(SettingsState settings) { var rootPanel = new JPanel(new GridBagLayout()); var gbc = new GridBagConstraints(); gbc.fill = GridBagConstraints.BOTH; @@ -277,154 +227,7 @@ public abstract class BaseChatToolWindowTabPanel implements ChatToolWindowTabPan gbc.weighty = 0; gbc.fill = GridBagConstraints.HORIZONTAL; gbc.gridy = 1; - rootPanel.add(createUserPromptPanel(), gbc); + rootPanel.add(createUserPromptPanel(settings), gbc); return rootPanel; } - - private class ChatToolWindowCompletionEventListener implements ToolWindowCompletionEventListener { - - private final Logger LOG = Logger.getInstance(ChatToolWindowCompletionEventListener.class); - - private final StringBuilder messageBuilder = new StringBuilder(); - private final ResponsePanel responsePanel; - private final ChatMessageResponseBody responseContainer; - - public ChatToolWindowCompletionEventListener(ResponsePanel responsePanel) { - this.responsePanel = responsePanel; - this.responseContainer = (ChatMessageResponseBody) responsePanel.getContent(); - } - - @Override - public void handleMessage(String partialMessage) { - try { - LOG.debug(partialMessage); - ApplicationManager.getApplication() - .invokeLater(() -> { - responseContainer.update(partialMessage); - messageBuilder.append(partialMessage); - - var ongoingTokens = encodingManager.countTokens(messageBuilder.toString()); - totalTokensPanel.update( - totalTokensPanel.getTokenDetails().getTotal() + ongoingTokens); - }); - } catch (Exception e) { - responseContainer.displayDefaultError(); - throw new RuntimeException("Error while updating the content", e); - } - } - - @Override - public void handleError(ErrorDetails error, Throwable ex) { - try { - if ("insufficient_quota".equals(error.getCode())) { - if (SettingsState.getInstance().getSelectedService() == ServiceType.OPENAI) { - OpenAISettingsState.getInstance().setOpenAIQuotaExceeded(true); - } - responseContainer.displayQuotaExceeded(); - } else { - responseContainer.displayError(error.getMessage()); - } - } finally { - LOG.error(error.getMessage(), ex); - responsePanel.enableActions(); - stopStreaming(responseContainer); - } - } - - @Override - public void handleTokensExceeded(Conversation conversation, Message message) { - var answer = OverlayUtils.showTokenLimitExceededDialog(); - if (answer == OK) { - TelemetryAction.IDE_ACTION.createActionMessage() - .property("action", "DISCARD_TOKEN_LIMIT") - .property("model", conversation.getModel()) - .send(); - - conversationService.discardTokenLimits(conversation); - call(conversation, message, responsePanel, true); - } else { - stopStreaming(responseContainer); - } - } - - @Override - public void handleCompleted( - String fullMessage, - Message message, - Conversation conversation, - boolean isRetry) { - try { - responsePanel.enableActions(); - conversationService.saveMessage(fullMessage, message, conversation, isRetry); - - var serpResults = serpResultsMapping.get(message.getId()); - var containsResults = serpResults != null && !serpResults.isEmpty(); - if (YouSettingsState.getInstance().isDisplayWebSearchResults() && containsResults) { - responseContainer.displaySerpResults(serpResults); - } - - if (containsResults) { - message.setSerpResults(serpResults); - } - - totalTokensPanel.updateUserPromptTokens(userPromptTextArea.getText()); - totalTokensPanel.updateConversationTokens(conversation); - } finally { - stopStreaming(responseContainer); - } - } - - @Override - public void handleSerpResults(List results, Message message) { - serpResultsMapping.put(message.getId(), results); - } - - private void stopStreaming(ChatMessageResponseBody responseContainer) { - streaming = false; - userPromptTextArea.setSubmitEnabled(true); - responseContainer.hideCarets(); - } - } - - private void addSelectionListeners() { - var editorFactory = EditorFactory.getInstance(); - for (var editor : editorFactory.getAllEditors()) { - editor.getSelectionModel().addSelectionListener(getSelectionListener()); - } - editorFactory.addEditorFactoryListener(new EditorFactoryListener() { - @Override - public void editorCreated(@NotNull EditorFactoryEvent event) { - event.getEditor().getSelectionModel().addSelectionListener(getSelectionListener()); - } - }, this); - } - - private SelectionListener getSelectionListener() { - return new SelectionListener() { - @Override - public void selectionChanged(@NotNull SelectionEvent e) { - var selectedText = e.getEditor().getDocument().getText(e.getNewRange()); - totalTokensPanel.updateHighlightedTokens(selectedText); - } - }; - } - - private DocumentAdapter getUserPromptDocumentAdapter() { - return new DocumentAdapter() { - @Override - protected void textChanged(@NotNull DocumentEvent event) { - try { - if (!streaming) { - var document = event.getDocument(); - var text = document.getText( - document.getStartPosition().getOffset(), - document.getEndPosition().getOffset() - 1); - totalTokensPanel.updateUserPromptTokens(text); - } - } catch (BadLocationException ex) { - LOG.error("Something went wrong while processing user input tokens", ex); - } - } - }; - } -} +} \ No newline at end of file diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowScrollablePanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowScrollablePanel.java index 41c04fd8..78ced557 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowScrollablePanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowScrollablePanel.java @@ -32,13 +32,12 @@ public class ChatToolWindowScrollablePanel extends ScrollablePanel { } public void displayLandingView(JComponent landingView) { - removeAll(); + clearAll(); add(landingView); if (settings.getSelectedService() == ServiceType.YOU && (!youUserManager.isAuthenticated() || !youUserManager.isSubscribed())) { add(new ResponsePanel().addContent(createYouCouponTextPane())); } - update(); } public ResponsePanel getMessageResponsePanel(UUID messageId) { 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 e9a21eff..d5c4d1b7 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java @@ -4,26 +4,17 @@ import com.intellij.openapi.Disposable; import ee.carlrobert.codegpt.conversations.Conversation; import ee.carlrobert.codegpt.conversations.message.Message; import javax.swing.JPanel; -import org.jetbrains.annotations.Nullable; public interface ChatToolWindowTabPanel extends Disposable { JPanel getContent(); - @Nullable Conversation getConversation(); + Conversation getConversation(); TokenDetails getTokenDetails(); - boolean isStreaming(); - - void setConversation(@Nullable Conversation conversation); - void displayLandingView(); - void displayConversation(Conversation conversation); - - void startNewConversation(Message message); - void sendMessage(Message message); void requestFocusForTextArea(); diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java new file mode 100644 index 00000000..874230d9 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java @@ -0,0 +1,152 @@ +package ee.carlrobert.codegpt.toolwindow.chat; + +import static com.intellij.openapi.ui.Messages.OK; + +import com.intellij.openapi.application.ApplicationManager; +import com.intellij.openapi.diagnostic.Logger; +import ee.carlrobert.codegpt.EncodingManager; +import ee.carlrobert.codegpt.completions.CompletionResponseEventListener; +import ee.carlrobert.codegpt.conversations.Conversation; +import ee.carlrobert.codegpt.conversations.ConversationService; +import ee.carlrobert.codegpt.conversations.message.Message; +import ee.carlrobert.codegpt.settings.service.ServiceType; +import ee.carlrobert.codegpt.settings.state.OpenAISettingsState; +import ee.carlrobert.codegpt.settings.state.SettingsState; +import ee.carlrobert.codegpt.settings.state.YouSettingsState; +import ee.carlrobert.codegpt.telemetry.TelemetryAction; +import ee.carlrobert.codegpt.toolwindow.chat.components.ChatMessageResponseBody; +import ee.carlrobert.codegpt.toolwindow.chat.components.ResponsePanel; +import ee.carlrobert.codegpt.toolwindow.chat.components.TotalTokensPanel; +import ee.carlrobert.codegpt.toolwindow.chat.components.UserPromptTextArea; +import ee.carlrobert.codegpt.util.OverlayUtils; +import ee.carlrobert.llm.client.openai.completion.ErrorDetails; +import ee.carlrobert.llm.client.you.completion.YouSerpResult; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +abstract class ToolWindowCompletionResponseEventListener implements + CompletionResponseEventListener { + + private static final Logger LOG = Logger.getInstance( + ToolWindowCompletionResponseEventListener.class); + + private final StringBuilder messageBuilder = new StringBuilder(); + private final Map> serpResultsMapping = new HashMap<>(); + private final EncodingManager encodingManager; + private final ConversationService conversationService; + private final ResponsePanel responsePanel; + private final ChatMessageResponseBody responseContainer; + private final TotalTokensPanel totalTokensPanel; + private final UserPromptTextArea userPromptTextArea; + + private volatile boolean completed; + + public ToolWindowCompletionResponseEventListener( + ConversationService conversationService, + ResponsePanel responsePanel, + TotalTokensPanel totalTokensPanel, + UserPromptTextArea userPromptTextArea) { + this.encodingManager = EncodingManager.getInstance(); + this.conversationService = conversationService; + this.responsePanel = responsePanel; + this.responseContainer = (ChatMessageResponseBody) responsePanel.getContent(); + this.totalTokensPanel = totalTokensPanel; + this.userPromptTextArea = userPromptTextArea; + } + + public abstract void handleTokensExceededPolicyAccepted(); + + @Override + public void handleMessage(String partialMessage) { + try { + ApplicationManager.getApplication() + .invokeLater(() -> { + responseContainer.update(partialMessage); + messageBuilder.append(partialMessage); + + if (!completed) { + var ongoingTokens = encodingManager.countTokens(messageBuilder.toString()); + totalTokensPanel.update( + totalTokensPanel.getTokenDetails().getTotal() + ongoingTokens); + } + }); + } catch (Exception e) { + responseContainer.displayDefaultError(); + throw new RuntimeException("Error while updating the content", e); + } + } + + @Override + public void handleError(ErrorDetails error, Throwable ex) { + try { + if ("insufficient_quota".equals(error.getCode())) { + if (SettingsState.getInstance().getSelectedService() == ServiceType.OPENAI) { + OpenAISettingsState.getInstance().setOpenAIQuotaExceeded(true); + } + responseContainer.displayQuotaExceeded(); + } else { + responseContainer.displayError(error.getMessage()); + } + } finally { + LOG.error(error.getMessage(), ex); + responsePanel.enableActions(); + stopStreaming(responseContainer); + } + } + + @Override + public void handleTokensExceeded(Conversation conversation, Message message) { + var answer = OverlayUtils.showTokenLimitExceededDialog(); + if (answer == OK) { + TelemetryAction.IDE_ACTION.createActionMessage() + .property("action", "DISCARD_TOKEN_LIMIT") + .property("model", conversation.getModel()) + .send(); + + conversationService.discardTokenLimits(conversation); + handleTokensExceededPolicyAccepted(); + } else { + stopStreaming(responseContainer); + } + } + + @Override + public void handleCompleted( + String fullMessage, + Message message, + Conversation conversation, + boolean retry) { + try { + responsePanel.enableActions(); + conversationService.saveMessage(fullMessage, message, conversation, retry); + + var serpResults = serpResultsMapping.get(message.getId()); + var containsResults = serpResults != null && !serpResults.isEmpty(); + if (YouSettingsState.getInstance().isDisplayWebSearchResults() && containsResults) { + responseContainer.displaySerpResults(serpResults); + } + + if (containsResults) { + message.setSerpResults(serpResults); + } + + totalTokensPanel.updateUserPromptTokens(userPromptTextArea.getText()); + totalTokensPanel.updateConversationTokens(conversation); + } finally { + stopStreaming(responseContainer); + } + } + + @Override + public void handleSerpResults(List results, Message message) { + serpResultsMapping.put(message.getId(), results); + } + + private void stopStreaming(ChatMessageResponseBody responseContainer) { + completed = true; + userPromptTextArea.setSubmitEnabled(true); + responseContainer.hideCarets(); + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/TotalTokensPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/TotalTokensPanel.java index a58a168c..5616502e 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/TotalTokensPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/TotalTokensPanel.java @@ -3,6 +3,12 @@ package ee.carlrobert.codegpt.toolwindow.chat.components; import static java.lang.String.format; import com.intellij.icons.AllIcons.General; +import com.intellij.openapi.Disposable; +import com.intellij.openapi.editor.EditorFactory; +import com.intellij.openapi.editor.event.EditorFactoryEvent; +import com.intellij.openapi.editor.event.EditorFactoryListener; +import com.intellij.openapi.editor.event.SelectionEvent; +import com.intellij.openapi.editor.event.SelectionListener; import com.intellij.ui.components.JBLabel; import ee.carlrobert.codegpt.EncodingManager; import ee.carlrobert.codegpt.conversations.Conversation; @@ -15,6 +21,7 @@ import java.util.Map; import java.util.stream.Collectors; import javax.swing.Box; import javax.swing.JPanel; +import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; public class TotalTokensPanel extends JPanel { @@ -25,17 +32,40 @@ public class TotalTokensPanel extends JPanel { public TotalTokensPanel( Conversation conversation, - String userPrompt, - @Nullable String highlightedText) { + @Nullable String highlightedText, + Disposable parentDisposable) { super(new FlowLayout(FlowLayout.LEADING, 0, 0)); this.encodingManager = EncodingManager.getInstance(); - this.tokenDetails = createTokenDetails(conversation, userPrompt, highlightedText); + this.tokenDetails = createTokenDetails(conversation, highlightedText); this.label = getLabel(tokenDetails); setOpaque(false); add(getContextHelpIcon(tokenDetails)); add(Box.createHorizontalStrut(4)); add(label); + addSelectionListeners(parentDisposable); + } + + private void addSelectionListeners(Disposable parentDisposable) { + var editorFactory = EditorFactory.getInstance(); + for (var editor : editorFactory.getAllEditors()) { + editor.getSelectionModel().addSelectionListener(getSelectionListener()); + } + editorFactory.addEditorFactoryListener(new EditorFactoryListener() { + @Override + public void editorCreated(@NotNull EditorFactoryEvent event) { + event.getEditor().getSelectionModel().addSelectionListener(getSelectionListener()); + } + }, parentDisposable); + } + + private SelectionListener getSelectionListener() { + return new SelectionListener() { + @Override + public void selectionChanged(@NotNull SelectionEvent e) { + updateHighlightedTokens(e.getEditor().getDocument().getText(e.getNewRange())); + } + }; } public TokenDetails getTokenDetails() { @@ -50,11 +80,15 @@ public class TotalTokensPanel extends JPanel { label.setText(getLabelHtml(total)); } - public void updateConversationTokens(Conversation conversation) { - tokenDetails.setConversationTokens(encodingManager.countConversationTokens(conversation)); + public void updateConversationTokens(int total) { + tokenDetails.setConversationTokens(total); update(); } + public void updateConversationTokens(Conversation conversation) { + updateConversationTokens(encodingManager.countConversationTokens(conversation)); + } + public void updateUserPromptTokens(String userPrompt) { tokenDetails.setUserPromptTokens(encodingManager.countTokens(userPrompt)); update(); @@ -67,11 +101,9 @@ public class TotalTokensPanel extends JPanel { private TokenDetails createTokenDetails( Conversation conversation, - String userPrompt, @Nullable String highlightedText) { var tokenDetails = new TokenDetails(encodingManager); tokenDetails.setConversationTokens(encodingManager.countConversationTokens(conversation)); - tokenDetails.setUserPromptTokens(encodingManager.countTokens(userPrompt)); if (highlightedText != null) { tokenDetails.setHighlightedTokens(encodingManager.countTokens(highlightedText)); } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/UserPromptTextArea.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/UserPromptTextArea.java index 126cc2d0..46dbb0ed 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/UserPromptTextArea.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/UserPromptTextArea.java @@ -1,6 +1,7 @@ package ee.carlrobert.codegpt.toolwindow.chat.components; import com.intellij.icons.AllIcons; +import com.intellij.openapi.diagnostic.Logger; import com.intellij.openapi.editor.ex.util.EditorUtil; import com.intellij.openapi.util.registry.Registry; import com.intellij.ui.DocumentAdapter; @@ -31,11 +32,14 @@ import javax.swing.JPanel; import javax.swing.KeyStroke; import javax.swing.UIManager; import javax.swing.event.DocumentEvent; +import javax.swing.text.BadLocationException; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; public class UserPromptTextArea extends JPanel { + private static final Logger LOG = Logger.getInstance(UserPromptTextArea.class); + private static final String TEXT_SUBMIT = "text-submit"; private static final String INSERT_BREAK = "insert-break"; private static final JBColor BACKGROUND_COLOR = JBColor.namedColor( @@ -49,11 +53,11 @@ public class UserPromptTextArea extends JPanel { private JPanel iconsPanel; private boolean submitEnabled = true; - public UserPromptTextArea(Consumer onSubmit, DocumentAdapter documentAdapter) { + public UserPromptTextArea(Consumer onSubmit, TotalTokensPanel totalTokensPanel) { this.onSubmit = onSubmit; textArea = new JBTextArea(); - textArea.getDocument().addDocumentListener(documentAdapter); + textArea.getDocument().addDocumentListener(getDocumentAdapter(totalTokensPanel)); textArea.setOpaque(false); textArea.setBackground(BACKGROUND_COLOR); textArea.setLineWrap(true); @@ -66,7 +70,11 @@ public class UserPromptTextArea extends JPanel { textArea.getActionMap().put(TEXT_SUBMIT, new AbstractAction() { @Override public void actionPerformed(ActionEvent e) { - handleSubmit(); + try { + handleSubmit(); + } finally { + totalTokensPanel.updateUserPromptTokens(""); + } } }); textArea.addFocusListener(new FocusListener() { @@ -90,6 +98,25 @@ public class UserPromptTextArea extends JPanel { init(); } + private DocumentAdapter getDocumentAdapter(TotalTokensPanel totalTokensPanel) { + return new DocumentAdapter() { + @Override + protected void textChanged(@NotNull DocumentEvent event) { + if (submitEnabled) { + try { + var document = event.getDocument(); + var text = document.getText( + document.getStartPosition().getOffset(), + document.getEndPosition().getOffset() - 1); + totalTokensPanel.updateUserPromptTokens(text); + } catch (BadLocationException ex) { + LOG.error("Something went wrong while processing user input tokens", ex); + } + } + } + }; + } + public String getText() { return textArea.getText().trim(); } @@ -133,7 +160,7 @@ public class UserPromptTextArea extends JPanel { if (submitEnabled && !textArea.getText().isEmpty()) { // Replacing each newline with two newlines to ensure proper Markdown formatting var text = textArea.getText().replace("\n", "\n\n"); - onSubmit.accept(text); + onSubmit.accept(text.trim()); textArea.setText(""); } } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/UserPromptTextAreaHeader.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/UserPromptTextAreaHeader.java index 60d5c1f6..5e3a1f8a 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/UserPromptTextAreaHeader.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/components/UserPromptTextAreaHeader.java @@ -21,8 +21,7 @@ public class UserPromptTextAreaHeader extends JPanel { public UserPromptTextAreaHeader( Project project, SettingsState settings, - TotalTokensPanel totalTokensPanel, - JBCheckBox gpt4CheckBox) { + TotalTokensPanel totalTokensPanel) { super(new BorderLayout()); setBackground(getPanelBackgroundColor()); setBorder(JBUI.Borders.emptyBottom(8)); @@ -32,6 +31,7 @@ public class UserPromptTextAreaHeader extends JPanel { add(totalTokensPanel, BorderLayout.LINE_START); break; case YOU: + JBCheckBox gpt4CheckBox = new YouProCheckbox(project); subscribeToYouTopics(project, gpt4CheckBox); add(gpt4CheckBox, BorderLayout.LINE_START); break; diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/contextual/ContextualChatToolWindowPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/contextual/ContextualChatToolWindowPanel.java index 52aed723..022e66c7 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/contextual/ContextualChatToolWindowPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/contextual/ContextualChatToolWindowPanel.java @@ -7,32 +7,28 @@ import com.intellij.openapi.actionSystem.DefaultActionGroup; import com.intellij.openapi.project.Project; import com.intellij.openapi.ui.SimpleToolWindowPanel; import com.intellij.openapi.util.Disposer; -import ee.carlrobert.codegpt.indexes.CodebaseIndexingAction; import ee.carlrobert.codegpt.actions.toolwindow.ClearChatWindowAction; +import ee.carlrobert.codegpt.conversations.Conversation; +import ee.carlrobert.codegpt.indexes.CodebaseIndexingAction; import org.jetbrains.annotations.NotNull; public class ContextualChatToolWindowPanel extends SimpleToolWindowPanel { - public ContextualChatToolWindowPanel(@NotNull Project project, @NotNull Disposable parentDisposable) { + public ContextualChatToolWindowPanel( + @NotNull Project project, + @NotNull Conversation conversation, + @NotNull Disposable parentDisposable) { super(true); - initialize(project, parentDisposable); - } - - private void initialize(Project project, Disposable parentDisposable) { - var tabPanel = new ContextualChatToolWindowTabPanel(project); + var tabPanel = new ContextualChatToolWindowTabPanel(project, conversation); setToolbar(createActionToolbar(tabPanel).getComponent()); setContent(tabPanel.getContent()); - Disposer.register(parentDisposable, tabPanel); } private ActionToolbar createActionToolbar(ContextualChatToolWindowTabPanel tabPanel) { var actionGroup = new DefaultActionGroup("TOOLBAR_ACTION_GROUP", false); - actionGroup.add(new ClearChatWindowAction(() -> { - tabPanel.displayLandingView(); - tabPanel.setConversation(null); - })); + actionGroup.add(new ClearChatWindowAction(tabPanel::displayLandingView)); actionGroup.addSeparator(); actionGroup.add(new CodebaseIndexingAction()); diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/contextual/ContextualChatToolWindowTabPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/contextual/ContextualChatToolWindowTabPanel.java index cc88a321..1b3448e6 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/contextual/ContextualChatToolWindowTabPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/contextual/ContextualChatToolWindowTabPanel.java @@ -9,24 +9,17 @@ import org.jetbrains.annotations.NotNull; public class ContextualChatToolWindowTabPanel extends BaseChatToolWindowTabPanel { - public ContextualChatToolWindowTabPanel(@NotNull Project project) { - super(project, true); + public ContextualChatToolWindowTabPanel( + @NotNull Project project, + @NotNull Conversation conversation) { + super(project, conversation, true); displayLandingView(); } @Override protected JComponent getLandingView() { - return new ContextualChatToolWindowLandingPanel(project, (prompt) -> { - var message = new Message(prompt); - if (conversation == null) { - startNewConversation(message); - } else { - sendMessage(message); - } - }); - } - - @Override - public void displayConversation(Conversation conversation) { + return new ContextualChatToolWindowLandingPanel( + project, + (prompt) -> sendMessage(new Message(prompt))); } } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowContentManager.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowContentManager.java index 3951d830..495d7382 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowContentManager.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowContentManager.java @@ -8,12 +8,14 @@ import com.intellij.openapi.wm.ToolWindow; import com.intellij.openapi.wm.ToolWindowManager; import com.intellij.ui.content.Content; import ee.carlrobert.codegpt.conversations.Conversation; +import ee.carlrobert.codegpt.conversations.ConversationService; import ee.carlrobert.codegpt.conversations.ConversationsState; import ee.carlrobert.codegpt.conversations.message.Message; import ee.carlrobert.codegpt.settings.configuration.ConfigurationState; import ee.carlrobert.codegpt.toolwindow.chat.ChatToolWindowTabPanel; import java.util.Arrays; import java.util.Optional; +import org.jetbrains.annotations.NotNull; @Service(Service.Level.PROJECT) public final class StandardChatToolWindowContentManager { @@ -42,27 +44,31 @@ public final class StandardChatToolWindowContentManager { toolWindow.show(); } - if (ConfigurationState.getInstance().isCreateNewChatOnEachAction() || ConversationsState.getCurrentConversation() == null) { - toolWindowTabPanel.startNewConversation(message); + if (ConfigurationState.getInstance().isCreateNewChatOnEachAction() + || ConversationsState.getCurrentConversation() == null) { + ConversationService.getInstance().startConversation(); } else { toolWindowTabPanel.sendMessage(message); } } - public void displayConversation(Conversation conversation) { + public void displayConversation(@NotNull Conversation conversation) { displayChatTab(); tryFindChatTabbedPane() - .ifPresent(tabbedPane -> tabbedPane.tryFindActiveConversationTitle(conversation.getId()) + .ifPresent(tabbedPane -> tabbedPane.tryFindTabTitle(conversation.getId()) .ifPresentOrElse( title -> tabbedPane.setSelectedIndex(tabbedPane.indexOfTab(title)), - () -> tabbedPane.addNewTab(new StandardChatToolWindowTabPanel(project, conversation)))); + () -> tabbedPane.addNewTab( + new StandardChatToolWindowTabPanel(project, conversation)))); } public StandardChatToolWindowTabPanel createNewTabPanel() { displayChatTab(); var tabbedPane = tryFindChatTabbedPane(); if (tabbedPane.isPresent()) { - var panel = new StandardChatToolWindowTabPanel(project); + var panel = new StandardChatToolWindowTabPanel( + project, + ConversationService.getInstance().startConversation()); tabbedPane.get().addNewTab(panel); return panel; } @@ -93,10 +99,12 @@ public final class StandardChatToolWindowContentManager { return Optional.empty(); } - public void resetActiveTab() { + public void resetAll() { tryFindChatTabbedPane().ifPresent(tabbedPane -> { tabbedPane.clearAll(); - tabbedPane.addNewTab(new StandardChatToolWindowTabPanel(project)); + tabbedPane.addNewTab(new StandardChatToolWindowTabPanel( + project, + ConversationService.getInstance().startConversation())); }); } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowPanel.java index 22dd8dd8..3b0eea95 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowPanel.java @@ -10,38 +10,48 @@ import com.intellij.openapi.util.Disposer; import ee.carlrobert.codegpt.actions.toolwindow.ClearChatWindowAction; import ee.carlrobert.codegpt.actions.toolwindow.CreateNewConversationAction; import ee.carlrobert.codegpt.actions.toolwindow.OpenInEditorAction; +import ee.carlrobert.codegpt.conversations.ConversationService; import ee.carlrobert.codegpt.conversations.ConversationsState; import java.awt.FlowLayout; import org.jetbrains.annotations.NotNull; public class StandardChatToolWindowPanel extends SimpleToolWindowPanel { - public StandardChatToolWindowPanel(@NotNull Project project, @NotNull Disposable parentDisposable) { + public StandardChatToolWindowPanel( + @NotNull Project project, + @NotNull Disposable parentDisposable) { super(true); - initialize(project, parentDisposable); + init(project, parentDisposable); } - private void initialize(Project project, Disposable parentDisposable) { + private void init(Project project, Disposable parentDisposable) { var conversation = ConversationsState.getCurrentConversation(); + if (conversation == null) { + conversation = ConversationService.getInstance().startConversation(); + } var tabPanel = new StandardChatToolWindowTabPanel(project, conversation); var tabbedPane = createTabbedPane(tabPanel, parentDisposable); var toolbarComponent = createActionToolbar(project, tabbedPane).getComponent(); toolbarComponent.setLayout(new FlowLayout()); - setToolbar(toolbarComponent); setContent(tabbedPane); Disposer.register(parentDisposable, tabPanel); } - private ActionToolbar createActionToolbar(Project project, StandardChatToolWindowTabbedPane tabbedPane) { + private ActionToolbar createActionToolbar( + Project project, + StandardChatToolWindowTabbedPane tabbedPane) { var actionGroup = new DefaultCompactActionGroup("TOOLBAR_ACTION_GROUP", false); actionGroup.add(new CreateNewConversationAction(() -> { - tabbedPane.addNewTab(new StandardChatToolWindowTabPanel(project)); + tabbedPane.addNewTab(new StandardChatToolWindowTabPanel( + project, + ConversationService.getInstance().startConversation())); repaint(); revalidate(); })); - actionGroup.add(new ClearChatWindowAction(tabbedPane::resetCurrentlyActiveTabPanel)); + actionGroup.add( + new ClearChatWindowAction(() -> tabbedPane.resetCurrentlyActiveTabPanel(project))); actionGroup.addSeparator(); actionGroup.add(new OpenInEditorAction()); @@ -51,7 +61,8 @@ public class StandardChatToolWindowPanel extends SimpleToolWindowPanel { return toolbar; } - private StandardChatToolWindowTabbedPane createTabbedPane(StandardChatToolWindowTabPanel tabPanel, Disposable parentDisposable) { + private StandardChatToolWindowTabbedPane createTabbedPane(StandardChatToolWindowTabPanel tabPanel, + Disposable parentDisposable) { var tabbedPane = new StandardChatToolWindowTabbedPane(parentDisposable); tabbedPane.addNewTab(tabPanel); return tabbedPane; diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowTabPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowTabPanel.java index 6b6dcbfb..5a49aadc 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowTabPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowTabPanel.java @@ -16,19 +16,14 @@ import ee.carlrobert.codegpt.util.OverlayUtils; import ee.carlrobert.codegpt.util.file.FileUtils; import javax.swing.JComponent; import org.jetbrains.annotations.NotNull; -import org.jetbrains.annotations.Nullable; public class StandardChatToolWindowTabPanel extends BaseChatToolWindowTabPanel { - public StandardChatToolWindowTabPanel(@NotNull Project project) { - this(project, null); - } - public StandardChatToolWindowTabPanel( @NotNull Project project, - @Nullable Conversation conversation) { - super(project, false); - if (conversation == null) { + @NotNull Conversation conversation) { + super(project, conversation, false); + if (conversation.getMessages().isEmpty()) { displayLandingView(); } else { displayConversation(conversation); @@ -55,16 +50,11 @@ public class StandardChatToolWindowTabPanel extends BaseChatToolWindowTabPanel { format("\n```%s\n%s\n```", fileExtension, editor.getSelectionModel().getSelectedText()))); message.setUserMessage(action.getUserMessage()); - if (conversation == null) { - startNewConversation(message); - } else { - sendMessage(message); - } + sendMessage(message); }); } - @Override - public void displayConversation(@NotNull Conversation conversation) { + private void displayConversation(@NotNull Conversation conversation) { clearWindow(); conversation.getMessages().forEach(message -> { var messageResponseBody = new ChatMessageResponseBody(project, this) @@ -83,6 +73,5 @@ public class StandardChatToolWindowTabPanel extends BaseChatToolWindowTabPanel { .withDeleteAction(() -> removeMessage(message.getId(), conversation)) .addContent(messageResponseBody)); }); - setConversation(conversation); } } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowTabbedPane.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowTabbedPane.java index be12e77f..de5233ee 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowTabbedPane.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/standard/StandardChatToolWindowTabbedPane.java @@ -2,11 +2,13 @@ package ee.carlrobert.codegpt.toolwindow.chat.standard; import com.intellij.icons.AllIcons; import com.intellij.openapi.Disposable; +import com.intellij.openapi.project.Project; import com.intellij.openapi.ui.JBMenuItem; import com.intellij.openapi.util.Disposer; import com.intellij.ui.components.JBLabel; import com.intellij.ui.components.JBTabbedPane; import com.intellij.util.ui.JBUI; +import ee.carlrobert.codegpt.conversations.ConversationService; import ee.carlrobert.codegpt.conversations.ConversationsState; import ee.carlrobert.codegpt.settings.state.SettingsState; import java.awt.Component; @@ -25,11 +27,12 @@ import javax.swing.SwingUtilities; public class StandardChatToolWindowTabbedPane extends JBTabbedPane { - private final Map activeTabMapping = new TreeMap<>((o1, o2) -> { - int n1 = Integer.parseInt(o1.replaceAll("\\D", "")); - int n2 = Integer.parseInt(o2.replaceAll("\\D", "")); - return Integer.compare(n1, n2); - }); + private final Map activeTabMapping = new TreeMap<>( + (o1, o2) -> { + int n1 = Integer.parseInt(o1.replaceAll("\\D", "")); + int n2 = Integer.parseInt(o2.replaceAll("\\D", "")); + return Integer.compare(n1, n2); + }); private final Disposable parentDisposable; public StandardChatToolWindowTabbedPane(Disposable parentDisposable) { @@ -68,7 +71,7 @@ public class StandardChatToolWindowTabbedPane extends JBTabbedPane { Disposer.register(parentDisposable, toolWindowPanel); } - public Optional tryFindActiveConversationTitle(UUID conversationId) { + public Optional tryFindTabTitle(UUID conversationId) { return activeTabMapping.entrySet().stream() .filter(entry -> { var panelConversation = entry.getValue().getConversation(); @@ -108,13 +111,17 @@ public class StandardChatToolWindowTabbedPane extends JBTabbedPane { } } - public void resetCurrentlyActiveTabPanel() { + public void resetCurrentlyActiveTabPanel(Project project) { tryFindActiveTabPanel().ifPresent(tabPanel -> { - tabPanel.displayLandingView(); - tabPanel.setConversation(null); - tabPanel.updateConversationTokens(); + Disposer.dispose(tabPanel); + activeTabMapping.remove(getTitleAt(getSelectedIndex())); + removeTabAt(getSelectedIndex()); + addNewTab(new StandardChatToolWindowTabPanel( + project, + ConversationService.getInstance().startConversation())); + repaint(); + revalidate(); }); - ConversationsState.getInstance().setCurrentConversation(null); } private JPanel createCloseableTabButtonPanel(String title) { @@ -174,7 +181,8 @@ public class StandardChatToolWindowTabbedPane extends JBTabbedPane { @Override public void show(Component invoker, int x, int y) { - selectedPopupTabIndex = StandardChatToolWindowTabbedPane.this.getUI().tabForCoordinate(StandardChatToolWindowTabbedPane.this, x, y); + selectedPopupTabIndex = StandardChatToolWindowTabbedPane.this.getUI() + .tabForCoordinate(StandardChatToolWindowTabbedPane.this, x, y); if (selectedPopupTabIndex > 0) { super.show(invoker, x, y); } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/conversations/ConversationPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/conversations/ConversationPanel.java index 7a184a1a..8bfec8fe 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/conversations/ConversationPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/conversations/ConversationPanel.java @@ -32,13 +32,20 @@ class ConversationPanel extends JPanel { @NotNull Conversation conversation, @NotNull Runnable onDelete) { super(new BorderLayout()); + var toolWindowContentManager = project.getService(StandardChatToolWindowContentManager.class); + init(toolWindowContentManager, conversation, onDelete); + } + + private void init( + StandardChatToolWindowContentManager toolWindowContentManager, + Conversation conversation, + Runnable onDelete) { setBackground(JBColor.background()); addMouseListener(new MouseAdapter() { @Override public void mouseClicked(MouseEvent e) { SettingsState.getInstance().sync(conversation); - project.getService(StandardChatToolWindowContentManager.class) - .displayConversation(conversation); + toolWindowContentManager.displayConversation(conversation); } }); addStyles(isSelected(conversation)); diff --git a/src/test/java/ee/carlrobert/codegpt/completions/DefaultCompletionRequestHandlerTest.java b/src/test/java/ee/carlrobert/codegpt/completions/DefaultCompletionRequestHandlerTest.java index 8ba13fb5..1592be68 100644 --- a/src/test/java/ee/carlrobert/codegpt/completions/DefaultCompletionRequestHandlerTest.java +++ b/src/test/java/ee/carlrobert/codegpt/completions/DefaultCompletionRequestHandlerTest.java @@ -171,14 +171,14 @@ public class DefaultCompletionRequestHandlerTest extends IntegrationTest { await().atMost(5, SECONDS).until(() -> "Hello!".equals(message.getResponse())); } - private ToolWindowCompletionEventListener getRequestEventListener(Message message) { - return new ToolWindowCompletionEventListener() { + private CompletionResponseEventListener getRequestEventListener(Message message) { + return new CompletionResponseEventListener() { @Override public void handleCompleted( String fullMessage, Message conversationMessage, Conversation conversation, - boolean isRetry) { + boolean retry) { message.setResponse(fullMessage); } }; diff --git a/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabPanelTest.java b/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabPanelTest.java index be9ff113..f0dd26e1 100644 --- a/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabPanelTest.java +++ b/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabPanelTest.java @@ -49,7 +49,11 @@ public class StandardChatToolWindowTabPanelTest extends IntegrationTest { panel.sendMessage(message); - await().atMost(5, SECONDS).until(() -> !panel.isStreaming()); + await().atMost(5, SECONDS) + .until(() -> { + var messages = conversation.getMessages(); + return !messages.isEmpty() && "Hello!".equals(messages.get(0).getResponse()); + }); var encodingManager = EncodingManager.getInstance(); assertThat(panel.getTokenDetails()).extracting( "systemPromptTokens", diff --git a/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabbedPaneTest.java b/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabbedPaneTest.java index 02416af9..17411989 100644 --- a/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabbedPaneTest.java +++ b/src/test/java/ee/carlrobert/codegpt/toolwindow/chat/StandardChatToolWindowTabbedPaneTest.java @@ -4,49 +4,49 @@ import static org.assertj.core.api.Assertions.assertThat; import com.intellij.openapi.util.Disposer; import com.intellij.testFramework.fixtures.BasePlatformTestCase; -import ee.carlrobert.codegpt.conversations.Conversation; +import ee.carlrobert.codegpt.conversations.ConversationService; import ee.carlrobert.codegpt.conversations.message.Message; import ee.carlrobert.codegpt.toolwindow.chat.standard.StandardChatToolWindowTabPanel; import ee.carlrobert.codegpt.toolwindow.chat.standard.StandardChatToolWindowTabbedPane; -import testsupport.IntegrationTest; public class StandardChatToolWindowTabbedPaneTest extends BasePlatformTestCase { public void testClearAllTabs() { var tabbedPane = new StandardChatToolWindowTabbedPane(Disposer.newDisposable()); - tabbedPane.addNewTab(new StandardChatToolWindowTabPanel(getProject())); + tabbedPane.addNewTab(createNewTabPanel()); tabbedPane.clearAll(); assertThat(tabbedPane.getActiveTabMapping()).isEmpty(); } + public void testAddingNewTabs() { var tabbedPane = new StandardChatToolWindowTabbedPane(Disposer.newDisposable()); - tabbedPane.addNewTab(new StandardChatToolWindowTabPanel(getProject())); - tabbedPane.addNewTab(new StandardChatToolWindowTabPanel(getProject())); - tabbedPane.addNewTab(new StandardChatToolWindowTabPanel(getProject())); + tabbedPane.addNewTab(createNewTabPanel()); + tabbedPane.addNewTab(createNewTabPanel()); + tabbedPane.addNewTab(createNewTabPanel()); - var tabMapping = tabbedPane.getActiveTabMapping(); - assertThat(tabMapping.keySet()).containsExactly("Chat 1", "Chat 2", "Chat 3"); - assertThat( - tabMapping.values().stream().allMatch(item -> item.getConversation() == null)).isTrue(); + assertThat(tabbedPane.getActiveTabMapping().keySet()) + .containsExactly("Chat 1", "Chat 2", "Chat 3"); } public void testResetCurrentlyActiveTabPanel() { var tabbedPane = new StandardChatToolWindowTabbedPane(Disposer.newDisposable()); - var tabPanel = new StandardChatToolWindowTabPanel(getProject()); - var message = new Message("TEST_PROMPT"); - message.setResponse("TEST_RESPONSE"); - var conversation = new Conversation(); - conversation.addMessage(message); - tabPanel.setConversation(conversation); - tabbedPane.addNewTab(tabPanel); + var conversation = ConversationService.getInstance().startConversation(); + conversation.addMessage(new Message("TEST_PROMPT", "TEST_RESPONSE")); + tabbedPane.addNewTab(new StandardChatToolWindowTabPanel(getProject(), conversation)); - tabbedPane.resetCurrentlyActiveTabPanel(); + tabbedPane.resetCurrentlyActiveTabPanel(getProject()); - var tabMapping = tabbedPane.getActiveTabMapping().get("Chat 1"); - assertThat(tabMapping.getConversation()).isNull(); + var tabPanel = tabbedPane.getActiveTabMapping().get("Chat 1"); + assertThat(tabPanel.getConversation().getMessages()).isEmpty(); + } + + private StandardChatToolWindowTabPanel createNewTabPanel() { + return new StandardChatToolWindowTabPanel( + getProject(), + ConversationService.getInstance().startConversation()); } }