hide services with missing api key & models not in users pricing plan

carlrobertoh · Sep 7, 2024 · 6e893be · 6e893be
1 parent 4ebe352
commit 6e893be
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 35 deletions.
diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/ModelComboBoxAction.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/textarea/ModelComboBoxAction.java
@@ -1,5 +1,10 @@
 package ee.carlrobert.codegpt.toolwindow.chat.ui.textarea;
 
+import static ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.ANTHROPIC_API_KEY;
+import static ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.AZURE_ACTIVE_DIRECTORY_TOKEN;
+import static ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.AZURE_OPENAI_API_KEY;
+import static ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.GOOGLE_API_KEY;
+import static ee.carlrobert.codegpt.credentials.CredentialsStore.CredentialKey.OPENAI_API_KEY;
 import static ee.carlrobert.codegpt.settings.service.ServiceType.ANTHROPIC;
 import static ee.carlrobert.codegpt.settings.service.ServiceType.AZURE;
 import static ee.carlrobert.codegpt.settings.service.ServiceType.CODEGPT;
@@ -22,6 +27,7 @@
 import ee.carlrobert.codegpt.CodeGPTKeys;
 import ee.carlrobert.codegpt.Icons;
 import ee.carlrobert.codegpt.completions.llama.LlamaModel;
+import ee.carlrobert.codegpt.credentials.CredentialsStore;
 import ee.carlrobert.codegpt.settings.GeneralSettings;
 import ee.carlrobert.codegpt.settings.service.ServiceType;
 import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTAvailableModels;
