From bc9778c035e7c09d2a242090dc318cb9f012e5f8 Mon Sep 17 00:00:00 2001
From: David Motsonashvili <davidmotson@gmail.com>
Date: Mon, 17 Jun 2024 12:18:10 -0700
Subject: [PATCH] fix int type, add long type, rename double type (#184)

Co-authored-by: David Motsonashvili <davidmotson@google.com>
---
 .../generativeai/type/FunctionDeclarations.kt    | 16 +++++++++++++---
 .../google/ai/client/generativeai/type/Type.kt   |  3 ++-
 2 files changed, 15 insertions(+), 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
index bf776b29..88247f3d 100644
--- 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
index c220031d..8486edab 100644
--- 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 =