From 3e1d36fcd735b8b6ace95db52ebf289975fb34a0 Mon Sep 17 00:00:00 2001 From: Stephan Janssen Date: Fri, 5 Jul 2024 08:28:33 +0200 Subject: [PATCH] Feat #157 Calc tokens for dir --- .../action/CalcTokensForDirectoryAction.java | 47 +++++++++++++ .../genie/service/ProjectContentService.java | 67 +++++++++++++++++-- .../ui/util/WindowContextFormatterUtil.java | 8 ++- src/main/resources/META-INF/plugin.xml | 8 +++ src/main/resources/application.properties | 2 +- 5 files changed, 122 insertions(+), 10 deletions(-) create mode 100644 src/main/java/com/devoxx/genie/action/CalcTokensForDirectoryAction.java diff --git a/src/main/java/com/devoxx/genie/action/CalcTokensForDirectoryAction.java b/src/main/java/com/devoxx/genie/action/CalcTokensForDirectoryAction.java new file mode 100644 index 00000000..5d324dba --- /dev/null +++ b/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()); + } +} diff --git a/src/main/java/com/devoxx/genie/service/ProjectContentService.java b/src/main/java/com/devoxx/genie/service/ProjectContentService.java index 32e26c4a..3d3fad65 100644 --- a/src/main/java/com/devoxx/genie/service/ProjectContentService.java +++ b/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 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; } diff --git a/src/main/java/com/devoxx/genie/ui/util/WindowContextFormatterUtil.java b/src/main/java/com/devoxx/genie/ui/util/WindowContextFormatterUtil.java index 4278bcd1..f977e163 100644 --- a/src/main/java/com/devoxx/genie/ui/util/WindowContextFormatterUtil.java +++ b/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); } diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml index f9a245df..30d0bba6 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -319,5 +319,13 @@ description="Copy the selected directory content to clipboard"> + + + + diff --git a/src/main/resources/application.properties b/src/main/resources/application.properties index 81bacecf..7022d475 100644 --- a/src/main/resources/application.properties +++ b/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