diff --git a/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextAction.java b/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextAction.java index 1455a079..a0cbc8c5 100644 --- a/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextAction.java +++ b/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextAction.java @@ -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()); } diff --git a/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextNotifier.java b/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextNotifier.java deleted file mode 100644 index c2f3b2f3..00000000 --- a/src/main/java/ee/carlrobert/codegpt/actions/IncludeFilesInContextNotifier.java +++ /dev/null @@ -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 FILES_INCLUDED_IN_CONTEXT_TOPIC = - Topic.create("filesIncludedInContext", IncludeFilesInContextNotifier.class); - - void filesIncluded(List includedFiles); -} 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 5493afe7..33e29294 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java @@ -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 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 { diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensDetails.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensDetails.java index d8533c49..42c81c52 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensDetails.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensDetails.java @@ -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; } diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensPanel.java index 084fb925..dc35dcdb 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/TotalTokensPanel.java @@ -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 includedFiles) { - totalTokensDetails.setReferencedFilesTokens(includedFiles.stream() - .mapToInt(file -> encodingManager.countTokens(file.fileContent())) - .sum()); + public void updateReferencedFilesTokens(List 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( "

%s: %d

", diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagDetailsComparator.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagDetailsComparator.kt index 14000f28..4b31be22 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagDetailsComparator.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagDetailsComparator.kt @@ -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 { @@ -9,14 +10,27 @@ internal class TagDetailsComparator : Comparator { 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 } } -} \ No newline at end of file +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt index e07473f8..1af5a88e 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt @@ -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) { + referencedFiles.forEach { userInputHeaderPanel.addTag(FileTagDetails(it)) } + } + override fun requestFocus() { invokeLater { promptTextField.requestFocusInWindow() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/UserInputHeaderPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/UserInputHeaderPanel.kt index 684e662d..e8b3daa1 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/UserInputHeaderPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/UserInputHeaderPanel.kt @@ -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) { + 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) { - includedFiles.forEach { tagManager.addTag(FileTagDetails(it)) } - } - } - private inner class TagPopupMenu : JBPopupMenu() { private val closeMenuItem = createPopupMenuItem(CodeGPTBundle.get("tagPopupMenuItem.close")) { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagDetails.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagDetails.kt index ec9e5a28..764fbecf 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagDetails.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagDetails.kt @@ -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 diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagManager.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagManager.kt index d8b50f62..656071e9 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagManager.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagManager.kt @@ -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) { diff --git a/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanelTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanelTest.kt index c01dcbcf..02df23a8 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanelTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanelTest.kt @@ -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(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(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")