feat: improve chat tags UX

This commit is contained in:
Carl-Robert Linnupuu 2025-03-19 23:47:33 +00:00
parent 7984d5211c
commit e60640f97f
11 changed files with 106 additions and 82 deletions

View file

@ -26,6 +26,7 @@ import ee.carlrobert.codegpt.CodeGPTBundle;
import ee.carlrobert.codegpt.EncodingManager;
import ee.carlrobert.codegpt.Icons;
import ee.carlrobert.codegpt.settings.IncludedFilesSettings;
import ee.carlrobert.codegpt.toolwindow.chat.ChatToolWindowContentManager;
import ee.carlrobert.codegpt.ui.UIUtil;
import ee.carlrobert.codegpt.ui.checkbox.FileCheckboxTree;
import ee.carlrobert.codegpt.ui.checkbox.VirtualFileCheckboxTree;
@ -80,9 +81,12 @@ public class IncludeFilesInContextAction extends AnAction {
totalTokensLabel,
checkboxTree);
if (show == OK_EXIT_CODE) {
project.getMessageBus()
.syncPublisher(IncludeFilesInContextNotifier.FILES_INCLUDED_IN_CONTEXT_TOPIC)
.filesIncluded(checkboxTree.getReferencedFiles());
project.getService(ChatToolWindowContentManager.class)
.tryFindActiveChatTabPanel()
.ifPresent(tabPanel -> {
tabPanel.includeFiles(checkboxTree.getReferencedFiles());
});
includedFilesSettings.setPromptTemplate(promptTemplateTextArea.getText());
includedFilesSettings.setRepeatableContext(repeatableContextTextArea.getText());
}

View file

@ -1,13 +0,0 @@
package ee.carlrobert.codegpt.actions;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.util.messages.Topic;
import java.util.List;
public interface IncludeFilesInContextNotifier {
Topic<IncludeFilesInContextNotifier> FILES_INCLUDED_IN_CONTEXT_TOPIC =
Topic.create("filesIncludedInContext", IncludeFilesInContextNotifier.class);
void filesIncluded(List<VirtualFile> includedFiles);
}

View file

