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 2025-03-11 20:57:55 +03:00 committed by GitHub
parent 97185109ad
commit 7d472e8cf1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
54 changed files with 1185 additions and 444 deletions

View file

@ -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);

View file

@ -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;
}

View file

@ -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);

View file

@ -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()
}

View file

@ -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()
}

View file

@ -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)
}
}

View file

@ -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;
}
}

View file

@ -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>",

View file

@ -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<String>? = listOf("\n\n")) {

View file

@ -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

View file

@ -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

View file

@ -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<PsiFile>): Set<ClassStructure> =
ReadAction.compute<Set<ClassStructure>, Throwable> {
val classStructureSet = mutableSetOf<ClassStructure>()
val processedPsiFiles = mutableSetOf<PsiFile?>()
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()
}
}

View file

@ -1,5 +0,0 @@
package ee.carlrobert.codegpt.codecompletions.psi.structure.models
enum class ClassLanguage {
KOTLIN,
}

View file

@ -1,4 +0,0 @@
package ee.carlrobert.codegpt.codecompletions.psi.structure.models
@JvmInline
value class ClassName(val value: String)

View file

@ -1,4 +0,0 @@
package ee.carlrobert.codegpt.codecompletions.psi.structure.models
@JvmInline
value class EnumEntryName(val value: String)

View file

@ -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<ReferencedFile>?,
var personaDetails: PersonaDetails?,
var psiStructure: Set<ClassStructure>?,
) : 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<ReferencedFile>? = null
private var personaDetails: PersonaDetails? = null
private var psiStructure: Set<ClassStructure>? = 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<ClassStructure>?) = apply { this.psiStructure = psiStructure }
fun build(): ChatCompletionParameters {
return ChatCompletionParameters(
conversation,
@ -73,7 +79,8 @@ class ChatCompletionParameters private constructor(
retry,
imageDetails,
referencedFiles,
personaDetails
personaDetails,
psiStructure,
)
}
}

View file

@ -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<PromptsSettings>().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

View file

@ -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<ReferencedFile>,
userPrompt: String?
userPrompt: String?,
psiStructure: Set<ClassStructure>?
): String {
val includedFilesSettings = service<IncludedFilesSettings>().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!!)
}
}

View file

@ -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<CodeGPTServiceSettings>().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()

View file

@ -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()))
)

View file

@ -15,7 +15,14 @@ class OllamaRequestFactory : BaseRequestFactory() {
val model = service<OllamaSettings>().state.model
val configuration = service<ConfigurationSettings>().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)

View file

