refactor: chat request message building

This commit is contained in:
Carl-Robert Linnupuu 2024-11-18 10:38:19 +00:00
parent 352c211323
commit d80a4a6556
28 changed files with 255 additions and 240 deletions

View file

@ -58,6 +58,10 @@ public final class EncodingManager {
}
public int countTokens(String text) {
if (text == null || text.isEmpty()) {
return 0;
}
try {
// #444: Cl100kParser.split() throws AssertionError "Input is not UTF-8: "
return encoding.countTokens(text.replaceAll("<|", "").replaceAll("|>", ""));

View file

@ -68,7 +68,6 @@ public class ProjectCompilationStatusListener implements CompilationStatusListen
message.setReferencedFilePaths(errorMapping.keySet().stream()
.map(ReferencedFile::getFilePath)
.toList());
message.setUserMessage(message.getPrompt());
message.setPrompt(CompletionRequestUtil.getPromptWithContext(
new ArrayList<>(errorMapping.keySet()),
prompt));

View file

@ -39,7 +39,6 @@ public class AskQuestionAction extends BaseEditorAction {
previousUserPrompt = dialog.getUserPrompt();
var message = new Message(
format("%s%n```%s%n%s%n```", previousUserPrompt, fileExtension, selectedText));
message.setUserMessage(previousUserPrompt);
SwingUtilities.invokeLater(() ->
project.getService(ChatToolWindowContentManager.class).sendMessage(message));
}

View file

@ -2,7 +2,6 @@ package ee.carlrobert.codegpt.actions.editor;
import static java.lang.String.format;
import com.intellij.icons.AllIcons.Actions;
import com.intellij.openapi.actionSystem.ActionManager;
import com.intellij.openapi.actionSystem.AnAction;
import com.intellij.openapi.actionSystem.DefaultActionGroup;
@ -12,7 +11,6 @@ import com.intellij.openapi.extensions.PluginId;
import com.intellij.openapi.project.Project;
import ee.carlrobert.codegpt.CodeGPTKeys;
import ee.carlrobert.codegpt.ReferencedFile;
import ee.carlrobert.codegpt.actions.IncludeFilesInContextAction;
import ee.carlrobert.codegpt.conversations.message.Message;
import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings;
import ee.carlrobert.codegpt.toolwindow.chat.ChatToolWindowContentManager;
@ -54,16 +52,15 @@ public class EditorActionsUtil {
var action = new BaseEditorAction(label, label) {
@Override
protected void actionPerformed(Project project, Editor editor, String selectedText) {
var toolWindowContentManager =
project.getService(ChatToolWindowContentManager.class);
toolWindowContentManager.getToolWindow().show();
var fileExtension = FileUtil.getFileExtension(
((EditorImpl) editor).getVirtualFile().getName());
var message = new Message(prompt.replace(
"{{selectedCode}}",
format("%n```%s%n%s%n```", fileExtension, selectedText)));
message.setUserMessage(prompt.replace("{{selectedCode}}", ""));
var toolWindowContentManager =
project.getService(ChatToolWindowContentManager.class);
toolWindowContentManager.getToolWindow().show();
message.setReferencedFilePaths(
Stream.ofNullable(project.getUserData(CodeGPTKeys.SELECTED_FILES))
.flatMap(Collection::stream)

View file

@ -45,23 +45,28 @@ public class ChatCompletionEventListener implements CompletionEventListener<Stri
@Override
public void onComplete(StringBuilder messageBuilder) {
eventListener.handleCompleted(messageBuilder.toString(), callParameters);
handleCompleted(messageBuilder);
}
@Override
public void onCancelled(StringBuilder messageBuilder) {
eventListener.handleCompleted(messageBuilder.toString(), callParameters);
handleCompleted(messageBuilder);
}
@Override
public void onError(ErrorDetails error, Throwable ex) {
try {
callParameters.getConversation().addMessage(callParameters.getMessage());
eventListener.handleError(error, ex);
} finally {
sendError(error, ex);
}
}
private void handleCompleted(StringBuilder messageBuilder) {
eventListener.handleCompleted(messageBuilder.toString(), callParameters);
}
private void sendError(ErrorDetails error, Throwable ex) {
var telemetryMessage = TelemetryAction.COMPLETION_ERROR.createActionMessage();
if ("insufficient_quota".equals(error.getCode())) {

View file

@ -1,22 +1,21 @@
package ee.carlrobert.codegpt.conversations.message;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
import ee.carlrobert.codegpt.settings.persona.PersonaDetails;
import ee.carlrobert.codegpt.ui.DocumentationDetails;
import ee.carlrobert.llm.client.you.completion.YouSerpResult;
import java.util.List;
import java.util.Objects;
import java.util.UUID;
import org.jetbrains.annotations.Nullable;
@JsonIgnoreProperties(ignoreUnknown = true)
public class Message {
private final UUID id;
private String prompt;
private String response;
private String userMessage;
private List<YouSerpResult> serpResults;
private List<String> referencedFilePaths;
private @Nullable String imageFilePath;
private boolean webSearchIncluded;
@ -54,22 +53,6 @@ public class Message {
this.response = response;
}
public String getUserMessage() {
return userMessage;
}
public void setUserMessage(String userMessage) {
this.userMessage = userMessage;
}
public List<YouSerpResult> getSerpResults() {
return serpResults;
}
public void setSerpResults(List<YouSerpResult> serpResults) {
this.serpResults = serpResults;
}
public List<String> getReferencedFilePaths() {
return referencedFilePaths;
}

View file

@ -7,9 +7,9 @@ import com.intellij.openapi.Disposable;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.editor.SelectionModel;
import com.intellij.openapi.editor.ex.EditorEx;
import com.intellij.openapi.editor.impl.EditorImpl;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.ui.JBColor;
import com.intellij.util.ui.JBUI;
import ee.carlrobert.codegpt.CodeGPTKeys;
@ -23,7 +23,6 @@ import ee.carlrobert.codegpt.conversations.Conversation;
import ee.carlrobert.codegpt.conversations.ConversationService;
import ee.carlrobert.codegpt.conversations.message.Message;
import ee.carlrobert.codegpt.telemetry.TelemetryAction;
import ee.carlrobert.codegpt.toolwindow.chat.actionprocessor.ActionProcessorFactory;
import ee.carlrobert.codegpt.toolwindow.chat.ui.ChatMessageResponseBody;
import ee.carlrobert.codegpt.toolwindow.chat.ui.ChatToolWindowScrollablePanel;
import ee.carlrobert.codegpt.toolwindow.chat.ui.ResponsePanel;
@ -39,9 +38,6 @@ import ee.carlrobert.codegpt.util.file.FileUtil;
import git4idea.GitCommit;
import java.awt.BorderLayout;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
@ -91,7 +87,7 @@ public class ChatToolWindowTabPanel implements Disposable {
if (conversation.getMessages().isEmpty()) {
displayLandingView();
} else {
displayConversation(conversation);
displayConversation();
}
}
@ -120,8 +116,8 @@ public class ChatToolWindowTabPanel implements Disposable {
totalTokensPanel.updateConversationTokens(conversation);
}
public void addSelection(String fileName, SelectionModel selectionModel) {
userInputPanel.addSelection(fileName, selectionModel);
public void addSelection(VirtualFile editorFile, SelectionModel selectionModel) {
userInputPanel.addSelection(editorFile, selectionModel);
}
public void addCommitReferences(List<GitCommit> gitCommits) {
@ -155,42 +151,27 @@ public class ChatToolWindowTabPanel implements Disposable {
}
public void sendMessage(Message message, ConversationType conversationType) {
sendMessage(message, conversationType, null);
}
public void sendMessage(
Message message,
ConversationType conversationType,
@Nullable String highlightedText) {
ApplicationManager.getApplication().invokeLater(() -> {
List<ReferencedFile> referencedFiles = getReferencedFiles();
if (!referencedFiles.isEmpty()) {
message.setReferencedFilePaths(referencedFiles.stream()
.map(ReferencedFile::getFilePath)
.toList());
message.setUserMessage(message.getPrompt());
}
var callParameters = ChatCompletionParameters.builder(conversation, message)
.sessionId(chatSession.getId())
.conversationType(conversationType)
.imageDetailsFromPath(CodeGPTKeys.IMAGE_ATTACHMENT_FILE_PATH.get(project))
.referencedFiles(getReferencedFiles())
.build();
String attachedImagePath = CodeGPTKeys.IMAGE_ATTACHMENT_FILE_PATH.get(project);
if (attachedImagePath != null) {
message.setImageFilePath(attachedImagePath);
}
totalTokensPanel.updateConversationTokens(conversation);
totalTokensPanel.updateReferencedFilesTokens(referencedFiles);
if (attachedImagePath != null || !referencedFiles.isEmpty()) {
var referencedFiles = callParameters.getReferencedFiles();
if ((referencedFiles != null && !referencedFiles.isEmpty())
|| callParameters.getImageDetails() != null) {
project.getService(ChatToolWindowContentManager.class)
.tryFindChatToolWindowPanel()
.ifPresent(panel -> panel.clearNotifications(project));
}
var callParameters = getCallParameters(
message,
conversationType,
referencedFiles,
highlightedText,
attachedImagePath);
totalTokensPanel.updateConversationTokens(conversation);
if (callParameters.getReferencedFiles() != null) {
totalTokensPanel.updateReferencedFilesTokens(callParameters.getReferencedFiles());
}
var responsePanel = createResponsePanel(callParameters);
var messagePanel = toolWindowScrollablePanel.addMessage(message.getId());
messagePanel.add(new UserMessagePanel(project, message, this));
@ -200,30 +181,6 @@ public class ChatToolWindowTabPanel implements Disposable {
});
}
private ChatCompletionParameters getCallParameters(
Message message,
ConversationType conversationType,
List<ReferencedFile> referencedFiles,
@Nullable String highlightedText,
@Nullable String attachedImagePath) {
var builder = ChatCompletionParameters.builder(conversation, message)
.sessionId(chatSession.getId())
.conversationType(conversationType)
.highlightedText(highlightedText)
.referencedFiles(referencedFiles);
if (attachedImagePath != null && !attachedImagePath.isEmpty()) {
try {
builder
.imageData(Files.readAllBytes(Path.of(attachedImagePath)))
.imageMediaType(FileUtil.getImageMediaType(attachedImagePath));
} catch (IOException e) {
throw new RuntimeException(e);
}
}
return builder.build();
}
private boolean hasReferencedFilePaths(Message message) {
return message.getReferencedFilePaths() != null && !message.getReferencedFilePaths().isEmpty();
}
@ -245,7 +202,6 @@ public class ChatToolWindowTabPanel implements Disposable {
.addContent(
new ChatMessageResponseBody(
project,
callParameters.getHighlightedText(),
true,
false,
message.isWebSearchIncluded(),
@ -320,39 +276,21 @@ public class ChatToolWindowTabPanel implements Disposable {
}
private Unit handleSubmit(String text, List<? extends AppliedActionInlay> appliedInlayActions) {
var message = new Message(text);
var editor = EditorUtil.getSelectedEditor(project);
var messageBuilder = new MessageBuilder(project, text)
.withSelectedEditorContent()
.withInlays(appliedInlayActions);
var remainingText = new StringBuilder(text);
var promptBuilder = new StringBuilder();
for (var actionInlay : appliedInlayActions) {
var inlayOffset = actionInlay.getInlay().getOffset();
promptBuilder.append(remainingText, 0, Math.min(inlayOffset, remainingText.length()))
.append("\n");
ActionProcessorFactory.getProcessor(actionInlay)
.process(message, actionInlay, editor, promptBuilder);
remainingText.delete(0, inlayOffset);
}
promptBuilder.append(remainingText);
String selectedText = "";
String selectedTextMd = "";
if (editor != null) {
var selectionModel = editor.getSelectionModel();
selectedText = selectionModel.getSelectedText();
if (selectedText != null && !selectedText.isEmpty()) {
var fileExtension = FileUtil.getFileExtension(
((EditorEx) editor).getVirtualFile().getName());
selectedTextMd = format("\n```%s\n%s\n```\n", fileExtension, selectedText);
selectionModel.removeSelection();
}
List<ReferencedFile> referencedFiles = getReferencedFiles();
if (!referencedFiles.isEmpty()) {
messageBuilder.withReferencedFiles(referencedFiles);
}
message.setUserMessage(selectedTextMd + promptBuilder);
message.setPrompt(selectedTextMd + promptBuilder);
String attachedImagePath = CodeGPTKeys.IMAGE_ATTACHMENT_FILE_PATH.get(project);
if (attachedImagePath != null) {
messageBuilder.withImage(attachedImagePath);
}
sendMessage(message, ConversationType.DEFAULT, selectedText);
sendMessage(messageBuilder.build(), ConversationType.DEFAULT);
return Unit.INSTANCE;
}
@ -390,18 +328,17 @@ public class ChatToolWindowTabPanel implements Disposable {
var message = new Message(action.getPrompt().replace(
"{{selectedCode}}",
format("%n```%s%n%s%n```", fileExtension, editor.getSelectionModel().getSelectedText())));
message.setUserMessage(action.getUserMessage());
sendMessage(message, ConversationType.DEFAULT);
return Unit.INSTANCE;
});
}
private void displayConversation(@NotNull Conversation conversation) {
private void displayConversation() {
clearWindow();
conversation.getMessages().forEach(message -> {
var response = message.getResponse() == null ? "" : message.getResponse();
var messageResponseBody =
new ChatMessageResponseBody(project, this).withResponse(message.getResponse());
new ChatMessageResponseBody(project, this).withResponse(response);
messageResponseBody.hideCaret();

View file

@ -1,11 +1,9 @@
package ee.carlrobert.codegpt.toolwindow.chat.actionprocessor;
import com.intellij.openapi.editor.Editor;
import ee.carlrobert.codegpt.conversations.message.Message;
import ee.carlrobert.codegpt.ui.textarea.AppliedActionInlay;
public interface ActionProcessor {
void process(Message message, AppliedActionInlay action, Editor editor,
StringBuilder promptBuilder);
void process(Message message, AppliedActionInlay action, StringBuilder promptBuilder);
}

View file

@ -1,14 +1,15 @@
package ee.carlrobert.codegpt.toolwindow.chat.actionprocessor;
import com.intellij.openapi.project.Project;
import ee.carlrobert.codegpt.ui.textarea.AppliedActionInlay;
import ee.carlrobert.codegpt.ui.textarea.AppliedCodeActionInlay;
import ee.carlrobert.codegpt.ui.textarea.AppliedSuggestionActionInlay;
public class ActionProcessorFactory {
public static ActionProcessor getProcessor(AppliedActionInlay action) {
public static ActionProcessor getProcessor(Project project, AppliedActionInlay action) {
if (action instanceof AppliedSuggestionActionInlay) {
return new SuggestionActionProcessor();
return new SuggestionActionProcessor(project);
} else if (action instanceof AppliedCodeActionInlay) {
return new CodeActionProcessor();
}

View file

@ -1,7 +1,5 @@
package ee.carlrobert.codegpt.toolwindow.chat.actionprocessor;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.editor.ex.EditorEx;
import ee.carlrobert.codegpt.conversations.message.Message;
import ee.carlrobert.codegpt.ui.textarea.AppliedActionInlay;
import ee.carlrobert.codegpt.ui.textarea.AppliedCodeActionInlay;
@ -10,19 +8,16 @@ import ee.carlrobert.codegpt.util.file.FileUtil;
public class CodeActionProcessor implements ActionProcessor {
@Override
public void process(Message message, AppliedActionInlay action, Editor editor,
StringBuilder promptBuilder) {
public void process(Message message, AppliedActionInlay action, StringBuilder promptBuilder) {
if (!(action instanceof AppliedCodeActionInlay codeAction)) {
throw new IllegalArgumentException("Invalid action type");
}
processCodeAction(codeAction, editor, promptBuilder);
processCodeAction(codeAction, promptBuilder);
}
private void processCodeAction(AppliedCodeActionInlay action, Editor editor,
StringBuilder promptBuilder) {
private void processCodeAction(AppliedCodeActionInlay action, StringBuilder promptBuilder) {
promptBuilder
.append("\n```%s\n".formatted(
FileUtil.getFileExtension(((EditorEx) editor).getVirtualFile().getName())))
.append("\n```%s\n".formatted(FileUtil.getFileExtension(action.getEditorFile().getName())))
.append(action.getCode())
.append("\n```\n");
}

View file

@ -1,6 +1,5 @@
package ee.carlrobert.codegpt.toolwindow.chat.actionprocessor;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.project.Project;
import ee.carlrobert.codegpt.CodeGPTKeys;
import ee.carlrobert.codegpt.conversations.message.Message;
@ -11,36 +10,38 @@ import ee.carlrobert.codegpt.ui.textarea.suggestion.item.DocumentationActionItem
import ee.carlrobert.codegpt.ui.textarea.suggestion.item.GitCommitActionItem;
import ee.carlrobert.codegpt.ui.textarea.suggestion.item.PersonaActionItem;
import ee.carlrobert.codegpt.ui.textarea.suggestion.item.WebSearchActionItem;
import org.jetbrains.annotations.Nullable;
public class SuggestionActionProcessor implements ActionProcessor {
private final Project project;
public SuggestionActionProcessor(Project project) {
this.project = project;
}
@Override
public void process(Message message, AppliedActionInlay action, Editor editor,
public void process(
Message message,
AppliedActionInlay action,
StringBuilder promptBuilder) {
if (!(action instanceof AppliedSuggestionActionInlay suggestionAction)) {
throw new IllegalArgumentException("Invalid action type");
}
processSuggestionAction(message, suggestionAction, editor, promptBuilder);
processSuggestionAction(message, suggestionAction, promptBuilder);
}
private void processSuggestionAction(
Message message,
AppliedSuggestionActionInlay action,
@Nullable Editor editor,
StringBuilder promptBuilder) {
message.setWebSearchIncluded(action.getSuggestion() instanceof WebSearchActionItem);
if (editor != null) {
processDocumentationAction(message, action, editor.getProject());
processPersonaAction(message, action, editor.getProject());
}
processDocumentationAction(message, action);
processPersonaAction(message, action);
processGitCommitAction(action, promptBuilder);
}
private void processDocumentationAction(
Message message,
AppliedSuggestionActionInlay action,
Project project) {
private void processDocumentationAction(Message message, AppliedSuggestionActionInlay action) {
var addedDocumentation = CodeGPTKeys.ADDED_DOCUMENTATION.get(project);
var appliedInlayExists = action.getSuggestion() instanceof DocumentationActionItem
|| action.getSuggestion() instanceof CreateDocumentationActionItem;
@ -51,10 +52,7 @@ public class SuggestionActionProcessor implements ActionProcessor {
}
}
private void processPersonaAction(
Message message,
AppliedSuggestionActionInlay action,
Project project) {
private void processPersonaAction(Message message, AppliedSuggestionActionInlay action) {
var addedPersona = CodeGPTKeys.ADDED_PERSONA.get(project);
var personaInlayExists = action.getSuggestion() instanceof PersonaActionItem;
if (addedPersona != null && personaInlayExists) {

View file

@ -38,7 +38,6 @@ import ee.carlrobert.codegpt.util.EditorUtil;
import java.awt.BorderLayout;
import javax.swing.JPanel;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
public class ResponseEditorPanel extends JPanel implements Disposable {
@ -49,7 +48,6 @@ public class ResponseEditorPanel extends JPanel implements Disposable {
String code,
String markdownLanguage,
boolean readOnly,
@Nullable String highlightedText,
Disposable disposableParent) {
super(new BorderLayout());
setBorder(JBUI.Borders.empty(8, 0));

View file

@ -43,7 +43,7 @@ import javax.swing.DefaultListModel;
import javax.swing.JEditorPane;
import javax.swing.JPanel;
import javax.swing.JTextPane;
import org.jetbrains.annotations.Nullable;
import org.jetbrains.annotations.NotNull;
public class ChatMessageResponseBody extends JPanel {
@ -56,25 +56,22 @@ public class ChatMessageResponseBody extends JPanel {
private final DefaultListModel<WebSearchEventDetails> webpageListModel = new DefaultListModel<>();
private final WebpageList webpageList = new WebpageList(webpageListModel);
private final ResponseBodyProgressPanel progressPanel = new ResponseBodyProgressPanel();
private final @Nullable String highlightedText;
private ResponseEditorPanel currentlyProcessedEditorPanel;
private JEditorPane currentlyProcessedTextPane;
private JPanel webpageListPanel;
public ChatMessageResponseBody(Project project, Disposable parentDisposable) {
this(project, null, false, false, false, false, parentDisposable);
this(project, false, false, false, false, parentDisposable);
}
public ChatMessageResponseBody(
Project project,
@Nullable String highlightedText,
boolean withGhostText,
boolean readOnly,
boolean webSearchIncluded,
boolean withProgress,
Disposable parentDisposable) {
this.project = project;
this.highlightedText = highlightedText;
this.parentDisposable = parentDisposable;
this.streamParser = new StreamParser();
this.readOnly = readOnly;
@ -97,7 +94,7 @@ public class ChatMessageResponseBody extends JPanel {
}
}
public ChatMessageResponseBody withResponse(String response) {
public ChatMessageResponseBody withResponse(@NotNull String response) {
try {
for (var message : MarkdownUtil.splitCodeBlocks(response)) {
processResponse(message, message.startsWith("```"), false);
@ -265,9 +262,8 @@ public class ChatMessageResponseBody extends JPanel {
private void prepareProcessingCode(String code, String markdownLanguage) {
hideCaret();
currentlyProcessedTextPane = null;
currentlyProcessedEditorPanel = new ResponseEditorPanel(project, code, markdownLanguage,
readOnly, highlightedText,
parentDisposable);
currentlyProcessedEditorPanel =
new ResponseEditorPanel(project, code, markdownLanguage, readOnly, parentDisposable);
add(currentlyProcessedEditorPanel);
}

View file

@ -47,15 +47,10 @@ public class UserMessagePanel extends JPanel {
displayImage(message.getImageFilePath());
}
var referencedFilePaths = message.getReferencedFilePaths();
if (referencedFilePaths != null && !referencedFilePaths.isEmpty()) {
add(createResponseBody(
project,
message.getUserMessage(),
parentDisposable), BorderLayout.SOUTH);
} else {
add(createResponseBody(project, message.getPrompt(), parentDisposable), BorderLayout.SOUTH);
}
add(createResponseBody(
project,
message.getPrompt(),
parentDisposable), BorderLayout.SOUTH);
}
public @Nullable JPanel getAdditionalContextPanel(Project project, Message message) {
@ -99,7 +94,7 @@ public class UserMessagePanel extends JPanel {
Project project,
String prompt,
Disposable parentDisposable) {
return new ChatMessageResponseBody(project, null, false, true, false, false, parentDisposable)
return new ChatMessageResponseBody(project, false, true, false, false, parentDisposable)
.withResponse(prompt);
}

View file

@ -1,8 +1,5 @@
package ee.carlrobert.codegpt.toolwindow.chat.ui.textarea;
import ee.carlrobert.codegpt.EncodingManager;
import ee.carlrobert.codegpt.settings.persona.PersonaSettings;
public class TotalTokensDetails {
private final int systemPromptTokens;
@ -11,8 +8,8 @@ public class TotalTokensDetails {
private int highlightedTokens;
private int referencedFilesTokens;
public TotalTokensDetails(EncodingManager encodingManager) {
systemPromptTokens = encodingManager.countTokens(PersonaSettings.getSystemPrompt());
public TotalTokensDetails(int systemPromptTokens) {
this.systemPromptTokens = systemPromptTokens;
}
public int getSystemPromptTokens() {

View file

@ -19,7 +19,9 @@ import ee.carlrobert.codegpt.EncodingManager;
import ee.carlrobert.codegpt.ReferencedFile;
import ee.carlrobert.codegpt.actions.IncludeFilesInContextNotifier;
import ee.carlrobert.codegpt.conversations.Conversation;
import ee.carlrobert.codegpt.conversations.message.Message;
import ee.carlrobert.codegpt.settings.GeneralSettings;
import ee.carlrobert.codegpt.settings.persona.PersonaSettings;
import ee.carlrobert.codegpt.settings.service.ServiceType;
import java.awt.FlowLayout;
import java.awt.event.MouseAdapter;
@ -105,6 +107,13 @@ public class TotalTokensPanel extends JPanel {
label.setText(getLabelHtml(total));
}
public void updateConversationTokens(Conversation conversation, Message message) {
totalTokensDetails.setConversationTokens(
encodingManager.countConversationTokens(conversation)
+ encodingManager.countMessageTokens("user", message.getPrompt()));
update();
}
public void updateConversationTokens(Conversation conversation) {
totalTokensDetails.setConversationTokens(encodingManager.countConversationTokens(conversation));
update();
@ -131,7 +140,8 @@ public class TotalTokensPanel extends JPanel {
Conversation conversation,
List<ReferencedFile> includedFiles,
@Nullable String highlightedText) {
var tokenDetails = new TotalTokensDetails(encodingManager);
var tokenDetails = new TotalTokensDetails(
encodingManager.countTokens(PersonaSettings.getSystemPrompt()));
tokenDetails.setConversationTokens(encodingManager.countConversationTokens(conversation));
if (includedFiles != null) {
tokenDetails.setReferencedFilesTokens(includedFiles.stream()

View file

@ -14,6 +14,6 @@ class AddSelectionToContextAction : BaseEditorAction(AllIcons.General.Add) {
val chatTabPanel = chatToolWindowContentManager
.tryFindActiveChatTabPanel()
.orElseThrow()
chatTabPanel.addSelection((editor as EditorEx).virtualFile.name, editor.selectionModel)
chatTabPanel.addSelection((editor as EditorEx).virtualFile, editor.selectionModel)
}
}

View file

@ -3,6 +3,9 @@ package ee.carlrobert.codegpt.completions
import ee.carlrobert.codegpt.ReferencedFile
import ee.carlrobert.codegpt.conversations.Conversation
import ee.carlrobert.codegpt.conversations.message.Message
import ee.carlrobert.codegpt.util.file.FileUtil
import java.nio.file.Files
import java.nio.file.Path
import java.util.*
interface CompletionParameters
@ -12,10 +15,8 @@ class ChatCompletionParameters private constructor(
val conversationType: ConversationType,
val message: Message,
var sessionId: UUID?,
var highlightedText: String?,
var retry: Boolean,
var imageMediaType: String?,
var imageData: ByteArray?,
var imageDetails: ImageDetails?,
var referencedFiles: List<ReferencedFile>?
) : CompletionParameters {
@ -23,10 +24,8 @@ class ChatCompletionParameters private constructor(
return Builder(conversation, message).apply {
sessionId(this@ChatCompletionParameters.sessionId)
conversationType(this@ChatCompletionParameters.conversationType)
highlightedText(this@ChatCompletionParameters.highlightedText)
retry(this@ChatCompletionParameters.retry)
imageMediaType(this@ChatCompletionParameters.imageMediaType)
imageData(this@ChatCompletionParameters.imageData)
imageDetails(this@ChatCompletionParameters.imageDetails)
referencedFiles(this@ChatCompletionParameters.referencedFiles)
}
}
@ -34,22 +33,25 @@ class ChatCompletionParameters private constructor(
class Builder(private val conversation: Conversation, private val message: Message) {
private var sessionId: UUID? = null
private var conversationType: ConversationType = ConversationType.DEFAULT
private var highlightedText: String? = null
private var retry: Boolean = false
private var imageMediaType: String? = null
private var imageData: ByteArray? = null
private var imageDetails: ImageDetails? = null
private var referencedFiles: List<ReferencedFile>? = null
fun sessionId(sessionId: UUID?) = apply { this.sessionId = sessionId }
fun conversationType(conversationType: ConversationType) =
apply { this.conversationType = conversationType }
fun highlightedText(highlightedText: String?) =
apply { this.highlightedText = highlightedText }
fun retry(retry: Boolean) = apply { this.retry = retry }
fun imageMediaType(imageMediaType: String?) = apply { this.imageMediaType = imageMediaType }
fun imageData(imageData: ByteArray?) = apply { this.imageData = imageData }
fun imageDetails(imageDetails: ImageDetails?) = apply { this.imageDetails = imageDetails }
fun imageDetailsFromPath(path: String?) = apply {
if (!path.isNullOrEmpty()) {
this.imageDetails = ImageDetails(
FileUtil.getImageMediaType(path),
Files.readAllBytes(Path.of(path))
)
}
}
fun referencedFiles(referencedFiles: List<ReferencedFile>?) =
apply { this.referencedFiles = referencedFiles }
@ -59,10 +61,8 @@ class ChatCompletionParameters private constructor(
conversationType,
message,
sessionId,
highlightedText,
retry,
imageMediaType,
imageData,
imageDetails,
referencedFiles
)
}
@ -84,4 +84,25 @@ data class LookupCompletionParameters(val prompt: String) : CompletionParameters
data class EditCodeCompletionParameters(
val prompt: String,
val selectedText: String
) : CompletionParameters
) : CompletionParameters
data class ImageDetails(
val mediaType: String,
val data: ByteArray
) {
override fun equals(other: Any?): Boolean {
if (this === other) return true
if (other !is ImageDetails) return false
if (mediaType != other.mediaType) return false
if (!data.contentEquals(other.data)) return false
return true
}
override fun hashCode(): Int {
var result = mediaType.hashCode()
result = 31 * result + data.contentHashCode()
return result
}
}

View file

@ -13,9 +13,7 @@ class AzureRequestFactory : BaseRequestFactory() {
override fun createChatRequest(params: ChatCompletionParameters): OpenAIChatCompletionRequest {
val configuration = service<ConfigurationSettings>().state
val requestBuilder: OpenAIChatCompletionRequest.Builder =
OpenAIChatCompletionRequest.Builder(
buildOpenAIMessages(null, params, params.referencedFiles)
)
OpenAIChatCompletionRequest.Builder(buildOpenAIMessages(null, params))
.setMaxTokens(configuration.maxTokens)
.setStream(true)
.setTemperature(configuration.temperature.toDouble())

View file

@ -28,13 +28,16 @@ class ClaudeRequestFactory : BaseRequestFactory() {
}
when {
params.imageMediaType != null && params.imageData != null -> {
params.imageDetails != null -> {
messages.add(
ClaudeCompletionDetailedMessage(
"user",
listOf(
ClaudeMessageImageContent(
ClaudeBase64Source(params.imageMediaType, params.imageData)
ClaudeBase64Source(
params.imageDetails!!.mediaType,
params.imageDetails!!.data
)
),
ClaudeMessageTextContent(params.message.prompt)
)

View file

@ -24,7 +24,7 @@ class CustomOpenAIRequestFactory : BaseRequestFactory() {
service<CustomServiceSettings>()
.state
.chatCompletionSettings,
OpenAIRequestFactory.buildOpenAIMessages(null, params, params.referencedFiles),
OpenAIRequestFactory.buildOpenAIMessages(null, params),
true,
getCredential(CredentialKey.CUSTOM_SERVICE_API_KEY)
)

View file

@ -2,8 +2,11 @@ package ee.carlrobert.codegpt.completions.factory
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.EncodingManager
import ee.carlrobert.codegpt.completions.*
import ee.carlrobert.codegpt.completions.BaseRequestFactory
import ee.carlrobert.codegpt.completions.ChatCompletionParameters
import ee.carlrobert.codegpt.completions.CompletionRequestUtil.FIX_COMPILE_ERRORS_SYSTEM_PROMPT
import ee.carlrobert.codegpt.completions.ConversationType
import ee.carlrobert.codegpt.completions.TotalUsageExceededException
import ee.carlrobert.codegpt.conversations.ConversationsState
import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings
import ee.carlrobert.codegpt.settings.persona.PersonaSettings
@ -142,15 +145,15 @@ class GoogleRequestFactory : BaseRequestFactory() {
messages.add(GoogleCompletionContent("model", listOf(prevMessage.response)))
}
if (params.imageMediaType != null && params.imageData != null) {
if (params.imageDetails != null) {
messages.add(
GoogleCompletionContent(
listOf(
GoogleContentPart(
null,
GoogleContentPart.Blob(
params.imageMediaType,
params.imageData
params.imageDetails!!.mediaType,
params.imageDetails!!.data
)
),
GoogleContentPart(message.prompt)

View file

@ -15,9 +15,7 @@ class OllamaRequestFactory : BaseRequestFactory() {
val model = service<OllamaSettings>().state.model
val configuration = service<ConfigurationSettings>().state
val requestBuilder: OpenAIChatCompletionRequest.Builder =
OpenAIChatCompletionRequest.Builder(
buildOpenAIMessages(model, params, params.referencedFiles)
)
OpenAIChatCompletionRequest.Builder(buildOpenAIMessages(model, params))
.setModel(model)
.setMaxTokens(configuration.maxTokens)
.setStream(true)

View file

@ -2,7 +2,6 @@ package ee.carlrobert.codegpt.completions.factory
import com.intellij.openapi.components.service
import ee.carlrobert.codegpt.EncodingManager
import ee.carlrobert.codegpt.ReferencedFile
import ee.carlrobert.codegpt.completions.*
import ee.carlrobert.codegpt.completions.CompletionRequestUtil.EDIT_CODE_SYSTEM_PROMPT
import ee.carlrobert.codegpt.completions.CompletionRequestUtil.FIX_COMPILE_ERRORS_SYSTEM_PROMPT
@ -25,9 +24,7 @@ class OpenAIRequestFactory : CompletionRequestFactory {
val model = service<OpenAISettings>().state.model
val configuration = service<ConfigurationSettings>().state
val requestBuilder: OpenAIChatCompletionRequest.Builder =
OpenAIChatCompletionRequest.Builder(
buildOpenAIMessages(model, params, params.referencedFiles)
)
OpenAIChatCompletionRequest.Builder(buildOpenAIMessages(model, params))
.setModel(model)
if ("o1-mini" == model || "o1-preview" == model) {
requestBuilder
@ -103,9 +100,8 @@ class OpenAIRequestFactory : CompletionRequestFactory {
fun buildOpenAIMessages(
model: String?,
callParameters: ChatCompletionParameters,
referencedFiles: List<ReferencedFile>? = mutableListOf()
): List<OpenAIChatCompletionMessage> {
val messages = buildOpenAIChatMessages(model, callParameters, referencedFiles)
val messages = buildOpenAIChatMessages(model, callParameters)
if (model == null) {
return messages
@ -140,7 +136,6 @@ class OpenAIRequestFactory : CompletionRequestFactory {
private fun buildOpenAIChatMessages(
model: String?,
callParameters: ChatCompletionParameters,
referencedFiles: List<ReferencedFile>? = mutableListOf()
): MutableList<OpenAIChatCompletionMessage> {
val message = callParameters.message
val messages = mutableListOf<OpenAIChatCompletionMessage>()
@ -202,15 +197,15 @@ class OpenAIRequestFactory : CompletionRequestFactory {
)
}
if (callParameters.imageMediaType != null && callParameters.imageData != null) {
if (callParameters.imageDetails != null) {
messages.add(
OpenAIChatCompletionDetailedMessage(
"user",
listOf(
OpenAIMessageImageURLContent(
OpenAIImageUrl(
callParameters.imageMediaType,
callParameters.imageData
callParameters.imageDetails!!.mediaType,
callParameters.imageDetails!!.data
)
),
OpenAIMessageTextContent(message.prompt)
@ -218,10 +213,13 @@ class OpenAIRequestFactory : CompletionRequestFactory {
)
)
} else {
val prompt = if (referencedFiles.isNullOrEmpty()) {
val prompt = if (callParameters.referencedFiles.isNullOrEmpty()) {
message.prompt
} else {
CompletionRequestUtil.getPromptWithContext(referencedFiles, message.prompt)
CompletionRequestUtil.getPromptWithContext(
callParameters.referencedFiles!!,
message.prompt
)
}
messages.add(OpenAIChatCompletionStandardMessage("user", prompt))
}

View file

@ -0,0 +1,80 @@
package ee.carlrobert.codegpt.toolwindow.chat
import com.intellij.openapi.editor.Editor
import com.intellij.openapi.editor.ex.EditorEx
import com.intellij.openapi.project.Project
import ee.carlrobert.codegpt.ReferencedFile
import ee.carlrobert.codegpt.conversations.message.Message
import ee.carlrobert.codegpt.toolwindow.chat.actionprocessor.ActionProcessorFactory
import ee.carlrobert.codegpt.ui.textarea.AppliedActionInlay
import ee.carlrobert.codegpt.util.EditorUtil.getSelectedEditor
class MessageBuilder(private val project: Project, private val text: String) {
private val message = Message("")
private var editorContent: String = ""
private var inlayContent: String = ""
fun withSelectedEditorContent(): MessageBuilder {
getSelectedEditor(project)?.let { editor ->
editorContent = processEditorSelectedText(editor)
}
return this
}
fun withInlays(inlays: List<AppliedActionInlay>): MessageBuilder {
if (inlays.isNotEmpty()) {
inlayContent = processInlays(message, inlays)
}
return this
}
fun withReferencedFiles(referencedFiles: List<ReferencedFile>): MessageBuilder {
if (referencedFiles.isNotEmpty()) {
message.referencedFilePaths = referencedFiles.map { it.filePath }
}
return this
}
fun withImage(attachedImagePath: String): MessageBuilder {
message.imageFilePath = attachedImagePath
return this
}
fun build(): Message {
message.prompt = buildString {
append(text)
if (editorContent.isNotBlank()) {
append("\n\n")
append(editorContent)
}
if (inlayContent.isNotBlank()) {
append("\n")
append(inlayContent)
}
}.trim()
return message
}
private fun processEditorSelectedText(editor: Editor): String {
return editor.selectionModel.selectedText?.let { selectedText ->
if (selectedText.isBlank()) return ""
val fileExtension = (editor as EditorEx).virtualFile?.name?.substringAfterLast('.', "") ?: ""
editor.selectionModel.removeSelection()
"```$fileExtension\n$selectedText\n```"
} ?: ""
}
private fun processInlays(
message: Message,
inlays: List<AppliedActionInlay>
): String = buildString {
inlays
.sortedBy { it.inlay.offset }
.forEach { actionInlay ->
ActionProcessorFactory.getProcessor(project, actionInlay)
.process(message, actionInlay, this)
}
}
}

View file

@ -12,6 +12,7 @@ import com.intellij.openapi.editor.ex.EditorEx
import com.intellij.openapi.fileTypes.FileTypes
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.TextRange
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.openapi.wm.ToolWindowManager
import com.intellij.ui.ComponentUtil.findParentByCondition
import com.intellij.ui.EditorTextField
@ -41,7 +42,8 @@ data class AppliedSuggestionActionInlay(
data class AppliedCodeActionInlay(
override val inlay: Inlay<PromptTextFieldInlayRenderer?>,
val code: String
val code: String,
val editorFile: VirtualFile
) : AppliedActionInlay
const val AT_CHAR = '@'
@ -85,7 +87,7 @@ class PromptTextField(
fun addInlayElement(
actionPrefix: String,
text: String,
fileName: String? = null,
editorFile: VirtualFile? = null,
tooltipText: String? = null
) {
editor?.let {
@ -93,7 +95,7 @@ class PromptTextField(
it.caretModel.offset,
actionPrefix,
text,
fileName = fileName,
editorFile = editorFile,
tooltipText = tooltipText
)
}
@ -104,7 +106,7 @@ class PromptTextField(
actionPrefix: String,
text: String?,
actionItem: SuggestionActionItem? = null,
fileName: String? = null,
editorFile: VirtualFile? = null,
tooltipText: String? = null
) {
runUndoTransparentWriteAction {
@ -117,7 +119,7 @@ class PromptTextField(
project,
actionPrefix,
text,
fileName ?: "",
editorFile?.name ?: "",
tooltipText
) { inlay ->
appliedInlays.removeIf { appliedInlay -> appliedInlay.inlay == inlay }
@ -125,10 +127,10 @@ class PromptTextField(
})
if (inlay != null) {
// TODO
if (tooltipText == null) {
appliedInlays.add(AppliedSuggestionActionInlay(inlay, actionItem))
if (tooltipText != null && editorFile != null) {
appliedInlays.add(AppliedCodeActionInlay(inlay, tooltipText, editorFile))
} else {
appliedInlays.add(AppliedCodeActionInlay(inlay, tooltipText))
appliedInlays.add(AppliedSuggestionActionInlay(inlay, actionItem))
}
editor?.caretModel?.moveToOffset(document.textLength)
}

View file

@ -10,6 +10,7 @@ import com.intellij.openapi.components.service
import com.intellij.openapi.editor.SelectionModel
import com.intellij.openapi.observable.properties.AtomicBooleanProperty
import com.intellij.openapi.project.Project
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.ui.components.AnActionLink
import com.intellij.ui.dsl.builder.AlignX
import com.intellij.ui.dsl.builder.RightGap
@ -209,11 +210,12 @@ class UserInputPanel(
}
}
fun addSelection(fileName: String, selectionModel: SelectionModel) {
fun addSelection(editorFile: VirtualFile, selectionModel: SelectionModel) {
val fileName = editorFile.name
promptTextField.addInlayElement(
"code",
"$fileName (${selectionModel.selectionStartPosition?.line}:${selectionModel.selectionEndPosition?.line})",
fileName = fileName,
editorFile = editorFile,
tooltipText = selectionModel.selectedText
)
promptTextField.requestFocusInWindow()

View file

@ -106,7 +106,6 @@ class ChatToolWindowTabPanelTest : IntegrationTest() {
useOpenAIService()
service<PersonaSettings>().state.selectedPersona.instructions = "TEST_SYSTEM_PROMPT"
val message = Message("TEST_MESSAGE")
message.userMessage = "TEST_MESSAGE"
message.referencedFilePaths =
listOf("TEST_FILE_PATH_1", "TEST_FILE_PATH_2", "TEST_FILE_PATH_3")
val conversation = ConversationService.getInstance().startConversation()
@ -202,7 +201,7 @@ class ChatToolWindowTabPanelTest : IntegrationTest() {
)
}
fun testSendingOpenAIMessageWithImage() {
fun testSendingOpenAIMessageWithImageInSession() {
val testImagePath =
Objects.requireNonNull(javaClass.getResource("/images/test-image.png")).path
project.putUserData(CodeGPTKeys.IMAGE_ATTACHMENT_FILE_PATH, testImagePath)
@ -302,7 +301,6 @@ class ChatToolWindowTabPanelTest : IntegrationTest() {
useOpenAIService()
service<PersonaSettings>().state.selectedPersona.instructions = "TEST_SYSTEM_PROMPT"
val message = Message("TEST_MESSAGE")
message.userMessage = "TEST_MESSAGE"
message.referencedFilePaths =
listOf("TEST_FILE_PATH_1", "TEST_FILE_PATH_2", "TEST_FILE_PATH_3")
val conversation = ConversationService.getInstance().startConversation()