diff --git a/src/main/java/ee/carlrobert/codegpt/ui/checkbox/PsiElementCheckboxTree.java b/src/main/java/ee/carlrobert/codegpt/ui/checkbox/PsiElementCheckboxTree.java
index 32140d0a..dd822278 100644
--- a/src/main/java/ee/carlrobert/codegpt/ui/checkbox/PsiElementCheckboxTree.java
+++ b/src/main/java/ee/carlrobert/codegpt/ui/checkbox/PsiElementCheckboxTree.java
@@ -40,6 +40,9 @@ public class PsiElementCheckboxTree extends FileCheckboxTree {
   }
 
   private static CheckedTreeNode createNode(PsiElement element) {
+    if (element instanceof com.intellij.psi.PsiClass) {
+      element = element.getContainingFile();
+    }
     if (!(element instanceof PsiDirectory || element instanceof PsiFile)) {
       return null;
     }