Skip to content

Commit

Permalink
Merge pull request #388 from devoxx/issue-387
Browse files Browse the repository at this point in the history 
Fix #387 : Use countTokensOrdinary and encodeOrdinary
  • Loading branch information
@stephanj
stephanj authored Dec 13, 2024
2 parents 4e6e068 + fa2c2e5 commit b33cd08
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 23 deletions.
4 changes: 3 additions & 1 deletion src/main/java/com/devoxx/genie/controller/ProjectContextController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,9 @@ public void addProjectContext() {
projectContext = "Project Context:\n" + projectContent.getContent();
isProjectContextAdded = true;
ApplicationManager.getApplication().invokeLater(() -> {
tokenCount = Encodings.newDefaultEncodingRegistry().getEncoding(EncodingType.CL100K_BASE).countTokens(projectContent.getContent());
tokenCount = Encodings.newDefaultEncodingRegistry()
.getEncoding(EncodingType.CL100K_BASE)
.countTokensOrdinary(projectContent.getContent());
actionButtonsPanel.updateAddProjectButton(isProjectContextAdded, tokenCount);
actionButtonsPanel.setAddProjectButtonEnabled(true);
actionButtonsPanel.updateTokenUsageBar(tokenCount, tokenLimit);
Expand Down
39 changes: 18 additions & 21 deletions src/main/java/com/devoxx/genie/service/projectscanner/ProjectScannerService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@
import java.util.Arrays;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicInteger;

public class ProjectScannerService {

Expand Down Expand Up @@ -60,20 +59,20 @@ public CompletableFuture<ScanContentResult> scanProject(Project project,

if (startDirectory == null) {
result.append("Directory Structure:\n");
fullContent = getContentFromModules(project, windowContextMaxTokens, result, scanContentResult);
fullContent = getContentFromModules(project, result, scanContentResult);
} else if (startDirectory.isDirectory()) {
result.append("Directory Structure:\n");
fullContent = processDirectory(project, startDirectory, result, scanContentResult, windowContextMaxTokens);
fullContent = processDirectory(project, startDirectory, result, scanContentResult);
} else {
// Handle the case where startDirectory is a file
result.append("File:\n");
fullContent = processSingleFile(project, startDirectory, result, scanContentResult);
fullContent = processSingleFile(startDirectory, result, scanContentResult);
}

String content = isTokenCalculation ? fullContent.toString() :
truncateToTokens(fullContent.toString(), windowContextMaxTokens, isTokenCalculation);

scanContentResult.setTokenCount(ENCODING.countTokens(content));
scanContentResult.setTokenCount(ENCODING.countTokensOrdinary(content));
scanContentResult.setContent(content);

return scanContentResult;
Expand All @@ -98,14 +97,14 @@ public ScanContentResult scanProjectSynchronously(Project project, VirtualFile s

if (startDirectory == null) {
result.append("Directory Structure:\n");
fullContent = getContentFromModules(project, windowContextMaxTokens, result, scanContentResult);
fullContent = getContentFromModules(project, result, scanContentResult);
} else if (startDirectory.isDirectory()) {
result.append("Directory Structure:\n");
fullContent = processDirectory(project, startDirectory, result, scanContentResult, windowContextMaxTokens);
fullContent = processDirectory(project, startDirectory, result, scanContentResult);
} else {
// Handle the case where startDirectory is a file
result.append("File:\n");
fullContent = processSingleFile(project, startDirectory, result, scanContentResult);
fullContent = processSingleFile(startDirectory, result, scanContentResult);
}

String content = isTokenCalculation ? fullContent.toString() : truncateToTokens(fullContent.toString(), windowContextMaxTokens, isTokenCalculation);
Expand Down Expand Up @@ -136,7 +135,6 @@ private void initGitignoreParser(Project project, VirtualFile startDirectory) {
}

private StringBuilder getContentFromModules(Project project,
int windowContextMaxTokens,
StringBuilder result,
ScanContentResult scanContentResult) {

Expand All @@ -157,31 +155,28 @@ private StringBuilder getContentFromModules(Project project,
.map(highestCommonRoot -> processDirectory(project,
highestCommonRoot,
result,
scanContentResult,
windowContextMaxTokens))
scanContentResult))
.orElseThrow();
}

private @NotNull StringBuilder processDirectory(Project project,
VirtualFile startDirectory,
@NotNull StringBuilder result,
ScanContentResult scanContentResult,
int windowContextMaxTokens) {
ScanContentResult scanContentResult) {
result.append(generateSourceTreeRecursive(startDirectory, 0));

result.append("\n\nFile Contents:\n");

ProjectFileIndex fileIndex = ProjectFileIndex.getInstance(project);

StringBuilder fullContent = new StringBuilder(result);
AtomicInteger currentTokens = new AtomicInteger(0);

walkThroughDirectory(startDirectory, fileIndex, fullContent, currentTokens, scanContentResult);
walkThroughDirectory(startDirectory, fileIndex, fullContent, scanContentResult);

return fullContent;
}

private @NotNull StringBuilder processSingleFile(Project project,
VirtualFile file,
private @NotNull StringBuilder processSingleFile(@NotNull VirtualFile file,
@NotNull StringBuilder result,
ScanContentResult scanContentResult) {
result.append(file.getName()).append("\n");
Expand All @@ -200,7 +195,6 @@ private StringBuilder getContentFromModules(Project project,
private void walkThroughDirectory(@NotNull VirtualFile directory,
@NotNull ProjectFileIndex fileIndex,
@NotNull StringBuilder fullContent,
@NotNull AtomicInteger currentTokens,
@NotNull ScanContentResult scanContentResult) {

VfsUtilCore.visitChildrenRecursively(directory, new VirtualFileVisitor<Void>() {
Expand All @@ -224,7 +218,9 @@ public boolean visitFile(@NotNull VirtualFile file) {
});
}

private void readFileContent(@NotNull VirtualFile file, @NotNull StringBuilder fullContent, @NotNull ScanContentResult scanContentResult) {
private void readFileContent(@NotNull VirtualFile file,
@NotNull StringBuilder fullContent,
@NotNull ScanContentResult scanContentResult) {
scanContentResult.incrementFileCount();
scanContentResult.addFile(Paths.get(file.getPath()));

Expand All @@ -249,10 +245,11 @@ private void readFileContent(@NotNull VirtualFile file, @NotNull StringBuilder f
}
}

private String truncateToTokens(String text,
private String truncateToTokens(@NotNull String text,
int windowContext,
boolean isTokenCalculation) {
IntArrayList tokens = ENCODING.encode(text);

IntArrayList tokens = ENCODING.encodeOrdinary(text);
if (tokens.size() <= windowContext) {
return text;
}
Expand Down
2 changes: 1 addition & 1 deletion src/main/resources/application.properties
View file
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
#Thu Dec 12 19:55:03 CET 2024
#Fri Dec 13 18:33:45 CET 2024
version=0.4.5

0 comments on commit b33cd08

Please sign in to comment.