KBaba7/llama.cpp
0
1#!/usr/bin/env python32from __future__ import annotations3 4import argparse5import itertools6import json7import re8import sys9from typing import Any, List, Optional, Set, Tuple, Union10 11def _build_repetition(item_rule, min_items, max_items, separator_rule=None):12 13 if min_items == 0 and max_items == 1:14 return f'{item_rule}?'15 16 if not separator_rule:17 if min_items == 1 and max_items is None:18 return f'{item_rule}+'19 elif min_items == 0 and max_items is None:20 return f'{item_rule}*'21 else:22 return f'{item_rule}{{{min_items},{max_items if max_items is not None else ""}}}'23 24 result = item_rule + ' ' + _build_repetition(f'({separator_rule} {item_rule})', min_items - 1 if min_items > 0 else 0, max_items - 1 if max_items is not None else None)25 return f'({result})?' if min_items == 0 else result26 27def _generate_min_max_int(min_value: Optional[int], max_value: Optional[int], out: list, decimals_left: int = 16, top_level: bool = True):28 has_min = min_value != None29 has_max = max_value != None30 31 def digit_range(from_char: str, to_char: str):32 out.append("[")33 if from_char == to_char:34 out.append(from_char)35 else:36 out.append(from_char)37 out.append("-")38 out.append(to_char)39 out.append("]")40 41 def more_digits(min_digits: int, max_digits: int):42 out.append("[0-9]")43 if min_digits == max_digits and min_digits == 1:44 return45 out.append("{")46 out.append(str(min_digits))47 if max_digits != min_digits:48 out.append(",")49 if max_digits != sys.maxsize:50 out.append(str(max_digits))51 out.append("}")52 53 def uniform_range(from_str: str, to_str: str):54 i = 055 while i < len(from_str) and from_str[i] == to_str[i]:56 i += 157 if i > 0:58 out.append("\"")59 out.append(from_str[:i])60 out.append("\"")61 if i < len(from_str):62 if i > 0:63 out.append(" ")64 sub_len = len(from_str) - i - 165 if sub_len > 0:66 from_sub = from_str[i+1:]67 to_sub = to_str[i+1:]68 sub_zeros = "0" * sub_len69 sub_nines = "9" * sub_len70 71 to_reached = False72 out.append("(")73 if from_sub == sub_zeros:74 digit_range(from_str[i], chr(ord(to_str[i]) - 1))75 out.append(" ")76 more_digits(sub_len, sub_len)77 else:78 out.append("[")79 out.append(from_str[i])80 out.append("] ")81 out.append("(")82 uniform_range(from_sub, sub_nines)83 out.append(")")84 if ord(from_str[i]) < ord(to_str[i]) - 1:85 out.append(" | ")86 if to_sub == sub_nines:87 digit_range(chr(ord(from_str[i]) + 1), to_str[i])88 to_reached = True89 else:90 digit_range(chr(ord(from_str[i]) + 1), chr(ord(to_str[i]) - 1))91 out.append(" ")92 more_digits(sub_len, sub_len)93 if not to_reached:94 out.append(" | ")95 digit_range(to_str[i], to_str[i])96 out.append(" ")97 uniform_range(sub_zeros, to_sub)98 out.append(")")99 else:100 out.append("[")101 out.append(from_str[i])102 out.append("-")103 out.append(to_str[i])104 out.append("]")105 106 if has_min and has_max:107 if min_value < 0 and max_value < 0:108 out.append("\"-\" (")109 _generate_min_max_int(-max_value, -min_value, out, decimals_left, top_level=True)110 out.append(")")111 return112 113 if min_value < 0:114 out.append("\"-\" (")115 _generate_min_max_int(0, -min_value, out, decimals_left, top_level=True)116 out.append(") | ")117 min_value = 0118 119 min_s = str(min_value)120 max_s = str(max_value)121 min_digits = len(min_s)122 max_digits = len(max_s)123 124 for digits in range(min_digits, max_digits):125 uniform_range(min_s, "9" * digits)126 min_s = "1" + "0" * digits127 out.append(" | ")128 uniform_range(min_s, max_s)129 return130 131 less_decimals = max(decimals_left - 1, 1)132 133 if has_min:134 if min_value < 0:135 out.append("\"-\" (")136 _generate_min_max_int(None, -min_value, out, decimals_left, top_level=False)137 out.append(") | [0] | [1-9] ")138 more_digits(0, decimals_left - 1)139 elif min_value == 0:140 if top_level:141 out.append("[0] | [1-9] ")142 more_digits(0, less_decimals)143 else:144 more_digits(1, decimals_left)145 elif min_value <= 9:146 c = str(min_value)147 range_start = '1' if top_level else '0'148 if c > range_start:149 digit_range(range_start, chr(ord(c) - 1))150 out.append(" ")151 more_digits(1, less_decimals)152 out.append(" | ")153 digit_range(c, "9")154 out.append(" ")155 more_digits(0, less_decimals)156 else:157 min_s = str(min_value)158 length = len(min_s)159 c = min_s[0]160 161 if c > "1":162 digit_range("1" if top_level else "0", chr(ord(c) - 1))163 out.append(" ")164 more_digits(length, less_decimals)165 out.append(" | ")166 digit_range(c, c)167 out.append(" (")168 _generate_min_max_int(int(min_s[1:]), None, out, less_decimals, top_level=False)169 out.append(")")170 if c < "9":171 out.append(" | ")172 digit_range(chr(ord(c) + 1), "9")173 out.append(" ")174 more_digits(length - 1, less_decimals)175 return176 177 if has_max:178 if max_value >= 0:179 if top_level:180 out.append("\"-\" [1-9] ")181 more_digits(0, less_decimals)182 out.append(" | ")183 _generate_min_max_int(0, max_value, out, decimals_left, top_level=True)184 else:185 out.append("\"-\" (")186 _generate_min_max_int(-max_value, None, out, decimals_left, top_level=False)187 out.append(")")188 return189 190 raise RuntimeError("At least one of min_value or max_value must be set")191 192class BuiltinRule:193 def __init__(self, content: str, deps: list | None = None):194 self.content = content195 self.deps = deps or []196 197# Constraining spaces to prevent model "running away".198SPACE_RULE = '| " " | "\\n" [ \\t]{0,20}'199 200PRIMITIVE_RULES = {201 'boolean' : BuiltinRule('("true" | "false") space', []),202 'decimal-part' : BuiltinRule('[0-9]{1,16}', []),203 'integral-part': BuiltinRule('[0] | [1-9] [0-9]{0,15}', []),204 'number' : BuiltinRule('("-"? integral-part) ("." decimal-part)? ([eE] [-+]? integral-part)? space', ['integral-part', 'decimal-part']),205 'integer' : BuiltinRule('("-"? integral-part) space', ['integral-part']),206 'value' : BuiltinRule('object | array | string | number | boolean | null', ['object', 'array', 'string', 'number', 'boolean', 'null']),207 'object' : BuiltinRule('"{" space ( string ":" space value ("," space string ":" space value)* )? "}" space', ['string', 'value']),208 'array' : BuiltinRule('"[" space ( value ("," space value)* )? "]" space', ['value']),209 'uuid' : BuiltinRule(r'"\"" [0-9a-fA-F]{8} "-" [0-9a-fA-F]{4} "-" [0-9a-fA-F]{4} "-" [0-9a-fA-F]{4} "-" [0-9a-fA-F]{12} "\"" space', []),210 'char' : BuiltinRule(r'[^"\\\x7F\x00-\x1F] | [\\] (["\\bfnrt] | "u" [0-9a-fA-F]{4})', []),211 'string' : BuiltinRule(r'"\"" char* "\"" space', ['char']),212 'null' : BuiltinRule('"null" space', []),213}214 215# TODO: support "uri", "email" string formats216STRING_FORMAT_RULES = {217 'date' : BuiltinRule('[0-9]{4} "-" ( "0" [1-9] | "1" [0-2] ) "-" ( \"0\" [1-9] | [1-2] [0-9] | "3" [0-1] )', []),218 'time' : BuiltinRule('([01] [0-9] | "2" [0-3]) ":" [0-5] [0-9] ":" [0-5] [0-9] ( "." [0-9]{3} )? ( "Z" | ( "+" | "-" ) ( [01] [0-9] | "2" [0-3] ) ":" [0-5] [0-9] )', []),219 'date-time' : BuiltinRule('date "T" time', ['date', 'time']),220 'date-string' : BuiltinRule('"\\"" date "\\"" space', ['date']),221 'time-string' : BuiltinRule('"\\"" time "\\"" space', ['time']),222 'date-time-string': BuiltinRule('"\\"" date-time "\\"" space', ['date-time']),223}224 225DOTALL = '[\\U00000000-\\U0010FFFF]'226DOT = '[^\\x0A\\x0D]'227 228RESERVED_NAMES = set(["root", "dot", *PRIMITIVE_RULES.keys(), *STRING_FORMAT_RULES.keys()])229 230INVALID_RULE_CHARS_RE = re.compile(r'[^a-zA-Z0-9-]+')231GRAMMAR_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"]')232GRAMMAR_RANGE_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"\]\-\\]')233GRAMMAR_LITERAL_ESCAPES = {'\r': '\\r', '\n': '\\n', '"': '\\"', '-': '\\-', ']': '\\]'}234 235NON_LITERAL_SET = set('|.()[]{}*+?')236ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = set('^$.[]()|{}*+?')237 238 239class SchemaConverter:240 def __init__(self, *, prop_order, allow_fetch, dotall, raw_pattern):241 self._prop_order = prop_order242 self._allow_fetch = allow_fetch243 self._dotall = dotall244 self._raw_pattern = raw_pattern245 self._rules = {246 'space': SPACE_RULE,247 }248 self._refs = {}249 self._refs_being_resolved = set()250 251 def _format_literal(self, literal):252 escaped = GRAMMAR_LITERAL_ESCAPE_RE.sub(253 lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)) or m.group(0), literal254 )255 return f'"{escaped}"'256 257 def not_literal(self, literal: str, dotall: bool = True, maybe_escaped_underscores = False) -> str:258 '''259 not_literal('a') -> '[^a]'260 not_literal('abc') -> '([^a] | "a" ([^b] | "b" ([^c])?)?)?'261 '''262 assert len(literal) > 0, 'Empty literal not supported'263 def recurse(i: int):264 c = literal[i]265 if maybe_escaped_underscores and c == '_':266 yield f'[^{c}\\\\]'267 yield ' | '268 yield f'"\\\\"? "{c}"'269 else:270 yield f'[^{c}]'271 if i < len(literal) - 1:272 yield ' | '273 yield self._format_literal(c)274 yield ' ('275 yield from recurse(i + 1)276 yield ')?'277 278 return ''.join(('(', *recurse(0), ')'))279 280 def _not_strings(self, strings):281 class TrieNode:282 def __init__(self):283 self.children = {}284 self.is_end_of_string = False285 286 def insert(self, string):287 node = self288 for c in string:289 node = node.children.setdefault(c, TrieNode())290 node.is_end_of_string = True291 292 trie = TrieNode()293 for s in strings:294 trie.insert(s)295 296 char_rule = self._add_primitive('char', PRIMITIVE_RULES['char'])297 out = ['["] ( ']298 299 def visit(node):300 rejects = []301 first = True302 for c in sorted(node.children.keys()):303 child = node.children[c]304 rejects.append(c)305 if first:306 first = False307 else:308 out.append(' | ')309 out.append(f'[{c}]')310 if child.children:311 out.append(f' (')312 visit(child)313 out.append(')')314 elif child.is_end_of_string:315 out.append(f' {char_rule}+')316 if node.children:317 if not first:318 out.append(' | ')319 out.append(f'[^"{"".join(rejects)}] {char_rule}*')320 visit(trie)321 322 out.append(f' ){"" if trie.is_end_of_string else "?"} ["] space')323 return ''.join(out)324 325 def _add_rule(self, name, rule):326 esc_name = INVALID_RULE_CHARS_RE.sub('-', name)327 if esc_name not in self._rules or self._rules[esc_name] == rule:328 key = esc_name329 else:330 i = 0331 while f'{esc_name}{i}' in self._rules and self._rules[f'{esc_name}{i}'] != rule:332 i += 1333 key = f'{esc_name}{i}'334 self._rules[key] = rule335 return key336 337 def resolve_refs(self, schema: dict, url: str):338 '''339 Resolves all $ref fields in the given schema, fetching any remote schemas,340 replacing $ref with absolute reference URL and populating self._refs with the341 respective referenced (sub)schema dictionaries.342 '''343 def visit(n: dict):344 if isinstance(n, list):345 return [visit(x) for x in n]346 elif isinstance(n, dict):347 ref = n.get('$ref')348 if ref is not None and ref not in self._refs:349 if ref.startswith('https://'):350 assert self._allow_fetch, 'Fetching remote schemas is not allowed (use --allow-fetch for force)'351 import requests352 353 frag_split = ref.split('#')354 base_url = frag_split[0]355 356 target = self._refs.get(base_url)357 if target is None:358 target = self.resolve_refs(requests.get(ref).json(), base_url)359 self._refs[base_url] = target360 361 if len(frag_split) == 1 or frag_split[-1] == '':362 return target363 elif ref.startswith('#/'):364 target = schema365 ref = f'{url}{ref}'366 n['$ref'] = ref367 else:368 raise ValueError(f'Unsupported ref {ref}')369 370 for sel in ref.split('#')[-1].split('/')[1:]:371 assert target is not None and sel in target, f'Error resolving ref {ref}: {sel} not in {target}'372 target = target[sel]373 374 self._refs[ref] = target375 else:376 for v in n.values():377 visit(v)378 379 return n380 return visit(schema)381 382 def _generate_union_rule(self, name, alt_schemas):383 return ' | '.join((384 self.visit(alt_schema, f'{name}{"-" if name else "alternative-"}{i}')385 for i, alt_schema in enumerate(alt_schemas)386 ))387 388 def _visit_pattern(self, pattern, name):389 '''390 Transforms a regular expression pattern into a GBNF rule.391 392 Input: https://json-schema.org/understanding-json-schema/reference/regular_expressions393 Output: https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md394 395 Unsupported features: negative/positive lookaheads, greedy/non-greedy modifiers.396 397 Mostly a 1:1 translation, except for {x} / {x,} / {x,y} quantifiers for which398 we define sub-rules to keep the output lean.399 '''400 401 assert pattern.startswith('^') and pattern.endswith('$'), 'Pattern must start with "^" and end with "$"'402 pattern = pattern[1:-1]403 sub_rule_ids = {}404 405 i = 0406 length = len(pattern)407 408 def to_rule(s: tuple[str, bool]) -> str:409 (txt, is_literal) = s410 return "\"" + txt + "\"" if is_literal else txt411 412 def transform() -> tuple[str, bool]:413 '''414 Parse a unit at index i (advancing it), and return its string representation + whether it's a literal.415 '''416 nonlocal i417 nonlocal pattern418 nonlocal sub_rule_ids419 420 start = i421 # For each component of this sequence, store its string representation and whether it's a literal.422 # We only need a flat structure here to apply repetition operators to the last item, and423 # to merge literals at the and (we're parsing grouped ( sequences ) recursively and don't treat '|' specially424 # (GBNF's syntax is luckily very close to regular expressions!)425 seq: list[tuple[str, bool]] = []426 427 def get_dot():428 if self._dotall:429 rule = DOTALL430 else:431 # Accept any character... except \n and \r line break chars (\x0A and \xOD)432 rule = DOT433 return self._add_rule(f'dot', rule)434 435 def join_seq():436 nonlocal seq437 ret = []438 for is_literal, g in itertools.groupby(seq, lambda x: x[1]):439 if is_literal:440 ret.append((''.join(x[0] for x in g), True))441 else:442 ret.extend(g)443 if len(ret) == 1:444 return ret[0]445 return (' '.join(to_rule(x) for x in seq), False)446 447 while i < length:448 c = pattern[i]449 if c == '.':450 seq.append((get_dot(), False))451 i += 1452 elif c == '(':453 i += 1454 if i < length:455 assert pattern[i] != '?', f'Unsupported pattern syntax "{pattern[i]}" at index {i} of /{pattern}/'456 seq.append((f'({to_rule(transform())})', False))457 elif c == ')':458 i += 1459 assert start > 0 and pattern[start-1] == '(', f'Unbalanced parentheses; start = {start}, i = {i}, pattern = {pattern}'460 return join_seq()461 elif c == '[':462 square_brackets = c463 i += 1464 while i < length and pattern[i] != ']':465 if pattern[i] == '\\':466 square_brackets += pattern[i:i+2]467 i += 2468 else:469 square_brackets += pattern[i]470 i += 1471 assert i < length, f'Unbalanced square brackets; start = {start}, i = {i}, pattern = {pattern}'472 square_brackets += ']'473 i += 1474 seq.append((square_brackets, False))475 elif c == '|':476 seq.append(('|', False))477 i += 1478 elif c in ('*', '+', '?'):479 seq[-1] = (to_rule(seq[-1]) + c, False)480 i += 1481 elif c == '{':482 curly_brackets = c483 i += 1484 while i < length and pattern[i] != '}':485 curly_brackets += pattern[i]486 i += 1487 assert i < length, f'Unbalanced curly brackets; start = {start}, i = {i}, pattern = {pattern}'488 curly_brackets += '}'489 i += 1490 nums = [s.strip() for s in curly_brackets[1:-1].split(',')]491 min_times = 0492 max_times = None493 try:494 if len(nums) == 1:495 min_times = int(nums[0])496 max_times = min_times497 else:498 assert len(nums) == 2499 min_times = int(nums[0]) if nums[0] else 0500 max_times = int(nums[1]) if nums[1] else None501 except ValueError:502 raise ValueError(f'Invalid quantifier {curly_brackets} in /{pattern}/')503 504 (sub, sub_is_literal) = seq[-1]505 506 if not sub_is_literal:507 id = sub_rule_ids.get(sub)508 if id is None:509 id = self._add_rule(f'{name}-{len(sub_rule_ids) + 1}', sub)510 sub_rule_ids[sub] = id511 sub = id512 513 seq[-1] = (_build_repetition(f'"{sub}"' if sub_is_literal else sub, min_times, max_times), False)514 else:515 literal = ''516 while i < length:517 if pattern[i] == '\\' and i < length - 1:518 next = pattern[i + 1]519 if next in ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS:520 i += 1521 literal += pattern[i]522 i += 1523 else:524 literal += pattern[i:i+2]525 i += 2526 elif pattern[i] == '"' and not self._raw_pattern:527 literal += '\\"'528 i += 1529 elif pattern[i] not in NON_LITERAL_SET and \530 (i == length - 1 or literal == '' or pattern[i+1] == '.' or pattern[i+1] not in NON_LITERAL_SET):531 literal += pattern[i]532 i += 1533 else:534 break535 if literal:536 seq.append((literal, True))537 538 return join_seq()539 540 return self._add_rule(541 name,542 to_rule(transform()) if self._raw_pattern \543 else "\"\\\"\" (" + to_rule(transform()) + ") \"\\\"\" space")544 545 546 def _resolve_ref(self, ref):547 ref_name = ref.split('/')[-1]548 if ref_name not in self._rules and ref not in self._refs_being_resolved:549 self._refs_being_resolved.add(ref)550 resolved = self._refs[ref]551 ref_name = self.visit(resolved, ref_name)552 self._refs_being_resolved.remove(ref)553 return ref_name554 555 def _generate_constant_rule(self, value):556 return self._format_literal(json.dumps(value))557 558 def visit(self, schema, name):559 schema_type = schema.get('type')560 schema_format = schema.get('format')561 rule_name = name + '-' if name in RESERVED_NAMES else name or 'root'562 563 if (ref := schema.get('$ref')) is not None:564 return self._add_rule(rule_name, self._resolve_ref(ref))565 566 elif 'oneOf' in schema or 'anyOf' in schema:567 return self._add_rule(rule_name, self._generate_union_rule(name, schema.get('oneOf') or schema['anyOf']))568 569 elif isinstance(schema_type, list):570 return self._add_rule(rule_name, self._generate_union_rule(name, [{**schema, 'type': t} for t in schema_type]))571 572 elif 'const' in schema:573 return self._add_rule(rule_name, self._generate_constant_rule(schema['const']) + ' space')574 575 elif 'enum' in schema:576 rule = '(' + ' | '.join((self._generate_constant_rule(v) for v in schema['enum'])) + ') space'577 return self._add_rule(rule_name, rule)578 579 elif schema_type in (None, 'object') and \580 ('properties' in schema or \581 ('additionalProperties' in schema and schema['additionalProperties'] is not True)):582 required = set(schema.get('required', []))583 properties = list(schema.get('properties', {}).items())584 return self._add_rule(rule_name, self._build_object_rule(properties, required, name, schema.get('additionalProperties')))585 586 elif schema_type in (None, 'object') and 'allOf' in schema:587 required = set()588 properties = []589 hybrid_name = name590 def add_component(comp_schema, is_required):591 if (ref := comp_schema.get('$ref')) is not None:592 comp_schema = self._refs[ref]593 594 if 'properties' in comp_schema:595 for prop_name, prop_schema in comp_schema['properties'].items():596 properties.append((prop_name, prop_schema))597 if is_required:598 required.add(prop_name)599 600 for t in schema['allOf']:601 if 'anyOf' in t:602 for tt in t['anyOf']:603 add_component(tt, is_required=False)604 else:605 add_component(t, is_required=True)606 607 return self._add_rule(rule_name, self._build_object_rule(properties, required, hybrid_name, additional_properties=None))608 609 elif schema_type in (None, 'array') and ('items' in schema or 'prefixItems' in schema):610 items = schema.get('items') or schema['prefixItems']611 if isinstance(items, list):612 return self._add_rule(613 rule_name,614 '"[" space ' +615 ' "," space '.join(616 self.visit(item, f'{name}{"-" if name else ""}tuple-{i}')617 for i, item in enumerate(items)) +618 ' "]" space')619 else:620 item_rule_name = self.visit(items, f'{name}{"-" if name else ""}item')621 min_items = schema.get("minItems", 0)622 max_items = schema.get("maxItems")623 return self._add_rule(rule_name, '"[" space ' + _build_repetition(item_rule_name, min_items, max_items, separator_rule='"," space') + ' "]" space')624 625 elif schema_type in (None, 'string') and 'pattern' in schema:626 return self._visit_pattern(schema['pattern'], rule_name)627 628 elif schema_type in (None, 'string') and re.match(r'^uuid[1-5]?$', schema_format or ''):629 return self._add_primitive(630 'root' if rule_name == 'root' else schema_format,631 PRIMITIVE_RULES['uuid']632 )633 634 elif schema_type in (None, 'string') and f'{schema_format}-string' in STRING_FORMAT_RULES:635 prim_name = f'{schema_format}-string'636 return self._add_rule(rule_name, self._add_primitive(prim_name, STRING_FORMAT_RULES[prim_name]))637 638 elif schema_type == 'string' and ('minLength' in schema or 'maxLength' in schema):639 char_rule = self._add_primitive('char', PRIMITIVE_RULES['char'])640 min_len = schema.get('minLength', 0)641 max_len = schema.get('maxLength')642 643 return self._add_rule(rule_name, r'"\"" ' + _build_repetition(char_rule, min_len, max_len) + r' "\"" space')644 645 elif schema_type in (None, 'integer') and \646 ('minimum' in schema or 'exclusiveMinimum' in schema or 'maximum' in schema or 'exclusiveMaximum' in schema):647 min_value = None648 max_value = None649 if 'minimum' in schema:650 min_value = schema['minimum']651 elif 'exclusiveMinimum' in schema:652 min_value = schema['exclusiveMinimum'] + 1653 if 'maximum' in schema:654 max_value = schema['maximum']655 elif 'exclusiveMaximum' in schema:656 max_value = schema['exclusiveMaximum'] - 1657 658 out = ["("]659 _generate_min_max_int(min_value, max_value, out)660 out.append(") space")661 return self._add_rule(rule_name, ''.join(out))662 663 elif (schema_type == 'object') or (len(schema) == 0):664 return self._add_rule(rule_name, self._add_primitive('object', PRIMITIVE_RULES['object']))665 666 else:667 assert schema_type in PRIMITIVE_RULES, f'Unrecognized schema: {schema}'668 # TODO: support minimum, maximum, exclusiveMinimum, exclusiveMaximum at least for zero669 return self._add_primitive('root' if rule_name == 'root' else schema_type, PRIMITIVE_RULES[schema_type])670 671 def _add_primitive(self, name: str, rule: BuiltinRule):672 n = self._add_rule(name, rule.content)673 674 for dep in rule.deps:675 dep_rule = PRIMITIVE_RULES.get(dep) or STRING_FORMAT_RULES.get(dep)676 assert dep_rule, f'Rule {dep} not known'677 if dep not in self._rules:678 self._add_primitive(dep, dep_rule)679 return n680 681 def _build_object_rule(self, properties: List[Tuple[str, Any]], required: Set[str], name: str, additional_properties: Optional[Union[bool, Any]]):682 prop_order = self._prop_order683 # sort by position in prop_order (if specified) then by original order684 sorted_props = [kv[0] for _, kv in sorted(enumerate(properties), key=lambda ikv: (prop_order.get(ikv[1][0], len(prop_order)), ikv[0]))]685 686 prop_kv_rule_names = {}687 for prop_name, prop_schema in properties:688 prop_rule_name = self.visit(prop_schema, f'{name}{"-" if name else ""}{prop_name}')689 prop_kv_rule_names[prop_name] = self._add_rule(690 f'{name}{"-" if name else ""}{prop_name}-kv',691 fr'{self._format_literal(json.dumps(prop_name))} space ":" space {prop_rule_name}'692 )693 required_props = [k for k in sorted_props if k in required]694 optional_props = [k for k in sorted_props if k not in required]695 696 if additional_properties is not None and additional_properties != False:697 sub_name = f'{name}{"-" if name else ""}additional'698 value_rule = self.visit(additional_properties, f'{sub_name}-value') if isinstance(additional_properties, dict) else \699 self._add_primitive('value', PRIMITIVE_RULES['value'])700 key_rule = self._add_primitive('string', PRIMITIVE_RULES['string']) if not sorted_props \701 else self._add_rule(f'{sub_name}-k', self._not_strings(sorted_props))702 703 prop_kv_rule_names["*"] = self._add_rule(704 f'{sub_name}-kv',705 f'{key_rule} ":" space {value_rule}'706 )707 optional_props.append("*")708 709 rule = '"{" space '710 rule += ' "," space '.join(prop_kv_rule_names[k] for k in required_props)711 712 if optional_props:713 rule += ' ('714 if required_props:715 rule += ' "," space ( '716 717 def get_recursive_refs(ks, first_is_optional):718 [k, *rest] = ks719 kv_rule_name = prop_kv_rule_names[k]720 comma_ref = f'( "," space {kv_rule_name} )'721 if first_is_optional:722 res = comma_ref + ('*' if k == '*' else '?')723 else:724 res = kv_rule_name + (' ' + comma_ref + "*" if k == '*' else '')725 if len(rest) > 0:726 res += ' ' + self._add_rule(727 f'{name}{"-" if name else ""}{k}-rest',728 get_recursive_refs(rest, first_is_optional=True)729 )730 return res731 732 rule += ' | '.join(733 get_recursive_refs(optional_props[i:], first_is_optional=False)734 for i in range(len(optional_props))735 )736 if required_props:737 rule += ' )'738 rule += ' )?'739 740 rule += ' "}" space'741 742 return rule743 744 def format_grammar(self):745 return '\n'.join(746 f'{name} ::= {rule}'747 for name, rule in sorted(self._rules.items(), key=lambda kv: kv[0])748 )749 750 751def main(args_in = None):752 parser = argparse.ArgumentParser(753 description='''754 Generates a grammar (suitable for use in ./llama-cli) that produces JSON conforming to a755 given JSON schema. Only a subset of JSON schema features are supported; more may be756 added in the future.757 ''',758 )759 parser.add_argument(760 '--prop-order',761 default=[],762 type=lambda s: s.split(','),763 help='''764 comma-separated property names defining the order of precedence for object properties;765 properties not specified here are given lower precedence than those that are, and766 are kept in their original order from the schema. Required properties are always767 given precedence over optional properties.768 '''769 )770 parser.add_argument(771 '--allow-fetch',772 action='store_true',773 default=False,774 help='Whether to allow fetching referenced schemas over HTTPS')775 parser.add_argument(776 '--dotall',777 action='store_true',778 default=False,779 help='Whether to treat dot (".") as matching all chars including line breaks in regular expression patterns')780 parser.add_argument(781 '--raw-pattern',782 action='store_true',783 default=False,784 help='Treats string patterns as raw patterns w/o quotes (or quote escapes)')785 786 parser.add_argument('schema', help='file containing JSON schema ("-" for stdin)')787 args = parser.parse_args(args_in)788 789 if args.schema.startswith('https://'):790 url = args.schema791 import requests792 schema = requests.get(url).json()793 elif args.schema == '-':794 url = 'stdin'795 schema = json.load(sys.stdin)796 else:797 url = f'file://{args.schema}'798 with open(args.schema) as f:799 schema = json.load(f)800 converter = SchemaConverter(801 prop_order={name: idx for idx, name in enumerate(args.prop_order)},802 allow_fetch=args.allow_fetch,803 dotall=args.dotall,804 raw_pattern=args.raw_pattern)805 schema = converter.resolve_refs(schema, url)806 converter.visit(schema, '')807 print(converter.format_grammar())808 809 810if __name__ == '__main__':811 main()812