diff --git a/src/main/java/ee/carlrobert/codegpt/Icons.java b/src/main/java/ee/carlrobert/codegpt/Icons.java index 596b3958..f952887e 100644 --- a/src/main/java/ee/carlrobert/codegpt/Icons.java +++ b/src/main/java/ee/carlrobert/codegpt/Icons.java @@ -25,6 +25,7 @@ public final class Icons { public static final Icon You = IconLoader.getIcon("/icons/you.svg", Icons.class); public static final Icon Ollama = IconLoader.getIcon("/icons/ollama.svg", Icons.class); public static final Icon User = IconLoader.getIcon("/icons/user.svg", Icons.class); + public static final Icon Tree = IconLoader.getIcon("/icons/tree.svg", Icons.class); public static final Icon Lightning = IconLoader.getIcon("/icons/lightning.svg", Icons.class); public static final Icon LightningDisabled = IconLoader.getIcon("/icons/lightning.svg", Icons.class); diff --git a/src/main/java/ee/carlrobert/codegpt/ProjectCompilationStatusListener.java b/src/main/java/ee/carlrobert/codegpt/ProjectCompilationStatusListener.java index 695700e8..963afe80 100644 --- a/src/main/java/ee/carlrobert/codegpt/ProjectCompilationStatusListener.java +++ b/src/main/java/ee/carlrobert/codegpt/ProjectCompilationStatusListener.java @@ -20,6 +20,7 @@ import ee.carlrobert.codegpt.ui.OverlayUtil; import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import org.jetbrains.annotations.NotNull; @@ -69,7 +70,8 @@ public class ProjectCompilationStatusListener implements CompilationStatusListen .toList()); message.setPrompt(CompletionRequestUtil.getPromptWithContext( new ArrayList<>(errorMapping.keySet()), - prompt)); + prompt, + new HashSet<>())); return message; } 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 dee3c219..5493afe7 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java @@ -21,10 +21,14 @@ import ee.carlrobert.codegpt.completions.ToolwindowChatCompletionRequestHandler; import ee.carlrobert.codegpt.conversations.Conversation; import ee.carlrobert.codegpt.conversations.ConversationService; import ee.carlrobert.codegpt.conversations.message.Message; +import ee.carlrobert.codegpt.psistructure.PsiStructureProvider; +import ee.carlrobert.codegpt.psistructure.models.ClassStructure; import ee.carlrobert.codegpt.settings.GeneralSettings; import ee.carlrobert.codegpt.settings.service.ServiceType; import ee.carlrobert.codegpt.telemetry.TelemetryAction; import ee.carlrobert.codegpt.toolwindow.chat.editor.actions.CopyAction; +import ee.carlrobert.codegpt.toolwindow.chat.structure.data.PsiStructureRepository; +import ee.carlrobert.codegpt.toolwindow.chat.structure.data.PsiStructureState; import ee.carlrobert.codegpt.toolwindow.chat.ui.ChatMessageResponseBody; import ee.carlrobert.codegpt.toolwindow.chat.ui.ChatToolWindowScrollablePanel; import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.TotalTokensDetails; @@ -34,17 +38,24 @@ import ee.carlrobert.codegpt.toolwindow.ui.ResponseMessagePanel; import ee.carlrobert.codegpt.toolwindow.ui.UserMessagePanel; import ee.carlrobert.codegpt.ui.OverlayUtil; import ee.carlrobert.codegpt.ui.textarea.UserInputPanel; +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.GitCommitTagDetails; import ee.carlrobert.codegpt.ui.textarea.header.tag.PersonaTagDetails; import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails; +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagManager; import ee.carlrobert.codegpt.util.EditorUtil; +import ee.carlrobert.codegpt.util.coroutines.CoroutineDispatchers; import ee.carlrobert.codegpt.util.file.FileUtil; import git4idea.GitCommit; import java.awt.BorderLayout; +import java.util.HashSet; import java.util.List; +import java.util.Objects; import java.util.Optional; +import java.util.Set; import java.util.UUID; +import java.util.stream.Collectors; import javax.swing.JComponent; import javax.swing.JPanel; import kotlin.Unit; @@ -64,6 +75,8 @@ public class ChatToolWindowTabPanel implements Disposable { private final ConversationService conversationService; private final TotalTokensPanel totalTokensPanel; private final ChatToolWindowScrollablePanel toolWindowScrollablePanel; + private final PsiStructureRepository psiStructureRepository; + private final TagManager tagManager; private @Nullable ToolwindowChatCompletionRequestHandler requestHandler; @@ -73,16 +86,27 @@ public class ChatToolWindowTabPanel implements Disposable { this.chatSession = new ChatSession(); conversationService = ConversationService.getInstance(); toolWindowScrollablePanel = new ChatToolWindowScrollablePanel(); + tagManager = new TagManager(this); + this.psiStructureRepository = new PsiStructureRepository( + this, + project, + tagManager, + new PsiStructureProvider(), + new CoroutineDispatchers() + ); + totalTokensPanel = new TotalTokensPanel( project, conversation, EditorUtil.getSelectedEditorSelectedText(project), - this); + this, + psiStructureRepository); userInputPanel = new UserInputPanel( project, conversation, totalTokensPanel, this, + tagManager, this::handleSubmit, this::handleCancel); userInputPanel.requestFocus(); @@ -134,13 +158,19 @@ public class ChatToolWindowTabPanel implements Disposable { private ChatCompletionParameters getCallParameters( Message message, - ConversationType conversationType) { - var selectedTags = userInputPanel.getSelectedTags(); + ConversationType conversationType, + Set psiStructure + ) { + final var selectedTags = tagManager.getTags().stream() + .filter(TagDetails::getSelected) + .collect(Collectors.toList()); + var builder = ChatCompletionParameters.builder(conversation, message) .sessionId(chatSession.getId()) .conversationType(conversationType) .imageDetailsFromPath(CodeGPTKeys.IMAGE_ATTACHMENT_FILE_PATH.get(project)) - .referencedFiles(getReferencedFiles(selectedTags)); + .referencedFiles(getReferencedFiles(selectedTags)) + .psiStructure(psiStructure); findTagOfType(selectedTags, PersonaTagDetails.class) .ifPresent(tag -> builder.personaDetails(tag.getPersonaDetails())); @@ -151,17 +181,27 @@ public class ChatToolWindowTabPanel implements Disposable { return builder.build(); } - private List getReferencedFiles() { - return getReferencedFiles(userInputPanel.getSelectedTags()); - } - private List getReferencedFiles(List tags) { return tags.stream() - .filter(FileTagDetails.class::isInstance) - .map(it -> ReferencedFile.from(((FileTagDetails) it).getVirtualFile())) + .map(this::getVirtualFile) + .filter(Objects::nonNull) + .map(ReferencedFile::from) .toList(); } + private VirtualFile getVirtualFile(TagDetails tag) { + VirtualFile virtualFile = null; + if (tag.getSelected()) { + if (tag instanceof FileTagDetails) { + virtualFile = ((FileTagDetails) tag).getVirtualFile(); + } else if (tag instanceof EditorTagDetails) { + virtualFile = ((EditorTagDetails) tag).getVirtualFile(); + } + + } + return virtualFile; + } + private Optional findTagOfType( List tags, Class tagClass) { @@ -172,8 +212,16 @@ public class ChatToolWindowTabPanel implements Disposable { } public void sendMessage(Message message, ConversationType conversationType) { + sendMessage(message, conversationType, new HashSet<>()); + } + + public void sendMessage( + Message message, + ConversationType conversationType, + Set psiStructure + ) { ApplicationManager.getApplication().invokeLater(() -> { - var callParameters = getCallParameters(message, conversationType); + var callParameters = getCallParameters(message, conversationType, psiStructure); if (callParameters.getImageDetails() != null) { project.getService(ChatToolWindowContentManager.class) .tryFindChatToolWindowPanel() @@ -233,7 +281,7 @@ public class ChatToolWindowTabPanel implements Disposable { } private void reloadMessage(ChatCompletionParameters prevParameters, - UserMessagePanel userMessagePanel) { + UserMessagePanel userMessagePanel) { var prevMessage = prevParameters.getMessage(); ResponseMessagePanel responsePanel = null; try { @@ -307,10 +355,22 @@ public class ChatToolWindowTabPanel implements Disposable { .executeOnPooledThread(() -> requestHandler.call(callParameters)); } - private Unit handleSubmit(String text, List appliedTags) { + private Unit handleSubmit(String text) { + final Set psiStructure; + if (psiStructureRepository.getStructureState().getValue() + instanceof PsiStructureState.Content content) { + psiStructure = content.getElements(); + } else { + psiStructure = new HashSet<>(); + } + + final var appliedTags = tagManager.getTags().stream() + .filter(TagDetails::getSelected) + .collect(Collectors.toList()); + var messageBuilder = new MessageBuilder(project, text).withInlays(appliedTags); - List referencedFiles = getReferencedFiles(); + List referencedFiles = getReferencedFiles(appliedTags); if (!referencedFiles.isEmpty()) { messageBuilder.withReferencedFiles(referencedFiles); } @@ -320,7 +380,7 @@ public class ChatToolWindowTabPanel implements Disposable { messageBuilder.withImage(attachedImagePath); } - sendMessage(messageBuilder.build(), ConversationType.DEFAULT); + sendMessage(messageBuilder.build(), ConversationType.DEFAULT, psiStructure); return Unit.INSTANCE; } @@ -336,6 +396,7 @@ public class ChatToolWindowTabPanel implements Disposable { panel.setBorder(JBUI.Borders.compound( JBUI.Borders.customLine(JBColor.border(), 1, 0, 0, 0), JBUI.Borders.empty(8))); + if (GeneralSettings.getSelectedService() != ServiceType.CODEGPT) { panel.add(JBUI.Panels.simplePanel(totalTokensPanel) .withBorder(JBUI.Borders.emptyBottom(8)), BorderLayout.NORTH); diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/structure/data/PsiStructureRepository.kt b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/structure/data/PsiStructureRepository.kt new file mode 100644 index 00000000..50a897f7 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/structure/data/PsiStructureRepository.kt @@ -0,0 +1,251 @@ +package ee.carlrobert.codegpt.toolwindow.chat.structure.data + +import com.intellij.openapi.Disposable +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.application.ReadAction +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.Disposer +import com.intellij.openapi.vfs.AsyncFileListener +import com.intellij.openapi.vfs.VirtualFile +import com.intellij.openapi.vfs.VirtualFileManager +import com.intellij.psi.PsiFile +import com.intellij.psi.PsiManager +import com.intellij.util.io.await +import ee.carlrobert.codegpt.psistructure.PsiStructureProvider +import ee.carlrobert.codegpt.settings.chat.ChatSettingsListener +import ee.carlrobert.codegpt.ui.textarea.header.tag.CurrentGitChangesTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.DocumentationTagDetails +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.EmptyTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.FileTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.FolderTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.GitCommitTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.PersonaTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.SelectionTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagManager +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagManagerListener +import ee.carlrobert.codegpt.ui.textarea.header.tag.WebTagDetails +import ee.carlrobert.codegpt.util.coroutines.CoroutineDispatchers +import ee.carlrobert.codegpt.util.coroutines.DisposableCoroutineScope +import kotlinx.coroutines.Job +import kotlinx.coroutines.asExecutor +import kotlinx.coroutines.currentCoroutineContext +import kotlinx.coroutines.ensureActive +import kotlinx.coroutines.flow.MutableStateFlow +import kotlinx.coroutines.flow.asStateFlow +import kotlinx.coroutines.sync.Mutex +import kotlinx.coroutines.sync.withLock + +class PsiStructureRepository( + parentDisposable: Disposable, + private val project: Project, + private val tagManager: TagManager, + private val psiStructureProvider: PsiStructureProvider, + private val dispatchers: CoroutineDispatchers, +) { + + private val mutex = Mutex() + private val coroutineScope = DisposableCoroutineScope(dispatchers.io()) + + @Volatile + private var updatePsiStructureJob: Job? = null + + private val tagsListener = object : TagManagerListener { + override fun onTagAdded(tag: TagDetails) { + updatePsiStructureIfNeeded() + } + + override fun onTagRemoved(tag: TagDetails) { + updatePsiStructureIfNeeded() + } + + override fun onTagSelectionChanged(tag: TagDetails) { + val tags = tagManager.getTags().getPsiAnalyzedTags() + update(tags) + } + + private fun updatePsiStructureIfNeeded() { + val tags = tagManager.getTags().getPsiAnalyzedTags() + if (isNeedUpdatePsiStructure(tags)) { + update(tags) + } + } + + private fun isNeedUpdatePsiStructure(tagsForAnalyze: Set): Boolean { + val currentlyAnalyzedTags = when (val currentState = _structureState.value) { + is PsiStructureState.Content -> currentState.currentlyAnalyzedTags + is PsiStructureState.UpdateInProgress -> currentState.currentlyAnalyzedTags + PsiStructureState.Disabled -> emptySet() + } + return tagsForAnalyze.toVirtualFilesSet() != currentlyAnalyzedTags.toVirtualFilesSet() + } + } + + private val asyncFileListener = AsyncFileListener { events -> + val currentlyAnalyzedTags = when (val currentState = _structureState.value) { + is PsiStructureState.Content -> currentState.currentlyAnalyzedTags + is PsiStructureState.UpdateInProgress -> currentState.currentlyAnalyzedTags + PsiStructureState.Disabled -> emptySet() + } + + val currentlyAnalyzedFiles = currentlyAnalyzedTags.toVirtualFilesSet() + + val hasRelevantChanges = events.any { event -> + event.file?.let { it in currentlyAnalyzedFiles } == true + } + + if (hasRelevantChanges) { + object : AsyncFileListener.ChangeApplier { + override fun afterVfsChange() { + update(currentlyAnalyzedTags) + } + } + } else { + null + } + } + + init { + Disposer.register(parentDisposable, coroutineScope) + tagManager.addListener(tagsListener) + Disposer.register(parentDisposable) { + tagManager.removeListener(tagsListener) + } + VirtualFileManager.getInstance().addAsyncFileListener(asyncFileListener, parentDisposable) + + val connection = ApplicationManager.getApplication().messageBus + .connect(parentDisposable) + + connection.subscribe( + ChatSettingsListener.TOPIC, + ChatSettingsListener { newState -> + if (newState.psiStructureEnabled) { + enable() + } else { + disable() + } + }) + } + + private val _structureState: MutableStateFlow = MutableStateFlow( + PsiStructureState.Content(emptySet(), emptySet()) + ) + val structureState = _structureState.asStateFlow() + + private fun disable() { + updatePsiStructureJob?.cancel() + _structureState.value = PsiStructureState.Disabled + } + + private fun enable() { + val tags = tagManager.getTags().getPsiAnalyzedTags() + _structureState.value = PsiStructureState.UpdateInProgress(tags) + update(tags) + } + + private fun update(tags: Set) { + updatePsiStructureJob?.cancel() + updatePsiStructureJob = coroutineScope.launch { + mutex.withLock { + if (_structureState.value == PsiStructureState.Disabled) return@launch + _structureState.value = PsiStructureState.UpdateInProgress(tags) + + val tagsVirtualFiles = tags.toVirtualFilesSet() + val coroutineContext = currentCoroutineContext() + + val psiFiles = ReadAction.nonBlocking> { + tagsVirtualFiles + .mapNotNull { virtualFile -> + coroutineContext.ensureActive() + PsiManager.getInstance(project).findFile(virtualFile) + } + } + .inSmartMode(project) + .submit(dispatchers.default().asExecutor()) + .await() + + val virtualFilesToRemoveFromStructure = tags.getExcludedVirtualFiles() + val result = psiStructureProvider.get(psiFiles) + .filter { classStructure -> + !virtualFilesToRemoveFromStructure.contains(classStructure.virtualFile) + } + .toSet() + + _structureState.value = PsiStructureState.Content(tags, result) + } + } + } + + private fun Set.getExcludedVirtualFiles(): Set = + mapNotNull { tagDetails -> + if (!tagDetails.selected) { + null + } else { + when (tagDetails) { + is SelectionTagDetails -> tagDetails.virtualFile + is FileTagDetails -> tagDetails.virtualFile + is EditorTagDetails -> tagDetails.virtualFile + + // Maybe need recursive find all files + is FolderTagDetails -> null + + is EditorSelectionTagDetails -> null + is DocumentationTagDetails -> null + is CurrentGitChangesTagDetails -> null + is GitCommitTagDetails -> null + is PersonaTagDetails -> null + is EmptyTagDetails -> null + is WebTagDetails -> null + } + } + } + .toSet() + + private fun Set.getPsiAnalyzedTags(): Set = + filter { tagDetails -> + when (tagDetails) { + is SelectionTagDetails -> tagDetails.selected + is FileTagDetails -> tagDetails.selected + is EditorSelectionTagDetails -> tagDetails.selected + is EditorTagDetails -> tagDetails.selected + + // Maybe need recursive find all files + is FolderTagDetails -> false + + is DocumentationTagDetails -> false + is CurrentGitChangesTagDetails -> false + is GitCommitTagDetails -> false + is PersonaTagDetails -> false + is EmptyTagDetails -> false + is WebTagDetails -> false + } + } + .toSet() + + private fun Set.toVirtualFilesSet(): Set = + mapNotNull { tagDetails -> + if (!tagDetails.selected) { + null + } else { + when (tagDetails) { + is SelectionTagDetails -> tagDetails.virtualFile + is FileTagDetails -> tagDetails.virtualFile + is EditorSelectionTagDetails -> tagDetails.virtualFile + is EditorTagDetails -> tagDetails.virtualFile + + // Maybe need recursive find all files + is FolderTagDetails -> null + + is DocumentationTagDetails -> null + is CurrentGitChangesTagDetails -> null + is GitCommitTagDetails -> null + is PersonaTagDetails -> null + is EmptyTagDetails -> null + is WebTagDetails -> null + } + } + } + .toSet() +} \ No newline at end of file diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/structure/data/PsiStructureState.kt b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/structure/data/PsiStructureState.kt new file mode 100644 index 00000000..1b74d1bf --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/structure/data/PsiStructureState.kt @@ -0,0 +1,18 @@ +package ee.carlrobert.codegpt.toolwindow.chat.structure.data + +import ee.carlrobert.codegpt.psistructure.models.ClassStructure +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails + +sealed class PsiStructureState { + + data class UpdateInProgress( + val currentlyAnalyzedTags: Set, + ) : PsiStructureState() + + data object Disabled : PsiStructureState() + + data class Content( + val currentlyAnalyzedTags: Set, + val elements: Set + ) : PsiStructureState() +} \ No newline at end of file diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/PsiStructureTotalTokenProvider.kt b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/PsiStructureTotalTokenProvider.kt new file mode 100644 index 00000000..aacd4225 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/PsiStructureTotalTokenProvider.kt @@ -0,0 +1,57 @@ +package ee.carlrobert.codegpt.toolwindow.chat.ui.textarea + +import com.intellij.openapi.Disposable +import com.intellij.openapi.util.Disposer +import ee.carlrobert.codegpt.EncodingManager +import ee.carlrobert.codegpt.psistructure.ClassStructureSerializer +import ee.carlrobert.codegpt.psistructure.models.ClassStructure +import ee.carlrobert.codegpt.toolwindow.chat.structure.data.PsiStructureRepository +import ee.carlrobert.codegpt.toolwindow.chat.structure.data.PsiStructureState +import ee.carlrobert.codegpt.util.coroutines.CoroutineDispatchers +import ee.carlrobert.codegpt.util.coroutines.DisposableCoroutineScope +import kotlinx.coroutines.flow.flowOn +import kotlinx.coroutines.flow.launchIn +import kotlinx.coroutines.flow.map +import kotlinx.coroutines.flow.onEach + +class PsiStructureTotalTokenProvider( + parentDisposable: Disposable, + private val classStructureSerializer: ClassStructureSerializer, + private val encodingManager: EncodingManager, + dispatchers: CoroutineDispatchers, + psiStructureRepository: PsiStructureRepository, + onPsiTokenHandled: (Int) -> Unit +) { + + private val coroutineScope = DisposableCoroutineScope() + + init { + Disposer.register(parentDisposable, coroutineScope) + psiStructureRepository.structureState + .map { structureState -> + when (structureState) { + is PsiStructureState.Content -> { + getPsiTokensCount(structureState.elements) + } + + PsiStructureState.Disabled -> 0 + + is PsiStructureState.UpdateInProgress -> 0 + } + } + .flowOn(dispatchers.io()) + .onEach { psiTokens -> + onPsiTokenHandled(psiTokens) + } + .launchIn(coroutineScope) + } + + private fun getPsiTokensCount(psiStructureSet: Set): Int = + psiStructureSet + .joinToString(separator = "\n\n") { psiStructure -> + classStructureSerializer.serialize(psiStructure) + } + .let { serializedPsiStructure -> + encodingManager.countTokens(serializedPsiStructure) + } +} \ No newline at end of file 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 c42b2368..d8533c49 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 @@ -7,6 +7,7 @@ public class TotalTokensDetails { private int userPromptTokens; private int highlightedTokens; private int referencedFilesTokens; + private int psiTokens; public TotalTokensDetails(int systemPromptTokens) { this.systemPromptTokens = systemPromptTokens; @@ -20,6 +21,10 @@ public class TotalTokensDetails { this.conversationTokens = conversationTokens; } + public void setPsiTokens(int psiTokens) { + this.psiTokens = psiTokens; + } + public int getConversationTokens() { return conversationTokens; } @@ -48,11 +53,16 @@ public class TotalTokensDetails { return referencedFilesTokens; } + public int getPsiTokens() { + return psiTokens; + } + public int getTotal() { return systemPromptTokens + conversationTokens + userPromptTokens + highlightedTokens - + referencedFilesTokens; + + referencedFilesTokens + + 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 c9a325d2..084fb925 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 @@ -18,9 +18,12 @@ 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.prompts.PromptsSettings; import ee.carlrobert.codegpt.settings.service.ServiceType; +import ee.carlrobert.codegpt.toolwindow.chat.structure.data.PsiStructureRepository; +import ee.carlrobert.codegpt.util.coroutines.CoroutineDispatchers; import java.awt.FlowLayout; import java.awt.event.MouseAdapter; import java.awt.event.MouseEvent; @@ -30,6 +33,7 @@ import java.util.Map; import java.util.stream.Collectors; import javax.swing.Box; import javax.swing.JPanel; +import kotlin.Unit; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -43,11 +47,25 @@ public class TotalTokensPanel extends JPanel { @NotNull Project project, Conversation conversation, @Nullable String highlightedText, - Disposable parentDisposable) { + Disposable parentDisposable, + PsiStructureRepository psiStructureRepository + ) { super(new FlowLayout(FlowLayout.LEADING, 0, 0)); this.totalTokensDetails = createTokenDetails(conversation, highlightedText); this.label = getLabel(totalTokensDetails); + new PsiStructureTotalTokenProvider( + parentDisposable, + ClassStructureSerializer.INSTANCE, + encodingManager, + new CoroutineDispatchers(), + psiStructureRepository, + psiTokens -> { + updatePsiTokenCount(psiTokens); + return Unit.INSTANCE; + } + ); + setBorder(JBUI.Borders.empty(4)); setOpaque(false); add(getContextHelpIcon(totalTokensDetails)); @@ -104,6 +122,11 @@ public class TotalTokensPanel extends JPanel { label.setText(getLabelHtml(total)); } + public void updatePsiTokenCount(int psiTokenCount) { + totalTokensDetails.setPsiTokens(psiTokenCount); + update(); + } + public void updateConversationTokens(Conversation conversation) { totalTokensDetails.setConversationTokens(encodingManager.countConversationTokens(conversation)); update(); @@ -148,7 +171,8 @@ public class TotalTokensPanel extends JPanel { "Conversation Tokens", totalTokensDetails.getConversationTokens(), "Input Tokens", totalTokensDetails.getUserPromptTokens(), "Highlighted Tokens", totalTokensDetails.getHighlightedTokens(), - "Referenced Files Tokens", totalTokensDetails.getReferencedFilesTokens())) + "Referenced Files Tokens", totalTokensDetails.getReferencedFilesTokens(), + "Dependency structure Tokens", totalTokensDetails.getPsiTokens())) .entrySet().stream() .map(entry -> format( "

%s: %d

", diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt index ae4eeab2..35864da4 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt @@ -1,6 +1,6 @@ package ee.carlrobert.codegpt.codecompletions -import ee.carlrobert.codegpt.codecompletions.psi.structure.ClassStructureSerializer +import ee.carlrobert.codegpt.psistructure.ClassStructureSerializer import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty enum class InfillPromptTemplate(val label: String, val stopTokens: List? = listOf("\n\n")) { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt index 320b1ade..9ed086a2 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt @@ -7,7 +7,7 @@ import com.intellij.psi.PsiElement import ee.carlrobert.codegpt.EncodingManager import ee.carlrobert.codegpt.codecompletions.psi.filePath import ee.carlrobert.codegpt.codecompletions.psi.readText -import ee.carlrobert.codegpt.codecompletions.psi.structure.models.ClassStructure +import ee.carlrobert.codegpt.psistructure.models.ClassStructure const val MAX_PROMPT_TOKENS = 256 diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt index 93782e3c..7c0d488b 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt @@ -7,7 +7,7 @@ import com.intellij.refactoring.suggested.startOffset import ee.carlrobert.codegpt.EncodingManager import ee.carlrobert.codegpt.codecompletions.psi.CompletionContextService import ee.carlrobert.codegpt.codecompletions.psi.readText -import ee.carlrobert.codegpt.codecompletions.psi.structure.PsiStructureProvider +import ee.carlrobert.codegpt.psistructure.PsiStructureProvider import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings import ee.carlrobert.codegpt.util.GitUtil diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/PsiStructureProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/PsiStructureProvider.kt deleted file mode 100644 index e4bf3136..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/PsiStructureProvider.kt +++ /dev/null @@ -1,36 +0,0 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure - -import com.intellij.openapi.application.ApplicationManager -import com.intellij.openapi.application.ReadAction -import com.intellij.psi.PsiFile -import ee.carlrobert.codegpt.codecompletions.psi.structure.models.ClassStructure -import org.jetbrains.kotlin.psi.KtFile - -class PsiStructureProvider { - - private val kotlinFileAnalyzerAvailable: Boolean = - ApplicationManager.getApplication().hasComponent(KotlinFileAnalyzer::class.java) - - fun get(psiFiles: List): Set = - ReadAction.compute, Throwable> { - val classStructureSet = mutableSetOf() - val processedPsiFiles = mutableSetOf() - val psiFileQueue = PsiFileQueue(psiFiles) - - while (true) { - val psiFile = psiFileQueue.pop() - when { - processedPsiFiles.contains(psiFile) -> Unit - - kotlinFileAnalyzerAvailable && psiFile is KtFile -> { - classStructureSet.addAll(KotlinFileAnalyzer(psiFileQueue, psiFile).analyze()) - processedPsiFiles.add(psiFile) - } - - psiFile == null -> break - } - } - - return@compute classStructureSet.toSet() - } -} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassLanguage.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassLanguage.kt deleted file mode 100644 index e8779496..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassLanguage.kt +++ /dev/null @@ -1,5 +0,0 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure.models - -enum class ClassLanguage { - KOTLIN, -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassName.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassName.kt deleted file mode 100644 index f5a7a5db..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassName.kt +++ /dev/null @@ -1,4 +0,0 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure.models - -@JvmInline -value class ClassName(val value: String) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/EnumEntryName.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/EnumEntryName.kt deleted file mode 100644 index 3221bc54..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/EnumEntryName.kt +++ /dev/null @@ -1,4 +0,0 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure.models - -@JvmInline -value class EnumEntryName(val value: String) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionParameters.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionParameters.kt index 97b0edad..dae658f9 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionParameters.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionParameters.kt @@ -3,6 +3,7 @@ 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.psistructure.models.ClassStructure import ee.carlrobert.codegpt.settings.prompts.PersonaDetails import ee.carlrobert.codegpt.util.file.FileUtil import java.nio.file.Files @@ -20,6 +21,7 @@ class ChatCompletionParameters private constructor( var imageDetails: ImageDetails?, var referencedFiles: List?, var personaDetails: PersonaDetails?, + var psiStructure: Set?, ) : CompletionParameters { fun toBuilder(): Builder { @@ -30,6 +32,7 @@ class ChatCompletionParameters private constructor( imageDetails(this@ChatCompletionParameters.imageDetails) referencedFiles(this@ChatCompletionParameters.referencedFiles) personaDetails(this@ChatCompletionParameters.personaDetails) + psiStructure(this@ChatCompletionParameters.psiStructure) } } @@ -40,6 +43,7 @@ class ChatCompletionParameters private constructor( private var imageDetails: ImageDetails? = null private var referencedFiles: List? = null private var personaDetails: PersonaDetails? = null + private var psiStructure: Set? = null private var gitDiff: String = "" fun sessionId(sessionId: UUID?) = apply { this.sessionId = sessionId } @@ -64,6 +68,8 @@ class ChatCompletionParameters private constructor( fun personaDetails(personaDetails: PersonaDetails?) = apply { this.personaDetails = personaDetails } + fun psiStructure(psiStructure: Set?) = apply { this.psiStructure = psiStructure } + fun build(): ChatCompletionParameters { return ChatCompletionParameters( conversation, @@ -73,7 +79,8 @@ class ChatCompletionParameters private constructor( retry, imageDetails, referencedFiles, - personaDetails + personaDetails, + psiStructure, ) } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt index 2bb12174..d5d5e26c 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt @@ -1,7 +1,15 @@ package ee.carlrobert.codegpt.completions import com.intellij.openapi.components.service -import ee.carlrobert.codegpt.completions.factory.* +import ee.carlrobert.codegpt.completions.factory.AzureRequestFactory +import ee.carlrobert.codegpt.completions.factory.ClaudeRequestFactory +import ee.carlrobert.codegpt.completions.factory.CodeGPTRequestFactory +import ee.carlrobert.codegpt.completions.factory.CustomOpenAIRequestFactory +import ee.carlrobert.codegpt.completions.factory.GoogleRequestFactory +import ee.carlrobert.codegpt.completions.factory.LlamaRequestFactory +import ee.carlrobert.codegpt.completions.factory.OllamaRequestFactory +import ee.carlrobert.codegpt.completions.factory.OpenAIRequestFactory +import ee.carlrobert.codegpt.psistructure.ClassStructureSerializer import ee.carlrobert.codegpt.settings.prompts.CoreActionsState import ee.carlrobert.codegpt.settings.prompts.PromptsSettings import ee.carlrobert.codegpt.settings.service.ServiceType @@ -17,7 +25,7 @@ interface CompletionRequestFactory { @JvmStatic fun getFactory(serviceType: ServiceType): CompletionRequestFactory { return when (serviceType) { - ServiceType.CODEGPT -> CodeGPTRequestFactory() + ServiceType.CODEGPT -> CodeGPTRequestFactory(ClassStructureSerializer) ServiceType.OPENAI -> OpenAIRequestFactory() ServiceType.CUSTOM_OPENAI -> CustomOpenAIRequestFactory() ServiceType.AZURE -> AzureRequestFactory() @@ -35,7 +43,8 @@ abstract class BaseRequestFactory : CompletionRequestFactory { val prompt = "Code to modify:\n${params.selectedText}\n\nInstructions: ${params.prompt}" return createBasicCompletionRequest( service().state.coreActions.editCode.instructions - ?: CoreActionsState.DEFAULT_EDIT_CODE_PROMPT, prompt, 8192, true) + ?: CoreActionsState.DEFAULT_EDIT_CODE_PROMPT, prompt, 8192, true + ) } override fun createCommitMessageRequest(params: CommitMessageCompletionParameters): CompletionRequest { @@ -65,7 +74,8 @@ abstract class BaseRequestFactory : CompletionRequestFactory { } else { CompletionRequestUtil.getPromptWithContext( it, - callParameters.message.prompt + callParameters.message.prompt, + callParameters.psiStructure, ) } } ?: return callParameters.message.prompt diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestUtil.kt index 7037ce8a..f77ffa91 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestUtil.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestUtil.kt @@ -2,19 +2,34 @@ package ee.carlrobert.codegpt.completions import com.intellij.openapi.components.service import ee.carlrobert.codegpt.ReferencedFile +import ee.carlrobert.codegpt.psistructure.ClassStructureSerializer +import ee.carlrobert.codegpt.psistructure.models.ClassStructure import ee.carlrobert.codegpt.settings.IncludedFilesSettings +import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty import java.util.stream.Collectors object CompletionRequestUtil { + + private val psiStructureSerializer = ClassStructureSerializer + + private val PSI_STRUCTURE_TITLE = """ + The following is the structure of the file dependencies that were attached above. + The structure contains a description of classes with their methods, method arguments, and return types. + If the type is specified as TypeUnknown, then the analyzer could not identify the type, + try to take it out of context, if necessary for the response. + """.trimIndent() + @JvmStatic fun getPromptWithContext( referencedFiles: List, - userPrompt: String? + userPrompt: String?, + psiStructure: Set? ): String { val includedFilesSettings = service().state - val repeatableContext = referencedFiles.stream() + val repeatableContext = includedFilesSettings.repeatableContext + val fileContext = referencedFiles.stream() .map { item: ReferencedFile -> - includedFilesSettings.repeatableContext + repeatableContext .replace("{FILE_PATH}", item.filePath()) .replace( "{FILE_CONTENT}", String.format( @@ -25,8 +40,25 @@ object CompletionRequestUtil { } .collect(Collectors.joining("\n\n")) + val structureContext = psiStructure + ?.map { structure: ClassStructure -> + val fileExtension = structure.virtualFile.extension ?: "" + repeatableContext + .replace("{FILE_PATH}", structure.virtualFile.path) + .replace( + "{FILE_CONTENT}", String.format( + "```%s%n%s%n```", + fileExtension, + psiStructureSerializer.serialize(structure) + ) + ) + } + ?.ifNotEmpty { + joinToString(prefix = "\n\n" + PSI_STRUCTURE_TITLE + "\n\n", separator = "\n\n") { it } + } + return includedFilesSettings.promptTemplate - .replace("{REPEATABLE_CONTEXT}", repeatableContext) + .replace("{REPEATABLE_CONTEXT}", fileContext + structureContext.orEmpty()) .replace("{QUESTION}", userPrompt!!) } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CodeGPTRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CodeGPTRequestFactory.kt index f56fa6c8..ea221622 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CodeGPTRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CodeGPTRequestFactory.kt @@ -6,12 +6,17 @@ import ee.carlrobert.codegpt.CodeGPTPlugin import ee.carlrobert.codegpt.completions.BaseRequestFactory import ee.carlrobert.codegpt.completions.ChatCompletionParameters import ee.carlrobert.codegpt.completions.factory.OpenAIRequestFactory.Companion.buildOpenAIMessages +import ee.carlrobert.codegpt.psistructure.ClassStructureSerializer import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings -import ee.carlrobert.llm.client.codegpt.request.chat.* +import ee.carlrobert.llm.client.codegpt.request.chat.AdditionalRequestContext +import ee.carlrobert.llm.client.codegpt.request.chat.ChatCompletionRequest +import ee.carlrobert.llm.client.codegpt.request.chat.ContextFile +import ee.carlrobert.llm.client.codegpt.request.chat.DocumentationDetails +import ee.carlrobert.llm.client.codegpt.request.chat.Metadata import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionStandardMessage -class CodeGPTRequestFactory : BaseRequestFactory() { +class CodeGPTRequestFactory(private val classStructureSerializer: ClassStructureSerializer) : BaseRequestFactory() { override fun createChatRequest(params: ChatCompletionParameters): ChatCompletionRequest { val model = service().state.chatCompletionSettings.model @@ -46,10 +51,18 @@ class CodeGPTRequestFactory : BaseRequestFactory() { DocumentationDetails(it.name, it.url) ) } - params.referencedFiles?.let { - requestBuilder.setContext(AdditionalRequestContext(it.map { file -> - ContextFile(file.fileName(), file.fileContent()) - })) + + val contextFiles = params.referencedFiles?.map { file -> + ContextFile(file.fileName(), file.fileContent()) + }.orEmpty() + + val psiContext = params.psiStructure?.map { classStructure -> + ContextFile(classStructure.virtualFile.name, classStructureSerializer.serialize(classStructure)) + }.orEmpty() + + val contextFilesWithPsi = contextFiles + psiContext + if (contextFilesWithPsi.isNotEmpty()) { + requestBuilder.setContext(AdditionalRequestContext(contextFilesWithPsi)) } return requestBuilder.build() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactory.kt index c7b7bafb..5813f9c9 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/CustomOpenAIRequestFactory.kt @@ -25,7 +25,7 @@ class CustomOpenAIRequestFactory : BaseRequestFactory() { .active val request = buildCustomOpenAIChatCompletionRequest( activeService.chatCompletionSettings, - OpenAIRequestFactory.buildOpenAIMessages(null, params), + OpenAIRequestFactory.buildOpenAIMessages(null, params, params.referencedFiles, params.psiStructure), true, getCredential(CredentialKey.CustomServiceApiKey(activeService.name.orEmpty())) ) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OllamaRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OllamaRequestFactory.kt index 1679cb03..840fd7da 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OllamaRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OllamaRequestFactory.kt @@ -15,7 +15,14 @@ class OllamaRequestFactory : BaseRequestFactory() { val model = service().state.model val configuration = service().state val requestBuilder: OpenAIChatCompletionRequest.Builder = - OpenAIChatCompletionRequest.Builder(buildOpenAIMessages(model, params)) + OpenAIChatCompletionRequest.Builder( + buildOpenAIMessages( + model = model, + callParameters = params, + referencedFiles = params.referencedFiles, + psiStructure = params.psiStructure, + ) + ) .setModel(model) .setMaxTokens(configuration.maxTokens) .setStream(true) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt index 6a2983df..775a4d66 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/factory/OpenAIRequestFactory.kt @@ -3,16 +3,33 @@ 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.ChatCompletionParameters +import ee.carlrobert.codegpt.completions.CommitMessageCompletionParameters +import ee.carlrobert.codegpt.completions.CompletionRequestFactory +import ee.carlrobert.codegpt.completions.CompletionRequestUtil +import ee.carlrobert.codegpt.completions.ConversationType +import ee.carlrobert.codegpt.completions.EditCodeCompletionParameters +import ee.carlrobert.codegpt.completions.LookupCompletionParameters +import ee.carlrobert.codegpt.completions.TotalUsageExceededException import ee.carlrobert.codegpt.conversations.ConversationsState +import ee.carlrobert.codegpt.psistructure.models.ClassStructure import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings.Companion.getState import ee.carlrobert.codegpt.settings.prompts.CoreActionsState import ee.carlrobert.codegpt.settings.prompts.PromptsSettings import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings import ee.carlrobert.codegpt.util.file.FileUtil.getImageMediaType -import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel.* -import ee.carlrobert.llm.client.openai.completion.request.* +import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel.O_1_MINI +import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel.O_1_PREVIEW +import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel.O_3_MINI +import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel.findByCode +import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionDetailedMessage +import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionMessage +import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionRequest +import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionStandardMessage +import ee.carlrobert.llm.client.openai.completion.request.OpenAIImageUrl +import ee.carlrobert.llm.client.openai.completion.request.OpenAIMessageImageURLContent +import ee.carlrobert.llm.client.openai.completion.request.OpenAIMessageTextContent import java.io.IOException import java.nio.file.Files import java.nio.file.Path @@ -110,12 +127,14 @@ class OpenAIRequestFactory : CompletionRequestFactory { fun buildOpenAIMessages( model: String?, callParameters: ChatCompletionParameters, - referencedFiles: List? = null + referencedFiles: List? = null, + psiStructure: Set? = null ): List { val messages = buildOpenAIChatMessages( - model, - callParameters, - referencedFiles ?: callParameters.referencedFiles + model = model, + callParameters = callParameters, + referencedFiles = referencedFiles ?: callParameters.referencedFiles, + psiStructure = psiStructure, ) if (model == null) { @@ -151,7 +170,8 @@ class OpenAIRequestFactory : CompletionRequestFactory { private fun buildOpenAIChatMessages( model: String?, callParameters: ChatCompletionParameters, - referencedFiles: List? = null + referencedFiles: List? = null, + psiStructure: Set? = null ): MutableList { val message = callParameters.message val messages = mutableListOf() @@ -254,7 +274,8 @@ class OpenAIRequestFactory : CompletionRequestFactory { } else { CompletionRequestUtil.getPromptWithContext( referencedFiles, - message.prompt + message.prompt, + psiStructure ) } messages.add(OpenAIChatCompletionStandardMessage("user", prompt)) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/ClassStructureSerializer.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/ClassStructureSerializer.kt similarity index 97% rename from src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/ClassStructureSerializer.kt rename to src/main/kotlin/ee/carlrobert/codegpt/psistructure/ClassStructureSerializer.kt index 765130fb..9f563dc0 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/ClassStructureSerializer.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/ClassStructureSerializer.kt @@ -1,6 +1,6 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure +package ee.carlrobert.codegpt.psistructure -import ee.carlrobert.codegpt.codecompletions.psi.structure.models.* +import ee.carlrobert.codegpt.psistructure.models.* import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty object ClassStructureSerializer { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/KotlinFileAnalyzer.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/KotlinFileAnalyzer.kt similarity index 90% rename from src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/KotlinFileAnalyzer.kt rename to src/main/kotlin/ee/carlrobert/codegpt/psistructure/KotlinFileAnalyzer.kt index 14abdd95..8a0e937e 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/KotlinFileAnalyzer.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/KotlinFileAnalyzer.kt @@ -1,4 +1,4 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure +package ee.carlrobert.codegpt.psistructure import com.intellij.openapi.roots.PackageIndex import com.intellij.psi.PsiElement @@ -7,10 +7,29 @@ import com.intellij.psi.PsiManager import com.intellij.psi.search.GlobalSearchScope import com.intellij.psi.search.PsiShortNamesCache import com.intellij.psi.util.PsiTreeUtil -import ee.carlrobert.codegpt.codecompletions.psi.structure.models.* +import ee.carlrobert.codegpt.psistructure.models.ClassName +import ee.carlrobert.codegpt.psistructure.models.ClassStructure +import ee.carlrobert.codegpt.psistructure.models.ClassType +import ee.carlrobert.codegpt.psistructure.models.ConstructorStructure +import ee.carlrobert.codegpt.psistructure.models.EnumEntryName +import ee.carlrobert.codegpt.psistructure.models.FieldStructure +import ee.carlrobert.codegpt.psistructure.models.MethodStructure +import ee.carlrobert.codegpt.psistructure.models.ParameterInfo import org.jetbrains.kotlin.asJava.classes.KtLightClass import org.jetbrains.kotlin.lexer.KtTokens -import org.jetbrains.kotlin.psi.* +import org.jetbrains.kotlin.psi.KtClass +import org.jetbrains.kotlin.psi.KtClassBody +import org.jetbrains.kotlin.psi.KtClassOrObject +import org.jetbrains.kotlin.psi.KtConstructor +import org.jetbrains.kotlin.psi.KtEnumEntry +import org.jetbrains.kotlin.psi.KtFile +import org.jetbrains.kotlin.psi.KtFunction +import org.jetbrains.kotlin.psi.KtModifierListOwner +import org.jetbrains.kotlin.psi.KtNamedFunction +import org.jetbrains.kotlin.psi.KtObjectDeclaration +import org.jetbrains.kotlin.psi.KtParameter +import org.jetbrains.kotlin.psi.KtProperty +import org.jetbrains.kotlin.psi.KtVariableDeclaration class KotlinFileAnalyzer( private val psiFileQueue: PsiFileQueue, @@ -112,6 +131,7 @@ class KotlinFileAnalyzer( modifierList = getModifiers(ktClass), packageName = ktClass.fqName?.parent()?.asString().orEmpty(), repositoryName = ktFile.project.name, + virtualFile = ktFile.virtualFile, ) analyzeSupertypes( diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/PsiFileQueue.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/PsiFileQueue.kt similarity index 82% rename from src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/PsiFileQueue.kt rename to src/main/kotlin/ee/carlrobert/codegpt/psistructure/PsiFileQueue.kt index 30763971..5609f545 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/PsiFileQueue.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/PsiFileQueue.kt @@ -1,4 +1,4 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure +package ee.carlrobert.codegpt.psistructure import com.intellij.psi.PsiFile diff --git a/src/main/kotlin/ee/carlrobert/codegpt/psistructure/PsiStructureProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/PsiStructureProvider.kt new file mode 100644 index 00000000..aa13e7cb --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/PsiStructureProvider.kt @@ -0,0 +1,76 @@ +package ee.carlrobert.codegpt.psistructure + +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.application.ReadAction +import com.intellij.psi.PsiFile +import com.intellij.util.io.await +import ee.carlrobert.codegpt.psistructure.models.ClassStructure +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.asExecutor +import kotlinx.coroutines.currentCoroutineContext +import kotlinx.coroutines.delay +import kotlinx.coroutines.ensureActive +import org.jetbrains.kotlin.psi.KtFile +import kotlin.coroutines.cancellation.CancellationException + +class PsiStructureProvider { + + suspend fun get(psiFiles: List): Set { + var result: Set? = null + var attempts = 0 + val maxAttempts = 5 + + val kotlinFileAnalyzerAvailable: Boolean = ApplicationManager + .getApplication() + .hasComponent(KotlinFileAnalyzer::class.java) + + while (result == null && attempts < maxAttempts) { + attempts++ + try { + val project = psiFiles + .map { it.project } + .firstOrNull { !it.isDisposed } ?: error("Project not available") + + val coroutineContext = currentCoroutineContext() + val future = ReadAction.nonBlocking> { + val classStructureSet = mutableSetOf() + val processedPsiFiles = mutableSetOf() + val psiFileQueue = PsiFileQueue(psiFiles) + + while (true) { + coroutineContext.ensureActive() + + val psiFile = psiFileQueue.pop() + when { + processedPsiFiles.contains(psiFile) -> Unit + + kotlinFileAnalyzerAvailable && psiFile is KtFile -> { + classStructureSet.addAll(KotlinFileAnalyzer(psiFileQueue, psiFile).analyze()) + processedPsiFiles.add(psiFile) + } + + psiFile == null -> break + } + } + + classStructureSet.toSet() + } + .inSmartMode(project) + .coalesceBy(this@PsiStructureProvider) + .submit(Dispatchers.Default.asExecutor()) + + result = future.await() + } catch (e: CancellationException) { + throw e + } catch (_: Exception) { + delay(DELAY_RESTART_READ_ACTION) + } + } + + return result ?: emptySet() + } + + private companion object { + const val DELAY_RESTART_READ_ACTION = 200L + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassLanguage.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassLanguage.kt new file mode 100644 index 00000000..f1c673f4 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassLanguage.kt @@ -0,0 +1,5 @@ +package ee.carlrobert.codegpt.psistructure.models + +enum class ClassLanguage { + KOTLIN, +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassName.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassName.kt new file mode 100644 index 00000000..e37faa68 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassName.kt @@ -0,0 +1,4 @@ +package ee.carlrobert.codegpt.psistructure.models + +@JvmInline +value class ClassName(val value: String) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassStructure.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassStructure.kt similarity index 83% rename from src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassStructure.kt rename to src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassStructure.kt index 26b3095d..e4e8a75e 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassStructure.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassStructure.kt @@ -1,7 +1,10 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure.models +package ee.carlrobert.codegpt.psistructure.models + +import com.intellij.openapi.vfs.VirtualFile data class ClassStructure( val name: ClassName, + val virtualFile: VirtualFile, val simpleName: ClassName, val classType: ClassType, val modifierList: List, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassType.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassType.kt similarity index 53% rename from src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassType.kt rename to src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassType.kt index ac9e5dc3..b1b13e8f 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ClassType.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ClassType.kt @@ -1,4 +1,4 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure.models +package ee.carlrobert.codegpt.psistructure.models enum class ClassType { ENUM, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ConstructorStructure.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ConstructorStructure.kt similarity index 62% rename from src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ConstructorStructure.kt rename to src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ConstructorStructure.kt index 8e91ac2e..afc0d1cc 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ConstructorStructure.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ConstructorStructure.kt @@ -1,4 +1,4 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure.models +package ee.carlrobert.codegpt.psistructure.models data class ConstructorStructure( val parameters: List, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/EnumEntryName.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/EnumEntryName.kt new file mode 100644 index 00000000..3cfdccb3 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/EnumEntryName.kt @@ -0,0 +1,4 @@ +package ee.carlrobert.codegpt.psistructure.models + +@JvmInline +value class EnumEntryName(val value: String) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/FieldStructure.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/FieldStructure.kt similarity index 62% rename from src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/FieldStructure.kt rename to src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/FieldStructure.kt index 21a9f6dc..f6236c91 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/FieldStructure.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/FieldStructure.kt @@ -1,4 +1,4 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure.models +package ee.carlrobert.codegpt.psistructure.models data class FieldStructure( val name: String, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/MethodStructure.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/MethodStructure.kt similarity index 70% rename from src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/MethodStructure.kt rename to src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/MethodStructure.kt index c53a748e..2bfd7aa2 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/MethodStructure.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/MethodStructure.kt @@ -1,4 +1,4 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure.models +package ee.carlrobert.codegpt.psistructure.models data class MethodStructure( val name: String, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ParameterInfo.kt b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ParameterInfo.kt similarity index 61% rename from src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ParameterInfo.kt rename to src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ParameterInfo.kt index 87a62a3c..305539ef 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/psi/structure/models/ParameterInfo.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/psistructure/models/ParameterInfo.kt @@ -1,4 +1,4 @@ -package ee.carlrobert.codegpt.codecompletions.psi.structure.models +package ee.carlrobert.codegpt.psistructure.models data class ParameterInfo( val name: String, diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/chat/ChatConfigurationConfigurable.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/chat/ChatConfigurationConfigurable.kt new file mode 100644 index 00000000..00fbdb63 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/chat/ChatConfigurationConfigurable.kt @@ -0,0 +1,62 @@ +package ee.carlrobert.codegpt.settings.chat + +import com.intellij.openapi.components.service +import com.intellij.openapi.options.Configurable +import com.intellij.openapi.ui.DialogPanel +import com.intellij.ui.components.JBCheckBox +import com.intellij.ui.dsl.builder.panel +import ee.carlrobert.codegpt.CodeGPTBundle +import javax.swing.JComponent + +class ChatConfigurationConfigurable : Configurable { + + private val editorContextTagCheckBox = JBCheckBox( + CodeGPTBundle.get("chatConfigurationConfigurable.editorContextTag.title"), + service().state.editorContextTagEnabled + ) + + private val psiStructureCheckBox = JBCheckBox( + CodeGPTBundle.get("chatConfigurationConfigurable.psiStructure.title"), + service().state.psiStructureEnabled + ) + + fun createPanel(): DialogPanel { + return panel { + row { + cell(editorContextTagCheckBox) + .comment(CodeGPTBundle.get("chatConfigurationConfigurable.editorContextTag.description")) + } + row { + cell(psiStructureCheckBox) + .comment(CodeGPTBundle.get("chatConfigurationConfigurable.psiStructure.description")) + } + } + } + + fun resetForm(prevState: ChatSettingsState) { + editorContextTagCheckBox.isSelected = prevState.editorContextTagEnabled + psiStructureCheckBox.isSelected = prevState.psiStructureEnabled + } + + override fun createComponent(): JComponent = createPanel() + + override fun isModified(): Boolean { + return ChatSettingsState().apply { + editorContextTagEnabled = editorContextTagCheckBox.isSelected + psiStructureEnabled = psiStructureCheckBox.isSelected + } != service().state + } + + override fun apply() { + service().loadState( + ChatSettingsState().apply { + editorContextTagEnabled = editorContextTagCheckBox.isSelected + psiStructureEnabled = psiStructureCheckBox.isSelected + } + ) + } + + override fun getDisplayName(): String = + CodeGPTBundle.get("chatConfigurationConfigurable.displayName") + +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/chat/ChatSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/chat/ChatSettings.kt new file mode 100644 index 00000000..c608b6b3 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/chat/ChatSettings.kt @@ -0,0 +1,29 @@ +package ee.carlrobert.codegpt.settings.chat + +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.components.BaseState +import com.intellij.openapi.components.Service +import com.intellij.openapi.components.SimplePersistentStateComponent +import com.intellij.openapi.components.State +import com.intellij.openapi.components.Storage + +@Service +@State( + name = "ProxyAI_ChatSettings", + storages = [Storage("ProxyAI_ChatSettings.xml")] +) +class ChatSettings : + SimplePersistentStateComponent(ChatSettingsState()) { + + override fun loadState(state: ChatSettingsState) { + super.loadState(state) + ApplicationManager.getApplication().messageBus + .syncPublisher(ChatSettingsListener.TOPIC) + .onChatSettingsChanged(state) + } +} + +class ChatSettingsState : BaseState() { + var editorContextTagEnabled by property(true) + var psiStructureEnabled by property(true) +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/chat/ChatSettingsListener.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/chat/ChatSettingsListener.kt new file mode 100644 index 00000000..2635d7f2 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/chat/ChatSettingsListener.kt @@ -0,0 +1,11 @@ +package ee.carlrobert.codegpt.settings.chat + +import com.intellij.util.messages.Topic + +fun interface ChatSettingsListener { + fun onChatSettingsChanged(newState: ChatSettingsState) + + companion object { + val TOPIC = Topic.create("Chat Settings Changed", ChatSettingsListener::class.java) + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceConfigurable.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceConfigurable.kt index 530aa646..2b3085c1 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceConfigurable.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/CustomServiceConfigurable.kt @@ -1,27 +1,17 @@ package ee.carlrobert.codegpt.settings.service.custom -import com.intellij.openapi.application.runInEdt import com.intellij.openapi.components.service import com.intellij.openapi.options.Configurable import ee.carlrobert.codegpt.settings.service.custom.form.CustomServiceListForm -import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.cancel import javax.swing.JComponent -import kotlin.coroutines.CoroutineContext - -object SwingDispatcher : CoroutineDispatcher() { - override fun dispatch(context: CoroutineContext, block: Runnable) { - runInEdt { - block.run() - } - } -} class CustomServiceConfigurable : Configurable { - private val coroutineScope = CoroutineScope(SupervisorJob() + SwingDispatcher) + private val coroutineScope = CoroutineScope(SupervisorJob() + Dispatchers.Main.immediate) private lateinit var component: CustomServiceListForm override fun getDisplayName(): String { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/MessageBuilder.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/MessageBuilder.kt index 32f5782f..39c7018a 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/MessageBuilder.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/chat/MessageBuilder.kt @@ -44,8 +44,10 @@ class MessageBuilder(private val project: Project, private val text: String) { message: Message, tags: List ): String = buildString { - tags.forEach { - TagProcessorFactory.getProcessor(project, it).process(message, it, this) - } + tags + .map { + TagProcessorFactory.getProcessor(project, it) + } + .forEach { it.process(message, this) } } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagDetailsComparator.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagDetailsComparator.kt new file mode 100644 index 00000000..14000f28 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagDetailsComparator.kt @@ -0,0 +1,22 @@ +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.TagDetails + +internal class TagDetailsComparator : Comparator { + override fun compare(o1: TagDetails, o2: TagDetails): Int { + val priority1 = getPriority(o1) + val priority2 = getPriority(o2) + + return priority1.compareTo(priority2) + } + + private fun getPriority(tag: TagDetails): Int { + return when (tag) { + is EditorSelectionTagDetails -> 0 + is EditorTagDetails -> 1 + else -> 2 + } + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagProcessor.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagProcessor.kt index 045d1cf7..1715b3a8 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagProcessor.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagProcessor.kt @@ -1,8 +1,7 @@ package ee.carlrobert.codegpt.ui.textarea import ee.carlrobert.codegpt.conversations.message.Message -import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails -interface TagProcessor { - fun process(message: Message, tagDetails: TagDetails, promptBuilder: StringBuilder) +fun interface TagProcessor { + fun process(message: Message, promptBuilder: StringBuilder) } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagProcessorFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagProcessorFactory.kt index 4e0dcb0e..9a444e30 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagProcessorFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/TagProcessorFactory.kt @@ -6,7 +6,18 @@ import com.intellij.openapi.project.Project import com.intellij.openapi.vfs.VirtualFile import ee.carlrobert.codegpt.EncodingManager import ee.carlrobert.codegpt.conversations.message.Message -import ee.carlrobert.codegpt.ui.textarea.header.tag.* +import ee.carlrobert.codegpt.ui.textarea.header.tag.CurrentGitChangesTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.DocumentationTagDetails +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.EmptyTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.FileTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.FolderTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.GitCommitTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.PersonaTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.SelectionTagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.WebTagDetails import ee.carlrobert.codegpt.util.GitUtil import git4idea.GitCommit @@ -14,28 +25,25 @@ object TagProcessorFactory { fun getProcessor(project: Project, tagDetails: TagDetails): TagProcessor { return when (tagDetails) { - is FileTagDetails -> FileTagProcessor() - is SelectionTagDetails -> SelectionTagProcessor() - is DocumentationTagDetails -> DocumentationTagProcessor() - is PersonaTagDetails -> PersonaTagProcessor() - is FolderTagDetails -> FolderTagProcessor() + is FileTagDetails -> FileTagProcessor(tagDetails) + is SelectionTagDetails -> SelectionTagProcessor(tagDetails) + is DocumentationTagDetails -> DocumentationTagProcessor(tagDetails) + is PersonaTagDetails -> PersonaTagProcessor(tagDetails) + is FolderTagDetails -> FolderTagProcessor(tagDetails) is WebTagDetails -> WebTagProcessor() - is GitCommitTagDetails -> GitCommitTagProcessor(project) + is GitCommitTagDetails -> GitCommitTagProcessor(project, tagDetails) is CurrentGitChangesTagDetails -> CurrentGitChangesTagProcessor(project) - else -> throw IllegalArgumentException("Unknown tag type: ${tagDetails::class.simpleName}") + is EditorSelectionTagDetails -> EditorSelectionTagProcessor(tagDetails) + is EditorTagDetails -> EditorTagProcessor(tagDetails) + is EmptyTagDetails -> TagProcessor { _, _ -> } } } } -class FileTagProcessor : TagProcessor { - override fun process( - message: Message, - tagDetails: TagDetails, - promptBuilder: StringBuilder - ) { - if (tagDetails !is FileTagDetails) { - return - } +class FileTagProcessor( + private val tagDetails: FileTagDetails, +) : TagProcessor { + override fun process(message: Message, promptBuilder: StringBuilder) { if (message.referencedFilePaths == null) { message.referencedFilePaths = mutableListOf() } @@ -43,23 +51,33 @@ class FileTagProcessor : TagProcessor { } } -class SelectionTagProcessor : TagProcessor { - override fun process( - message: Message, - tagDetails: TagDetails, - promptBuilder: StringBuilder - ) { - val selectionTagDetails = tagDetails as? SelectionTagDetails ?: return - if (selectionTagDetails.selectedText.isNullOrEmpty()) { +class EditorTagProcessor( + private val tagDetails: EditorTagDetails, +) : TagProcessor { + + override fun process(message: Message, promptBuilder: StringBuilder) { + if (message.referencedFilePaths == null) { + message.referencedFilePaths = mutableListOf() + } + message.referencedFilePaths?.add(tagDetails.virtualFile.path) + } +} + +class SelectionTagProcessor( + private val tagDetails: SelectionTagDetails, +) : TagProcessor { + + override fun process(message: Message, promptBuilder: StringBuilder) { + if (tagDetails.selectedText.isNullOrEmpty()) { return } promptBuilder .append("\n```${tagDetails.virtualFile.extension}\n") - .append(selectionTagDetails.selectedText) + .append(tagDetails.selectedText) .append("\n```\n") - selectionTagDetails.selectionModel.let { + tagDetails.selectionModel.let { if (it.hasSelection()) { it.removeSelection() } @@ -67,42 +85,44 @@ class SelectionTagProcessor : TagProcessor { } } -class DocumentationTagProcessor : TagProcessor { - override fun process( - message: Message, - tagDetails: TagDetails, - promptBuilder: StringBuilder - ) { - if (tagDetails !is DocumentationTagDetails) { +class EditorSelectionTagProcessor( + private val tagDetails: EditorSelectionTagDetails, +) : TagProcessor { + override fun process(message: Message, promptBuilder: StringBuilder) { + if (tagDetails.selectedText.isNullOrEmpty()) { return } + + promptBuilder + .append("\n```${tagDetails.virtualFile.extension}\n") + .append(tagDetails.selectedText) + .append("\n```\n") + } +} + +class DocumentationTagProcessor( + private val tagDetails: DocumentationTagDetails, +) : TagProcessor { + override fun process(message: Message, promptBuilder: StringBuilder) { message.documentationDetails = tagDetails.documentationDetails } } -class PersonaTagProcessor : TagProcessor { - override fun process( - message: Message, - tagDetails: TagDetails, - promptBuilder: StringBuilder - ) { - if (tagDetails !is PersonaTagDetails) { - return - } +class PersonaTagProcessor( + private val tagDetails: PersonaTagDetails, +) : TagProcessor { + override fun process(message: Message, promptBuilder: StringBuilder) { message.personaName = tagDetails.personaDetails.name } } -class FolderTagProcessor : TagProcessor { +class FolderTagProcessor( + private val tagDetails: FolderTagDetails, +) : TagProcessor { override fun process( message: Message, - tagDetails: TagDetails, promptBuilder: StringBuilder ) { - if (tagDetails !is FolderTagDetails) { - return - } - if (message.referencedFilePaths == null) { message.referencedFilePaths = mutableListOf() } @@ -123,25 +143,17 @@ class FolderTagProcessor : TagProcessor { class WebTagProcessor : TagProcessor { override fun process( message: Message, - tagDetails: TagDetails, promptBuilder: StringBuilder ) { - if (tagDetails !is WebTagDetails) { - return - } message.isWebSearchIncluded = true } } -class GitCommitTagProcessor(private val project: Project) : TagProcessor { - override fun process( - message: Message, - tagDetails: TagDetails, - promptBuilder: StringBuilder - ) { - if (tagDetails !is GitCommitTagDetails) { - return - } +class GitCommitTagProcessor( + private val project: Project, + private val tagDetails: GitCommitTagDetails, +) : TagProcessor { + override fun process(message: Message, promptBuilder: StringBuilder) { promptBuilder .append("\n```shell\n") .append(getDiffString(project, tagDetails.gitCommit)) @@ -167,16 +179,14 @@ class GitCommitTagProcessor(private val project: Project) : TagProcessor { } } -class CurrentGitChangesTagProcessor(private val project: Project) : TagProcessor { +class CurrentGitChangesTagProcessor( + private val project: Project, +) : TagProcessor { + override fun process( message: Message, - tagDetails: TagDetails, promptBuilder: StringBuilder ) { - if (tagDetails !is CurrentGitChangesTagDetails) { - return - } - ProgressManager.getInstance().runProcessWithProgressSynchronously( { GitUtil.getCurrentChanges(project)?.let { 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 84c3ea03..e07473f8 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt @@ -37,10 +37,17 @@ import ee.carlrobert.codegpt.ui.textarea.header.UserInputHeaderPanel 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 +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagManager import ee.carlrobert.codegpt.ui.textarea.suggestion.SuggestionsPopupManager +import ee.carlrobert.codegpt.util.coroutines.DisposableCoroutineScope import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel import git4idea.GitCommit -import java.awt.* +import java.awt.BasicStroke +import java.awt.BorderLayout +import java.awt.Graphics +import java.awt.Graphics2D +import java.awt.Insets +import java.awt.RenderingHints import java.awt.geom.Area import java.awt.geom.Rectangle2D import java.awt.geom.RoundRectangle2D @@ -51,7 +58,8 @@ class UserInputPanel( private val conversation: Conversation, private val totalTokensPanel: TotalTokensPanel, parentDisposable: Disposable, - private val onSubmit: (String, List) -> Unit, + tagManager: TagManager, + private val onSubmit: (String) -> Unit, private val onStop: () -> Unit ) : JPanel(BorderLayout()) { @@ -59,11 +67,12 @@ class UserInputPanel( private const val CORNER_RADIUS = 16 } + private val disposableCoroutineScope = DisposableCoroutineScope() private val suggestionsPopupManager = SuggestionsPopupManager(project, this) private val promptTextField = PromptTextField(project, suggestionsPopupManager, ::updateUserTokens, ::handleSubmit) private val userInputHeaderPanel = - UserInputHeaderPanel(project, suggestionsPopupManager, promptTextField) + UserInputHeaderPanel(project, tagManager, suggestionsPopupManager, promptTextField) private val submitButton = IconActionButton( object : AnAction( CodeGPTBundle.get("smartTextPane.submitButton.title"), @@ -71,7 +80,7 @@ class UserInputPanel( IconUtil.scale(Icons.Send, null, 0.85f) ) { override fun actionPerformed(e: AnActionEvent) { - handleSubmit(promptTextField.text, userInputHeaderPanel.getSelectedTags()) + handleSubmit(promptTextField.text) } }, "SUBMIT" @@ -94,6 +103,7 @@ class UserInputPanel( get() = promptTextField.text init { + Disposer.register(parentDisposable, disposableCoroutineScope) background = service().globalScheme.defaultBackground add(userInputHeaderPanel, BorderLayout.NORTH) add(promptTextField, BorderLayout.CENTER) @@ -187,12 +197,8 @@ class UserInputPanel( override fun getInsets(): Insets = JBUI.insets(4) private fun handleSubmit(text: String) { - handleSubmit(text, userInputHeaderPanel.getSelectedTags()) - } - - private fun handleSubmit(text: String, appliedTags: List = emptyList()) { if (text.isNotEmpty() && submitButton.isEnabled) { - onSubmit(text, appliedTags) + onSubmit(text) promptTextField.clear() } } 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 8262bed9..34a6dd76 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 @@ -1,11 +1,9 @@ package ee.carlrobert.codegpt.ui.textarea.header import com.intellij.icons.AllIcons -import com.intellij.openapi.components.service import com.intellij.openapi.editor.Editor import com.intellij.openapi.editor.EditorKind import com.intellij.openapi.editor.SelectionModel -import com.intellij.openapi.fileEditor.FileEditorManager import com.intellij.openapi.fileEditor.FileEditorManagerEvent import com.intellij.openapi.fileEditor.FileEditorManagerListener import com.intellij.openapi.project.Project @@ -17,39 +15,36 @@ import ee.carlrobert.codegpt.EditorNotifier import ee.carlrobert.codegpt.actions.IncludeFilesInContextNotifier import ee.carlrobert.codegpt.ui.WrapLayout import ee.carlrobert.codegpt.ui.textarea.PromptTextField -import ee.carlrobert.codegpt.ui.textarea.header.tag.* +import ee.carlrobert.codegpt.ui.textarea.TagDetailsComparator +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.SelectionTagPanel +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagManager +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagManagerListener +import ee.carlrobert.codegpt.ui.textarea.header.tag.TagPanel import ee.carlrobert.codegpt.ui.textarea.suggestion.SuggestionsPopupManager import ee.carlrobert.codegpt.util.EditorUtil import ee.carlrobert.codegpt.util.EditorUtil.getSelectedEditor -import ee.carlrobert.codegpt.util.EditorUtil.getSelectedEditorFile import java.awt.Cursor import java.awt.Dimension import java.awt.FlowLayout import java.awt.Graphics import javax.swing.JButton import javax.swing.JPanel -import kotlin.math.min class UserInputHeaderPanel( private val project: Project, + private val tagManager: TagManager, suggestionsPopupManager: SuggestionsPopupManager, private val promptTextField: PromptTextField ) : JPanel(WrapLayout(FlowLayout.LEFT, 4, 4)), TagManagerListener { companion object { private const val INITIAL_VISIBLE_FILES = 2 - private const val TAG_INSERTION_OFFSET = 1 } - private val tagManager = TagManager() - private val selectedFileTagPanel = object : SelectedFileTagPanel(project, promptTextField) { - override fun onClose() { - this.isVisible = false - if (tagManager.getTags().isEmpty()) { - emptyText.isVisible = true - } - } - } private val emptyText = JBLabel("No context included").apply { foreground = JBUI.CurrentTheme.Label.disabledForeground() font = JBUI.Fonts.smallFont() @@ -57,7 +52,8 @@ class UserInputHeaderPanel( preferredSize = Dimension(preferredSize.width, 20) verticalAlignment = JBLabel.CENTER } - private val selectionTagPanel = SelectionTagPanel(project, promptTextField) + + // private val selectionTagPanel = SelectionTagPanel(project, tagManager, promptTextField) private val defaultHeaderTagsPanel = CustomFlowPanel().apply { add(AddButton { if (suggestionsPopupManager.isPopupVisible()) { @@ -66,9 +62,8 @@ class UserInputHeaderPanel( suggestionsPopupManager.showPopup(this) } }) + add(RemoveAllButton()) add(emptyText) - add(selectionTagPanel) - add(selectedFileTagPanel) } init { @@ -80,17 +75,6 @@ class UserInputHeaderPanel( fun getSelectedTags(): List { val selectedTags = tagManager.getTags().filter { it.selected }.toMutableList() - val selectedFile = getSelectedFile() - if (selectedFileTagPanel.isVisible && selectedFileTagPanel.isSelected && selectedFile != null) { - selectedTags.add(FileTagDetails(selectedFile)) - } - - (selectionTagPanel.tagDetails as? SelectionTagDetails)?.let { - if (!it.selectedText.isNullOrEmpty()) { - selectedTags.add(it) - } - } - return selectedTags } @@ -99,113 +83,68 @@ class UserInputHeaderPanel( } override fun onTagAdded(tag: TagDetails) { - emptyText.isVisible = false - add(createTag(tag), getInsertionIndex()) - - if (tagManager.getTags().filter { !it.selected }.size > 2) { - components - .lastOrNull { it is TagPanel && it.tagDetails is FileTagDetails && !it.isSelected } - ?.let { tagManager.removeTag((it as TagPanel).tagDetails.id) } - } - - promptTextField.requestFocus() - - revalidate() - repaint() + onTagsChanged() } override fun onTagRemoved(tag: TagDetails) { - val componentToRemove = - components.find { it is TagPanel && it.id == tag.id } ?: return - remove(componentToRemove) - - if (getSelectedEditorFile(project) == null) { - selectedFileTagPanel.isVisible = false - } - - if (tagManager.getTags().isEmpty() && !selectedFileTagPanel.isVisible) { - emptyText.isVisible = true - } - promptTextField.requestFocus() - revalidate() - repaint() + onTagsChanged() } override fun onTagSelectionChanged(tag: TagDetails) { - val existingTagComponent = - components.filterIsInstance().find { it.id == tag.id } - ?: return - existingTagComponent.update(tag.name, tag.icon) - updateTagPosition(existingTagComponent) + onTagsChanged() } - private fun createTag(tagDetails: TagDetails) = - object : TagPanel(tagDetails, true) { + private fun onTagsChanged() { + components.filterIsInstance() + .forEach { remove(it) } - init { - cursor = - if (tagDetails is FileTagDetails) Cursor(Cursor.HAND_CURSOR) else Cursor(Cursor.DEFAULT_CURSOR) - } + val allTags = tagManager.getTags() - override fun onSelect(tagDetails: TagDetails) { - if (tagDetails is FileTagDetails) { - if (tagDetails.selected) { - project.service().openFile(tagDetails.virtualFile) - return - } + val editorVirtualFilesSet = allTags + .filterIsInstance() + .map { it.virtualFile } + .toSet() - tagDetails.selected = true - - val canAddNewTag = tagManager.getTags() - .filterIsInstance() - .count { !it.selected } < 2 - if (canAddNewTag) { - addNextOpenFile() - } - update(tagDetails.name, tagDetails.icon) - } - } - - override fun onClose() { - tagManager.removeTag(tagDetails.id) + /** + * Filter the tags collection to prioritize EditorTagDetails over FileTagDetails + * Keep all tags except FileTagDetails that have a corresponding EditorTagDetails + */ + val tags = allTags.filter { tag -> + if (tag is FileTagDetails) { + !editorVirtualFilesSet.contains(tag.virtualFile) + } else { + true } } + .sortedWith(TagDetailsComparator()) + .toSet() + + emptyText.isVisible = tags.none { it.selected } + + tags.forEach { add(createTagPanel(it)) } - private fun updateTagPosition(tag: TagPanel) { - remove(tag) - add(tag, getInsertionIndex()) revalidate() repaint() } - private fun getInsertionIndex(): Int { - val lastSelectionTagIndex = getLastSelectedTagIndex() - return if (lastSelectionTagIndex != -1) { - min(lastSelectionTagIndex + TAG_INSERTION_OFFSET + 1, components.size) + private fun createTagPanel(tagDetails: TagDetails) = + if (tagDetails is EditorSelectionTagDetails) { + SelectionTagPanel(tagDetails, tagManager, promptTextField) } else { - TAG_INSERTION_OFFSET - } - } + object : TagPanel(tagDetails, tagManager, false) { - private fun getLastSelectedTagIndex(): Int = - components - .filter { it !is SelectedFileTagPanel && it !is SelectionTagPanel } - .filterIsInstance() - .indexOfLast { it.tagDetails.selected } + init { + cursor = + if (tagDetails is FileTagDetails) Cursor(Cursor.HAND_CURSOR) else Cursor(Cursor.DEFAULT_CURSOR) + } - private fun getSortedOpenFileTags(): MutableList = - EditorUtil.getOpenLocalFiles(project) - .filterNot { tagManager.isFileTagExists(it) } - .map { FileTagDetails(it) } - .toMutableList() + override fun onSelect(tagDetails: TagDetails) = Unit - private fun addNextOpenFile() { - getSortedOpenFileTags() - .firstOrNull { it.virtualFile != getSelectedFile() } - ?.let { - tagManager.addTag(it.apply { selected = false }) + override fun onClose() { + tagManager.remove(tagDetails) + } } - } + } private fun initializeUI() { isOpaque = false @@ -217,7 +156,12 @@ class UserInputHeaderPanel( private fun addInitialTags() { val selectedFile = getSelectedEditor(project)?.virtualFile - getSortedOpenFileTags() + if (selectedFile != null) { + tagManager.addTag(EditorTagDetails(selectedFile)) + } + + EditorUtil.getOpenLocalFiles(project) + .map { EditorTagDetails(it) } .filterNot { it.virtualFile == selectedFile } .take(INITIAL_VISIBLE_FILES) .forEach { @@ -260,6 +204,29 @@ class UserInputHeaderPanel( } } + private inner class RemoveAllButton : JButton() { + init { + addActionListener { + tagManager.clear() + } + + cursor = Cursor(Cursor.HAND_CURSOR) + preferredSize = Dimension(20, 20) + isContentAreaFilled = false + isOpaque = false + border = null + toolTipText = "Remove All Context" + icon = IconUtil.scale(AllIcons.Actions.Close, null, 0.75f) + rolloverIcon = IconUtil.scale(AllIcons.Actions.CloseHovered, null, 0.75f) + pressedIcon = IconUtil.scale(AllIcons.Actions.CloseHovered, null, 0.75f) + } + + override fun paintComponent(g: Graphics) { + PaintUtil.drawRoundedBackground(g, this, true) + super.paintComponent(g) + } + } + private inner class EditorSelectionChangeListener : EditorNotifier.SelectionChange { override fun selectionChanged(selectionModel: SelectionModel, virtualFile: VirtualFile) { handleSelectionChange(selectionModel, virtualFile) @@ -269,37 +236,28 @@ class UserInputHeaderPanel( selectionModel: SelectionModel, virtualFile: VirtualFile ) { - selectionTagPanel.update(virtualFile, selectionModel) + if (selectionModel.hasSelection()) { + tagManager.addTag(EditorSelectionTagDetails(virtualFile, selectionModel)) + } else { + tagManager.remove(EditorSelectionTagDetails(virtualFile, selectionModel)) + } + } } private inner class EditorReleasedListener : EditorNotifier.Released { override fun editorReleased(editor: Editor) { if (editor.editorKind == EditorKind.MAIN_EDITOR && !editor.isDisposed && editor.virtualFile != null) { - tagManager.removeFileTag(editor.virtualFile) + tagManager.remove(EditorTagDetails(editor.virtualFile)) } } } - private fun getSelectedFile(): VirtualFile? { - return (selectedFileTagPanel.tagDetails as? FileTagDetails)?.virtualFile - } - private inner class FileSelectionListener : FileEditorManagerListener { override fun selectionChanged(event: FileEditorManagerEvent) { event.newFile?.let { newFile -> - var existingFileTag = tagManager.getFileTag(newFile) - if (existingFileTag != null) { - tagManager.removeTag(existingFileTag.id) - } else { - existingFileTag = FileTagDetails(newFile).apply { selected = false } - } - - if (selectedFileTagPanel.tagDetails !is EmptyTagDetails) { - tagManager.addTag(selectedFileTagPanel.tagDetails) - } - - selectedFileTagPanel.update(existingFileTag) + val editorTagDetails = EditorTagDetails(newFile) + tagManager.addTag(editorTagDetails) emptyText.isVisible = false } } @@ -307,9 +265,7 @@ class UserInputHeaderPanel( private inner class IncludedFilesListener : IncludeFilesInContextNotifier { override fun filesIncluded(includedFiles: MutableList) { - includedFiles - .filterNot { tagManager.isFileTagExists(it) || getSelectedFile() == it } - .forEach { tagManager.addTag(FileTagDetails(it)) } + includedFiles.forEach { tagManager.addTag(FileTagDetails(it)) } } } } \ No newline at end of file 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 b4e892f6..ec9e5a28 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 @@ -10,12 +10,13 @@ import git4idea.GitCommit import java.util.* import javax.swing.Icon -open class TagDetails( - open val name: String, +sealed class TagDetails( + val name: String, val icon: Icon? = null, - open var selected: Boolean = true -) { val id: UUID = UUID.randomUUID() +) { + + var selected: Boolean = true override fun equals(other: Any?): Boolean { if (this === other) return true @@ -28,8 +29,46 @@ open class TagDetails( } } -data class FileTagDetails(var virtualFile: VirtualFile) : - TagDetails(virtualFile.name, virtualFile.fileType.icon) +class EditorTagDetails(val virtualFile: VirtualFile) : TagDetails(virtualFile.name, virtualFile.fileType.icon) { + + private val type: String = "EditorTagDetails" + + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as EditorTagDetails + + if (virtualFile != other.virtualFile) return false + if (type != other.type) return false + + return true + } + + override fun hashCode(): Int = + 31 * virtualFile.hashCode() + type.hashCode() + +} + +class FileTagDetails(val virtualFile: VirtualFile) : TagDetails(virtualFile.name, virtualFile.fileType.icon) { + + private val type: String = "FileTagDetails" + + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (javaClass != other?.javaClass) return false + + other as FileTagDetails + + if (virtualFile != other.virtualFile) return false + if (type != other.type) return false + + return true + } + + override fun hashCode(): Int = + 31 * virtualFile.hashCode() + type.hashCode() +} data class SelectionTagDetails( var virtualFile: VirtualFile, @@ -42,6 +81,26 @@ data class SelectionTagDetails( private set } +class EditorSelectionTagDetails( + val virtualFile: VirtualFile, + val selectionModel: SelectionModel +) : TagDetails( + "${virtualFile.name} (${selectionModel.selectionStartPosition?.line}:${selectionModel.selectionEndPosition?.line})", + virtualFile.fileType.icon +) { + var selectedText: String? = selectionModel.selectedText + private set + + override fun equals(other: Any?): Boolean { + if (other === null) return false + return other::class == this::class + } + + override fun hashCode(): Int { + return this::class.hashCode() + } +} + data class DocumentationTagDetails(var documentationDetails: DocumentationDetails) : TagDetails(documentationDetails.name, AllIcons.Toolwindows.Documentation) 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 93822347..cc8b501c 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 @@ -1,52 +1,84 @@ package ee.carlrobert.codegpt.ui.textarea.header.tag -import com.intellij.openapi.vfs.VirtualFile -import java.util.* +import com.intellij.openapi.Disposable +import com.intellij.openapi.application.ApplicationManager +import ee.carlrobert.codegpt.settings.chat.ChatSettingsListener +import java.util.concurrent.CopyOnWriteArraySet -class TagManager { +class TagManager(parentDisposable: Disposable) { private val tags = mutableSetOf() - private val listeners = mutableListOf() + private val listeners = CopyOnWriteArraySet() + + @Volatile + private var enabled: Boolean = true + + init { + val connection = ApplicationManager.getApplication().messageBus + .connect(parentDisposable) + + connection.subscribe( + ChatSettingsListener.TOPIC, + ChatSettingsListener { newState -> + if (newState.editorContextTagEnabled) { + enabled = true + } else { + enabled = false + clear() + } + }) + } fun addListener(listener: TagManagerListener) { listeners.add(listener) } - fun getTags(): Set = tags.toSet() + fun removeListener(listener: TagManagerListener) { + listeners.remove(listener) + } + + fun getTags(): Set = synchronized(this) { tags.toSet() } fun addTag(tagDetails: TagDetails) { - if (tags.add(tagDetails)) { + val wasAdded = synchronized(this) { + if (!enabled && isEditorTag(tagDetails)) return + + if (tagDetails is EditorSelectionTagDetails) { + tags.remove(tagDetails) + } + + tags.add(tagDetails) + } + if (wasAdded) { listeners.forEach { it.onTagAdded(tagDetails) } } } - fun removeFileTag(virtualFile: VirtualFile) { - getFileTag(virtualFile)?.let { - removeTag(it.id) + fun notifySelectionChanged(tagDetails: TagDetails) { + val containsTag = synchronized(this) { tags.contains(tagDetails) } + if (containsTag) { + listeners.forEach { it.onTagSelectionChanged(tagDetails) } } } - fun removeTag(id: UUID) { - tags.find { it.id == id } - ?.let { tag -> - if (tags.removeIf { it.id == tag.id }) { - listeners.forEach { it.onTagRemoved(tag) } - } - } + fun remove(tagDetails: TagDetails) { + val wasRemoved = synchronized(this) { tags.remove(tagDetails) } + if (wasRemoved) { + listeners.forEach { it.onTagRemoved(tagDetails) } + } } - fun getTag(id: UUID): TagDetails? = tags.find { it.id == id } - - fun getFileTag(file: VirtualFile): FileTagDetails? = - tags.filterIsInstance().find { it.virtualFile == file } - - fun isFileTagExists(file: VirtualFile): Boolean = getFileTag(file) != null - fun clear() { - val tagsToRemove = tags.toList() - tags.clear() - tagsToRemove.forEach { tag -> + val removedTags = mutableListOf() + synchronized(this) { + removedTags.addAll(tags) + tags.clear() + } + removedTags.forEach { tag -> listeners.forEach { it.onTagRemoved(tag) } } } + + private fun isEditorTag(tagDetails: TagDetails): Boolean = + tagDetails is EditorSelectionTagDetails || tagDetails is EditorTagDetails } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagPanel.kt index f9ea6ec5..92d5ceee 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagPanel.kt @@ -3,35 +3,23 @@ package ee.carlrobert.codegpt.ui.textarea.header.tag import com.intellij.icons.AllIcons import com.intellij.icons.AllIcons.Actions.Close import com.intellij.openapi.components.service -import com.intellij.openapi.editor.SelectionModel import com.intellij.openapi.editor.colors.EditorColorsManager -import com.intellij.openapi.project.Project -import com.intellij.openapi.vfs.VirtualFile import com.intellij.ui.components.JBLabel import com.intellij.util.IconUtil import com.intellij.util.ui.JBUI -import com.jetbrains.rd.util.UUID import ee.carlrobert.codegpt.ui.textarea.PromptTextField import ee.carlrobert.codegpt.ui.textarea.header.PaintUtil -import ee.carlrobert.codegpt.util.EditorUtil.getSelectedEditor -import ee.carlrobert.codegpt.util.EditorUtil.getSelectedEditorFile -import java.awt.Cursor -import java.awt.Dimension -import java.awt.GridBagConstraints -import java.awt.GridBagLayout -import java.awt.Graphics +import java.awt.* import javax.swing.Icon import javax.swing.JButton import javax.swing.JToggleButton -import javax.swing.JPanel abstract class TagPanel( var tagDetails: TagDetails, + private val tagManager: TagManager, private val shouldPreventDeselection: Boolean = true, ) : JToggleButton() { - val id: UUID = tagDetails.id - private val label = TagLabel(tagDetails.name, tagDetails.icon, tagDetails.selected) private val closeButton = CloseButton { isVisible = isSelected @@ -99,6 +87,7 @@ abstract class TagPanel( closeButton.isVisible = isSelected tagDetails.selected = isSelected + tagManager.notifySelectionChanged(tagDetails) label.update(isSelected) } @@ -154,58 +143,18 @@ abstract class TagPanel( } } -abstract class SelectedFileTagPanel( - private val project: Project, - private val promptTextField: PromptTextField, - virtualFile: VirtualFile? = getSelectedEditorFile(project) -) : TagPanel( - (if (virtualFile == null) EmptyTagDetails() - else FileTagDetails(virtualFile)).apply { selected = true }, - false -) { - - init { - isVisible = getSelectedEditorFile(project) != null - } - - override fun onSelect(tagDetails: TagDetails) { - if (tagDetails is FileTagDetails) { - update(tagDetails.virtualFile.name, tagDetails.virtualFile.fileType.icon) - } - promptTextField.requestFocus() - } - - fun update(tagDetails: FileTagDetails) { - this.tagDetails = tagDetails - isVisible = true - isSelected = true - update(tagDetails.name, tagDetails.virtualFile.fileType.icon) - } -} - class SelectionTagPanel( - project: Project, - private val promptTextField: PromptTextField -) : TagPanel(getDefaultSelectionTagDetails(project), true) { + tagDetails: EditorSelectionTagDetails, + tagManager: TagManager, + private val promptTextField: PromptTextField, +) : TagPanel(tagDetails, tagManager, true) { init { cursor = Cursor(Cursor.DEFAULT_CURSOR) - } - - companion object { - fun getDefaultSelectionTagDetails(project: Project): TagDetails { - val editor = getSelectedEditor(project) - val selectionModel = editor?.selectionModel - return if (selectionModel?.hasSelection() == true) { - SelectionTagDetails(editor.virtualFile, selectionModel) - } else { - EmptyTagDetails() - } - } - } - - init { - isVisible = tagDetails !is EmptyTagDetails + update( + "${tagDetails.virtualFile.name}:${tagDetails.selectionModel.selectionStart}-${tagDetails.selectionModel.selectionEnd}", + tagDetails.virtualFile.fileType.icon + ) } override fun onSelect(tagDetails: TagDetails) { @@ -213,15 +162,6 @@ class SelectionTagPanel( } override fun onClose() { - (tagDetails as? SelectionTagDetails)?.selectionModel?.removeSelection() - } - - fun update(virtualFile: VirtualFile, selectionModel: SelectionModel) { - tagDetails = SelectionTagDetails(virtualFile, selectionModel) - isVisible = selectionModel.hasSelection() - update( - "${virtualFile.name}:${selectionModel.selectionStart}-${selectionModel.selectionEnd}", - virtualFile.fileType.icon - ) + (tagDetails as? EditorSelectionTagDetails)?.selectionModel?.removeSelection() } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/coroutines/CoroutineDispatchers.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/coroutines/CoroutineDispatchers.kt new file mode 100644 index 00000000..f9ef72ee --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/coroutines/CoroutineDispatchers.kt @@ -0,0 +1,15 @@ +package ee.carlrobert.codegpt.util.coroutines + +import kotlinx.coroutines.CoroutineDispatcher +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.MainCoroutineDispatcher + +class CoroutineDispatchers { + fun default(): CoroutineDispatcher = Dispatchers.Default + + fun main(): MainCoroutineDispatcher = Dispatchers.Main + + fun unconfined(): CoroutineDispatcher = Dispatchers.Unconfined + + fun io(): CoroutineDispatcher = Dispatchers.IO +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/coroutines/DisposableCoroutineScope.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/coroutines/DisposableCoroutineScope.kt new file mode 100644 index 00000000..2c800390 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/util/coroutines/DisposableCoroutineScope.kt @@ -0,0 +1,23 @@ +package ee.carlrobert.codegpt.util.coroutines + +import com.intellij.openapi.Disposable +import kotlinx.coroutines.* +import kotlin.coroutines.CoroutineContext + +internal class DisposableCoroutineScope( + scopeDispatcher: CoroutineDispatcher = Dispatchers.Main.immediate +) : Disposable, CoroutineScope { + + private val coroutineScope = CoroutineScope(SupervisorJob() + scopeDispatcher) + + fun launch(block: suspend CoroutineScope.() -> Unit): Job = + coroutineScope.launch { block() } + + + override fun dispose() { + coroutineScope.cancel() + } + + override val coroutineContext: CoroutineContext + get() = coroutineScope.coroutineContext +} \ No newline at end of file diff --git a/src/main/resources/META-INF/plugin-kotlin.xml b/src/main/resources/META-INF/plugin-kotlin.xml index c28e1e2c..6fca3917 100644 --- a/src/main/resources/META-INF/plugin-kotlin.xml +++ b/src/main/resources/META-INF/plugin-kotlin.xml @@ -1,6 +1,6 @@ + serviceImplementation="ee.carlrobert.codegpt.psistructure.KotlinFileAnalyzer"/> diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml index 91419ee9..82ed152e 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -54,6 +54,8 @@ instance="ee.carlrobert.codegpt.settings.documentation.DocumentationsConfigurable"/> + \ No newline at end of file diff --git a/src/main/resources/messages/codegpt.properties b/src/main/resources/messages/codegpt.properties index a7a4eabb..134e8a52 100644 --- a/src/main/resources/messages/codegpt.properties +++ b/src/main/resources/messages/codegpt.properties @@ -147,6 +147,11 @@ configurationConfigurable.section.codeCompletion.gitDiff.title=Enable git diff c configurationConfigurable.section.codeCompletion.collectDependencyStructure.title=Enable dependency analyzer configurationConfigurable.section.codeCompletion.collectDependencyStructure.description=Enabling the setting allows the plugin to collect the dependency structure, which increases the accuracy of the proposed data, but consumes more tokens per request. Currently, it is implemented only for the Kotlin language. configurationConfigurable.section.codeCompletion.gitDiff.description=If checked, the user's most recent unstaged git diff will be included when requesting completion. +chatConfigurationConfigurable.displayName=ProxyAI: Chat Settings +chatConfigurationConfigurable.editorContextTag.title=Enable editor context tags +chatConfigurationConfigurable.editorContextTag.description=If enabled, open files in the editor will be added to the chat tags. +chatConfigurationConfigurable.psiStructure.title=Enable dependency structure analysis of attached files. +chatConfigurationConfigurable.psiStructure.description=If enabled, the class structure that is present in the imports of the attached files will be added in the context of the dialog. A structure refers to the source code in files that include constructors, fields, and methods, with all modifiers, arguments, and return types, but without an implementation. The implementation of dependencies is intentionally excluded in order to find a balance between a high-quality chat context and saving tokens. settingsConfigurable.service.llama.topK.label=Top K: settingsConfigurable.service.llama.topK.comment=Limit the next token selection to the K most probable tokens (default: 40) settingsConfigurable.service.llama.topP.label=Top P: