add changes from #171

google-gemini · Jun 4, 2024 · eda2f8f · eda2f8f
1 parent 54d52cc
commit eda2f8f
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/common/src/main/kotlin/com/google/ai/client/generativeai/common/Request.kt b/common/src/main/kotlin/com/google/ai/client/generativeai/common/Request.kt
@@ -38,4 +38,10 @@ data class GenerateContentRequest(
 ) : Request
 
 @Serializable
-data class CountTokensRequest(val generateContentRequest: GenerateContentRequest) : Request
+data class CountTokensRequest(
+  val generateContentRequest: GenerateContentRequest,
+  val model: String? = null,
+  val contents: List<Content>? = null,
+  @SerialName("tool_config") var toolConfig: ToolConfig? = null,
+  @SerialName("system_instruction") val systemInstruction: Content? = null,
+) : Request
diff --git a/common/src/test/java/com/google/ai/client/generativeai/common/APIControllerTests.kt b/common/src/test/java/com/google/ai/client/generativeai/common/APIControllerTests.kt
@@ -311,4 +311,4 @@ fun textGenerateContentRequest(prompt: String) =
     contents = listOf(Content(parts = listOf(TextPart(prompt)))),
   )
 
-fun textCountTokenRequest(prompt: String) = CountTokensRequest(textGenerateContentRequest(prompt))
+fun textCountTokenRequest(prompt: String) = CountTokensRequest(generateContentRequest = textGenerateContentRequest(prompt))
diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/GenerativeModel.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/GenerativeModel.kt
@@ -240,7 +240,7 @@ internal constructor(
     )
 
   private fun constructCountTokensRequest(vararg prompt: Content) =
-    CountTokensRequest(constructRequest(*prompt))
+    CountTokensRequest(generateContentRequest = constructRequest(*prompt))
 
   private fun GenerateContentResponse.validate() = apply {
     if (candidates.isEmpty() && promptFeedback == null) {