devoxx / DevoxxGenieIDEAPlugin

DevoxxGenie is a plugin for IntelliJ IDEA that uses local LLM's (Ollama, LMStudio, GPT4All, Llama.cpp and Exo) and Cloud based LLMs to help review, test, explain your project code.
https://devoxx.com
MIT License
156 stars 24 forks source link

Incorporate code into the prompt window context based on the selected AST options #216

Open stephanj opened 1 month ago

stephanj commented 1 month ago

This AST featured used PSIClass which is only supported in IntelliJ IDEA but not in PyCharm, CLion etc. So removed this feature for now.

However it would still be nice to have something similar, maybe using AST instead?

The AST options:

  1. Include Project Parent Class(es)
  2. Include Class References
  3. Include Field References

The "old" PSI analyzer:

MessageCreationService.class

private @NotNull UserMessage constructUserMessageWithEditorContent(@NotNull ChatMessageContext chatMessageContext) {
        StringBuilder stringBuilder = new StringBuilder();

        // The user prompt is always added
        appendIfNotEmpty(stringBuilder, chatMessageContext.getUserPrompt());

        // Add the editor content or selected text
        String editorContent = getEditorContentOrSelectedText(chatMessageContext);

        if (!editorContent.isEmpty()) {
            // Add the context prompt if it is not empty
            appendIfNotEmpty(stringBuilder, CONTEXT_PROMPT);
            appendIfNotEmpty(stringBuilder, editorContent);
        }

        if (DevoxxGenieSettingsServiceProvider.getInstance().getAstMode()) {
>>>            addASTContext(chatMessageContext, stringBuilder);
        }

        UserMessage userMessage = new UserMessage(stringBuilder.toString());
        chatMessageContext.setUserMessage(userMessage);
        return userMessage;
    }

 /**
     * Add AST prompt context (selected code snippet) to the chat message.
     * @param chatMessageContext the chat message context
     * @param sb                 the string builder
     */
    private void addASTContext(@NotNull ChatMessageContext chatMessageContext,
                                      @NotNull StringBuilder sb) {
        sb.append("\n\nRelated classes:\n\n");
        List<VirtualFile> tempFiles = new ArrayList<>();

        chatMessageContext.getEditorInfo().getSelectedFiles().forEach(file ->
            PSIAnalyzerService.getInstance().analyze(chatMessageContext.getProject(), file)
                .ifPresent(psiClasses ->
                    psiClasses.forEach(psiClass -> {
                        tempFiles.add(psiClass.getContainingFile().getVirtualFile());
                        sb.append(psiClass.getText()).append("\n");
                    })));

        chatMessageContext.getEditorInfo().getSelectedFiles().addAll(tempFiles);
    }

The PSIAnalyzerService:

package com.devoxx.genie.service;

import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.*;
import com.intellij.psi.util.InheritanceUtil;
import com.intellij.psi.util.PsiTreeUtil;

import org.jetbrains.annotations.NotNull;

import java.util.*;

import static com.intellij.openapi.application.ActionsKt.runReadAction;

public class PSIAnalyzerService {

    public static PSIAnalyzerService getInstance() {
        return ApplicationManager.getApplication().getService(PSIAnalyzerService.class);
    }

    /**
     * Analyze the PSI tree of a file and return a list of related classes.
     * @param project the project
     * @param file the virtual file
     * @return list of PSI classes
     */
    public Optional<Set<PsiClass>> analyze(Project project, VirtualFile file) {

        Set<PsiClass> relatedClasses = new HashSet<>();

        PsiFile psiFile = runReadAction(() -> PsiManager.getInstance(project).findFile(file));
        if (!(psiFile instanceof PsiJavaFile javaFile)) {
            return Optional.empty();
        }

        for (PsiClass psiClass : runReadAction(javaFile::getClasses)) {
            if (DevoxxGenieSettingsServiceProvider.getInstance().getAstParentClass()) {
                extractBaseClass(psiClass, relatedClasses);
            }
            if (DevoxxGenieSettingsServiceProvider.getInstance().getAstClassReference()) {
                extractReferenceClasses(psiClass, relatedClasses);
            }
            if (DevoxxGenieSettingsServiceProvider.getInstance().getAstFieldReference()) {
                PsiField[] fields = runReadAction(psiClass::getFields);
                extractPSIFields(fields, relatedClasses);
            }
        }

        return Optional.of(relatedClasses);
    }

    /**
     * Extract the base class of a PsiClass and add it to the list of related classes.
     * @param psiClass the PsiClass
     * @param relatedClasses the list of related classes
     */
    private void extractBaseClass(@NotNull PsiClass psiClass, Set<PsiClass> relatedClasses) {
        runReadAction(() -> {
            InheritanceUtil.getSuperClasses(psiClass, relatedClasses, false);
            return null;
        });
    }

    /**
     * Extract all referenced classes in a PsiClass and add them to the list of related classes.
     * @param psiClass the PsiClass
     * @param relatedClasses the list of related classes
     */
    private void extractReferenceClasses(PsiClass psiClass, Set<PsiClass> relatedClasses) {
        PsiClass[] referencedClasses = PsiTreeUtil.getChildrenOfType(psiClass, PsiClass.class);
        if (referencedClasses != null) {
            Collections.addAll(relatedClasses, referencedClasses);
        }
    }

    /**
     * Extract all fields in a PsiClass and add the corresponding classes to the list of related classes.
     * @param psiFields     the PsiFields
     * @param relatedClasses the list of related classes
     */
    private void extractPSIFields(@NotNull PsiField @NotNull [] psiFields,
                                  Set<PsiClass> relatedClasses) {
        for (PsiField field : psiFields) {
            // Get the type of the field
            PsiType fieldType = runReadAction(field::getType);
            // If the type is a class type, add the corresponding PsiClass to the list
            if (fieldType instanceof PsiClassType) {
                PsiClass fieldClass = runReadAction(() -> ((PsiClassType) fieldType).resolve());
                if (fieldClass != null) {
                    relatedClasses.add(fieldClass);
                }
            }
        }
    }
}