Merge pull request #46 from nadheesh/main

Restrict Agent's input schema to use types supported by JSON Schema
ballerina-platform · Feb 15, 2024 · e288bc1 · e288bc1
2 parents 937367e + 183f9fa
commit e288bc1
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 23 deletions.
diff --git a/ballerina/Ballerina.toml b/ballerina/Ballerina.toml
@@ -2,7 +2,7 @@
 distribution = "2201.8.4"
 org = "ballerinax"
 name = "ai.agent"
-version = "0.7.5"
+version = "0.7.6"
 license = ["Apache-2.0"]
 authors = ["Ballerina"]
 keywords = ["AI/Agent", "Cost/Freemium"]

diff --git a/ballerina/openapi_utils.bal b/ballerina/openapi_utils.bal
@@ -433,8 +433,9 @@ class OpenApiSpecVisitor {
     }
 
     private isolated function visitPrimitiveTypeSchema(PrimitiveTypeSchema schema, boolean isXml) returns PrimitiveInputSchema|ObjectInputSchema|error {
+        INTEGER|NUMBER|FLOAT|STRING|BOOLEAN 'type = schema.'type;
         PrimitiveInputSchema inputSchema = {
-            'type: schema.'type
+            'type: 'type is FLOAT ? NUMBER : 'type
         };
 
         if self.additionalInfoFlags.extractDescription {
@@ -460,9 +461,7 @@ class OpenApiSpecVisitor {
             inputSchema.pattern = pattern;
             inputSchema.'enum = schema.'enum;
         }
-        if schema is NumberSchema {
-            inputSchema.'type = FLOAT;
-        }
+
         if isXml {
             string? xmlNamespace = schema.'xml?.namespace;
             string? xmlPrefix = schema.'xml?.prefix;

diff --git a/ballerina/tests/openapi_utils_test.bal b/ballerina/tests/openapi_utils_test.bal
@@ -250,15 +250,15 @@ function testExtractToolsFromOpenAPISpecJSONFile2() returns error? {
                             },
                             suffix: {'type: "string"},
                             max_tokens: {'type: "integer"},
-                            temperature: {'type: "float"},
-                            top_p: {'type: "float"},
+                            temperature: {'type: "number"},
+                            top_p: {'type: "number"},
                             n: {'type: "integer"},
                             'stream: {'type: "boolean"},
                             logprobs: {'type: "integer"},
                             echo: {'type: "boolean"},
                             stop: {oneOf: [{'type: "string"}, {'type: "array", items: {'type: "string"}}]},
-                            presence_penalty: {'type: "float"},
-                            frequency_penalty: {'type: "float"},
+                            presence_penalty: {'type: "number"},
+                            frequency_penalty: {'type: "number"},
                             best_of: {'type: "integer"},
                             logit_bias: {'type: "object", properties: {}},
                             user: {'type: "string"}
@@ -289,14 +289,14 @@ function testExtractToolsFromOpenAPISpecJSONFile2() returns error? {
                                     properties: {role: {'type: "string", 'enum: ["system", "user", "assistant"]}, content: {'type: "string"}, name: {'type: "string"}}
                                 }
                             },
-                            temperature: {'type: "float"},
-                            top_p: {'type: "float"},
+                            temperature: {'type: "number"},
+                            top_p: {'type: "number"},
                             n: {'type: "integer"},
                             'stream: {'type: "boolean"},
                             stop: {"oneOf": [{'type: "string"}, {'type: "array", items: {'type: "string"}}]},
                             max_tokens: {'type: "integer"},
-                            presence_penalty: {'type: "float"},
-                            frequency_penalty: {'type: "float"},
+                            presence_penalty: {'type: "number"},
+                            frequency_penalty: {'type: "number"},
                             logit_bias: {'type: "object", properties: {}},
                             user: {'type: "string"}
                         },
@@ -378,10 +378,10 @@ function testExtractToolsFromOpenAPISpecJSONFile3() returns error? {
                                         "type": "integer"
                                     },
                                     "temperature": {
-                                        "type": "float"
+                                        "type": "number"
                                     },
                                     "top_p": {
-                                        "type": "float"
+                                        "type": "number"
                                     },
                                     "n": {
                                         "type": "integer"
@@ -417,10 +417,10 @@ function testExtractToolsFromOpenAPISpecJSONFile3() returns error? {
                                         ]
                                     },
                                     "presence_penalty": {
-                                        "type": "float"
+                                        "type": "number"
                                     },
                                     "frequency_penalty": {
-                                        "type": "float"
+                                        "type": "number"
                                     },
                                     "best_of": {
                                         "type": "integer"
@@ -491,7 +491,7 @@ function testExtractToolsFromOpenAPISpecJSONFile3() returns error? {
                                         }
                                     },
                                     "@temperature": {
-                                        "type": "float"
+                                        "type": "number"
                                     },
                                     "top:top_p": {
                                         "type": "object",
@@ -500,7 +500,7 @@ function testExtractToolsFromOpenAPISpecJSONFile3() returns error? {
                                                 "const": "http://openai.com/docs/1.0/parameters"
                                             },
                                             "#content": {
-                                                "type": "float"
+                                                "type": "number"
                                             }
                                         }
                                     },
@@ -527,10 +527,10 @@ function testExtractToolsFromOpenAPISpecJSONFile3() returns error? {
                                         "type": "integer"
                                     },
                                     "presence_penalty": {
-                                        "type": "float"
+                                        "type": "number"
                                     },
                                     "frequency_penalty": {
-                                        "type": "float"
+                                        "type": "number"
                                     },
                                     "logit_bias": {
                                         "type": "object",

diff --git a/ballerina/tool_types.bal b/ballerina/tool_types.bal
@@ -52,8 +52,8 @@ public type BaseInputTypeSchema record {|
 # Defines a primitive input field in the schema.
 public type PrimitiveInputSchema record {|
     *BaseInputTypeSchema;
-    # Input data type. Should be one of `STRING`, `INTEGER`, `NUMBER`, `FLOAT`, or `BOOLEAN`.
-    STRING|INTEGER|NUMBER|FLOAT|BOOLEAN 'type;
+    # Input data type. Should be one of `STRING`, `INTEGER`, `NUMBER`, or `BOOLEAN`.
+    STRING|INTEGER|NUMBER|BOOLEAN 'type;
     # Format of the input. This is not applicable for `BOOLEAN` type.
     string format?;
     # Pattern of the input. This is only applicable for `STRING` type.