From 309e070d23674174320c925e37a93a1339a048ee Mon Sep 17 00:00:00 2001 From: David Motsonashvili Date: Mon, 3 Jun 2024 14:45:04 -0700 Subject: [PATCH] add extra schema to the test --- .../generativeai/FunctionCallingTests.kt | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/generativeai/src/test/java/com/google/ai/client/generativeai/FunctionCallingTests.kt b/generativeai/src/test/java/com/google/ai/client/generativeai/FunctionCallingTests.kt index 78109bc7..5f38fc91 100644 --- a/generativeai/src/test/java/com/google/ai/client/generativeai/FunctionCallingTests.kt +++ b/generativeai/src/test/java/com/google/ai/client/generativeai/FunctionCallingTests.kt @@ -30,12 +30,28 @@ internal class FunctionCallingTests { @Test fun `function calls with valid args should succeed`() = doBlocking { val functionDeclaration = - defineFunction("name", "description", Schema.str("param1", "description")) { param1 -> + defineFunction( + "name", + "description", + Schema.str("param1", "description"), + Schema.num("param2", "description"), + Schema.bool("param3", "description"), + Schema.int("param4", "description"), + ) { param1, param2, param3, param4 -> JSONObject(mapOf("result" to "success")) } val model = GenerativeModel("model", "key", tools = listOf(Tool(listOf(functionDeclaration)))) - val functionCall = FunctionCallPart("name", mapOf("param1" to "valid parameter")) + val functionCall = + FunctionCallPart( + "name", + mapOf( + ("param1" to "valid parameter"), + ("param2" to "2.2"), + ("param3" to "false"), + ("param4" to "2"), + ) + ) val result = model.executeFunction(functionCall)