From 2ddcde136373475f10cb80c48668f99bd1513534 Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Mon, 19 Aug 2024 01:17:52 +0300
Subject: [PATCH] feat: image support for Custom OpenAI provider (#649)

---
 .../codegpt/settings/service/ServiceType.java | 38 -------------------
 .../codegpt/ui/textarea/UserInputPanel.kt     | 34 ++++++++++++++++-
 2 files changed, 33 insertions(+), 39 deletions(-)

diff --git a/src/main/java/ee/carlrobert/codegpt/settings/service/ServiceType.java b/src/main/java/ee/carlrobert/codegpt/settings/service/ServiceType.java
index 0a95adad..e8c8216f 100644
--- a/src/main/java/ee/carlrobert/codegpt/settings/service/ServiceType.java
+++ b/src/main/java/ee/carlrobert/codegpt/settings/service/ServiceType.java
@@ -1,15 +1,7 @@
 package ee.carlrobert.codegpt.settings.service;
 
-import static ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel.GPT_4_O;
-import static ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel.GPT_4_O_MINI;
-import static ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel.GPT_4_VISION_PREVIEW;
-
-import com.intellij.openapi.application.ApplicationManager;
 import ee.carlrobert.codegpt.CodeGPTBundle;
-import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings;
-import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
 public enum ServiceType {
@@ -64,34 +56,4 @@ public enum ServiceType {
     }
     return serviceType;
   }
-
-  public boolean isImageActionSupported() {
-    return switch (this) {
-      case ANTHROPIC:
-      case OLLAMA:
-        yield true;
-      case CODEGPT:
-        var codegptModel = ApplicationManager.getApplication()
-            .getService(CodeGPTServiceSettings.class)
-            .getState()
-            .getChatCompletionSettings()
-            .getModel();
-        yield List.of(
-            "gpt-4o",
-            "gpt-4o-mini",
-            "claude-3-opus",
-            "claude-3.5-sonnet"
-        ).contains(codegptModel);
-      case OPENAI:
-        var openaiModel = ApplicationManager.getApplication().getService(OpenAISettings.class)
-            .getState()
-            .getModel();
-        yield List.of(
-            GPT_4_VISION_PREVIEW.getCode(),
-            GPT_4_O.getCode(),
-            GPT_4_O_MINI.getCode()).contains(openaiModel);
-      default:
-        yield false;
-    };
-  }
 }
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt
index 12d1e815..af1f2dfc 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/UserInputPanel.kt
@@ -18,10 +18,14 @@ import ee.carlrobert.codegpt.actions.AttachImageAction
 import ee.carlrobert.codegpt.conversations.ConversationService
 import ee.carlrobert.codegpt.conversations.ConversationsState
 import ee.carlrobert.codegpt.settings.GeneralSettings
+import ee.carlrobert.codegpt.settings.service.ServiceType
+import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings
+import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings
 import ee.carlrobert.codegpt.toolwindow.chat.ChatToolWindowContentManager
 import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.ModelComboBoxAction
 import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.TotalTokensPanel
 import ee.carlrobert.codegpt.ui.IconActionButton
+import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel
 import java.awt.*
 import javax.swing.JPanel
 
@@ -143,6 +147,34 @@ class UserInputPanel(
     }
 
     private fun isImageActionSupported(): Boolean {
-        return service<GeneralSettings>().state.selectedService.isImageActionSupported
+        return when (service<GeneralSettings>().state.selectedService) {
+            ServiceType.CUSTOM_OPENAI,
+            ServiceType.ANTHROPIC,
+            ServiceType.OLLAMA -> true
+
+            ServiceType.CODEGPT -> {
+                listOf(
+                    "gpt-4o",
+                    "gpt-4o-mini",
+                    "claude-3-opus",
+                    "claude-3.5-sonnet"
+                ).contains(
+                    service<CodeGPTServiceSettings>()
+                        .state
+                        .chatCompletionSettings
+                        .model
+                )
+            }
+
+            ServiceType.OPENAI -> {
+                listOf(
+                    OpenAIChatCompletionModel.GPT_4_VISION_PREVIEW.code,
+                    OpenAIChatCompletionModel.GPT_4_O.code,
+                    OpenAIChatCompletionModel.GPT_4_O_MINI.code
+                ).contains(service<OpenAISettings>().state.model)
+            }
+
+            else -> false
+        }
     }
 }