mirror of
https://github.com/carlrobertoh/ProxyAI.git
synced 2026-05-20 17:52:23 +00:00
feat: psi structure for chat (#897)
* Change SwingDispatcher to Dispatchers.Main.immediate * Change PSI analyzer package * Added a blank for PSI analysis in the chat * Add support for code structure analysis and improve tag management - Refactor TagManager to use thread-safe collections - Add support for new tag types (EditorTagDetails, FileTagDetails) - Update UI components to handle structure analysis - Add new icon for structure tags * Refactoring tags v2 * Add PSI structure to chat settings * Add VirtualFile to ClassStructure and improve PSI token tracking * Support passing PSI structure to completion requests * Add removeListener method to TagManager and fix memory leak * Update buildOpenAIMessages to support PSI structure for all providers * Add selected editor tag when initializing user input header * Add chat settings configuration screen * Remove unused editor tag and PSI structure settings panels --------- Co-authored-by: a.iudin <a.iudin@vk.team>
This commit is contained in:
Aleksandr
GitHub
parent
97185109ad
commit
7d472e8cf1
54 changed files with 1185 additions and 444 deletions
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<ClassStructure> 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<ReferencedFile> getReferencedFiles() {
|
||||
return getReferencedFiles(userInputPanel.getSelectedTags());
|
||||
}
|
||||
|
||||
private List<ReferencedFile> getReferencedFiles(List<? extends TagDetails> 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 <T extends TagDetails> Optional<T> findTagOfType(
|
||||
List<? extends TagDetails> tags,
|
||||
Class<T> 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<ClassStructure> 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<? extends TagDetails> appliedTags) {
|
||||
private Unit handleSubmit(String text) {
|
||||
final Set<ClassStructure> 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<ReferencedFile> referencedFiles = getReferencedFiles();
|
||||
List<ReferencedFile> 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);
|
||||
|
|
|
|||
|
|
@ -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<TagDetails>): 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<PsiStructureState> = 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<TagDetails>) {
|
||||
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<List<PsiFile>> {
|
||||
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<TagDetails>.getExcludedVirtualFiles(): Set<VirtualFile> =
|
||||
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<TagDetails>.getPsiAnalyzedTags(): Set<TagDetails> =
|
||||
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<TagDetails>.toVirtualFilesSet(): Set<VirtualFile> =
|
||||
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()
|
||||
}
|
||||
|
|
@ -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<TagDetails>,
|
||||
) : PsiStructureState()
|
||||
|
||||
data object Disabled : PsiStructureState()
|
||||
|
||||
data class Content(
|
||||
val currentlyAnalyzedTags: Set<TagDetails>,
|
||||
val elements: Set<ClassStructure>
|
||||
) : PsiStructureState()
|
||||
}
|
||||
|
|
@ -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<ClassStructure>): Int =
|
||||
psiStructureSet
|
||||
.joinToString(separator = "\n\n") { psiStructure ->
|
||||
classStructureSerializer.serialize(psiStructure)
|
||||
}
|
||||
.let { serializedPsiStructure ->
|
||||
encodingManager.countTokens(serializedPsiStructure)
|
||||
}
|
||||
}
|
||||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
"<p style=\"margin: 0; padding: 0;\"><small>%s: <strong>%d</strong></small></p>",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue