Skip to content

Commit d018ad4

Browse files
This commit adds a new test to ToolTests.kt to verify the handling of enum parameters in tool function schemas.
The new test, `testTools_enumParameter`, defines a function with a parameter that is constrained to a list of enum values. It then verifies that the model correctly calls this function with one of the allowed enum values when prompted. This increases the test coverage of the tool functionality in the Firebase AI library.
1 parent 64b488f commit d018ad4

File tree

1 file changed

+27
-0
lines changed
  • firebase-ai/src/androidTest/kotlin/com/google/firebase/ai

1 file changed

+27
-0
lines changed

firebase-ai/src/androidTest/kotlin/com/google/firebase/ai/ToolTests.kt

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,33 @@ class ToolTests {
293293
}
294294
}
295295

296+
@Test
297+
fun testTools_enumParameter() {
298+
val enumValues = listOf("red", "green", "blue")
299+
val schema =
300+
mapOf(
301+
Pair("color", Schema.enumeration(enumValues, "A color from the provided list"))
302+
)
303+
val model =
304+
setupModel(
305+
FunctionDeclaration(
306+
name = "getHexCode",
307+
description = "returns the hex code for a given color",
308+
parameters = schema
309+
),
310+
)
311+
runBlocking {
312+
val response = model.generateContent("what is the hex code for green?")
313+
validator.validateResponse((response))
314+
assert(response.functionCalls.size == 1)
315+
val call = response.functionCalls[0]
316+
assert(call.name == "getHexCode")
317+
validateSchema(schema, call.args)
318+
val color = call.args["color"]!!.jsonPrimitive.content
319+
assert(color in enumValues)
320+
}
321+
}
322+
296323
companion object {
297324
@JvmStatic
298325
fun setupModel(vararg functions: FunctionDeclaration): GenerativeModel {

0 commit comments

Comments
 (0)