@ -96,7 +96,6 @@ public class ChatToolWindowTabPanel implements Disposable {
);
totalTokensPanel = new TotalTokensPanel(
project,
conversation,
EditorUtil.getSelectedEditorSelectedText(project),
this,
@ -230,7 +229,8 @@ public class ChatToolWindowTabPanel implements Disposable {
totalTokensPanel.updateConversationTokens(conversation);
if (callParameters.getReferencedFiles() != null) {
totalTokensPanel.updateReferencedFilesTokens(callParameters.getReferencedFiles());
totalTokensPanel.updateReferencedFilesTokens(
callParameters.getReferencedFiles().stream().map(ReferencedFile::fileContent).toList());
}
var userMessagePanel = createUserMessagePanel(message, callParameters);
@ -244,6 +244,12 @@ public class ChatToolWindowTabPanel implements Disposable {
});
}
public void includeFiles(List<VirtualFile> referencedFiles) {
userInputPanel.includeFiles(referencedFiles);
totalTokensPanel.updateReferencedFilesTokens(
referencedFiles.stream().map(it -> ReferencedFile.from(it).fileContent()).toList());
}
private boolean hasReferencedFilePaths(Message message) {
return message.getReferencedFilePaths() != null && !message.getReferencedFilePaths().isEmpty();
}
@ -280,8 +286,9 @@ public class ChatToolWindowTabPanel implements Disposable {
return panel;
}
private void reloadMessage(ChatCompletionParameters prevParameters,
UserMessagePanel userMessagePanel) {
private void reloadMessage(
ChatCompletionParameters prevParameters,
UserMessagePanel userMessagePanel) {
var prevMessage = prevParameters.getMessage();
ResponseMessagePanel responsePanel = null;
try {

View file

@ -21,10 +21,6 @@ public class TotalTokensDetails {
this.conversationTokens = conversationTokens;
}
public void setPsiTokens(int psiTokens) {
this.psiTokens = psiTokens;
}
public int getConversationTokens() {
return conversationTokens;
}
@ -53,6 +49,10 @@ public class TotalTokensDetails {
return referencedFilesTokens;
}
public void setPsiTokens(int psiTokens) {
this.psiTokens = psiTokens;
}
public int getPsiTokens() {
return psiTokens;
}

View file

@ -12,14 +12,14 @@ 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.project.Project;
import com.intellij.openapi.util.text.Strings;
import com.intellij.ui.components.JBLabel;
import com.intellij.util.ui.JBUI;
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.psistructure.ClassStructureSerializer;
import ee.carlrobert.codegpt.settings.GeneralSettings;
import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings;
import ee.carlrobert.codegpt.settings.prompts.PromptsSettings;
import ee.carlrobert.codegpt.settings.service.ServiceType;
import ee.carlrobert.codegpt.toolwindow.chat.structure.data.PsiStructureRepository;
@ -44,7 +44,6 @@ public class TotalTokensPanel extends JPanel {
private final JBLabel label;
public TotalTokensPanel(
@NotNull Project project,
Conversation conversation,
@Nullable String highlightedText,
Disposable parentDisposable,
@ -61,7 +60,10 @@ public class TotalTokensPanel extends JPanel {
new CoroutineDispatchers(),
psiStructureRepository,
psiTokens -> {
updatePsiTokenCount(psiTokens);
if (ConfigurationSettings.getState().getChatCompletionSettings()
.getPsiStructureEnabled()) {
updatePsiTokenCount(psiTokens);
}
return Unit.INSTANCE;
}
);
@ -72,13 +74,6 @@ public class TotalTokensPanel extends JPanel {
add(Box.createHorizontalStrut(4));
add(label);
addSelectionListeners(parentDisposable);
project.getMessageBus()
.connect()
.subscribe(IncludeFilesInContextNotifier.FILES_INCLUDED_IN_CONTEXT_TOPIC,
(IncludeFilesInContextNotifier) includedFiles ->
updateReferencedFilesTokens(
includedFiles.stream().map(ReferencedFile::from).toList()));
}
private void addSelectionListeners(Disposable parentDisposable) {
@ -142,10 +137,9 @@ public class TotalTokensPanel extends JPanel {
update();
}
public void updateReferencedFilesTokens(List<ReferencedFile> includedFiles) {
totalTokensDetails.setReferencedFilesTokens(includedFiles.stream()
.mapToInt(file -> encodingManager.countTokens(file.fileContent()))
.sum());
public void updateReferencedFilesTokens(List<String> includedFileContents) {
totalTokensDetails.setReferencedFilesTokens(
encodingManager.countTokens(Strings.join(includedFileContents, "\n")));
update();
}
@ -172,7 +166,7 @@ public class TotalTokensPanel extends JPanel {
"Input Tokens", totalTokensDetails.getUserPromptTokens(),
"Highlighted Tokens", totalTokensDetails.getHighlightedTokens(),
"Referenced Files Tokens", totalTokensDetails.getReferencedFilesTokens(),
"Dependency structure Tokens", totalTokensDetails.getPsiTokens()))
"Dependency Structure Tokens", totalTokensDetails.getPsiTokens()))
.entrySet().stream()
.map(entry -> format(
"<p style=\"margin: 0; padding: 0;\"><small>%s: <strong>%d</strong></small></p>",

View file

@ -2,6 +2,7 @@ package ee.carlrobert.codegpt.ui.textarea
import ee.carlrobert.codegpt.ui.textarea.header.tag.EditorSelectionTagDetails
import ee.carlrobert.codegpt.ui.textarea.header.tag.EditorTagDetails
import ee.carlrobert.codegpt.ui.textarea.header.tag.FileTagDetails
import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails
internal class TagDetailsComparator : Comparator<TagDetails> {
@ -9,14 +10,27 @@ internal class TagDetailsComparator : Comparator<TagDetails> {
val priority1 = getPriority(o1)
val priority2 = getPriority(o2)
return priority1.compareTo(priority2)
if (priority1 != priority2) {
return priority1.compareTo(priority2)
}
if (priority1 == 2) {
return o2.createdOn.compareTo(o1.createdOn)
}
return 0
}
private fun getPriority(tag: TagDetails): Int {
return when (tag) {
is EditorSelectionTagDetails -> 0
is EditorTagDetails -> 1
else -> 2
is EditorTagDetails -> {
if (tag.selected) 1 else 2
}
is FileTagDetails -> {
if (tag.selected) 3 else 4
}
else -> 5
}
}
}
}

View file

@ -34,6 +34,7 @@ import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.ModelComboBoxAction
import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.TotalTokensPanel
import ee.carlrobert.codegpt.ui.IconActionButton
import ee.carlrobert.codegpt.ui.textarea.header.UserInputHeaderPanel
import ee.carlrobert.codegpt.ui.textarea.header.tag.FileTagDetails
import ee.carlrobert.codegpt.ui.textarea.header.tag.GitCommitTagDetails
import ee.carlrobert.codegpt.ui.textarea.header.tag.SelectionTagDetails
import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails
@ -72,7 +73,13 @@ class UserInputPanel(
private val promptTextField =
PromptTextField(project, suggestionsPopupManager, ::updateUserTokens, ::handleSubmit)
private val userInputHeaderPanel =
UserInputHeaderPanel(project, tagManager, suggestionsPopupManager, promptTextField)
UserInputHeaderPanel(
project,
tagManager,
totalTokensPanel,
suggestionsPopupManager,
promptTextField
)
private val submitButton = IconActionButton(
object : AnAction(
CodeGPTBundle.get("smartTextPane.submitButton.title"),
@ -153,6 +160,10 @@ class UserInputPanel(
}
}
fun includeFiles(referencedFiles: MutableList<VirtualFile>) {
referencedFiles.forEach { userInputHeaderPanel.addTag(FileTagDetails(it)) }
}
override fun requestFocus() {
invokeLater {
promptTextField.requestFocusInWindow()

View file

@ -10,12 +10,13 @@ import com.intellij.openapi.project.Project
import com.intellij.openapi.ui.JBMenuItem
import com.intellij.openapi.ui.JBPopupMenu
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.openapi.vfs.readText
import com.intellij.ui.components.JBLabel
import com.intellij.util.IconUtil
import com.intellij.util.ui.JBUI
import ee.carlrobert.codegpt.CodeGPTBundle
import ee.carlrobert.codegpt.EditorNotifier
import ee.carlrobert.codegpt.actions.IncludeFilesInContextNotifier
import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.TotalTokensPanel
import ee.carlrobert.codegpt.ui.WrapLayout
import ee.carlrobert.codegpt.ui.textarea.PromptTextField
import ee.carlrobert.codegpt.ui.textarea.TagDetailsComparator
@ -31,6 +32,7 @@ import javax.swing.JPanel
class UserInputHeaderPanel(
private val project: Project,
private val tagManager: TagManager,
private val totalTokensPanel: TotalTokensPanel,
suggestionsPopupManager: SuggestionsPopupManager,
private val promptTextField: PromptTextField
) : JPanel(WrapLayout(FlowLayout.LEFT, 4, 4)), TagManagerListener {
@ -111,7 +113,8 @@ class UserInputHeaderPanel(
.sortedWith(TagDetailsComparator())
.toSet()
emptyText.isVisible = tags.none { it.selected }
updateReferencedFilesTokens(tags)
emptyText.isVisible = tags.isEmpty()
tags.forEach { add(createTagPanel(it)) }
@ -157,23 +160,32 @@ class UserInputHeaderPanel(
}
EditorUtil.getOpenLocalFiles(project)
.map { EditorTagDetails(it) }
.filterNot { it.virtualFile == selectedFile }
.filterNot { it == selectedFile }
.take(INITIAL_VISIBLE_FILES)
.forEach {
tagManager.addTag(it.apply { selected = false })
tagManager.addTag(EditorTagDetails(it).apply { selected = false })
}
}
private fun updateReferencedFilesTokens(tags: Set<TagDetails>) {
val referencedFileContents = tags.asSequence()
.filter { it.selected }
.mapNotNull { tag ->
when (tag) {
is FileTagDetails -> tag.virtualFile.readText()
is EditorTagDetails -> tag.virtualFile.readText()
else -> null
}
}
.toList()
totalTokensPanel.updateReferencedFilesTokens(referencedFileContents)
}
private fun initializeEventListeners() {
project.messageBus.connect().apply {
subscribe(EditorNotifier.SelectionChange.TOPIC, EditorSelectionChangeListener())
subscribe(EditorNotifier.Released.TOPIC, EditorReleasedListener())
subscribe(FileEditorManagerListener.FILE_EDITOR_MANAGER, FileSelectionListener())
subscribe(
IncludeFilesInContextNotifier.FILES_INCLUDED_IN_CONTEXT_TOPIC,
IncludedFilesListener()
)
}
}
@ -229,19 +241,17 @@ class UserInputHeaderPanel(
private inner class FileSelectionListener : FileEditorManagerListener {
override fun selectionChanged(event: FileEditorManagerEvent) {
event.newFile?.let { newFile ->
val editorTagDetails = EditorTagDetails(newFile)
tagManager.addTag(editorTagDetails)
val containsTag = tagManager.getTags()
.none { it is EditorTagDetails && it.virtualFile == newFile }
if (containsTag) {
tagManager.addTag(EditorTagDetails(newFile).apply { selected = false })
}
emptyText.isVisible = false
}
}
}
private inner class IncludedFilesListener : IncludeFilesInContextNotifier {
override fun filesIncluded(includedFiles: MutableList<VirtualFile>) {
includedFiles.forEach { tagManager.addTag(FileTagDetails(it)) }
}
}
private inner class TagPopupMenu : JBPopupMenu() {
private val closeMenuItem =
createPopupMenuItem(CodeGPTBundle.get("tagPopupMenuItem.close")) {

View file

@ -13,7 +13,8 @@ import javax.swing.Icon
sealed class TagDetails(
val name: String,
val icon: Icon? = null,
val id: UUID = UUID.randomUUID()
val id: UUID = UUID.randomUUID(),
val createdOn: Long = System.currentTimeMillis()
) {
var selected: Boolean = true

View file

@ -45,6 +45,10 @@ class TagManager(parentDisposable: Disposable) {
tags.remove(tagDetails)
}
if (tags.count { !it.selected } == 2) {
tags.remove(tags.sortedBy { it.createdOn }.first { !it.selected })
}
tags.add(tagDetails)
}
if (wasAdded) {

View file

@ -4,8 +4,6 @@ import com.intellij.openapi.components.service
import com.intellij.testFramework.LightVirtualFile
import ee.carlrobert.codegpt.CodeGPTKeys
import ee.carlrobert.codegpt.EncodingManager
import ee.carlrobert.codegpt.actions.IncludeFilesInContextNotifier
import ee.carlrobert.codegpt.actions.IncludeFilesInContextNotifier.FILES_INCLUDED_IN_CONTEXT_TOPIC
import ee.carlrobert.codegpt.completions.ConversationType
import ee.carlrobert.codegpt.completions.HuggingFaceModel
import ee.carlrobert.codegpt.completions.llama.PromptTemplate.LLAMA
@ -106,15 +104,11 @@ class ChatToolWindowTabPanelTest : IntegrationTest() {
listOf("TEST_FILE_PATH_1", "TEST_FILE_PATH_2", "TEST_FILE_PATH_3")
val conversation = ConversationService.getInstance().startConversation()
val panel = ChatToolWindowTabPanel(project, conversation)
project.messageBus
.syncPublisher<IncludeFilesInContextNotifier>(FILES_INCLUDED_IN_CONTEXT_TOPIC)
.filesIncluded(
listOf(
LightVirtualFile("TEST_FILE_NAME_1", "TEST_FILE_CONTENT_1"),
LightVirtualFile("TEST_FILE_NAME_2", "TEST_FILE_CONTENT_2"),
LightVirtualFile("TEST_FILE_NAME_3", "TEST_FILE_CONTENT_3"),
)
)
panel.includeFiles(listOf(
LightVirtualFile("TEST_FILE_NAME_1", "TEST_FILE_CONTENT_1"),
LightVirtualFile("TEST_FILE_NAME_2", "TEST_FILE_CONTENT_2"),
LightVirtualFile("TEST_FILE_NAME_3", "TEST_FILE_CONTENT_3"),
))
expectOpenAI(StreamHttpExchange { request: RequestEntity ->
assertThat(request.uri.path).isEqualTo("/v1/chat/completions")
assertThat(request.method).isEqualTo("POST")
@ -306,15 +300,13 @@ class ChatToolWindowTabPanelTest : IntegrationTest() {
listOf("TEST_FILE_PATH_1", "TEST_FILE_PATH_2", "TEST_FILE_PATH_3")
val conversation = ConversationService.getInstance().startConversation()
val panel = ChatToolWindowTabPanel(project, conversation)
project.messageBus
.syncPublisher<IncludeFilesInContextNotifier>(FILES_INCLUDED_IN_CONTEXT_TOPIC)
.filesIncluded(
listOf(
LightVirtualFile("TEST_FILE_NAME_1", "TEST_FILE_CONTENT_1"),
LightVirtualFile("TEST_FILE_NAME_2", "TEST_FILE_CONTENT_2"),
LightVirtualFile("TEST_FILE_NAME_3", "TEST_FILE_CONTENT_3"),
)
panel.includeFiles(
listOf(
LightVirtualFile("TEST_FILE_NAME_1", "TEST_FILE_CONTENT_1"),
LightVirtualFile("TEST_FILE_NAME_2", "TEST_FILE_CONTENT_2"),
LightVirtualFile("TEST_FILE_NAME_3", "TEST_FILE_CONTENT_3"),
)
)
expectOpenAI(StreamHttpExchange { request: RequestEntity ->
assertThat(request.uri.path).isEqualTo("/v1/chat/completions")
assertThat(request.method).isEqualTo("POST")