Merge pull request #158 from devoxx/issue-157

Feat #157 Calc tokens for dir

src/main/java/com/devoxx/genie/action/CalcTokensForDirectoryAction.java

@@ -0,0 +1,47 @@
+package com.devoxx.genie.action;
+
+import com.devoxx.genie.service.ProjectContentService;
+import com.devoxx.genie.ui.util.NotificationUtil;
+import com.devoxx.genie.ui.util.WindowContextFormatterUtil;
+import com.intellij.openapi.actionSystem.AnAction;
+import com.intellij.openapi.actionSystem.AnActionEvent;
+import com.intellij.openapi.actionSystem.CommonDataKeys;
+import com.intellij.openapi.progress.ProgressIndicator;
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.openapi.progress.Task;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.vfs.VirtualFile;
+import org.jetbrains.annotations.NotNull;
+
+public class CalcTokensForDirectoryAction extends AnAction {
+
+    @Override
+    public void actionPerformed(@NotNull AnActionEvent e) {
+        Project project = e.getProject();
+        VirtualFile selectedDir = e.getData(CommonDataKeys.VIRTUAL_FILE);
+
+        if (project == null || selectedDir == null || !selectedDir.isDirectory()) {
+            return;
+        }
+
+        ProgressManager.getInstance().run(new Task.Backgroundable(project, "Calculating Tokens", false) {
+            @Override
+            public void run(@NotNull ProgressIndicator indicator) {
+                ProjectContentService.getInstance()
+                    .getDirectoryContentAndTokens(project, selectedDir, Integer.MAX_VALUE, true)
+                    .thenAccept(result -> {
+                        String message = String.format("Directory '%s' contains approximately %s tokens",
+                            selectedDir.getName(),
+                            WindowContextFormatterUtil.format(result.getTokenCount()));
+                        NotificationUtil.sendNotification(project, message);
+                    });
+            }
+        });
+    }
+
+    @Override
+    public void update(@NotNull AnActionEvent e) {
+        VirtualFile file = e.getData(CommonDataKeys.VIRTUAL_FILE);
+        e.getPresentation().setEnabledAndVisible(file != null && file.isDirectory());
+    }
+}

src/main/java/com/devoxx/genie/service/ProjectContentService.java

@@ -13,11 +13,13 @@
 import com.knuddels.jtokkit.Encodings;
 import com.knuddels.jtokkit.api.Encoding;
 import com.knuddels.jtokkit.api.EncodingType;
+import org.jetbrains.annotations.NotNull;
 
 import java.awt.*;
 import java.awt.datatransfer.Clipboard;
 import java.awt.datatransfer.StringSelection;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * The ProjectContentService class provides methods to retrieve and
@@ -87,14 +89,23 @@ public CompletableFuture<ContentResult> getDirectoryContentAndTokens(Project pro
                                                                          VirtualFile directory,
                                                                          int tokenLimit,
                                                                          boolean isTokenCalculation) {
-        return ProjectScannerService.getInstance()
-            .scanProject(project, directory, tokenLimit, isTokenCalculation)
-            .thenApply(content -> {
-                int tokenCount = ENCODING.countTokens(content);
-                return new ContentResult(content, tokenCount);
-            });
+        return CompletableFuture.supplyAsync(() -> {
+            AtomicLong totalTokens = new AtomicLong(0);
+            StringBuilder content = new StringBuilder();
+
+            processDirectoryRecursively(project, directory, content, totalTokens, isTokenCalculation);
+
+            return new ContentResult(content.toString(), totalTokens.intValue());
+        });
     }
 
+    /**
+     * Calculates the number of tokens and estimated cost for a specified Project.
+     * @param project The Project to scan for content
+     * @param windowContext Integer representing the desired Window Context (ignored in this implementation)
+     * @param provider ModelProvider enum value representing the provider to use for cost calculation
+     * @param languageModel LanguageModel object representing the model to use for cost calculation
+     */
     public void calculateTokensAndCost(Project project,
                                        int windowContext,
                                        ModelProvider provider,
@@ -125,6 +136,50 @@ public void calculateTokensAndCost(Project project,
             });
     }
 
