Generalize naming support

google-gemini · Jan 31, 2024 · 6430fa8 · 6430fa8
1 parent 5a24830
commit 6430fa8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/internal/api/APIController.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/internal/api/APIController.kt
@@ -114,7 +114,7 @@ internal class APIController(
  * Models must be prepended with the `models/` prefix when communicating with the backend.
  */
 private fun fullModelName(name: String): String =
-  name.takeIf { it.startsWith("models/") || it.startsWith("tunedModels/") } ?: "models/$name"
+  name.takeIf { it.contains("/") } ?: "models/$name"
 
 /**
  * Makes a POST request to the specified [url] and returns a [Flow] of deserialized response objects

diff --git a/generativeai/src/test/java/com/google/ai/client/generativeai/GenerativeModelTests.kt b/generativeai/src/test/java/com/google/ai/client/generativeai/GenerativeModelTests.kt
@@ -85,10 +85,10 @@ internal class ModelNamingTests(private val modelName: String, private val actua
     fun data() =
       listOf(
         arrayOf("gemini-pro", "models/gemini-pro"),
-        arrayOf("x/gemini-pro", "models/x/gemini-pro"),
+        arrayOf("x/gemini-pro", "x/gemini-pro"),
         arrayOf("models/gemini-pro", "models/gemini-pro"),
-        arrayOf("tunedModels/mymodel", "tunedModels/mymodel"),
-        arrayOf("tuneModels/mymodel", "models/tuneModels/mymodel"),
+        arrayOf("/modelname", "/modelname"),
+        arrayOf("modifiedNaming/mymodel", "modifiedNaming/mymodel"),
       )
   }
 }