fix xlam's schema issues (#1917)

fix xlam

Signed-off-by: dafnapension <dafnashein@yahoo.com>
Co-authored-by: Elron Bandel <elronbandel@gmail.com>

Author: dafnapension

prepare/cards/xlam_function_calling.py

@@ -4,6 +4,7 @@
 from unitxt.operators import (
     Copy,
     ExecuteExpression,
+    FixJsonSchemaOfParameterTypes,
     Move,
     Set,
 )
@@ -24,22 +25,22 @@
         LoadJson(field="answers", to_field="reference_calls"),
         LoadJson(field="tools"),
         Move(field="tools/*/parameters", to_field="properties"),
-        Set(fields={"tools/*/parameters": {"type": "object"}}, use_deepcopy=True),
         Copy(
             field="properties",
             to_field="tools/*/parameters/properties",
             set_every_value=True,
         ),
+        Set(fields={"tools/*/parameters/type": "object"}, use_deepcopy=True),
         ExecuteExpression(
             to_field="required",
-            expression="[[p for p, c in tool['parameters']['properties'].items() if 'optional' not in c['type']] for tool in tools]",
+            expression="[[p for p, c in tool['parameters']['properties'].items() if 'optional' not in c['type'].lower()] for tool in tools]",
         ),
         Copy(
             field="required",
             to_field="tools/*/parameters/required",
             set_every_value=True,
         ),
-        "operators.fix_json_schema",
+        FixJsonSchemaOfParameterTypes(main_field="tools"),
     ],
     task="tasks.tool_calling.multi_turn",
     templates=["templates.tool_calling.multi_turn"],

src/unitxt/catalog/cards/xlam_function_calling_60k.json

@@ -45,33 +45,34 @@
             "field": "tools/*/parameters",
             "to_field": "properties"
         },
-        {
-            "__type__": "set",
-            "fields": {
-                "tools/*/parameters": {
-                    "type": "object"
-                }
-            },
-            "use_deepcopy": true
-        },
         {
             "__type__": "copy",
             "field": "properties",
             "to_field": "tools/*/parameters/properties",
             "set_every_value": true
         },
+        {
+            "__type__": "set",
+            "fields": {
+                "tools/*/parameters/type": "object"
+            },
+            "use_deepcopy": true
+        },
         {
             "__type__": "execute_expression",
             "to_field": "required",
-            "expression": "[[p for p, c in tool['parameters']['properties'].items() if 'optional' not in c['type']] for tool in tools]"
+            "expression": "[[p for p, c in tool['parameters']['properties'].items() if 'optional' not in c['type'].lower()] for tool in tools]"
         },
         {
             "__type__": "copy",
             "field": "required",
             "to_field": "tools/*/parameters/required",
             "set_every_value": true
         },
-        "operators.fix_json_schema"
+        {
+            "__type__": "fix_json_schema_of_parameter_types",
+            "main_field": "tools"
+        }
     ],
     "task": "tasks.tool_calling.multi_turn",
     "templates": [

src/unitxt/operators.py

@@ -40,6 +40,7 @@
 """
 
 import operator
+import re
 import uuid
 import warnings
 import zipfile
@@ -2611,3 +2612,105 @@ def process_value(self, value: str) -> str:
         # Read from local file
         with open(value, encoding=self.encoding) as f:
             return f.read()
+
+
+class FixJsonSchemaOfParameterTypes(InstanceOperator):
+    main_field: str
+
+    def prepare(self):
+        self.simple_mapping = {
+            "": "object",
+            "any": "object",
+            "Any": "object",
+            "Array": "array",
+            "ArrayList": "array",
+            "Bigint": "integer",
+            "bool": "boolean",
+            "Boolean": "boolean",
+            "byte": "integer",
+            "char": "string",
+            "dict": "object",
+            "Dict": "object",
+            "double": "number",
+            "float": "number",
+            "HashMap": "object",
+            "Hashtable": "object",
+            "int": "integer",
+            "list": "array",
+            "List": "array",
+            "long": "integer",
+            "Queue": "array",
+            "short": "integer",
+            "Stack": "array",
+            "tuple": "array",
+            "Set": "array",
+            "set": "array",
+            "str": "string",
+            "String": "string",
+        }
+
+    def dict_type_of(self, type_str: str) -> dict:
+        return {"type": type_str}
+
+    def recursive_trace_for_type_fields(self, containing_element):
+        if isinstance(containing_element, dict):
+            keys = list(containing_element.keys())
+            for key in keys:
+                if key == "type" and isinstance(containing_element["type"], str):
+                    jsonschema_dict = self.type_str_to_jsonschema_dict(
+                        containing_element["type"]
+                    )
+                    containing_element.pop("type")
+                    containing_element.update(jsonschema_dict)
+                else:
+                    self.recursive_trace_for_type_fields(containing_element[key])
+        elif isinstance(containing_element, list):
+            for list_element in containing_element:
+                self.recursive_trace_for_type_fields(list_element)
+
+    def type_str_to_jsonschema_dict(self, type_str: str) -> dict:
+        if type_str in self.simple_mapping:
+            return self.dict_type_of(self.simple_mapping[type_str])
+        m = re.match(r"^(List|Tuple)\[(.*?)\]$", type_str)
+        if m:
+            basic_type = self.dict_type_of("array")
+            basic_type["items"] = self.type_str_to_jsonschema_dict(
+                m.group(2) if m.group(1) == "List" else m.group(2).split(",")[0].strip()
+            )
+            return basic_type
+
+        m = re.match(r"^(Union)\[(.*?)\]$", type_str)
+        if m:
+            args = m.group(2).split(",")
+            for i in range(len(args)):
+                args[i] = args[i].strip()
+            return {"anyOf": [self.type_str_to_jsonschema_dict(arg) for arg in args]}
+        if re.match(r"^(Callable)\[(.*?)\]$", type_str):
+            return self.dict_type_of("object")
+        if "," in type_str:
+            sub_types = type_str.split(",")
+            for i in range(len(sub_types)):
+                sub_types[i] = sub_types[i].strip()
+            assert len(sub_types) in [
+                2,
+                3,
+            ], f"num of subtypes should be 2 or 3, got {type_str}"
+            basic_type = self.type_str_to_jsonschema_dict(sub_types[0])
+            for sub_type in sub_types[1:]:
+                if sub_type.lower().startswith("default"):
+                    basic_type["default"] = re.split(r"[= ]", sub_type, maxsplit=1)[1]
+            for sub_type in sub_types[1:]:
+                if sub_type.lower().startswith("optional"):
+                    return {"anyOf": [basic_type, self.dict_type_of("null")]}
+            return basic_type
+
+        return self.dict_type_of(type_str)  # otherwise - return what arrived
+
+    def process(
+        self, instance: Dict[str, Any], stream_name: Optional[str] = None
+    ) -> Dict[str, Any]:
+        assert (
+            self.main_field in instance
+        ), f"field '{self.main_field}' must reside in instance in order to verify its jsonschema correctness. got {instance}"
+        self.recursive_trace_for_type_fields(instance[self.main_field])
+        return instance