+    /**
+     * Processes a directory recursively, calculating the number of tokens and building a content string.
+     * @param project The Project containing the directory to scan
+     * @param directory VirtualFile representing the directory to scan
+     * @param content StringBuilder object to hold the content of the scanned files
+     * @param totalTokens AtomicLong object to hold the total token count
+     * @param isTokenCalculation Boolean flag indicating whether to calculate tokens or not
+     */
+    private void processDirectoryRecursively(Project project,
+                                             VirtualFile directory,
+                                             StringBuilder content,
+                                             AtomicLong totalTokens,
+                                             boolean isTokenCalculation) {
+        DevoxxGenieStateService settings = DevoxxGenieStateService.getInstance();
+
+        for (VirtualFile child : directory.getChildren()) {
+            if (child.isDirectory()) {
+                if (!settings.getExcludedDirectories().contains(child.getName())) {
+                    processDirectoryRecursively(project, child, content, totalTokens, isTokenCalculation);
+                }
+            } else if (shouldIncludeFile(child, settings)) {
+                String fileContent = readFileContent(child);
+                if (!isTokenCalculation) {
+                    content.append("File: ").append(child.getPath()).append("\n");
+                    content.append(fileContent).append("\n\n");
+                }
+                totalTokens.addAndGet(ENCODING.countTokens(fileContent));
+            }
+        }
+    }
+
+    private boolean shouldIncludeFile(@NotNull VirtualFile file, DevoxxGenieStateService settings) {
+        String extension = file.getExtension();
+        return extension != null && settings.getIncludedFileExtensions().contains(extension.toLowerCase());
+    }
+
+    private @NotNull String readFileContent(VirtualFile file) {
+        try {
+            return new String(file.contentsToByteArray());
+        } catch (Exception e) {
+            return "Error reading file: " + e.getMessage();
+        }
+    }
+
     private double calculateCost(int tokenCount, double tokenCost) {
         return (tokenCount / 1_000_000.0) * tokenCost;
     }

src/main/java/com/devoxx/genie/ui/util/WindowContextFormatterUtil.java

@@ -5,10 +5,12 @@
 public class WindowContextFormatterUtil {
 
     public static @NotNull String format(int tokens, String suffix) {
-        if (tokens >= 1_000_000) {
-            return String.format("%dM %s", tokens / 1_000_000, suffix);
+        if (tokens >= 1_000_000_000) {
+            return String.format("%.2fB %s", tokens / 1_000_000_000.0, suffix);
+        } else if (tokens >= 1_000_000) {
+            return String.format("%.2fM %s", tokens / 1_000_000.0, suffix);
         } else if (tokens >= 1_000) {
-            return String.format("%dK %s", tokens / 1_000, suffix);
+            return String.format("%.2fK %s", tokens / 1_000.0, suffix);
         } else {
             return String.format("%d %s", tokens, suffix);
         }

src/main/resources/META-INF/plugin.xml

@@ -319,5 +319,13 @@
                 description="Copy the selected directory content to clipboard">
             <add-to-group group-id="ProjectViewPopupMenu" anchor="after" relative-to-action="AddDirectoryToContextWindow"/>
         </action>
+
+        <action id="CalcTokenDirectory"
+                class="com.devoxx.genie.action.CalcTokensForDirectoryAction"
+                text="Calc Tokens for Directory"
+                icon="/icons/pluginIcon.svg"
+                description="Calculate the tokens for selected directory">
+            <add-to-group group-id="ProjectViewPopupMenu" anchor="after" relative-to-action="CopyDirectoryToClipboard"/>
+        </action>
     </actions>
 </idea-plugin>

src/main/resources/application.properties

@@ -1,2 +1,2 @@
-#Thu Jul 04 22:09:48 CEST 2024
+#Fri Jul 05 08:27:25 CEST 2024
 version=0.2.3