fix(server): resolve JSON Schema refs and union types in tool arguments

python/freetoken/server/function_call_parser.py

@@ -347,17 +347,102 @@ def _ends_with_partial_token(self, buffer: str, bot_token: str) -> int:
                 return i
         return 0
 
-    def _get_param_config(self, func_name: str, tools: List[Tool]) -> Dict:
-        """Extract the parameter properties (JSON schema) for one tool."""
+    def _resolve_local_ref(self, schema: Dict, root_schema: Dict, _seen: Optional[set[str]] = None,) -> Dict:
+        """
+        Resolve local JSON Schema references.
+
+        Unresolvable or cyclic references are returned unresolved. Their values
+        therefore follow the parser's existing loose/untyped conversion behavior.
+        """
+        if not isinstance(schema, dict):
+            return schema
+        ref = schema.get("$ref")
+        if not isinstance(ref, str) or not ref.startswith("#/"):
+            return schema
+        if _seen is None:
+            _seen = set()
+        if ref in _seen:
+            # Cyclic ref: stop resolution and preserve the unresolved schema.
+            return schema
+        _seen.add(ref)
+        node = root_schema
+        try:
+            for part in ref[2:].split("/"):
+                part = part.replace("~1", "/").replace("~0", "~")
+                node = node[part]
+        except (KeyError, TypeError):
+            return schema
+        if not isinstance(node, dict):
+            return schema
+        merged = dict(node)
+        merged.update({k: v for k, v in schema.items() if k != "$ref"})
+        if "$ref" in merged:
+            return self._resolve_local_ref(merged, root_schema, _seen)
+        return merged
+
+    def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict:
+        """Resolve the schema parts needed for tool argument typing."""
+        if not isinstance(schema, dict):
+            return {"type": "string"}
+        schema = self._resolve_local_ref(schema, root_schema)
+        param_type = schema.get("type")
+        if isinstance(param_type, list):
+            non_null = [t for t in param_type if t != "null"]
+            result = dict(schema)
+            result["type"] = non_null[0] if len(non_null) == 1 else "loose"
+            return result
+        if isinstance(param_type, str):
+            return schema
+        for keyword in ("oneOf", "anyOf"):
+            branches = schema.get(keyword)
+            if not isinstance(branches, list):
+                continue
+            types = []
+            for branch in branches:
+                if not isinstance(branch, dict):
+                    continue
+                branch = self._resolve_local_ref(branch, root_schema)
+                branch_type = branch.get("type")
+                if isinstance(branch_type, str) and branch_type != "null":
+                    types.append(branch_type)
+            types = list(dict.fromkeys(types))
+            result = dict(schema)
+            result["type"] = types[0] if len(types) == 1 else "loose"
+            return result
+        if isinstance(schema.get("properties"), dict):
+            result = dict(schema)
+            result["type"] = "object"
+            return result
+        if isinstance(schema.get("items"), (dict, list)):
+            result = dict(schema)
+            result["type"] = "array"
+            return result
+        result = dict(schema)
+        result["type"] = "loose"
+        return result
+
+    def _get_tool_schema(self, func_name: str, tools: List[Tool]) -> Dict:
+        """Return the root JSON schema for one tool."""
         for tool in tools:
-            if tool.function.name == func_name and tool.function.parameters:
-                params = tool.function.parameters
-                if isinstance(params, dict) and "properties" in params:
-                    return params["properties"]
-                elif isinstance(params, dict):
-                    return params
+            if tool.function.name != func_name or not tool.function.parameters:
+                continue
+            params = tool.function.parameters
+            return params if isinstance(params, dict) else {}
         return {}
 
+    def _get_param_config(self, func_name: str, tools: List[Tool]) -> Dict:
+        """Extract and normalize parameter properties for one tool."""
+        params = self._get_tool_schema(func_name, tools)
+        if not params:
+            return {}
+        properties = params.get("properties")
+        if not isinstance(properties, dict):
+            return params
+        return {
+            name: self._normalize_param_schema(prop, params)
+            for name, prop in properties.items()
+        }
+
     def _convert_param_value(self, value: str, param_name: str, param_config: Dict, func_name: str) -> Any:
         """Convert parameter value based on schema type. Safe alternative to eval()."""
         if value.lower() == "null":
