factory methods and removing fullModelName duplication

common/src/main/kotlin/com/google/ai/client/generativeai/common/APIController.kt

@@ -19,6 +19,7 @@ package com.google.ai.client.generativeai.common
 import android.util.Log
 import com.google.ai.client.generativeai.common.server.FinishReason
 import com.google.ai.client.generativeai.common.util.decodeToFlow
+import com.google.ai.client.generativeai.common.util.fullModelName
 import io.ktor.client.HttpClient
 import io.ktor.client.call.body
 import io.ktor.client.engine.HttpClientEngine
@@ -213,12 +214,6 @@ interface HeaderProvider {
   suspend fun generateHeaders(): Map<String, String>
 }
 
-/**
- * Ensures the model name provided has a `models/` prefix
- *
- * Models must be prepended with the `models/` prefix when communicating with the backend.
- */
-private fun fullModelName(name: String): String = name.takeIf { it.contains("/") } ?: "models/$name"
 
 private suspend fun validateResponse(response: HttpResponse) {
   if (response.status == HttpStatusCode.OK) return

common/src/main/kotlin/com/google/ai/client/generativeai/common/Request.kt

@@ -21,6 +21,7 @@ import com.google.ai.client.generativeai.common.client.Tool
 import com.google.ai.client.generativeai.common.client.ToolConfig
 import com.google.ai.client.generativeai.common.shared.Content
 import com.google.ai.client.generativeai.common.shared.SafetySetting
+import com.google.ai.client.generativeai.common.util.fullModelName
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 
@@ -39,9 +40,27 @@ data class GenerateContentRequest(
 
 @Serializable
 data class CountTokensRequest(
-  val generateContentRequest: GenerateContentRequest,
+  val generateContentRequest: GenerateContentRequest? = null,
   val model: String? = null,
   val contents: List<Content>? = null,
-  @SerialName("tool_config") var toolConfig: ToolConfig? = null,
+  val tools: List<Tool>? = null,
   @SerialName("system_instruction") val systemInstruction: Content? = null,
-) : Request
+) : Request {
+  companion object {
+    fun forGenAI(generateContentRequest: GenerateContentRequest) =
+      CountTokensRequest(
+        generateContentRequest =
+          generateContentRequest.model?.let {
+            generateContentRequest.copy(model = fullModelName(it))
+          } ?: generateContentRequest
+      )
+
+    fun forVertexAI(generateContentRequest: GenerateContentRequest) =
+      CountTokensRequest(
+        model = generateContentRequest.model?.let { fullModelName(it) },
+        contents = generateContentRequest.contents,
+        tools = generateContentRequest.tools,
+        systemInstruction = generateContentRequest.systemInstruction,
+      )
+  }
+}

common/src/main/kotlin/com/google/ai/client/generativeai/common/util/util.kt

@@ -0,0 +1,8 @@
+package com.google.ai.client.generativeai.common.util
+
+/**
+ * Ensures the model name provided has a `models/` prefix
+ *
+ * Models must be prepended with the `models/` prefix when communicating with the backend.
+ */
+fun fullModelName(name: String): String = name.takeIf { it.contains("/") } ?: "models/$name"

generativeai/src/main/java/com/google/ai/client/generativeai/GenerativeModel.kt

@@ -240,7 +240,7 @@ internal constructor(
     )
 
   private fun constructCountTokensRequest(vararg prompt: Content) =
-    CountTokensRequest(generateContentRequest = constructRequest(*prompt))
+    CountTokensRequest.forGenAI(constructRequest(*prompt))
 
   private fun GenerateContentResponse.validate() = apply {
     if (candidates.isEmpty() && promptFeedback == null) {