fix int type, add long type, rename double type (google-gemini#184)

Co-authored-by: David Motsonashvili <davidmotson@google.com>
AVidhanR · Jun 17, 2024 · bc9778c · bc9778c
1 parent 12274f4
commit bc9778c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
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
@@ -78,15 +78,25 @@ class Schema<T>(
   fun fromString(value: String?) = type.parse(value)
 
   companion object {
-    /** Registers a schema for an integer number */
+    /** Registers a schema for a 32 bit integer number */
     fun int(name: String, description: String) =
-      Schema<Long>(
+      Schema<Int>(
         name = name,
         description = description,
+        format = "int32",
         type = FunctionType.INTEGER,
         nullable = false,
       )
 
+    /** Registers a schema for a 64 bit integer number */
+    fun long(name: String, description: String) =
+      Schema<Long>(
+        name = name,
+        description = description,
+        type = FunctionType.LONG,
+        nullable = false,
+      )
+
     /** Registers a schema for a string */
     fun str(name: String, description: String) =
       Schema<String>(
@@ -106,7 +116,7 @@ class Schema<T>(
       )
 
     /** Registers a schema for a floating point number */
-    fun num(name: String, description: String) =
+    fun double(name: String, description: String) =
       Schema<Double>(
         name = name,
         description = description,

diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/type/Type.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/type/Type.kt
@@ -30,7 +30,8 @@ import org.json.JSONObject
 class FunctionType<T>(val name: String, val parse: (String?) -> T?) {
   companion object {
     val STRING = FunctionType<String>("STRING") { it }
-    val INTEGER = FunctionType<Long>("INTEGER") { it?.toLongOrNull() }
+    val INTEGER = FunctionType<Int>("INTEGER") { it?.toIntOrNull() }
+    val LONG = FunctionType<Long>("INTEGER") { it?.toLongOrNull() }
     val NUMBER = FunctionType<Double>("NUMBER") { it?.toDoubleOrNull() }
     val BOOLEAN = FunctionType<Boolean>("BOOLEAN") { it?.toBoolean() }
     val ARRAY =