@@ -392,6 +477,8 @@ def _convert_param_value(self, value: str, param_name: str, param_config: Dict,
                     return ast.literal_eval(value)
                 except (ValueError, SyntaxError, TypeError):
                     return value
+        elif param_type == "loose":
+            return _parse_loose_json_value(value)
         return value
 
     def _schema_param_type(self, param_name: str, param_config: Dict, missing: str = "string") -> str:
@@ -2591,18 +2678,23 @@ def _union_leaf(self, raw: str, subs: list) -> Any:
                     return v
         return raw
 
-    def _nested_value(self, raw: str, schema: Any = None) -> Any:
+    def _nested_value(self, raw: str, schema: Any = None, root_schema: Any = None) -> Any:
+        if isinstance(schema, dict) and isinstance(root_schema, dict):
+            schema = self._normalize_param_schema(schema, root_schema)
         items, stray = self._scan_elements(raw)
         if not items:
             return self._typed_leaf(raw, schema)
-        return self._structure(items, schema, stray=stray)
+        return self._structure(items, schema, root_schema=root_schema, stray=stray)
 
-    def _structure(self, items: List[tuple], schema: Any = None, stray: str = "") -> Any:
+    def _structure(self, items: List[tuple], schema: Any = None, root_schema: Any = None, stray: str = "") -> Any:
         """Composite value from scanned elements, threading the JSON schema down
         (``properties`` / ``items`` / ``additionalProperties``). Only ``<item>``
         children render as a bare array -- the template's array convention;
         repeated siblings under any other tag stay an object with an array-valued
         key. Mixed text+children keeps the text under ``"$text"``."""
+        if isinstance(schema, dict) and isinstance(root_schema, dict):
+            schema = self._normalize_param_schema(schema, root_schema)
+
         props = schema.get("properties") if isinstance(schema, dict) else None
         item_schema = schema.get("items") if isinstance(schema, dict) else None
 
@@ -2612,11 +2704,13 @@ def _sub_schema(key: str) -> Any:
                 ap = schema.get("additionalProperties")
                 if isinstance(ap, dict):
                     sub = ap
+            if isinstance(sub, dict) and isinstance(root_schema, dict):
+                sub = self._normalize_param_schema(sub, root_schema)
             return sub
 
         names = {k for k, _ in items}
         if names == {"item"}:
-            return [self._nested_value(raw, item_schema) for _, raw in items]
+            return [self._nested_value(raw, item_schema, root_schema=root_schema) for _, raw in items]
         counts: Dict[str, int] = {}
         for k, _ in items:
             counts[k] = counts.get(k, 0) + 1
@@ -2632,7 +2726,7 @@ def _sub_schema(key: str) -> Any:
                 and str(sub.get("type", "")).lower() == "array"
             ):
                 sub = sub.get("items")
-            value = self._nested_value(raw, sub)
+            value = self._nested_value(raw, sub, root_schema=root_schema)
             if key in out:
                 prev = out[key]
                 out[key] = (prev if isinstance(prev, list) else [prev]) + [value]
@@ -2643,16 +2737,19 @@ def _sub_schema(key: str) -> Any:
         return out
 
     def _args_from_items(
-        self, func_name: str, items: List[tuple], tools: List[Tool]
+        self, func_name: str, items: List[tuple], tools: List[Tool],
     ) -> Dict:
+        root_schema = self._get_tool_schema(func_name, tools)
         config = self._get_param_config(func_name, tools)
         args: Dict[str, Any] = {}
         for key, raw in items:
             prop = config.get(key) if isinstance(config, dict) and key in config else None
             nested, stray = self._scan_elements(raw)
             if nested:
-                value: Any = self._structure(nested, prop, stray=stray)
+                value: Any = self._structure(nested, prop, root_schema=root_schema, stray=stray)
             else:
+                if isinstance(prop, dict) and isinstance(root_schema, dict):
+                    prop = self._normalize_param_schema(prop, root_schema)
                 value = self._typed_leaf(raw, prop)
             if key in args:
                 prev = args[key]