diff --git a/src/main/java/com/devoxx/genie/action/AddDirectoryAction.java b/src/main/java/com/devoxx/genie/action/AddDirectoryAction.java index 468430ec..20d85787 100644 --- a/src/main/java/com/devoxx/genie/action/AddDirectoryAction.java +++ b/src/main/java/com/devoxx/genie/action/AddDirectoryAction.java @@ -1,5 +1,6 @@ package com.devoxx.genie.action; +import com.devoxx.genie.model.enumarations.ModelProvider; import com.devoxx.genie.service.FileListManager; import com.devoxx.genie.service.ProjectContentService; import com.devoxx.genie.ui.settings.DevoxxGenieStateService; @@ -53,14 +54,16 @@ private void addDirectoryToContext(Project project, @NotNull VirtualFile directo if (!filesToAdd.isEmpty()) { fileListManager.addFiles(filesToAdd); - // Get the content and token count - ProjectContentService.getInstance().getDirectoryContentAndTokens(project, directory, Integer.MAX_VALUE, false) + ModelProvider selectedProvider = ModelProvider.valueOf(settings.getSelectedProvider()); + + ProjectContentService.getInstance() + .getDirectoryContentAndTokens(project, directory, false, selectedProvider) .thenAccept(result -> { int fileCount = filesToAdd.size(); int tokenCount = result.getTokenCount(); NotificationUtil.sendNotification(project, - String.format("Added %d files from directory: %s (Approximately %s tokens)", - fileCount, directory.getName(), WindowContextFormatterUtil.format(tokenCount))); + String.format("Added %d files from directory: %s (Approximately %s tokens using %s tokenizer)", + fileCount, directory.getName(), WindowContextFormatterUtil.format(tokenCount), selectedProvider.getName())); }); } } diff --git a/src/main/java/com/devoxx/genie/action/CalcTokensForDirectoryAction.java b/src/main/java/com/devoxx/genie/action/CalcTokensForDirectoryAction.java index 5d324dba..da996b50 100644 --- a/src/main/java/com/devoxx/genie/action/CalcTokensForDirectoryAction.java +++ b/src/main/java/com/devoxx/genie/action/CalcTokensForDirectoryAction.java @@ -1,6 +1,8 @@ package com.devoxx.genie.action; +import com.devoxx.genie.model.enumarations.ModelProvider; import com.devoxx.genie.service.ProjectContentService; +import com.devoxx.genie.ui.settings.DevoxxGenieStateService; import com.devoxx.genie.ui.util.NotificationUtil; import com.devoxx.genie.ui.util.WindowContextFormatterUtil; import com.intellij.openapi.actionSystem.AnAction; @@ -24,15 +26,19 @@ public void actionPerformed(@NotNull AnActionEvent e) { return; } + DevoxxGenieStateService stateService = DevoxxGenieStateService.getInstance(); + ModelProvider selectedProvider = ModelProvider.valueOf(stateService.getSelectedProvider()); + 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) + .getDirectoryContentAndTokens(project, selectedDir, true, selectedProvider) .thenAccept(result -> { - String message = String.format("Directory '%s' contains approximately %s tokens", + String message = String.format("Directory '%s' contains approximately %s tokens (using %s tokenizer)", selectedDir.getName(), - WindowContextFormatterUtil.format(result.getTokenCount())); + WindowContextFormatterUtil.format(result.getTokenCount()), + selectedProvider.getName()); NotificationUtil.sendNotification(project, message); }); } diff --git a/src/main/java/com/devoxx/genie/service/ProjectContentService.java b/src/main/java/com/devoxx/genie/service/ProjectContentService.java index 3d3fad65..51fdfeea 100644 --- a/src/main/java/com/devoxx/genie/service/ProjectContentService.java +++ b/src/main/java/com/devoxx/genie/service/ProjectContentService.java @@ -78,22 +78,16 @@ public CompletableFuture getDirectoryContent(Project project, }); } - /** - * Retrieves and processes the content of a specified directory within a Project. - * Also calculates number of tokens in this content if required by user settings or provider configurations. - * @param project The Project containing the directory to be scanned - * @param directory VirtualFile representing the directory to scan for content - * @return ContentResult object holding both content and token count, optionally copied to clipboard based on flag - */ public CompletableFuture getDirectoryContentAndTokens(Project project, VirtualFile directory, - int tokenLimit, - boolean isTokenCalculation) { + boolean isTokenCalculation, + ModelProvider modelProvider) { return CompletableFuture.supplyAsync(() -> { AtomicLong totalTokens = new AtomicLong(0); StringBuilder content = new StringBuilder(); - processDirectoryRecursively(project, directory, content, totalTokens, isTokenCalculation); + Encoding encoding = getEncodingForProvider(modelProvider); + processDirectoryRecursively(project, directory, content, totalTokens, isTokenCalculation, encoding); return new ContentResult(content.toString(), totalTokens.intValue()); }); @@ -136,6 +130,19 @@ public void calculateTokensAndCost(Project project, }); } + private Encoding getEncodingForProvider(@NotNull ModelProvider provider) { + return switch (provider) { + case OpenAI, Anthropic, Gemini -> + Encodings.newDefaultEncodingRegistry().getEncoding(EncodingType.CL100K_BASE); + case Mistral, DeepInfra, Groq -> + // These often use the Llama tokenizer or similar + Encodings.newDefaultEncodingRegistry().getEncoding(EncodingType.R50K_BASE); + default -> + // Default to cl100k_base for unknown providers + Encodings.newDefaultEncodingRegistry().getEncoding(EncodingType.CL100K_BASE); + }; + } + /** * Processes a directory recursively, calculating the number of tokens and building a content string. * @param project The Project containing the directory to scan @@ -145,16 +152,17 @@ public void calculateTokensAndCost(Project project, * @param isTokenCalculation Boolean flag indicating whether to calculate tokens or not */ private void processDirectoryRecursively(Project project, - VirtualFile directory, + @NotNull VirtualFile directory, StringBuilder content, AtomicLong totalTokens, - boolean isTokenCalculation) { + boolean isTokenCalculation, + Encoding encoding) { DevoxxGenieStateService settings = DevoxxGenieStateService.getInstance(); for (VirtualFile child : directory.getChildren()) { if (child.isDirectory()) { if (!settings.getExcludedDirectories().contains(child.getName())) { - processDirectoryRecursively(project, child, content, totalTokens, isTokenCalculation); + processDirectoryRecursively(project, child, content, totalTokens, isTokenCalculation, encoding); } } else if (shouldIncludeFile(child, settings)) { String fileContent = readFileContent(child); @@ -162,7 +170,7 @@ private void processDirectoryRecursively(Project project, content.append("File: ").append(child.getPath()).append("\n"); content.append(fileContent).append("\n\n"); } - totalTokens.addAndGet(ENCODING.countTokens(fileContent)); + totalTokens.addAndGet(encoding.countTokens(fileContent)); } } } diff --git a/src/main/java/com/devoxx/genie/ui/DevoxxGenieToolWindowContent.java b/src/main/java/com/devoxx/genie/ui/DevoxxGenieToolWindowContent.java index 675f62b1..40cafcb9 100644 --- a/src/main/java/com/devoxx/genie/ui/DevoxxGenieToolWindowContent.java +++ b/src/main/java/com/devoxx/genie/ui/DevoxxGenieToolWindowContent.java @@ -321,8 +321,7 @@ private void processModelNameSelection(@NotNull ActionEvent e) { * Set the model provider and update the model names. */ private void handleModelProviderSelectionChange(@NotNull ActionEvent e) { - if (!e.getActionCommand() - .equals(Constant.COMBO_BOX_CHANGED) || !isInitializationComplete || isUpdatingModelNames) return; + if (!e.getActionCommand().equals(Constant.COMBO_BOX_CHANGED) || !isInitializationComplete || isUpdatingModelNames) return; isUpdatingModelNames = true; @@ -330,6 +329,9 @@ private void handleModelProviderSelectionChange(@NotNull ActionEvent e) { JComboBox comboBox = (JComboBox) e.getSource(); ModelProvider modelProvider = (ModelProvider) comboBox.getSelectedItem(); if (modelProvider != null) { + // Update the selectedProvider in DevoxxGenieStateService + DevoxxGenieStateService.getInstance().setSelectedProvider(modelProvider.getName()); + updateModelNamesComboBox(modelProvider.getName()); modelNameComboBox.setRenderer(new ModelInfoRenderer()); // Re-apply the renderer modelNameComboBox.revalidate(); diff --git a/src/main/java/com/devoxx/genie/ui/settings/DevoxxGenieStateService.java b/src/main/java/com/devoxx/genie/ui/settings/DevoxxGenieStateService.java index 8899cdd3..1136003c 100644 --- a/src/main/java/com/devoxx/genie/ui/settings/DevoxxGenieStateService.java +++ b/src/main/java/com/devoxx/genie/ui/settings/DevoxxGenieStateService.java @@ -56,6 +56,8 @@ public static DevoxxGenieStateService getInstance() { private Integer maxSearchResults = MAX_SEARCH_RESULTS; // Last selected language model + @Getter + @Setter private String selectedProvider; private String selectedLanguageModel; @@ -160,4 +162,5 @@ public void setModelWindowContext(ModelProvider provider, String modelName, int public void setLanguageModels(List models) { this.languageModels = new ArrayList<>(models); } + } diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml index af4cd48e..cbd2f28f 100644 --- a/src/main/resources/META-INF/plugin.xml +++ b/src/main/resources/META-INF/plugin.xml @@ -40,6 +40,7 @@
  • Feat #148: Create custom commands
  • Feat #157: Calc tokens for directory
  • Fix #153: Use the "Copy Project" settings when using "Add Directory to Context Window"
    • +
  • Feat #159: Introduce variable TokenCalculator based on selected LLM Provider

    • v0.2.2