Fix #133 Use display model name

devoxx · Jul 1, 2024 · c10fa77 · c10fa77
1 parent 2f53f88
commit c10fa77
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/src/main/java/com/devoxx/genie/service/ProjectContentService.java b/src/main/java/com/devoxx/genie/service/ProjectContentService.java
@@ -1,5 +1,6 @@
 package com.devoxx.genie.service;
 
+import com.devoxx.genie.model.LanguageModel;
 import com.devoxx.genie.model.enumarations.ModelProvider;
 import com.devoxx.genie.ui.settings.DevoxxGenieStateService;
 import com.devoxx.genie.ui.util.NotificationUtil;
@@ -55,7 +56,7 @@ public CompletableFuture<String> getDirectoryContent(Project project,
     public void calculateTokensAndCost(Project project,
                                        int windowContext,
                                        ModelProvider provider,
-                                       String modelName) {
+                                       LanguageModel languageModel) {
         if (!DefaultLLMSettings.isApiBasedProvider(provider)) {
             getProjectContent(project, windowContext, true)
                 .thenAccept(projectContent -> {
@@ -68,15 +69,15 @@ public void calculateTokensAndCost(Project project,
         }
 
         DevoxxGenieStateService settings = DevoxxGenieStateService.getInstance();
-        double inputCost = settings.getModelInputCost(provider, modelName);
+        double inputCost = settings.getModelInputCost(provider, languageModel.getModelName());
 
         getProjectContent(project, windowContext, true)
             .thenAccept(projectContent -> {
                 int tokenCount = ENCODING.countTokens(projectContent);
                 double estimatedInputCost = calculateCost(tokenCount, inputCost);
                 String message = String.format("Project contains %s. Estimated min. cost using %s is $%.6f",
                     WindowContextFormatterUtil.format(tokenCount, "tokens"),
-                    modelName,
+                    languageModel.getDisplayName(),
                     estimatedInputCost);
                 NotificationUtil.sendNotification(project, message);
             });

diff --git a/src/main/java/com/devoxx/genie/ui/component/TokenUsageBar.java b/src/main/java/com/devoxx/genie/ui/component/TokenUsageBar.java
@@ -6,7 +6,6 @@
 import javax.swing.*;
 import java.awt.*;
 
-@Setter
 public class TokenUsageBar extends JComponent {
     private int maxTokens = 100;
     private int usedTokens;
@@ -15,9 +14,9 @@ public TokenUsageBar() {
         setPreferredSize(new Dimension(200, 20));
     }
 
-    public void reset() {
-        this.usedTokens = 0;
-        this.maxTokens = 100;
+    public void setMaxTokens(int maxTokens) {
+        this.maxTokens = maxTokens;
+        repaint();
     }
 
     public void setTokens(int usedTokens, int maxTokens) {
@@ -26,6 +25,12 @@ public void setTokens(int usedTokens, int maxTokens) {
         repaint();
     }
 
+    public void reset() {
+        this.usedTokens = 0;
+        this.maxTokens = 100;
+        repaint();
+    }
+
     @Override
     protected void paintComponent(Graphics g) {
         super.paintComponent(g);

diff --git a/src/main/java/com/devoxx/genie/ui/panel/ActionButtonsPanel.java b/src/main/java/com/devoxx/genie/ui/panel/ActionButtonsPanel.java
@@ -600,7 +600,7 @@ private void calculateTokensAndCost() {
             project,
             getWindowContext(),
             selectedProvider,
-            selectedModel.getModelName()
+            selectedModel
         );
     }