@ -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<ReferencedFile>? = null
referencedFiles: List<ReferencedFile>? = null,
psiStructure: Set<ClassStructure>? = null
): List<OpenAIChatCompletionMessage> {
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<ReferencedFile>? = null
referencedFiles: List<ReferencedFile>? = null,
psiStructure: Set<ClassStructure>? = null
): MutableList<OpenAIChatCompletionMessage> {
val message = callParameters.message
val messages = mutableListOf<OpenAIChatCompletionMessage>()
@ -254,7 +274,8 @@ class OpenAIRequestFactory : CompletionRequestFactory {
} else {
CompletionRequestUtil.getPromptWithContext(
referencedFiles,
message.prompt
message.prompt,
psiStructure
)
}
messages.add(OpenAIChatCompletionStandardMessage("user", prompt))

View file

@ -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 {

View file

@ -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(

View file

@ -1,4 +1,4 @@
package ee.carlrobert.codegpt.codecompletions.psi.structure
package ee.carlrobert.codegpt.psistructure
import com.intellij.psi.PsiFile

View file

@ -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<PsiFile>): Set<ClassStructure> {
var result: Set<ClassStructure>? = 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<Set<ClassStructure>> {
val classStructureSet = mutableSetOf<ClassStructure>()
val processedPsiFiles = mutableSetOf<PsiFile?>()
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
}
}

View file

@ -0,0 +1,5 @@
package ee.carlrobert.codegpt.psistructure.models
enum class ClassLanguage {
KOTLIN,
}

View file

@ -0,0 +1,4 @@
package ee.carlrobert.codegpt.psistructure.models
@JvmInline
value class ClassName(val value: String)

View file

@ -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<String>,

View file

@ -1,4 +1,4 @@
package ee.carlrobert.codegpt.codecompletions.psi.structure.models
package ee.carlrobert.codegpt.psistructure.models
enum class ClassType {
ENUM,

View file

@ -1,4 +1,4 @@
package ee.carlrobert.codegpt.codecompletions.psi.structure.models
package ee.carlrobert.codegpt.psistructure.models
data class ConstructorStructure(
val parameters: List<ParameterInfo>,

View file

@ -0,0 +1,4 @@
package ee.carlrobert.codegpt.psistructure.models
@JvmInline
value class EnumEntryName(val value: String)

View file

@ -1,4 +1,4 @@
package ee.carlrobert.codegpt.codecompletions.psi.structure.models
package ee.carlrobert.codegpt.psistructure.models
data class FieldStructure(
val name: String,

View file

@ -1,4 +1,4 @@
package ee.carlrobert.codegpt.codecompletions.psi.structure.models
package ee.carlrobert.codegpt.psistructure.models
data class MethodStructure(
val name: String,

View file

@ -1,4 +1,4 @@
package ee.carlrobert.codegpt.codecompletions.psi.structure.models
package ee.carlrobert.codegpt.psistructure.models
data class ParameterInfo(
val name: String,

View file

@ -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<ChatSettings>().state.editorContextTagEnabled
)
private val psiStructureCheckBox = JBCheckBox(
CodeGPTBundle.get("chatConfigurationConfigurable.psiStructure.title"),
service<ChatSettings>().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<ChatSettings>().state
}
override fun apply() {
service<ChatSettings>().loadState(
ChatSettingsState().apply {
editorContextTagEnabled = editorContextTagCheckBox.isSelected
psiStructureEnabled = psiStructureCheckBox.isSelected
}
)
}
override fun getDisplayName(): String =
CodeGPTBundle.get("chatConfigurationConfigurable.displayName")
}

View file

@ -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>(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)
}

View file

@ -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)
}
}

View file

@ -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 {

View file

@ -44,8 +44,10 @@ class MessageBuilder(private val project: Project, private val text: String) {
message: Message,
tags: List<TagDetails>
): String = buildString {
tags.forEach {
TagProcessorFactory.getProcessor(project, it).process(message, it, this)
}
tags
.map {
TagProcessorFactory.getProcessor(project, it)
}
.forEach { it.process(message, this) }
}
}

View file

@ -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<TagDetails> {
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
}
}
}

View file

@ -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)
}

View file

@ -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<Unit, Exception>(
{
GitUtil.getCurrentChanges(project)?.let {

View file

@ -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<TagDetails>) -> 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<EditorColorsManager>().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<TagDetails> = emptyList()) {
if (text.isNotEmpty() && submitButton.isEnabled) {
onSubmit(text, appliedTags)
onSubmit(text)
promptTextField.clear()
}
}

View file

@ -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<TagDetails> {
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<TagPanel>().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<TagPanel>()
.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<FileEditorManager>().openFile(tagDetails.virtualFile)
return
}
val editorVirtualFilesSet = allTags
.filterIsInstance<EditorTagDetails>()
.map { it.virtualFile }
.toSet()
tagDetails.selected = true
val canAddNewTag = tagManager.getTags()
.filterIsInstance<FileTagDetails>()
.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<TagPanel>()
.indexOfLast { it.tagDetails.selected }
init {
cursor =
if (tagDetails is FileTagDetails) Cursor(Cursor.HAND_CURSOR) else Cursor(Cursor.DEFAULT_CURSOR)
}
private fun getSortedOpenFileTags(): MutableList<FileTagDetails> =
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<VirtualFile>) {
includedFiles
.filterNot { tagManager.isFileTagExists(it) || getSelectedFile() == it }
.forEach { tagManager.addTag(FileTagDetails(it)) }
includedFiles.forEach { tagManager.addTag(FileTagDetails(it)) }
}
}
}

View file

@ -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)

View file

