unsloth/Kimi-K2.7-Code

1"""2Encode structured tool declaration to typescript style string.3"""4import dataclasses5import json6import logging7from collections.abc import Sequence8from typing import Any9 10logger = logging.getLogger(__name__)11 12_TS_INDENT = "  "13_TS_FIELD_DELIMITER = ",\n"14 15 16class _SchemaRegistry:17    """Registry for schema definitions to handle $ref resolution"""18 19    def __init__(self):20        self.definitions = {}21        self.has_self_ref = False22 23    def register_definitions(self, defs: dict[str, Any]):24        """Register schema definitions from $defs section"""25        if not defs:26            return27        for def_name, def_schema in defs.items():28            self.definitions[def_name] = def_schema29 30    def resolve_ref(self, ref: str) -> dict[str, Any]:31        """Resolve a reference to its schema definition"""32        if ref == "#":33            self.has_self_ref = True34            return {"$self_ref": True}35        elif ref.startswith("#/$defs/"):36            def_name = ref.split("/")[-1]37            if def_name not in self.definitions:38                raise ValueError(f"Reference not found: {ref}")39            return self.definitions[def_name]40        else:41            raise ValueError(f"Unsupported reference format: {ref}")42 43 44def _format_description(description: str, indent: str = "") -> str:45    return "\n".join([46        f"{indent}// {line}" if line else ""47        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: v64            for k, v in extra_props.items() if k in allowed_constraint_keys65        }66 67    def to_typescript_style(self, indent: str = "") -> str:68        raise NotImplementedError69 70    def format_docstring(self, indent: str) -> str:71        lines = []72        if self.description:73            lines.append(_format_description(self.description, indent))74        if self.constraints:75            constraints_str = ", ".join(f"{k}: {v}" for k, v in sorted(76                self.constraints.items(), key=lambda kv: kv[0]))77            lines.append(f"{indent}// {constraints_str}")78 79        return "".join(x + "\n" for x in lines)80 81 82class _ParameterTypeScalar(_BaseType):83    type: str84 85    def __init__(self, type: str, extra_props: dict[str, Any] | None = None):86        self.type = type87 88        allowed_constraint_keys: list[str] = []89        if self.type == "string":90            allowed_constraint_keys = ["maxLength", "minLength", "pattern"]91        elif self.type in ("number", "integer"):92            allowed_constraint_keys = ["maximum", "minimum"]93 94        super().__init__(extra_props or {},95                         allowed_constraint_keys=allowed_constraint_keys)96 97    def to_typescript_style(self, indent: str = "") -> str:98        # Map integer to number in TypeScript99        if self.type == "integer":100            return "number"101        return self.type102 103 104class _ParameterTypeObject(_BaseType):105    properties: list["_Parameter"]106    additional_properties: Any | None = None107 108    def __init__(self,109                 json_schema_object: dict[str, Any],110                 registry: _SchemaRegistry | None = None):111        super().__init__(json_schema_object)112 113        self.properties = []114        self.additional_properties = None115 116        if not json_schema_object:117            return118 119        if "$defs" in json_schema_object and registry:120            registry.register_definitions(json_schema_object["$defs"])121 122        self.additional_properties = json_schema_object.get(123            "additionalProperties")124        if isinstance(self.additional_properties, dict):125            self.additional_properties = _parse_parameter_type(126                self.additional_properties, registry)127 128        if "properties" not in json_schema_object:129            return130 131        required_parameters = json_schema_object.get("required", [])132        optional_parameters = set(133            json_schema_object["properties"].keys()) - set(required_parameters)134 135        self.properties = [136            _Parameter(137                name=name,138                type=_parse_parameter_type(prop, registry),139                optional=name in optional_parameters,140                default=prop.get("default")141                if isinstance(prop, dict) else None,142            ) for name, prop in json_schema_object["properties"].items()143        ]144 145    def to_typescript_style(self, indent: str = "") -> str:146        # sort by optional, make the required parameters first147        parameters = [p for p in self.properties if not p.optional]148        opt_params = [p for p in self.properties if p.optional]149 150        parameters = sorted(parameters, key=lambda p: p.name)151        parameters.extend(sorted(opt_params, key=lambda p: p.name))152 153        param_strs = []154        for p in parameters:155            one = p.to_typescript_style(indent=indent + _TS_INDENT)156            param_strs.append(one)157 158        if self.additional_properties is not None:159            ap_type_str = "any"160            if self.additional_properties is True:161                ap_type_str = "any"162            elif self.additional_properties is False:163                ap_type_str = "never"164            elif isinstance(self.additional_properties, _ParameterType):165                ap_type_str = self.additional_properties.to_typescript_style(166                    indent=indent + _TS_INDENT)167            else:168                raise ValueError(169                    f"Unknown additionalProperties: {self.additional_properties}"170                )171            param_strs.append(172                f"{indent + _TS_INDENT}[k: string]: {ap_type_str}")173 174        if not param_strs:175            return "{}"176 177        params_str = _TS_FIELD_DELIMITER.join(param_strs)178        if params_str:179            # add new line before and after180            params_str = f"\n{params_str}\n"181        # always wrap with object182        return f"{{{params_str}{indent}}}"183 184 185class _ParameterTypeArray(_BaseType):186    item: "_ParameterType"187 188    def __init__(self,189                 json_schema_object: dict[str, Any],190                 registry: _SchemaRegistry | None = None):191        super().__init__(json_schema_object,192                         allowed_constraint_keys=("minItems", "maxItems"))193        if json_schema_object.get("items"):194            self.item = _parse_parameter_type(json_schema_object["items"],195                                              registry)196        else:197            self.item = _ParameterTypeScalar(type="any")198 199    def to_typescript_style(self, indent: str = "") -> str:200        item_docstring = self.item.format_docstring(indent + _TS_INDENT)201        if item_docstring:202            return ("Array<\n" + item_docstring + indent + _TS_INDENT +203                    self.item.to_typescript_style(indent=indent + _TS_INDENT) +204                    "\n" + indent + ">")205        else:206            return f"Array<{self.item.to_typescript_style(indent=indent)}>"207 208 209class _ParameterTypeEnum(_BaseType):210    # support scalar types only211    enum: list[str | int | float | bool | None]212 213    def __init__(self, json_schema_object: dict[str, Any]):214        super().__init__(json_schema_object)215        self.enum = json_schema_object["enum"]216 217        # Validate enum values against declared type if present218        if "type" in json_schema_object:219            typ = json_schema_object["type"]220            if isinstance(typ, list):221                if len(typ) == 1:222                    typ = typ[0]223                elif len(typ) == 2:224                    if "null" not in typ:225                        raise ValueError(f"Enum type {typ} is not supported")226                    else:227                        typ = typ[0] if typ[0] != "null" else typ[1]228                else:229                    raise ValueError(f"Enum type {typ} is not supported")230            for val in self.enum:231                if val is None:232                    continue233                if typ == "string" and not isinstance(val, str):234                    raise ValueError(f"Enum value {val} is not a string")235                elif typ == "number" and not isinstance(val, (int, float)):236                    raise ValueError(f"Enum value {val} is not a number")237                elif typ == "integer" and not isinstance(val, int):238                    raise ValueError(f"Enum value {val} is not an integer")239                elif typ == "boolean" and not isinstance(val, bool):240                    raise ValueError(f"Enum value {val} is not a boolean")241 242    def to_typescript_style(self, indent: str = "") -> str:243        return " | ".join(244            [f'"{e}"' if isinstance(e, str) else str(e) for e in self.enum])245 246 247class _ParameterTypeAnyOf(_BaseType):248    types: list["_ParameterType"]249 250    def __init__(251        self,252        json_schema_object: dict[str, Any],253        registry: _SchemaRegistry | None = None,254    ):255        super().__init__(json_schema_object)256        self.types = [257            _parse_parameter_type(t, registry)258            for t in json_schema_object["anyOf"]259        ]260 261    def to_typescript_style(self, indent: str = "") -> str:262        return " | ".join(263            [t.to_typescript_style(indent=indent) for t in self.types])264 265 266class _ParameterTypeUnion(_BaseType):267    types: list[str]268 269    def __init__(self, json_schema_object: dict[str, Any]):270        super().__init__(json_schema_object)271 272        mapping = {273            "string": "string",274            "number": "number",275            "integer": "number",276            "boolean": "boolean",277            "null": "null",278            "object": "{}",279            "array": "Array<any>",280        }281        self.types = [mapping[t] for t in json_schema_object["type"]]282 283    def to_typescript_style(self, indent: str = "") -> str:284        return " | ".join(self.types)285 286 287class _ParameterTypeRef(_BaseType):288    ref_name: str289    is_self_ref: bool = False290 291    def __init__(self, json_schema_object: dict[str, Any],292                 registry: _SchemaRegistry):293        super().__init__(json_schema_object)294 295        ref = json_schema_object["$ref"]296        resolved_schema = registry.resolve_ref(ref)297 298        if resolved_schema.get("$self_ref", False):299            self.ref_name = "parameters"300            self.is_self_ref = True301        else:302            self.ref_name = ref.split("/")[-1]303 304    def to_typescript_style(self, indent: str = "") -> str:305        return self.ref_name306 307 308_ParameterType = (_ParameterTypeScalar309                  | _ParameterTypeObject310                  | _ParameterTypeArray311                  | _ParameterTypeEnum312                  | _ParameterTypeAnyOf313                  | _ParameterTypeUnion314                  | _ParameterTypeRef)315 316 317@dataclasses.dataclass318class _Parameter:319    """320    A parameter in a function, or a field in a object.321    It consists of the type as well as the name.322    """323 324    type: _ParameterType325    name: str = "_"326    optional: bool = True327    default: Any | None = None328 329    @classmethod330    def parse_extended(cls, attributes: dict[str, Any]) -> "_Parameter":331        if not attributes:332            raise ValueError("attributes is empty")333 334        return cls(335            name=attributes.get("name", "_"),336            type=_parse_parameter_type(attributes),337            optional=attributes.get("optional", False),338            default=attributes.get("default"),339        )340 341    def to_typescript_style(self, indent: str = "") -> str:342        comments = self.type.format_docstring(indent)343 344        if self.default is not None:345            default_repr = (json.dumps(self.default, ensure_ascii=False)346                            if not isinstance(self.default, (int, float, bool))347                            else repr(self.default))348            comments += f"{indent}// Default: {default_repr}\n"349 350        return (351            comments +352            f"{indent}{self.name}{'?' if self.optional else ''}: {self.type.to_typescript_style(indent=indent)}"353        )354 355 356def _parse_parameter_type(357        json_schema_object: dict[str, Any] | bool,358        registry: _SchemaRegistry | None = None) -> _ParameterType:359    if isinstance(json_schema_object, bool):360        if json_schema_object:361            return _ParameterTypeScalar(type="any")362        else:363            logger.warning(364                f"Warning: Boolean value {json_schema_object} is not supported, use null instead."365            )366            return _ParameterTypeScalar(type="null")367 368    if "$ref" in json_schema_object and registry:369        return _ParameterTypeRef(json_schema_object, registry)370 371    if "anyOf" in json_schema_object:372        return _ParameterTypeAnyOf(json_schema_object, registry)373    elif "enum" in json_schema_object:374        return _ParameterTypeEnum(json_schema_object)375    elif "type" in json_schema_object:376        typ = json_schema_object["type"]377        if isinstance(typ, list):378            return _ParameterTypeUnion(json_schema_object)379        elif typ == "object":380            return _ParameterTypeObject(json_schema_object, registry)381        elif typ == "array":382            return _ParameterTypeArray(json_schema_object, registry)383        else:384            return _ParameterTypeScalar(typ, json_schema_object)385    elif json_schema_object == {}:386        return _ParameterTypeScalar(type="any")387    else:388        raise ValueError(f"Invalid JSON Schema object: {json_schema_object}")389 390 391def _openai_function_to_typescript_style(function: dict[str, Any], ) -> str:392    """Convert OpenAI function definition (dict) to TypeScript style string."""393    registry = _SchemaRegistry()394    parameters = function.get("parameters") or {}395    parsed = _ParameterTypeObject(parameters, registry)396 397    interfaces = []398    root_interface_name = None399    if registry.has_self_ref:400        root_interface_name = "parameters"401        params_str = _TS_FIELD_DELIMITER.join([402            p.to_typescript_style(indent=_TS_INDENT) for p in parsed.properties403        ])404        params_str = f"\n{params_str}\n" if params_str else ""405        interface_def = f"interface {root_interface_name} {{{params_str}}}"406        interfaces.append(interface_def)407 408    definitions_copy = dict(registry.definitions)409    for def_name, def_schema in definitions_copy.items():410        obj_type = _parse_parameter_type(def_schema, registry)411        params_str = obj_type.to_typescript_style()412 413        description_part = ""414        if obj_description := def_schema.get("description", ""):415            description_part = _format_description(obj_description) + "\n"416 417        interface_def = f"{description_part}interface {def_name} {params_str}"418        interfaces.append(interface_def)419 420    interface_str = "\n".join(interfaces)421    function_name = function.get("name", "function")422    if root_interface_name:423        type_def = f"type {function_name} = (_: {root_interface_name}) => any;"424    else:425        params_str = parsed.to_typescript_style()426        type_def = f"type {function_name} = (_: {params_str}) => any;"427 428    description = function.get("description")429    return "\n".join(430        filter(431            bool,432            [433                interface_str,434                ((description and _format_description(description)) or ""),435                type_def,436            ],437        ))438 439 440def encode_tools_to_typescript_style(tools: list[dict[str, Any]], ) -> str:441    """442    Convert tools (list of dict) to TypeScript style string.443    444    Supports OpenAI format: {"type": "function", "function": {...}}445    446    Args:447        tools: List of tool definitions in dict format448        449    Returns:450        TypeScript style string representation of the tools451    """452    if not tools:453        return ""454 455    functions = []456 457    for tool in tools:458        tool_type = tool.get("type")459        if tool_type == "function":460            func_def = tool.get("function", {})461            if func_def:462                functions.append(463                    _openai_function_to_typescript_style(func_def))464        else:465            # Skip unsupported tool types (like "_plugin")466            continue467 468    if not functions:469        return ""470 471    functions_str = "\n".join(functions)472    result = "# Tools\n\n"473 474    if functions_str:475        result += "## functions\nnamespace functions {\n"476        result += functions_str + "\n"477        result += "}\n"478 479    return result480