tool_declaration_ts.py
16.1 KB · 480 lines · python Raw
1 """
2 Encode structured tool declaration to typescript style string.
3 """
4 import dataclasses
5 import json
6 import logging
7 from collections.abc import Sequence
8 from typing import Any
9
10 logger = logging.getLogger(__name__)
11
12 _TS_INDENT = " "
13 _TS_FIELD_DELIMITER = ",\n"
14
15
16 class _SchemaRegistry:
17 """Registry for schema definitions to handle $ref resolution"""
18
19 def __init__(self):
20 self.definitions = {}
21 self.has_self_ref = False
22
23 def register_definitions(self, defs: dict[str, Any]):
24 """Register schema definitions from $defs section"""
25 if not defs:
26 return
27 for def_name, def_schema in defs.items():
28 self.definitions[def_name] = def_schema
29
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 = True
34 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
44 def _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
51 class _BaseType:
52 description: str
53 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
64 for k, v in extra_props.items() if k in allowed_constraint_keys
65 }
66
67 def to_typescript_style(self, indent: str = "") -> str:
68 raise NotImplementedError
69
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
82 class _ParameterTypeScalar(_BaseType):
83 type: str
84
85 def __init__(self, type: str, extra_props: dict[str, Any] | None = None):
86 self.type = type
87
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 TypeScript
99 if self.type == "integer":
100 return "number"
101 return self.type
102
103
104 class _ParameterTypeObject(_BaseType):
105 properties: list["_Parameter"]
106 additional_properties: Any | None = None
107
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 = None
115
116 if not json_schema_object:
117 return
118
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 return
130
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 first
147 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 after
180 params_str = f"\n{params_str}\n"
181 # always wrap with object
182 return f"{{{params_str}{indent}}}"
183
184
185 class _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
209 class _ParameterTypeEnum(_BaseType):
210 # support scalar types only
211 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 present
218 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 continue
233 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
247 class _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
266 class _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
287 class _ParameterTypeRef(_BaseType):
288 ref_name: str
289 is_self_ref: bool = False
290
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 = True
301 else:
302 self.ref_name = ref.split("/")[-1]
303
304 def to_typescript_style(self, indent: str = "") -> str:
305 return self.ref_name
306
307
308 _ParameterType = (_ParameterTypeScalar
309 | _ParameterTypeObject
310 | _ParameterTypeArray
311 | _ParameterTypeEnum
312 | _ParameterTypeAnyOf
313 | _ParameterTypeUnion
314 | _ParameterTypeRef)
315
316
317 @dataclasses.dataclass
318 class _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: _ParameterType
325 name: str = "_"
326 optional: bool = True
327 default: Any | None = None
328
329 @classmethod
330 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
356 def _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
391 def _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 = None
399 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.properties
403 ])
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
440 def 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 format
448
449 Returns:
450 TypeScript style string representation of the tools
451 """
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 continue
467
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 result
480