From ce69eba6ef060c1d93f08f7dd602ba3ccf58b719 Mon Sep 17 00:00:00 2001 From: Stephan Lukasczyk Date: Mon, 25 Nov 2024 11:51:49 +0100 Subject: [PATCH] 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. --- .../org/jetbrains/research/testspark/helpers/LLMHelper.kt | 8 ++++---- 1 file changed, 4 insertions(+), 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 index 0f00c9fa6..1dd86ab28 100644 --- a/src/main/kotlin/org/jetbrains/research/testspark/helpers/LLMHelper.kt +++ b/src/main/kotlin/org/jetbrains/research/testspark/helpers/LLMHelper.kt @@ -378,9 +378,9 @@ object LLMHelper { val models = mutableListOf() 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 @@ -388,7 +388,7 @@ object LLMHelper { .getAsJsonArray("supportedGenerationMethods") .map { method -> method.asString } if (methods.contains("generateContent")) { - models.add(id) + models.add(id.removePrefix("models/")) } } }