Support for Long, Double, List, and data classes as tool arguments (#210)

Author: Rizzen

agents/agents-tools/src/jvmMain/kotlin/ai/koog/agents/core/tools/reflect/util.kt

@@ -14,6 +14,7 @@ import kotlin.reflect.KType
 import kotlin.reflect.full.findAnnotation
 import kotlin.reflect.full.functions
 import kotlin.reflect.full.instanceParameter
+import kotlin.reflect.full.memberProperties
 import kotlin.reflect.jvm.javaMethod
 import kotlin.reflect.jvm.kotlinFunction
 
@@ -197,6 +198,15 @@ public fun KType.asToolType(): ToolParameterType {
         Int::class -> ToolParameterType.Integer
         Float::class -> ToolParameterType.Float
         Boolean::class -> ToolParameterType.Boolean
+        Long::class -> ToolParameterType.Integer
+        Double::class -> ToolParameterType.Float
+
+        List::class -> {
+            val listItemType = this.arguments[0].type ?: error("List item type is null")
+            val listItemToolType = listItemType.asToolType()
+            ToolParameterType.List(listItemToolType)
+        }
+
         is KClass<*> -> {
             val classJava = classifier.java
             when {
@@ -211,6 +221,18 @@ public fun KType.asToolType(): ToolParameterType {
                     ToolParameterType.List(arrayItemToolType)
                 }
 
+                classifier.isData -> {
+                    val properties = classifier.memberProperties.map { prop ->
+                        val description = prop.findAnnotation<LLMDescription>()?.description ?: prop.name
+                        ToolParameterDescriptor(
+                            name = prop.name,
+                            description = description,
+                            type = prop.returnType.asToolType() // Recursive call
+                        )
+                    }
+                    ToolParameterType.Object(properties)
+                }
+
                 else -> throw kotlin.IllegalArgumentException("Unsupported type $classifier")
             }
         }

agents/agents-tools/src/jvmTest/kotlin/ai/koog/agents/core/tools/reflect/ReflectionArgsSerializerTest.kt

@@ -1,12 +1,22 @@
+@file:OptIn(InternalAgentToolsApi::class)
+
 package ai.koog.agents.core.tools.reflect
 
+import ai.koog.agents.core.tools.annotations.InternalAgentToolsApi
+import ai.koog.agents.core.tools.annotations.LLMDescription
+import ai.koog.agents.core.tools.annotations.Tool
 import ai.koog.agents.core.tools.reflect.ToolFromCallable.VarArgsSerializer
+import ai.koog.agents.core.tools.reflect.ToolsFromCallableTest.Companion.ToolsEnabler
+import kotlinx.coroutines.runBlocking
+import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.Json
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.Arguments
 import org.junit.jupiter.params.provider.MethodSource
 import kotlin.reflect.KCallable
+import kotlin.test.Test
 import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
 
 fun foo(i: Int, s: String, b: Boolean = true) = println("$i $s")
 fun Any.fooEx(i: Int, s: String, b: Boolean = true) = println("$i $s")
@@ -33,4 +43,97 @@ class ReflectionArgsSerializerTest {
             decodedArguments.asNamedValues().sortedBy { it.first }.toList()
         )
     }
+
+
+    @Serializable
+    data class MySpecificToolArgs(
+        @LLMDescription("arg Long") val argLong: Long,
+        @LLMDescription("arg Double") val argDouble: Double)
+
+    class MySpecificTool() {
+        @Tool
+        @LLMDescription("Specific tool")
+        suspend fun execute(@LLMDescription("arg Long") argLong: Long): String {
+            return "Specific tool called with $argLong"
+        }
+
+        @Tool
+        @LLMDescription("Specific tool without tool annotation")
+        suspend fun executeDouble(@LLMDescription("arg Long") argDouble: Double): String {
+            return "Specific tool called with $argDouble"
+        }
+
+        @Tool
+        @LLMDescription("Specific tool without tool annotation")
+        suspend fun executeWithArgs(@LLMDescription("args Tool") args: MySpecificToolArgs,
+                                    @LLMDescription("args Tool") args2: MySpecificToolArgs): String {
+            return "Specific tool called with ${args.argLong} and ${args.argDouble}"
+        }
+
+        suspend fun executeWithListArg(@LLMDescription("args Tool") args: List<MySpecificToolArgs>): String {
+            return "Specific tool called with ${args.joinToString(", ") { "${it.argLong} and ${it.argDouble}" }}"
+        }
+    }
+
+    @Test
+    fun testToolLongArg() {
+        val toolClass = MySpecificTool()
+        val tool = toolClass::execute.asTool(ToolsFromCallableTest.Companion.json)
+
+        assertEquals(
+            "Specific tool called with 42",
+            runBlocking {
+                val args = tool.decodeArgsFromString("""{ "argLong": 42 }""")
+                val (rawResult, _) = tool.executeAndSerialize(args, ToolsEnabler)
+                rawResult.result
+            },
+        )
+    }
+
+    @Test
+    fun testToolDoubleArg() {
+        val toolClass = MySpecificTool()
+        val tool: ToolFromCallable = toolClass::executeDouble.asTool(ToolsFromCallableTest.Companion.json)
+
+        assertEquals(
+            "Specific tool called with 42.0",
+            runBlocking {
+                val args = tool.decodeArgsFromString("""{ "argDouble": 42.0 }""")
+                val (rawResult, _) = tool.executeAndSerialize(args, ToolsEnabler)
+                rawResult.result
+            },
+        )
+    }
+
+    @Test
+    fun testToolWithArgs() {
+        val toolClass = MySpecificTool()
+        val tool = toolClass::executeWithArgs.asTool(ToolsFromCallableTest.Companion.json)
+
+        assertEquals(
+            "Specific tool called with 42 and 3.14",
+            runBlocking {
+                val args: ToolFromCallable.VarArgs = tool.decodeArgsFromString("""{ "args": {"argLong": 42, "argDouble": 3.14 }, "args2": {"argLong": 22, "argDouble": 3.14 }}""")
+                val (rawResult, _) = tool.executeAndSerialize(args, ToolsEnabler)
+                rawResult.result
+            }
+        )
+    }
+
+    @Test
+    fun testToolWithListArg() {
+        val toolClass = MySpecificTool()
+        val tool = toolClass::executeWithListArg.asTool(ToolsFromCallableTest.Companion.json)
+
+        assertEquals(
+            "Specific tool called with 42 and 3.14, 22 and 3.14",
+            runBlocking {
+                val args: ToolFromCallable.VarArgs = tool.decodeArgsFromString(
+                    """{ "args": [{"argLong": 42, "argDouble": 3.14 }, {"argLong": 22, "argDouble": 3.14 }] }"""
+                )
+                val (rawResult, _) = tool.executeAndSerialize(args, ToolsEnabler)
+                rawResult.result
+            }
+        )
+    }
 }

agents/agents-tools/src/jvmTest/kotlin/ai/koog/agents/core/tools/reflect/ToolsFromCallableTest.kt

@@ -285,4 +285,5 @@ class ToolsFromCallableTest {
         }.trim()
         assertEquals(expectedDescription, rendered)
     }
+
 }