Modelname.feature

.changes/condition-company-cloth-distribution.json

@@ -0,0 +1 @@
+{"type":"MAJOR","changes":["Support a general model naming schema"]}

generativeai/build.gradle.kts

@@ -94,6 +94,7 @@ dependencies {
     testImplementation("junit:junit:4.13.2")
     testImplementation("io.kotest:kotest-assertions-core:4.0.7")
     testImplementation("io.kotest:kotest-assertions-jvm:4.0.7")
+    testImplementation("io.kotest:kotest-assertions-json:4.0.7")
     testImplementation("io.ktor:ktor-client-mock:$ktorVersion")
     androidTestImplementation("androidx.test.ext:junit:1.1.5")
     androidTestImplementation("androidx.test.espresso:espresso-core:3.5.1")

generativeai/src/main/java/com/google/ai/client/generativeai/internal/api/APIController.kt

@@ -118,8 +118,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/") } ?: "models/$name"
+private fun fullModelName(name: String): String = name.takeIf { it.contains("/") } ?: "models/$name"
 
 /**
  * Makes a POST request to the specified [url] and returns a [Flow] of deserialized response objects

generativeai/src/test/java/com/google/ai/client/generativeai/GenerativeModelTests.kt

@@ -19,16 +19,26 @@ package com.google.ai.client.generativeai
 import com.google.ai.client.generativeai.type.RequestOptions
 import com.google.ai.client.generativeai.type.RequestTimeoutException
 import com.google.ai.client.generativeai.util.commonTest
+import com.google.ai.client.generativeai.util.createGenerativeModel
 import com.google.ai.client.generativeai.util.createResponses
+import com.google.ai.client.generativeai.util.doBlocking
 import com.google.ai.client.generativeai.util.prepareStreamingResponse
 import io.kotest.assertions.throwables.shouldThrow
 import io.kotest.matchers.shouldBe
+import io.kotest.matchers.string.shouldContain
+import io.ktor.client.engine.mock.MockEngine
+import io.ktor.client.engine.mock.respond
+import io.ktor.http.HttpHeaders
+import io.ktor.http.HttpStatusCode
+import io.ktor.http.headersOf
+import io.ktor.utils.io.ByteChannel
 import io.ktor.utils.io.close
 import io.ktor.utils.io.writeFully
 import kotlin.time.Duration.Companion.seconds
-import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.withTimeout
 import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
 
 internal class GenerativeModelTests {
   private val testTimeout = 5.seconds
@@ -57,3 +67,40 @@ internal class GenerativeModelTests {
       }
     }
 }
+
+@RunWith(Parameterized::class)
+internal class ModelNamingTests(private val modelName: String, private val actualName: String) {
+
+  @Test
+  fun `request should include right model name`() = doBlocking {
+    val channel = ByteChannel(autoFlush = true)
+    val mockEngine = MockEngine {
+      respond(channel, HttpStatusCode.OK, headersOf(HttpHeaders.ContentType, "application/json"))
+    }
+    prepareStreamingResponse(createResponses("Random")).forEach { channel.writeFully(it) }
+    val model =
+      createGenerativeModel(modelName, "super_cool_test_key", RequestOptions(), mockEngine)
+
+    withTimeout(5.seconds) {
+      model.generateContentStream().collect {
+        it.candidates.isEmpty() shouldBe false
+        channel.close()
+      }
+    }
+
+    mockEngine.requestHistory.first().url.encodedPath shouldContain actualName
+  }
+
+  companion object {
+    @JvmStatic
+    @Parameterized.Parameters
+    fun data() =
+      listOf(
+        arrayOf("gemini-pro", "models/gemini-pro"),
+        arrayOf("x/gemini-pro", "x/gemini-pro"),
+        arrayOf("models/gemini-pro", "models/gemini-pro"),
+        arrayOf("/modelname", "/modelname"),
+        arrayOf("modifiedNaming/mymodel", "modifiedNaming/mymodel"),
+      )
+  }
+}

generativeai/src/test/java/com/google/ai/client/generativeai/util/tests.kt

@@ -107,18 +107,30 @@ internal fun commonTest(
   val mockEngine = MockEngine {
     respond(channel, status, headersOf(HttpHeaders.ContentType, "application/json"))
   }
-  val controller =
-    APIController(
-      "super_cool_test_key",
-      "gemini-pro",
-      requestOptions.apiVersion,
-      requestOptions.timeout,
-      mockEngine
-    )
-  val model = GenerativeModel("gemini-pro", "super_cool_test_key", controller = controller)
+  val model = createGenerativeModel("gemini-pro", "super_cool_test_key", requestOptions, mockEngine)
   CommonTestScope(channel, model).block()
 }
 
+/** Simple wrapper that guarantees the model and APIController are created using the same data */
+internal fun createGenerativeModel(
+  name: String,
+  apikey: String,
+  requestOptions: RequestOptions = RequestOptions(),
+  engine: MockEngine
+) =
+  GenerativeModel(
+    name,
+    apikey,
+    controller =
+      APIController(
+        "super_cool_test_key",
+        name,
+        requestOptions.apiVersion,
+        requestOptions.timeout,
+        engine
+      )
+  )
+
 /**
  * A variant of [commonTest] for performing *streaming-based* snapshot tests.
  *