@ -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<TagDetails>()
private val listeners = mutableListOf<TagManagerListener>()
private val listeners = CopyOnWriteArraySet<TagManagerListener>()
@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<TagDetails> = tags.toSet()
fun removeListener(listener: TagManagerListener) {
listeners.remove(listener)
}
fun getTags(): Set<TagDetails> = 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<FileTagDetails>().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<TagDetails>()
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
}

View file

@ -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()
}
}

View file

@ -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
}

View file

@ -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
}

View file

@ -1,6 +1,6 @@
<idea-plugin>
<extensions defaultExtensionNs="com.intellij">
<applicationService
serviceImplementation="ee.carlrobert.codegpt.codecompletions.psi.structure.KotlinFileAnalyzer"/>
serviceImplementation="ee.carlrobert.codegpt.psistructure.KotlinFileAnalyzer"/>
</extensions>
</idea-plugin>

View file

@ -54,6 +54,8 @@
instance="ee.carlrobert.codegpt.settings.documentation.DocumentationsConfigurable"/>
<applicationConfigurable id="settings.codegpt.configuration" parentId="settings.codegpt" displayName="Configuration"
instance="ee.carlrobert.codegpt.settings.configuration.ConfigurationConfigurable"/>
<applicationConfigurable id="settings.codegpt.chat" parentId="settings.codegpt" displayName="Chat Settings"
instance="ee.carlrobert.codegpt.settings.chat.ChatConfigurationConfigurable"/>
<applicationConfigurable id="settings.codegpt.advanced" parentId="settings.codegpt" displayName="Advanced Settings"
instance="ee.carlrobert.codegpt.settings.advanced.AdvancedSettingsConfigurable"/>
<applicationConfigurable

View file

@ -0,0 +1 @@
<svg fill="#000000" height="16" width="16" version="1.2" baseProfile="tiny" id="Layer_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" viewBox="-63 65 128 128" xml:space="preserve"><path d="M52.4,133.6H30l15.9-19.8c0.7-0.8,0.7-2,0-2.7c-0.8-0.7-1.9-0.6-2.6,0.1l-10.5,12.7v-13.3c0-0.8-0.7-1.4-1.5-1.4 s-1.4,0.6-1.4,1.4v13.7l-15.7,20l-10.3,8v-6.1v-7.6v-14.1l20.9-20.9c0.9-0.9,1-2.5,0-3.5c-0.9-0.9-2.5-0.8-3.4,0.1l-10.2,10.2 l-1.1-12.6L27.9,80c0.7-0.7,0.7-1.8,0-2.5c-0.7-0.7-1.8-0.7-2.5,0L9.5,93.4L6.7,68.2c0-1.1-0.9-1.9-1.9-1.9c-1.1,0-1.8,0.9-1.8,1.9 l1.6,16.4l-9.4-9.4c-0.6-0.6-1.5-0.5-2.1,0c-0.6,0.6-0.5,1.4,0,2l9.7,9.7l2.9,25.2l-9.2,9.2c0,0-1.6,1.7-1.6,3.7c0,1,0,9.1,0,17.5 c0,4.5,0,9.1,0,12.6c0,3.7,0,6.2,0,6.2l-2.9-3.1l-6.3-6.3l-7-7v-38.7l8.5-8.5c0.6-0.6,0.6-1.5,0-2.1c-0.6-0.6-1.5-0.6-2.1,0 l-8.4,8.4V87.5c0-1.1-0.9-2-2-2c-1.1,0.1-2,0.9-2,2l0.1,32.5L-38,109.2c-0.9-0.9-2.3-0.9-3.2,0c-0.8,0.8-0.8,2.2,0,3.1l13.9,13.9 v15.2l-15-15.1c-1.1-1.1-3-1.1-4.1,0.1c-1.1,1.1-1.1,2.9,0,4l26.7,26.7l6.3,6.3l8.3,8.3v20h8.9v-32.1l10.3-8h19.5 c1.3,0,2.5-1.1,2.5-2.5c0-1.3-1.2-2.4-2.5-2.4H19.2l8.2-9.7h25.1c1,0,1.7-0.8,1.7-1.8C54.2,134.4,53.4,133.6,52.4,133.6z" /></svg>

After

Width:  |  Height:  |  Size: 1.2 KiB

View file

@ -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: