Support integer and float in tools (#68)

* Support integer and float in tools

Signed-off-by: Ira <[email protected]>

* Fixed indentation

Signed-off-by: Ira <[email protected]>

---------

Signed-off-by: Ira <[email protected]>
Signed-off-by: Ira Rosen <[email protected]>

Author: irar2

pkg/llm-d-inference-sim/tools_test.go

@@ -218,7 +218,7 @@ var toolWithObjects = []openai.ChatCompletionToolParam{
 								"type": "string",
 							},
 							"quantity": map[string]string{
-								"type": "number",
+								"type": "integer",
 							},
 							"address": map[string]interface{}{
 								"type": "object",
@@ -227,7 +227,7 @@ var toolWithObjects = []openai.ChatCompletionToolParam{
 										"type": "string",
 									},
 									"number": map[string]interface{}{
-										"type": "number",
+										"type": "integer",
 									},
 									"home": map[string]interface{}{
 										"type": "boolean",
@@ -264,7 +264,7 @@ var toolWithObjectAndArray = []openai.ChatCompletionToolParam{
 								"description": "The user's name",
 							},
 							"age": map[string]string{
-								"type":        "number",
+								"type":        "integer",
 								"description": "The user's age",
 							},
 							"hobbies": map[string]interface{}{
@@ -496,7 +496,7 @@ var _ = Describe("Simulator for request with tools", func() {
 			Expect(tc.Function.Name).To(Equal("multiply_numbers"))
 			Expect(tc.ID).NotTo(BeEmpty())
 			Expect(string(tc.Type)).To(Equal("function"))
-			args := make(map[string][]int)
+			args := make(map[string][]float64)
 			err = json.Unmarshal([]byte(tc.Function.Arguments), &args)
 			Expect(err).NotTo(HaveOccurred())
 			Expect(args["numbers"]).ToNot(BeEmpty())

pkg/llm-d-inference-sim/tools_utils.go

@@ -142,8 +142,10 @@ func createArgument(property any) (any, error) {
 	switch paramType {
 	case "string":
 		return getStringArgument(), nil
-	case "number":
+	case "integer":
 		return randomInt(0, 100), nil
+	case "number":
+		return randomFloat(0, 100), nil
 	case "boolean":
 		return flipCoin(), nil
 	case "array":
@@ -250,6 +252,7 @@ const schema = `{
             "array",
             "string",
             "number",
+            "integer",
             "boolean",
             "null"
           ]
@@ -308,6 +311,7 @@ const schema = `{
           "enum": [
             "string",
             "number",
+            "integer",
             "boolean",
             "array",
             "object",
@@ -323,6 +327,7 @@ const schema = `{
             "type": [
               "string",
               "number",
+              "integer",
               "boolean"
             ]
           }
@@ -404,6 +409,25 @@ const schema = `{
             }
           }
         },
+        {
+          "if": {
+            "properties": {
+              "type": {
+                "const": "integer"
+              }
+            }
+          },
+          "then": {
+            "properties": {
+              "enum": {
+                "type": "array",
+                "items": {
+                  "type": "integer"
+                }
+              }
+            }
+          }
+        },
         {
           "if": {
             "properties": {

pkg/llm-d-inference-sim/utils.go

@@ -123,6 +123,13 @@ func flipCoin() bool {
 	return randomInt(0, 1) != 0
 }
 
+// Returns a random float64 in the range [min, max)
+func randomFloat(min float64, max float64) float64 {
+	src := rand.NewSource(time.Now().UnixNano())
+	r := rand.New(src)
+	return r.Float64()*(max-min) + min
+}
+
 // Regular expression for the response tokenization
 var re *regexp.Regexp