mirror of
https://github.com/carlrobertoh/ProxyAI.git
synced 2026-05-16 11:15:34 +00:00
Fix minor issues related to total tokens calculation
This commit is contained in:
Carl-Robert Linnupuu
parent
346218b512
commit
318dd4286a
24 changed files with 424 additions and 397 deletions
|
|
@ -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<UUID, List<YouSerpResult>> 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<YouSerpResult> 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);
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue