Skip to content

Commit

Permalink
Feat #157 Calc tokens for dir
Browse files Browse the repository at this point in the history
  • Loading branch information
stephanj committed Jul 5, 2024
1 parent dcea01c commit 3e1d36f
Show file tree
Hide file tree
Showing 5 changed files with 122 additions and 10 deletions.
Original file line number Diff line number Diff line change
@@ -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());
}
}
67 changes: 61 additions & 6 deletions src/main/java/com/devoxx/genie/service/ProjectContentService.java
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
8 changes: 8 additions & 0 deletions src/main/resources/META-INF/plugin.xml
Original file line number Diff line number Diff line change
Expand Up @@ -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>
2 changes: 1 addition & 1 deletion src/main/resources/application.properties
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
#Thu Jul 04 22:09:48 CEST 2024
#Fri Jul 05 08:27:25 CEST 2024
version=0.2.3

0 comments on commit 3e1d36f

Please sign in to comment.