Remove the models prefix for Google AI

Every model retrieved from the API is prefixed by `models/` which does not add any value to the user, thus we remove it.
JetBrains-Research · Nov 25, 2024 · ce69eba · ce69eba
1 parent c415f4b
commit ce69eba
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/main/kotlin/org/jetbrains/research/testspark/helpers/LLMHelper.kt b/src/main/kotlin/org/jetbrains/research/testspark/helpers/LLMHelper.kt
@@ -378,17 +378,17 @@ object LLMHelper {
         val models = mutableListOf<String>()
 
         try {
-            httpRequest.connect { it ->
-                if ((it.connection as HttpURLConnection).responseCode == HttpURLConnection.HTTP_OK) {
-                    val jsonObject = JsonParser.parseString(it.readString()).asJsonObject
+            httpRequest.connect { request ->
+                if ((request.connection as HttpURLConnection).responseCode == HttpURLConnection.HTTP_OK) {
+                    val jsonObject = JsonParser.parseString(request.readString()).asJsonObject
                     val dataArray = jsonObject.getAsJsonArray("models")
                     for (dataObject in dataArray) {
                         val id = dataObject.asJsonObject.getAsJsonPrimitive("name").asString
                         val methods = dataObject.asJsonObject
                             .getAsJsonArray("supportedGenerationMethods")
                             .map { method -> method.asString }
                         if (methods.contains("generateContent")) {
-                            models.add(id)
+                            models.add(id.removePrefix("models/"))
                         }
                     }
                 }