Skip to content

Commit 6479604

Browse files
carlrobertohcarlrobertoh
committed
feat: improve chat UI performance
1 parent 658e78f commit 6479604

File tree

6 files changed

+143
-154
lines changed

6 files changed

+143
-154
lines changed

‎src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestHandler.java

+43-64
Original file line numberDiff line numberDiff line change
@@ -4,35 +4,37 @@
44
import com.fasterxml.jackson.databind.ObjectMapper;
55
import ee.carlrobert.codegpt.events.CodeGPTEvent;
66
import ee.carlrobert.codegpt.settings.GeneralSettings;
7-
import ee.carlrobert.codegpt.settings.GeneralSettingsState;
87
import ee.carlrobert.codegpt.telemetry.TelemetryAction;
98
import ee.carlrobert.llm.client.openai.completion.ErrorDetails;
109
import ee.carlrobert.llm.completion.CompletionEventListener;
11-
import java.util.List;
12-
import javax.swing.SwingWorker;
1310
import okhttp3.sse.EventSource;
1411

1512
public class CompletionRequestHandler {
1613

1714
private final StringBuilder messageBuilder = new StringBuilder();
1815
private final CompletionResponseEventListener completionResponseEventListener;
19-
private SwingWorker<Void, String> swingWorker;
2016
private EventSource eventSource;
2117

2218
public CompletionRequestHandler(CompletionResponseEventListener completionResponseEventListener) {
2319
this.completionResponseEventListener = completionResponseEventListener;
2420
}
2521

2622
public void call(CallParameters callParameters) {
27-
swingWorker = new CompletionRequestWorker(callParameters);
28-
swingWorker.execute();
23+
try {
24+
eventSource = startCall(callParameters, new RequestCompletionEventListener(callParameters));
25+
} catch (TotalUsageExceededException e) {
26+
completionResponseEventListener.handleTokensExceeded(
27+
callParameters.getConversation(),
28+
callParameters.getMessage());
29+
} finally {
30+
sendInfo(callParameters);
31+
}
2932
}
3033

3134
public void cancel() {
3235
if (eventSource != null) {
3336
eventSource.cancel();
3437
}
35-
swingWorker.cancel(true);
3638
}
3739

3840
private EventSource startCall(
@@ -57,79 +59,48 @@ private void handleCallException(Throwable ex) {
5759
completionResponseEventListener.handleError(new ErrorDetails(errorMessage), ex);
5860
}
5961

60-
private class CompletionRequestWorker extends SwingWorker<Void, String> {
62+
class RequestCompletionEventListener implements CompletionEventListener<String> {
6163

6264
private final CallParameters callParameters;
6365

64-
public CompletionRequestWorker(CallParameters callParameters) {
66+
public RequestCompletionEventListener(CallParameters callParameters) {
6567
this.callParameters = callParameters;
6668
}
6769

68-
protected Void doInBackground() {
69-
var settings = GeneralSettings.getCurrentState();
70+
@Override
71+
public void onEvent(String data) {
7072
try {
71-
eventSource = startCall(callParameters, new RequestCompletionEventListener());
72-
} catch (TotalUsageExceededException e) {
73-
completionResponseEventListener.handleTokensExceeded(
74-
callParameters.getConversation(),
75-
callParameters.getMessage());
76-
} finally {
77-
sendInfo(settings);
73+
var event = new ObjectMapper().readValue(data, CodeGPTEvent.class);
74+
completionResponseEventListener.handleCodeGPTEvent(event);
75+
} catch (JsonProcessingException e) {
76+
// ignore
7877
}
79-
return null;
8078
}
8179

82-
protected void process(List<String> chunks) {
80+
@Override
81+
public void onMessage(String message, EventSource eventSource) {
82+
messageBuilder.append(message);
8383
callParameters.getMessage().setResponse(messageBuilder.toString());
84-
for (String text : chunks) {
85-
messageBuilder.append(text);
86-
completionResponseEventListener.handleMessage(text);
87-
}
84+
completionResponseEventListener.handleMessage(message);
8885
}
8986

90-
class RequestCompletionEventListener implements CompletionEventListener<String> {
91-
92-
@Override
93-
public void onEvent(String data) {
94-
try {
95-
var event = new ObjectMapper().readValue(data, CodeGPTEvent.class);
96-
completionResponseEventListener.handleCodeGPTEvent(event);
97-
} catch (JsonProcessingException e) {
98-
// ignore
99-
}
100-
}
101-
102-
@Override
103-
public void onMessage(String message, EventSource eventSource) {
104-
publish(message);
105-
}
106-
107-
@Override
108-
public void onComplete(StringBuilder messageBuilder) {
109-
completionResponseEventListener.handleCompleted(messageBuilder.toString(), callParameters);
110-
}
111-
112-
@Override
113-
public void onCancelled(StringBuilder messageBuilder) {
114-
completionResponseEventListener.handleCompleted(messageBuilder.toString(), callParameters);
115-
}
87+
@Override
88+
public void onComplete(StringBuilder messageBuilder) {
89+
completionResponseEventListener.handleCompleted(messageBuilder.toString(), callParameters);
90+
}
11691

117-
@Override
118-
public void onError(ErrorDetails error, Throwable ex) {
119-
try {
120-
completionResponseEventListener.handleError(error, ex);
121-
} finally {
122-
sendError(error, ex);
123-
}
124-
}
92+
@Override
93+
public void onCancelled(StringBuilder messageBuilder) {
94+
completionResponseEventListener.handleCompleted(messageBuilder.toString(), callParameters);
12595
}
12696

127-
private void sendInfo(GeneralSettingsState settings) {
128-
TelemetryAction.COMPLETION.createActionMessage()
129-
.property("conversationId", callParameters.getConversation().getId().toString())
130-
.property("model", callParameters.getConversation().getModel())
131-
.property("service", settings.getSelectedService().getCode().toLowerCase())
132-
.send();
97+
@Override
98+
public void onError(ErrorDetails error, Throwable ex) {
99+
try {
100+
completionResponseEventListener.handleError(error, ex);
101+
} finally {
102+
sendError(error, ex);
103+
}
133104
}
134105

135106
private void sendError(ErrorDetails error, Throwable ex) {
@@ -147,4 +118,12 @@ private void sendError(ErrorDetails error, Throwable ex) {
147118
telemetryMessage.send();
148119
}
149120
}
121+
122+
private void sendInfo(CallParameters callParameters) {
123+
TelemetryAction.COMPLETION.createActionMessage()
124+
.property("conversationId", callParameters.getConversation().getId().toString())
125+
.property("model", callParameters.getConversation().getModel())
126+
.property("service", GeneralSettings.getSelectedService().getCode().toLowerCase())
127+
.send();
128+
}
150129
}

‎src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java

+3-10
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,12 @@
55
import static java.lang.String.format;
66

77
import com.intellij.openapi.Disposable;
8+
import com.intellij.openapi.application.ApplicationManager;
89
import com.intellij.openapi.diagnostic.Logger;
910
import com.intellij.openapi.project.Project;
1011
import com.intellij.ui.JBColor;
1112
import com.intellij.util.ui.JBUI;
1213
import ee.carlrobert.codegpt.CodeGPTKeys;
13-
import ee.carlrobert.codegpt.EncodingManager;
1414
import ee.carlrobert.codegpt.ReferencedFile;
1515
import ee.carlrobert.codegpt.actions.ActionType;
1616
import ee.carlrobert.codegpt.completions.CallParameters;
@@ -41,7 +41,6 @@
4141
import java.util.UUID;
4242
import javax.swing.JComponent;
4343
import javax.swing.JPanel;
44-
import javax.swing.SwingUtilities;
4544
import kotlin.Unit;
4645
import org.jetbrains.annotations.NotNull;
4746
import org.jetbrains.annotations.Nullable;
@@ -111,7 +110,7 @@ public void sendMessage(Message message) {
111110
}
112111

113112
public void sendMessage(Message message, ConversationType conversationType) {
114-
SwingUtilities.invokeLater(() -> {
113+
ApplicationManager.getApplication().invokeLater(() -> {
115114
var referencedFiles = project.getUserData(CodeGPTKeys.SELECTED_FILES);
116115
var chatToolWindowPanel = project.getService(ChatToolWindowContentManager.class)
117116
.tryFindChatToolWindowPanel();
@@ -127,6 +126,7 @@ public void sendMessage(Message message, ConversationType conversationType) {
127126

128127
chatToolWindowPanel.ifPresent(panel -> panel.clearNotifications(project));
129128
}
129+
totalTokensPanel.updateConversationTokens(conversation);
130130

131131
var userMessagePanel = new UserMessagePanel(project, message, this);
132132
var attachedFilePath = CodeGPTKeys.IMAGE_ATTACHMENT_FILE_PATH.get(project);
@@ -142,7 +142,6 @@ public void sendMessage(Message message, ConversationType conversationType) {
142142

143143
var responsePanel = createResponsePanel(message, conversationType);
144144
messagePanel.add(responsePanel);
145-
updateTotalTokens(message);
146145
call(callParameters, responsePanel);
147146
});
148147
}
@@ -163,12 +162,6 @@ private CallParameters getCallParameters(
163162
return callParameters;
164163
}
165164

166-
private void updateTotalTokens(Message message) {
167-
int userPromptTokens = EncodingManager.getInstance().countTokens(message.getPrompt());
168-
int conversationTokens = EncodingManager.getInstance().countConversationTokens(conversation);
169-
totalTokensPanel.updateConversationTokens(conversationTokens + userPromptTokens);
170-
}
171-
172165
private ResponsePanel createResponsePanel(Message message, ConversationType conversationType) {
173166
return new ResponsePanel()
174167
.withReloadAction(() -> reloadMessage(message, conversation, conversationType))

‎src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ToolWindowCompletionResponseEventListener.java

+15-16
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
import ee.carlrobert.codegpt.ui.OverlayUtil;
1919
import ee.carlrobert.codegpt.ui.textarea.UserInputPanel;
2020
import ee.carlrobert.llm.client.openai.completion.ErrorDetails;
21-
import javax.swing.SwingUtilities;
2221

2322
abstract class ToolWindowCompletionResponseEventListener implements
2423
CompletionResponseEventListener {
@@ -54,17 +53,16 @@ public ToolWindowCompletionResponseEventListener(
5453
@Override
5554
public void handleMessage(String partialMessage) {
5655
try {
57-
ApplicationManager.getApplication()
58-
.invokeLater(() -> {
59-
responseContainer.update(partialMessage);
60-
messageBuilder.append(partialMessage);
61-
62-
if (!completed) {
63-
var ongoingTokens = encodingManager.countTokens(messageBuilder.toString());
64-
totalTokensPanel.update(
65-
totalTokensPanel.getTokenDetails().getTotal() + ongoingTokens);
66-
}
67-
});
56+
responseContainer.update(partialMessage);
57+
messageBuilder.append(partialMessage);
58+
59+
if (!completed) {
60+
var ongoingTokens = encodingManager.countTokens(messageBuilder.toString());
61+
ApplicationManager.getApplication().invokeLater(() -> {
62+
totalTokensPanel.update(
63+
totalTokensPanel.getTokenDetails().getTotal() + ongoingTokens);
64+
});
65+
}
6866
} catch (Exception e) {
6967
responseContainer.displayError("Something went wrong.");
7068
throw new RuntimeException("Error while updating the content", e);
@@ -73,7 +71,7 @@ public void handleMessage(String partialMessage) {
7371

7472
@Override
7573
public void handleError(ErrorDetails error, Throwable ex) {
76-
SwingUtilities.invokeLater(() -> {
74+
ApplicationManager.getApplication().invokeLater(() -> {
7775
try {
7876
if ("insufficient_quota".equals(error.getCode())) {
7977
responseContainer.displayQuotaExceeded();
@@ -90,7 +88,7 @@ public void handleError(ErrorDetails error, Throwable ex) {
9088

9189
@Override
9290
public void handleTokensExceeded(Conversation conversation, Message message) {
93-
SwingUtilities.invokeLater(() -> {
91+
ApplicationManager.getApplication().invokeLater(() -> {
9492
var answer = OverlayUtil.showTokenLimitExceededDialog();
9593
if (answer == OK) {
9694
TelemetryAction.IDE_ACTION.createActionMessage()
@@ -110,7 +108,7 @@ public void handleTokensExceeded(Conversation conversation, Message message) {
110108
public void handleCompleted(String fullMessage, CallParameters callParameters) {
111109
conversationService.saveMessage(fullMessage, callParameters);
112110

113-
SwingUtilities.invokeLater(() -> {
111+
ApplicationManager.getApplication().invokeLater(() -> {
114112
try {
115113
responsePanel.enableActions();
116114
totalTokensPanel.updateUserPromptTokens(textArea.getText());
@@ -123,7 +121,8 @@ public void handleCompleted(String fullMessage, CallParameters callParameters) {
123121

124122
@Override
125123
public void handleCodeGPTEvent(CodeGPTEvent event) {
126-
responseContainer.displayWebSearchItem(event.getEvent().getDetails());
124+
ApplicationManager.getApplication().invokeLater(() ->
125+
responseContainer.displayWebSearchItem(event.getEvent().getDetails()));
127126
}
128127

129128
private void stopStreaming(ChatMessageResponseBody responseContainer) {

0 commit comments

Comments
 (0)