clean up serialization

google-gemini · Mar 15, 2024 · 90f2863 · 90f2863
1 parent 470594a
commit 90f2863
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 31 deletions.
diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/internal/api/client/Types.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/internal/api/client/Types.kt
@@ -18,7 +18,6 @@ package com.google.ai.client.generativeai.internal.api.client
 
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
-import kotlinx.serialization.json.JsonObject
 
 @Serializable
 internal data class GenerationConfig(
@@ -41,7 +40,15 @@ internal data class FunctionDeclaration(
 
 @Serializable
 internal data class FunctionParameters(
-  val properties: JsonObject,
+  val properties: Map<String, FunctionParameterProperties>,
   val required: List<String>,
   val type: String,
 )
+
+@Serializable
+internal data class FunctionParameterProperties(
+  val type: String,
+  val description: String,
+  val format: String?,
+  val enum: List<String>?
+)
diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/internal/util/conversions.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/internal/util/conversions.kt
@@ -21,6 +21,7 @@ import android.graphics.BitmapFactory
 import android.util.Base64
 import com.google.ai.client.generativeai.internal.api.CountTokensResponse
 import com.google.ai.client.generativeai.internal.api.GenerateContentResponse
+import com.google.ai.client.generativeai.internal.api.client.FunctionParameterProperties
 import com.google.ai.client.generativeai.internal.api.client.GenerationConfig
 import com.google.ai.client.generativeai.internal.api.server.BlockReason
 import com.google.ai.client.generativeai.internal.api.server.Candidate
@@ -46,16 +47,14 @@ import com.google.ai.client.generativeai.type.CitationMetadata
 import com.google.ai.client.generativeai.type.FunctionDeclaration
 import com.google.ai.client.generativeai.type.GenerativeBeta
 import com.google.ai.client.generativeai.type.ImagePart
+import com.google.ai.client.generativeai.type.ParameterDeclaration
 import com.google.ai.client.generativeai.type.SerializationException
 import com.google.ai.client.generativeai.type.Tool
 import com.google.ai.client.generativeai.type.content
 import java.io.ByteArrayOutputStream
 import kotlinx.serialization.decodeFromString
 import kotlinx.serialization.json.Json
-import kotlinx.serialization.json.JsonArray
 import kotlinx.serialization.json.JsonObject
-import kotlinx.serialization.json.JsonPrimitive
-import kotlinx.serialization.json.buildJsonObject
 import org.json.JSONObject
 
 private const val BASE_64_FLAGS = Base64.NO_WRAP
@@ -120,33 +119,19 @@ internal fun Tool.toInternal() =
   )
 
 @GenerativeBeta
-internal fun FunctionDeclaration.toInternal():
-  com.google.ai.client.generativeai.internal.api.client.FunctionDeclaration {
-  val convertedParams = buildJsonObject {
-    getParameters().forEach { pramDeclaration ->
-      put(
-        pramDeclaration.name,
-        buildJsonObject {
-          put("type", JsonPrimitive(pramDeclaration.type.name))
-          put("description", JsonPrimitive(pramDeclaration.description))
-          pramDeclaration.format?.let { format -> put("format", JsonPrimitive(format)) }
-          pramDeclaration.enum?.let { enum ->
-            put("enum", JsonArray(enum.map { JsonPrimitive(it) }))
-          }
-        },
-      )
-    }
-  }
-  return com.google.ai.client.generativeai.internal.api.client.FunctionDeclaration(
+internal fun FunctionDeclaration.toInternal() =
+  com.google.ai.client.generativeai.internal.api.client.FunctionDeclaration(
     name,
     description,
     com.google.ai.client.generativeai.internal.api.client.FunctionParameters(
-      convertedParams,
+      getParameters().associate { it.name to it.toInternal() },
       getParameters().map { it.name },
       "OBJECT",
     ),
   )
-}
+
+internal fun <T> ParameterDeclaration<T>.toInternal() =
+  FunctionParameterProperties(type.name, description, format, enum)
 
 internal fun JSONObject.toInternal() = Json.decodeFromString<JsonObject>(toString())
 

diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/type/FunctionDeclarations.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/type/FunctionDeclarations.kt
@@ -181,22 +181,46 @@ class ParameterDeclaration<T>(
 
   companion object {
     fun int(name: String, description: String) =
-      ParameterDeclaration<Long>(name, description, null, null, FunctionType.INTEGER)
+      ParameterDeclaration<Long>(
+        name = name,
+        description = description,
+        type = FunctionType.INTEGER,
+      )
 
     fun str(name: String, description: String) =
-      ParameterDeclaration<String>(name, description, null, null, FunctionType.STRING)
+      ParameterDeclaration<String>(
+        name = name,
+        description = description,
+        type = FunctionType.STRING,
+      )
 
     fun bool(name: String, description: String) =
-      ParameterDeclaration<Boolean>(name, description, null, null, FunctionType.BOOLEAN)
+      ParameterDeclaration<Boolean>(
+        name = name,
+        description = description,
+        type = FunctionType.BOOLEAN,
+      )
 
     fun num(name: String, description: String) =
-      ParameterDeclaration<Double>(name, description, null, null, FunctionType.NUMBER)
+      ParameterDeclaration<Double>(
+        name = name,
+        description = description,
+        type = FunctionType.NUMBER,
+      )
 
     fun obj(name: String, description: String) =
-      ParameterDeclaration<JSONObject>(name, description, null, null, FunctionType.OBJECT)
+      ParameterDeclaration<JSONObject>(
+        name = name,
+        description = description,
+        type = FunctionType.OBJECT,
+      )
 
     fun arr(name: String, description: String) =
-      ParameterDeclaration<List<String>>(name, description, null, null, FunctionType.ARRAY)
+      ParameterDeclaration<List<String>>(
+        name = name,
+        description = description,
+        type = FunctionType.ARRAY,
+      )
 
     fun enum(name: String, description: String, values: List<String>) =
       ParameterDeclaration<String>(name, description, "enum", values, FunctionType.STRING)