Add missing settings storage for Google AI

JetBrains-Research · Nov 25, 2024 · c415f4b · c415f4b
1 parent e2e6706
commit c415f4b
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 1 deletion.
diff --git a/src/main/kotlin/org/jetbrains/research/testspark/actions/llm/LLMSetupPanelBuilder.kt b/src/main/kotlin/org/jetbrains/research/testspark/actions/llm/LLMSetupPanelBuilder.kt
@@ -146,6 +146,10 @@ class LLMSetupPanelBuilder(e: AnActionEvent, private val project: Project) : Pan
                 llmSettingsState.huggingFaceToken = llmPlatforms[index].token
                 llmSettingsState.huggingFaceModel = llmPlatforms[index].model
             }
+            if (llmPlatforms[index].name == llmSettingsState.geminiName) {
+                llmSettingsState.geminiToken = llmPlatforms[index].token
+                llmSettingsState.geminiModel = llmPlatforms[index].model
+            }
         }
         llmSettingsState.junitVersion = junitSelector.selectedItem!! as JUnitVersion
 

diff --git a/src/main/kotlin/org/jetbrains/research/testspark/helpers/LLMHelper.kt b/src/main/kotlin/org/jetbrains/research/testspark/helpers/LLMHelper.kt
@@ -98,6 +98,12 @@ object LLMHelper {
                     modelSelector.selectedItem = settingsState.huggingFaceModel
                     llmPlatforms[index].model = modelSelector.selectedItem!!.toString()
                 }
+                if (llmPlatforms[index].name == settingsState.geminiName &&
+                    llmPlatforms[index].name == platformSelector.selectedItem!!.toString()
+                ) {
+                    modelSelector.selectedItem = settingsState.geminiModel
+                    llmPlatforms[index].model = modelSelector.selectedItem!!.toString()
+                }
             }
             modelSelector.isEnabled = true
             if (models.contentEquals(arrayOf(""))) modelSelector.isEnabled = false
@@ -135,6 +141,12 @@ object LLMHelper {
                 llmUserTokenField.text = settingsState.huggingFaceToken
                 llmPlatforms[index].token = settingsState.huggingFaceToken
             }
+            if (llmPlatforms[index].name == settingsState.geminiName &&
+                llmPlatforms[index].name == platformSelector.selectedItem!!.toString()
+            ) {
+                llmUserTokenField.text = settingsState.geminiToken
+                llmPlatforms[index].token = settingsState.geminiToken
+            }
         }
     }
 

diff --git a/src/main/kotlin/org/jetbrains/research/testspark/settings/llm/LLMSettingsComponent.kt b/src/main/kotlin/org/jetbrains/research/testspark/settings/llm/LLMSettingsComponent.kt
@@ -45,7 +45,7 @@ class LLMSettingsComponent(private val project: Project) : SettingsComponent {
 
     // Models
     private var modelSelector = ComboBox(arrayOf(""))
-    private var platformSelector = ComboBox(arrayOf(llmSettingsState.openAIName, llmSettingsState.huggingFaceName))
+    private var platformSelector = ComboBox(arrayOf(llmSettingsState.openAIName, llmSettingsState.huggingFaceName, llmSettingsState.geminiName))
 
     // Default LLM Requests
     private var defaultLLMRequestsSeparator =

diff --git a/src/main/kotlin/org/jetbrains/research/testspark/settings/llm/LLMSettingsConfigurable.kt b/src/main/kotlin/org/jetbrains/research/testspark/settings/llm/LLMSettingsConfigurable.kt
@@ -46,6 +46,10 @@ class LLMSettingsConfigurable(private val project: Project) : SettingsConfigurab
                 settingsComponent!!.llmPlatforms[index].token = llmSettingsState.huggingFaceToken
                 settingsComponent!!.llmPlatforms[index].model = llmSettingsState.huggingFaceModel
             }
+            if (settingsComponent!!.llmPlatforms[index].name == llmSettingsState.geminiName) {
+                settingsComponent!!.llmPlatforms[index].token = llmSettingsState.geminiToken
+                settingsComponent!!.llmPlatforms[index].model = llmSettingsState.geminiModel
+            }
         }
         settingsComponent!!.currentLLMPlatformName = llmSettingsState.currentLLMPlatformName
         settingsComponent!!.maxLLMRequest = llmSettingsState.maxLLMRequest
@@ -89,6 +93,10 @@ class LLMSettingsConfigurable(private val project: Project) : SettingsConfigurab
                 modified = modified or (settingsComponent!!.llmPlatforms[index].token != llmSettingsState.huggingFaceToken)
                 modified = modified or (settingsComponent!!.llmPlatforms[index].model != llmSettingsState.huggingFaceModel)
             }
+            if (settingsComponent!!.llmPlatforms[index].name == llmSettingsState.geminiName) {
+                modified = modified or (settingsComponent!!.llmPlatforms[index].token != llmSettingsState.geminiToken)
+                modified = modified or (settingsComponent!!.llmPlatforms[index].model != llmSettingsState.geminiModel)
+            }
         }
         modified = modified or (settingsComponent!!.currentLLMPlatformName != llmSettingsState.currentLLMPlatformName)
         modified = modified or (settingsComponent!!.maxLLMRequest != llmSettingsState.maxLLMRequest)
@@ -150,6 +158,10 @@ class LLMSettingsConfigurable(private val project: Project) : SettingsConfigurab
                 llmSettingsState.huggingFaceToken = settingsComponent!!.llmPlatforms[index].token
                 llmSettingsState.huggingFaceModel = settingsComponent!!.llmPlatforms[index].model
             }
+            if (settingsComponent!!.llmPlatforms[index].name == llmSettingsState.geminiName) {
+                llmSettingsState.geminiToken = settingsComponent!!.llmPlatforms[index].token
+                llmSettingsState.geminiModel = settingsComponent!!.llmPlatforms[index].model
+            }
         }
         llmSettingsState.currentLLMPlatformName = settingsComponent!!.currentLLMPlatformName
         llmSettingsState.maxLLMRequest = settingsComponent!!.maxLLMRequest