@@ -89,12 +95,13 @@ private AnAction[] getCodeGPTModelActions(Project project, Presentation presenta
   protected @NotNull DefaultActionGroup createPopupActionGroup(JComponent button) {
     var presentation = ((ComboBoxButton) button).getPresentation();
     var actionGroup = new DefaultActionGroup();
+    var credentialStore = CredentialsStore.INSTANCE;
 
     if (availableProviders.contains(CODEGPT)) {
       actionGroup.addSeparator("CodeGPT");
       actionGroup.addAll(getCodeGPTModelActions(project, presentation));
     }
-    if (availableProviders.contains(OPENAI)) {
+    if (availableProviders.contains(OPENAI) && credentialStore.isCredentialSet(OPENAI_API_KEY)) {
       actionGroup.addSeparator("OpenAI");
       List.of(
               OpenAIChatCompletionModel.GPT_4_O,
@@ -114,20 +121,22 @@ private AnAction[] getCodeGPTModelActions(Project project, Presentation presenta
           Icons.OpenAI,
           presentation));
     }
-    if (availableProviders.contains(ANTHROPIC)) {
+    if (availableProviders.contains(ANTHROPIC)
+            && credentialStore.isCredentialSet(ANTHROPIC_API_KEY)) {
       actionGroup.addSeparator("Anthropic");
       actionGroup.add(createModelAction(
           ANTHROPIC,
           "Anthropic (Claude)",
           Icons.Anthropic,
           presentation));
     }
-    if (availableProviders.contains(AZURE)) {
+    if (availableProviders.contains(AZURE) && (credentialStore.isCredentialSet(AZURE_OPENAI_API_KEY)
+            || credentialStore.isCredentialSet(AZURE_ACTIVE_DIRECTORY_TOKEN))) {
       actionGroup.addSeparator("Azure");
       actionGroup.add(
           createModelAction(AZURE, "Azure OpenAI", Icons.Azure, presentation));
     }
-    if (availableProviders.contains(GOOGLE)) {
+    if (availableProviders.contains(GOOGLE) && credentialStore.isCredentialSet(GOOGLE_API_KEY)) {
       actionGroup.addSeparator("Google");
       actionGroup.add(createModelAction(
           GOOGLE,

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTAvailableModels.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTAvailableModels.kt
@@ -9,48 +9,24 @@ object CodeGPTAvailableModels {
 
     @JvmStatic
     fun getToolWindowModels(pricingPlan: PricingPlan?): List<CodeGPTModel> {
-        val anonymousModels = listOf(
-            CodeGPTModel("GPT-4o", "gpt-4o", Icons.OpenAI, INDIVIDUAL),
-            CodeGPTModel("Claude 3.5 Sonnet", "claude-3.5-sonnet", Icons.Anthropic, INDIVIDUAL),
-            CodeGPTModel("Llama 3.1 (405B)", "llama-3.1-405b", Icons.Meta, INDIVIDUAL),
-            CodeGPTModel("DeepSeek Coder V2", "deepseek-coder-v2", Icons.DeepSeek, INDIVIDUAL),
-            CodeGPTModel("GPT-4o mini - FREE", "gpt-4o-mini", Icons.OpenAI, ANONYMOUS),
-            CodeGPTModel("Llama 3 (8B) - FREE", "llama-3-8b", Icons.Meta, ANONYMOUS)
-        )
-        if (pricingPlan == null) {
-            return anonymousModels
-        }
         return when (pricingPlan) {
-            ANONYMOUS -> anonymousModels
-
-            FREE -> listOf(
-                CodeGPTModel("GPT-4o", "gpt-4o", Icons.OpenAI, INDIVIDUAL),
-                CodeGPTModel("Claude 3.5 Sonnet", "claude-3.5-sonnet", Icons.Anthropic, INDIVIDUAL),
-                CodeGPTModel("GPT-4o mini", "gpt-4o-mini", Icons.OpenAI, ANONYMOUS),
-                CodeGPTModel("Llama 3 (70B)", "llama-3-70b", Icons.Meta, FREE),
-                CodeGPTModel("Mixtral (8x22B)", "mixtral-8x22b", Icons.CodeGPTModel, FREE),
-                CodeGPTModel("Code Llama (70B)", "codellama:chat", Icons.Meta, FREE),
-            )
-
-            else -> BASE_CHAT_MODELS
+            null, ANONYMOUS -> ALL_CHAT_MODELS.filter { it.individual == ANONYMOUS }
+            FREE -> ALL_CHAT_MODELS.filter { it.individual != INDIVIDUAL }
+            else -> ALL_CHAT_MODELS.filter { it.individual == INDIVIDUAL }
         }
     }
 
     @JvmStatic
-    val BASE_CHAT_MODELS: List<CodeGPTModel> = listOf(
+    val ALL_CHAT_MODELS: List<CodeGPTModel> = listOf(
         CodeGPTModel("GPT-4o", "gpt-4o", Icons.OpenAI, INDIVIDUAL),
         CodeGPTModel("GPT-4o mini", "gpt-4o-mini", Icons.OpenAI, ANONYMOUS),
         CodeGPTModel("Claude 3 Opus", "claude-3-opus", Icons.Anthropic, INDIVIDUAL),
         CodeGPTModel("Claude 3.5 Sonnet", "claude-3.5-sonnet", Icons.Anthropic, INDIVIDUAL),
         CodeGPTModel("Llama 3.1 (405B)", "llama-3.1-405b", Icons.Meta, INDIVIDUAL),
         CodeGPTModel("Llama 3 (70B)", "llama-3-70b", Icons.Meta, FREE),
+        CodeGPTModel("Llama 3 (8B)", "llama-3-8b", Icons.Meta, ANONYMOUS),
         CodeGPTModel("DeepSeek Coder V2", "deepseek-coder-v2", Icons.DeepSeek, INDIVIDUAL),
         CodeGPTModel("DBRX", "dbrx", Icons.Databricks, INDIVIDUAL),
-    )
-
-    @JvmStatic
-    val ALL_CHAT_MODELS: List<CodeGPTModel> = BASE_CHAT_MODELS + listOf(
-        CodeGPTModel("Llama 3 (8B) - FREE", "llama-3-8b", Icons.Meta, ANONYMOUS),
         CodeGPTModel("Code Llama (70B)", "codellama:chat", Icons.Meta, FREE),
         CodeGPTModel("Mixtral (8x22B)", "mixtral-8x22b", Icons.CodeGPTModel, FREE),
         CodeGPTModel("DeepSeek Coder (33B)", "deepseek-coder-33b", Icons.CodeGPTModel, FREE),
@@ -61,7 +37,7 @@ object CodeGPTAvailableModels {
     val CODE_MODELS: List<CodeGPTModel> = listOf(
         CodeGPTModel("GPT-3.5 Turbo Instruct", "gpt-3.5-turbo-instruct", Icons.OpenAI, INDIVIDUAL),
         CodeGPTModel("StarCoder (16B)", "starcoder-16b", Icons.CodeGPTModel, FREE),
-        CodeGPTModel("StarCoder (7B) - FREE", "starcoder-7b", Icons.CodeGPTModel, FREE),
+        CodeGPTModel("StarCoder (7B)", "starcoder-7b", Icons.CodeGPTModel, FREE),
         CodeGPTModel("WizardCoder Python (34B)", "wizardcoder-python", Icons.CodeGPTModel, FREE),
         CodeGPTModel("Phind Code LLaMA v2 (34B)", "phind-codellama", Icons.CodeGPTModel, FREE)
     )
@@ -77,4 +53,4 @@ data class CodeGPTModel(
     val code: String,
     val icon: Icon,
     val individual: PricingPlan
-)
+)