inclusionAI/LLaDA2.2-mini

1"""2Encode structured tool declaration to typescript style string.3"""4 5import dataclasses6import json7import logging8from collections.abc import Sequence9from typing import Any10 11logger = logging.getLogger(__name__)12 13_TS_INDENT = "  "14_TS_FIELD_DELIMITER = ",\n"15 16 17class _SchemaRegistry:18    """Registry for schema definitions to handle $ref resolution"""19 20    def __init__(self):21        self.definitions = {}22        self.has_self_ref = False23 24    def register_definitions(self, defs: dict[str, Any]):25        """Register schema definitions from $defs section"""26        if not defs:27            return28        for def_name, def_schema in defs.items():29            self.definitions[def_name] = def_schema30 31    def resolve_ref(self, ref: str) -> dict[str, Any]:32        """Resolve a reference to its schema definition"""33        if ref == "#":34            self.has_self_ref = True35            return {"$self_ref": True}36        elif ref.startswith("#/$defs/"):37            def_name = ref.split("/")[-1]38            if def_name not in self.definitions:39                raise ValueError(f"Reference not found: {ref}")40            return self.definitions[def_name]41        else:42            raise ValueError(f"Unsupported reference format: {ref}")43 44 45def _format_description(description: str, indent: str = "") -> str:46    return "\n".join(47        [f"{indent}// {line}" if line else "" for line in description.split("\n")]48    )49 50 51class _BaseType:52    description: str53    constraints: dict[str, Any]54 55    def __init__(56        self,57        extra_props: dict[str, Any],58        *,59        allowed_constraint_keys: Sequence[str] = (),60    ):61        self.description = extra_props.get("description", "")62        self.constraints = {63            k: v for k, v in extra_props.items() if k in allowed_constraint_keys64        }65 66    def to_typescript_style(self, indent: str = "") -> str:67        raise NotImplementedError68 69    def format_docstring(self, indent: str) -> str:70        lines = []71        if self.description:72            lines.append(_format_description(self.description, indent))73        if self.constraints:74            constraints_str = ", ".join(75                f"{k}: {v}"76                for k, v in sorted(self.constraints.items(), key=lambda kv: kv[0])77            )78            lines.append(f"{indent}// {constraints_str}")79 80        return "".join(x + "\n" for x in lines)81 82 83class _ParameterTypeScalar(_BaseType):84    type: str85 86    def __init__(self, type: str, extra_props: dict[str, Any] | None = None):87        self.type = type88 89        allowed_constraint_keys: list[str] = []90        if self.type == "string":91            allowed_constraint_keys = ["maxLength", "minLength", "pattern"]92        elif self.type in ("number", "integer"):93            allowed_constraint_keys = ["maximum", "minimum"]94 95        super().__init__(96            extra_props or {}, allowed_constraint_keys=allowed_constraint_keys97        )98 99    def to_typescript_style(self, indent: str = "") -> str:100        # Map integer to number in TypeScript101        if self.type == "integer":102            return "number"103        return self.type104 105 106class _ParameterTypeObject(_BaseType):107    properties: list["_Parameter"]108    additional_properties: Any | None = None109 110    def __init__(111        self,112        json_schema_object: dict[str, Any],113        registry: _SchemaRegistry | None = None,114    ):115        super().__init__(json_schema_object)116 117        self.properties = []118        self.additional_properties = None119 120        if not json_schema_object:121            return122 123        if "$defs" in json_schema_object and registry:124            registry.register_definitions(json_schema_object["$defs"])125 126        self.additional_properties = json_schema_object.get("additionalProperties")127        if isinstance(self.additional_properties, dict):128            self.additional_properties = _parse_parameter_type(129                self.additional_properties, registry130            )131 132        if "properties" not in json_schema_object:133            return134 135        required_parameters = json_schema_object.get("required", [])136        optional_parameters = set(json_schema_object["properties"].keys()) - set(137            required_parameters138        )139 140        self.properties = [141            _Parameter(142                name=name,143                type=_parse_parameter_type(prop, registry),144                optional=name in optional_parameters,145                default=prop.get("default") if isinstance(prop, dict) else None,146            )147            for name, prop in json_schema_object["properties"].items()148        ]149 150    def to_typescript_style(self, indent: str = "") -> str:151        # sort by optional, make the required parameters first152        parameters = [p for p in self.properties if not p.optional]153        opt_params = [p for p in self.properties if p.optional]154 155        parameters = sorted(parameters, key=lambda p: p.name)156        parameters.extend(sorted(opt_params, key=lambda p: p.name))157 158        param_strs = []159        for p in parameters:160            one = p.to_typescript_style(indent=indent + _TS_INDENT)161            param_strs.append(one)162 163        if self.additional_properties is not None:164            ap_type_str = "any"165            if self.additional_properties is True:166                ap_type_str = "any"167            elif self.additional_properties is False:168                ap_type_str = "never"169            elif isinstance(self.additional_properties, _ParameterType):170                ap_type_str = self.additional_properties.to_typescript_style(171                    indent=indent + _TS_INDENT172                )173            else:174                raise ValueError(175                    f"Unknown additionalProperties: {self.additional_properties}"176                )177            param_strs.append(f"{indent + _TS_INDENT}[k: string]: {ap_type_str}")178 179        if not param_strs:180            return "{}"181 182        params_str = _TS_FIELD_DELIMITER.join(param_strs)183        if params_str:184            # add new line before and after185            params_str = f"\n{params_str}\n"186        # always wrap with object187        return f"{{{params_str}{indent}}}"188 189 190class _ParameterTypeArray(_BaseType):191    item: "_ParameterType"192 193    def __init__(194        self,195        json_schema_object: dict[str, Any],196        registry: _SchemaRegistry | None = None,197    ):198        super().__init__(199            json_schema_object, allowed_constraint_keys=("minItems", "maxItems")200        )201        if json_schema_object.get("items"):202            self.item = _parse_parameter_type(json_schema_object["items"], registry)203        else:204            self.item = _ParameterTypeScalar(type="any")205 206    def to_typescript_style(self, indent: str = "") -> str:207        item_docstring = self.item.format_docstring(indent + _TS_INDENT)208        if item_docstring:209            return (210                "Array<\n"211                + item_docstring212                + indent213                + _TS_INDENT214                + self.item.to_typescript_style(indent=indent + _TS_INDENT)215                + "\n"216                + indent217                + ">"218            )219        else:220            return f"Array<{self.item.to_typescript_style(indent=indent)}>"221 222 223class _ParameterTypeEnum(_BaseType):224    # support scalar types only225    enum: list[str | int | float | bool | None]226 227    def __init__(self, json_schema_object: dict[str, Any]):228        super().__init__(json_schema_object)229        self.enum = json_schema_object["enum"]230 231        # Validate enum values against declared type if present232        if "type" in json_schema_object:233            typ = json_schema_object["type"]234            if isinstance(typ, list):235                if len(typ) == 1:236                    typ = typ[0]237                elif len(typ) == 2:238                    if "null" not in typ:239                        raise ValueError(f"Enum type {typ} is not supported")240                    else:241                        typ = typ[0] if typ[0] != "null" else typ[1]242                else:243                    raise ValueError(f"Enum type {typ} is not supported")244            for val in self.enum:245                if val is None:246                    continue247                if typ == "string" and not isinstance(val, str):248                    raise ValueError(f"Enum value {val} is not a string")249                elif typ == "number" and not isinstance(val, (int, float)):250                    raise ValueError(f"Enum value {val} is not a number")251                elif typ == "integer" and not isinstance(val, int):252                    raise ValueError(f"Enum value {val} is not an integer")253                elif typ == "boolean" and not isinstance(val, bool):254                    raise ValueError(f"Enum value {val} is not a boolean")255 256    def to_typescript_style(self, indent: str = "") -> str:257        return " | ".join(258            [f'"{e}"' if isinstance(e, str) else str(e) for e in self.enum]259        )260 261 262class _ParameterTypeAnyOf(_BaseType):263    types: list["_ParameterType"]264 265    def __init__(266        self,267        json_schema_object: dict[str, Any],268        registry: _SchemaRegistry | None = None,269    ):270        super().__init__(json_schema_object)271        self.types = [272            _parse_parameter_type(t, registry) for t in json_schema_object["anyOf"]273        ]274 275    def to_typescript_style(self, indent: str = "") -> str:276        return " | ".join([t.to_typescript_style(indent=indent) for t in self.types])277 278 279class _ParameterTypeUnion(_BaseType):280    types: list[str]281 282    def __init__(self, json_schema_object: dict[str, Any]):283        super().__init__(json_schema_object)284 285        mapping = {286            "string": "string",287            "number": "number",288            "integer": "number",289            "boolean": "boolean",290            "null": "null",291            "object": "{}",292            "array": "Array<any>",293        }294        self.types = [mapping[t] for t in json_schema_object["type"]]295 296    def to_typescript_style(self, indent: str = "") -> str:297        return " | ".join(self.types)298 299 300class _ParameterTypeRef(_BaseType):301    ref_name: str302    is_self_ref: bool = False303 304    def __init__(self, json_schema_object: dict[str, Any], registry: _SchemaRegistry):305        super().__init__(json_schema_object)306 307        ref = json_schema_object["$ref"]308        resolved_schema = registry.resolve_ref(ref)309 310        if resolved_schema.get("$self_ref", False):311            self.ref_name = "parameters"312            self.is_self_ref = True313        else:314            self.ref_name = ref.split("/")[-1]315 316    def to_typescript_style(self, indent: str = "") -> str:317        return self.ref_name318 319 320_ParameterType = (321    _ParameterTypeScalar322    | _ParameterTypeObject323    | _ParameterTypeArray324    | _ParameterTypeEnum325    | _ParameterTypeAnyOf326    | _ParameterTypeUnion327    | _ParameterTypeRef328)329 330 331@dataclasses.dataclass332class _Parameter:333    """334    A parameter in a function, or a field in a object.335    It consists of the type as well as the name.336    """337 338    type: _ParameterType339    name: str = "_"340    optional: bool = True341    default: Any | None = None342 343    @classmethod344    def parse_extended(cls, attributes: dict[str, Any]) -> "_Parameter":345        if not attributes:346            raise ValueError("attributes is empty")347 348        return cls(349            name=attributes.get("name", "_"),350            type=_parse_parameter_type(attributes),351            optional=attributes.get("optional", False),352            default=attributes.get("default"),353        )354 355    def to_typescript_style(self, indent: str = "") -> str:356        comments = self.type.format_docstring(indent)357 358        if self.default is not None:359            default_repr = (360                json.dumps(self.default, ensure_ascii=False)361                if not isinstance(self.default, (int, float, bool))362                else repr(self.default)363            )364            comments += f"{indent}// Default: {default_repr}\n"365 366        return (367            comments368            + f"{indent}{self.name}{'?' if self.optional else ''}: {self.type.to_typescript_style(indent=indent)}"369        )370 371 372def _parse_parameter_type(373    json_schema_object: dict[str, Any] | bool, registry: _SchemaRegistry | None = None374) -> _ParameterType:375    if isinstance(json_schema_object, bool):376        if json_schema_object:377            return _ParameterTypeScalar(type="any")378        else:379            logger.warning(380                f"Warning: Boolean value {json_schema_object} is not supported, use null instead."381            )382            return _ParameterTypeScalar(type="null")383 384    if "$ref" in json_schema_object and registry:385        return _ParameterTypeRef(json_schema_object, registry)386 387    if "anyOf" in json_schema_object:388        return _ParameterTypeAnyOf(json_schema_object, registry)389    elif "enum" in json_schema_object:390        return _ParameterTypeEnum(json_schema_object)391    elif "type" in json_schema_object:392        typ = json_schema_object["type"]393        if isinstance(typ, list):394            return _ParameterTypeUnion(json_schema_object)395        elif typ == "object":396            return _ParameterTypeObject(json_schema_object, registry)397        elif typ == "array":398            return _ParameterTypeArray(json_schema_object, registry)399        else:400            return _ParameterTypeScalar(typ, json_schema_object)401    elif json_schema_object == {}:402        return _ParameterTypeScalar(type="any")403    else:404        raise ValueError(f"Invalid JSON Schema object: {json_schema_object}")405 406 407def _openai_function_to_typescript_style(408    function: dict[str, Any],409) -> str:410    """Convert OpenAI function definition (dict) to TypeScript style string."""411    registry = _SchemaRegistry()412    parameters = function.get("parameters") or {}413    parsed = _ParameterTypeObject(parameters, registry)414 415    interfaces = []416    root_interface_name = None417    if registry.has_self_ref:418        root_interface_name = "parameters"419        params_str = _TS_FIELD_DELIMITER.join(420            [p.to_typescript_style(indent=_TS_INDENT) for p in parsed.properties]421        )422        params_str = f"\n{params_str}\n" if params_str else ""423        interface_def = f"interface {root_interface_name} {{{params_str}}}"424        interfaces.append(interface_def)425 426    definitions_copy = dict(registry.definitions)427    for def_name, def_schema in definitions_copy.items():428        obj_type = _parse_parameter_type(def_schema, registry)429        params_str = obj_type.to_typescript_style()430 431        description_part = ""432        if obj_description := def_schema.get("description", ""):433            description_part = _format_description(obj_description) + "\n"434 435        interface_def = f"{description_part}interface {def_name} {params_str}"436        interfaces.append(interface_def)437 438    interface_str = "\n".join(interfaces)439    function_name = function.get("name", "function")440    if root_interface_name:441        type_def = f"type {function_name} = (_: {root_interface_name}) => any;"442    else:443        params_str = parsed.to_typescript_style()444        type_def = f"type {function_name} = (_: {params_str}) => any;"445 446    description = function.get("description")447    return "\n".join(448        filter(449            bool,450            [451                interface_str,452                ((description and _format_description(description)) or ""),453                type_def,454            ],455        )456    )457 458 459def encode_tools_to_typescript_style(460    tools: list[dict[str, Any]],461) -> str:462    """463    Convert tools (list of dict) to TypeScript style string.464 465    Supports OpenAI format: {"type": "function", "function": {...}}466 467    Args:468        tools: List of tool definitions in dict format469 470    Returns:471        TypeScript style string representation of the tools472    """473    if not tools:474        return ""475 476    functions = []477 478    for tool in tools:479        tool_type = tool.get("type")480        if tool_type == "function":481            func_def = tool.get("function", {})482            if func_def:483                functions.append(_openai_function_to_typescript_style(func_def))484        else:485            # Skip unsupported tool types (like "_plugin")486            continue487 488    if not functions:489        return ""490 491    functions_str = "\n".join(functions)492    result = "# Tools\n\n"493 494    if functions_str:495        result += "## functions\nnamespace functions {\n"496        result += functions_str + "\n"497        result += "}\n"498 499    return result500