char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Aldehir Rojascommon : implement common_schema internal representation for JSON schemas (#28736)acecd5603

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
20.4 KiB514 linesraw
1#include "json-schema.h"
2#include "common.h"
3
4#include <cmath>
5#include <map>
6#include <stdexcept>
7#include <string>
8#include <unordered_set>
9#include <utility>
10#include <vector>
11
12class common_chat_schema_builder {
13    const common_json &      root_;
14    common_chat_schema_document & doc_;
15
16    // the targets built here, moved into doc_ once the whole schema is built
17    std::map<std::string, common_chat_schema_ptr> refs_;
18
19    // ref nodes get their target once every $ref is built, a cycle would otherwise need it too early
20    std::vector<common_chat_schema_ref *> pending_;
21
22    [[noreturn]] static void fail(const std::string & path, const std::string & msg) {
23        throw std::runtime_error("JSON schema error at " + path + ": " + msg);
24    }
25
26    static int get_count(const common_json & schema, const std::string & key, const std::string & path, int def) {
27        if (!schema.contains(key)) {
28            return def;
29        }
30        const common_json & value = schema.at(key);
31        if (!value.is_number_integer() || value.get<int>() < 0) {
32            fail(path, key + " must be a non-negative integer");
33        }
34        return value.get<int>();
35    }
36
37    // a fractional bound is rounded inwards, towards the integers it still admits
38    static int64_t get_bound(const common_json & schema, const std::string & key, const std::string & path, bool round_up) {
39        const common_json & value = schema.at(key);
40        if (value.is_number_integer()) {
41            return value.get<int64_t>();
42        }
43        if (!value.is_number()) {
44            fail(path, key + " must be a number");
45        }
46        double d = value.get<double>();
47        return (int64_t) (round_up ? std::ceil(d) : std::floor(d));
48    }
49
50    static common_chat_schema::string_format get_format(const common_json & schema, const std::string & path) {
51        if (!schema.contains("format")) {
52            return common_chat_schema::FORMAT_NONE;
53        }
54        const common_json & value = schema.at("format");
55        if (!value.is_string()) {
56            fail(path, "format must be a string");
57        }
58        std::string format = value.get<std::string>();
59        if (format == "date") {
60            return common_chat_schema::FORMAT_DATE;
61        }
62        if (format == "time") {
63            return common_chat_schema::FORMAT_TIME;
64        }
65        if (format == "date-time") {
66            return common_chat_schema::FORMAT_DATE_TIME;
67        }
68        if (format == "uuid" || (format.size() == 5 && format.compare(0, 4, "uuid") == 0 && format[4] >= '1' && format[4] <= '5')) {
69            return common_chat_schema::FORMAT_UUID;
70        }
71        return common_chat_schema::FORMAT_NONE;
72    }
73
74    const common_json & resolve_ref(const std::string & ref, const std::string & path) {
75        const common_json * target = &root_;
76        auto tokens = string_split(ref.substr(1), "/");
77        for (size_t i = 1; i < tokens.size(); i++) {
78            const std::string & sel = tokens[i];
79            if (target->is_object() && target->contains(sel)) {
80                target = &target->at(sel);
81            } else if (target->is_array()) {
82                size_t idx;
83                try {
84                    idx = std::stoull(sel);
85                } catch (const std::logic_error &) {
86                    idx = target->size();
87                }
88                if (idx >= target->size()) {
89                    fail(path, "cannot resolve $ref " + ref + ", " + sel + " is out of range");
90                }
91                target = &target->at(idx);
92            } else {
93                fail(path, "cannot resolve $ref " + ref + ", " + sel + " not found");
94            }
95        }
96        return *target;
97    }
98
99    common_chat_schema_ptr build_ref(const common_json & value, const std::string & path) {
100        if (!value.is_string()) {
101            fail(path, "$ref must be a string");
102        }
103        std::string ref = value.get<std::string>();
104        if (ref.compare(0, 2, "#/") != 0) {
105            fail(path, "unsupported $ref " + ref + ", only references into the same document are supported");
106        }
107        if (refs_.find(ref) == refs_.end()) {
108            // reserve the key first, so that a cycle back to this $ref stops here
109            refs_[ref] = nullptr;
110            refs_[ref] = build_node(resolve_ref(ref, path), ref);
111        }
112        auto node = std::make_unique<common_chat_schema_ref>(ref);
113        pending_.push_back(node.get());
114        return node;
115    }
116
117    template <typename T>
118    common_chat_schema_ptr build_alternatives(const common_json & alts, const std::string & path) {
119        if (!alts.is_array()) {
120            fail(path, "must be an array of schemas");
121        }
122        if (alts.empty()) {
123            fail(path, "must not be empty");
124        }
125        auto node = std::make_unique<T>();
126        size_t i = 0;
127        for (const auto & alt : alts) {
128            node->children.push_back(build_node(alt, path + "/" + std::to_string(i++)));
129        }
130        return node;
131    }
132
133    common_chat_schema_ptr build_object(const common_json & schema, const std::string & path) {
134        auto node = std::make_unique<common_chat_schema_object>();
135
136        std::unordered_set<std::string> required;
137        if (schema.contains("required") && schema.at("required").is_array()) {
138            for (const auto & name : schema.at("required")) {
139                if (name.is_string()) {
140                    required.insert(name.get<std::string>());
141                }
142            }
143        }
144
145        if (schema.contains("properties")) {
146            const common_json & properties = schema.at("properties");
147            if (!properties.is_object()) {
148                fail(path, "properties must be an object");
149            }
150            for (const auto & [name, prop] : properties.items()) {
151                node->properties.push_back({name, build_node(prop, path + "/properties/" + name), required.count(name) > 0});
152            }
153        }
154
155        if (schema.contains("additionalProperties")) {
156            const common_json & additional = schema.at("additionalProperties");
157            if (additional.is_boolean()) {
158                if (additional.get<bool>()) {
159                    node->additional_properties = std::make_unique<common_chat_schema_any>();
160                }
161            } else if (additional.is_object()) {
162                node->additional_properties = build_node(additional, path + "/additionalProperties");
163            } else {
164                fail(path, "additionalProperties must be a boolean or a schema");
165            }
166        } else if (!schema.contains("properties")) {
167            // {"type": "object"} on its own accepts any object
168            node->additional_properties = std::make_unique<common_chat_schema_any>();
169        }
170
171        return node;
172    }
173
174    common_chat_schema_ptr build_array(const common_json & schema, const std::string & path) {
175        auto node = std::make_unique<common_chat_schema_array>();
176        if (schema.contains("items") || schema.contains("prefixItems")) {
177            // "items" wins when both are present; as in the converter, a schema instead of an array is the item schema
178            const std::string key = schema.contains("items") ? "items" : "prefixItems";
179            const common_json & items = schema.at(key);
180            if (items.is_array()) {
181                auto tuple = std::make_unique<common_chat_schema_tuple>();
182                size_t i = 0;
183                for (const auto & item : items) {
184                    tuple->items.push_back(build_node(item, path + "/" + key + "/" + std::to_string(i++)));
185                }
186                return tuple;
187            }
188            node->items = build_node(items, path + "/" + key);
189        } else {
190            node->items = std::make_unique<common_chat_schema_any>();
191        }
192        node->min_items = get_count(schema, "minItems", path, 0);
193        node->max_items = get_count(schema, "maxItems", path, -1);
194        return node;
195    }
196
197    common_chat_schema_ptr build_string(const common_json & schema, const std::string & path) {
198        auto node = std::make_unique<common_chat_schema_string>();
199        if (schema.contains("pattern")) {
200            const common_json & pattern = schema.at("pattern");
201            if (!pattern.is_string()) {
202                fail(path, "pattern must be a string");
203            }
204            node->pattern = pattern.get<std::string>();
205        }
206        node->format     = get_format(schema, path);
207        node->min_length = get_count(schema, "minLength", path, 0);
208        node->max_length = get_count(schema, "maxLength", path, -1);
209        return node;
210    }
211
212    common_chat_schema_ptr build_integer(const common_json & schema, const std::string & path) {
213        auto node = std::make_unique<common_chat_schema_integer>();
214        if (schema.contains("minimum")) {
215            node->minimum = get_bound(schema, "minimum", path, /* round_up */ true);
216        } else if (schema.contains("exclusiveMinimum")) {
217            node->minimum = get_bound(schema, "exclusiveMinimum", path, /* round_up */ false) + 1;
218        }
219        if (schema.contains("maximum")) {
220            node->maximum = get_bound(schema, "maximum", path, /* round_up */ false);
221        } else if (schema.contains("exclusiveMaximum")) {
222            node->maximum = get_bound(schema, "exclusiveMaximum", path, /* round_up */ true) - 1;
223        }
224        return node;
225    }
226
227    common_chat_schema_ptr build_node(const common_json & schema, const std::string & path) {
228        if (!schema.is_object()) {
229            fail(path, "schema must be an object");
230        }
231        if (schema.contains("$ref")) {
232            return build_ref(schema.at("$ref"), path);
233        }
234        if (schema.contains("oneOf") || schema.contains("anyOf")) {
235            const std::string key = schema.contains("oneOf") ? "oneOf" : "anyOf";
236            return build_alternatives<common_chat_schema_any_of>(schema.at(key), path + "/" + key);
237        }
238
239        common_json type;
240        if (schema.contains("type")) {
241            type = schema.at("type");
242        }
243        if (type.is_array()) {
244            // {"type": ["a", "b"], ...} is {"anyOf": [{"type": "a", ...}, {"type": "b", ...}]}
245            if (type.empty()) {
246                fail(path, "type must not be empty");
247            }
248            auto node = std::make_unique<common_chat_schema_any_of>();
249            size_t i = 0;
250            for (const auto & t : type) {
251                common_json alt = schema;
252                alt["type"] = t;
253                node->children.push_back(build_node(alt, path + "/type/" + std::to_string(i++)));
254            }
255            return node;
256        }
257        if (schema.contains("const")) {
258            return std::make_unique<common_chat_schema_const>(schema.at("const"));
259        }
260        if (schema.contains("enum")) {
261            const common_json & values = schema.at("enum");
262            if (!values.is_array() || values.empty()) {
263                fail(path, "enum must be a non-empty array");
264            }
265            auto node = std::make_unique<common_chat_schema_enum>();
266            for (const auto & value : values) {
267                node->values.push_back(value);
268            }
269            return node;
270        }
271        if (!type.is_null() && !type.is_string()) {
272            fail(path, "type must be a string or an array of strings");
273        }
274
275        const std::string type_name = type.is_string() ? type.get<std::string>() : "";
276        const bool has_properties = schema.contains("properties") ||
277            (schema.contains("additionalProperties") && schema.at("additionalProperties") != true);
278
279        if (type_name.empty()) {
280            // without a type the structural keywords decide, in the same order as the converter
281            if (has_properties) {
282                return build_object(schema, path);
283            }
284            if (schema.contains("allOf")) {
285                return build_alternatives<common_chat_schema_all_of>(schema.at("allOf"), path + "/allOf");
286            }
287            if (schema.contains("items") || schema.contains("prefixItems")) {
288                return build_array(schema, path);
289            }
290            if (schema.contains("pattern") || schema.contains("minLength") || schema.contains("maxLength") || get_format(schema, path) != common_chat_schema::FORMAT_NONE) {
291                return build_string(schema, path);
292            }
293            return std::make_unique<common_chat_schema_any>();
294        }
295        if (type_name == "object") {
296            if (!has_properties && schema.contains("allOf")) {
297                return build_alternatives<common_chat_schema_all_of>(schema.at("allOf"), path + "/allOf");
298            }
299            return build_object(schema, path);
300        }
301        if (type_name == "string") {
302            if (schema.contains("allOf")) {
303                return build_alternatives<common_chat_schema_all_of>(schema.at("allOf"), path + "/allOf");
304            }
305            return build_string(schema, path);
306        }
307        if (type_name == "array") {
308            return build_array(schema, path);
309        }
310        if (type_name == "integer") {
311            return build_integer(schema, path);
312        }
313        if (type_name == "number") {
314            return std::make_unique<common_chat_schema_number>();
315        }
316        if (type_name == "boolean") {
317            return std::make_unique<common_chat_schema_boolean>();
318        }
319        if (type_name == "null") {
320            return std::make_unique<common_chat_schema_null>();
321        }
322        fail(path, "unrecognized type " + type_name);
323    }
324
325  public:
326    common_chat_schema_builder(const common_json & root, common_chat_schema_document & doc) : root_(root), doc_(doc) {}
327
328    common_chat_schema_ptr build() {
329        auto node = build_node(root_, "#");
330        for (auto & entry : refs_) {
331            doc_.refs[entry.first] = std::move(entry.second);
332        }
333        for (auto * ref : pending_) {
334            ref->target = doc_.refs.at(ref->ref).get();
335        }
336        return node;
337    }
338};
339
340common_chat_schema_document common_chat_schema_from_json(const common_json & schema) {
341    common_chat_schema_document doc;
342    doc.root = common_chat_schema_builder(schema, doc).build();
343    return doc;
344}
345
346static common_chat_schema::value_type json_type(const common_json & value) {
347    if (value.is_null()) {
348        return common_chat_schema::TYPE_NULL;
349    }
350    if (value.is_boolean()) {
351        return common_chat_schema::TYPE_BOOLEAN;
352    }
353    if (value.is_number_integer()) {
354        return common_chat_schema::TYPE_INTEGER;
355    }
356    if (value.is_number()) {
357        return common_chat_schema::TYPE_NUMBER;
358    }
359    if (value.is_string()) {
360        return common_chat_schema::TYPE_STRING;
361    }
362    if (value.is_array()) {
363        return common_chat_schema::TYPE_ARRAY;
364    }
365    return common_chat_schema::TYPE_OBJECT;
366}
367
368static common_chat_schema::type_set value_types_impl(const common_chat_schema & s, std::unordered_set<const common_chat_schema *> & visited) {
369    switch (s.kind()) {
370        case common_chat_schema::KIND_ANY:
371            return common_chat_schema::type_set::all();
372        case common_chat_schema::KIND_NULL:
373            return { common_chat_schema::TYPE_NULL };
374        case common_chat_schema::KIND_BOOLEAN:
375            return { common_chat_schema::TYPE_BOOLEAN };
376        case common_chat_schema::KIND_NUMBER:
377            return { common_chat_schema::TYPE_NUMBER, common_chat_schema::TYPE_INTEGER };
378        case common_chat_schema::KIND_INTEGER:
379            return { common_chat_schema::TYPE_INTEGER };
380        case common_chat_schema::KIND_STRING:
381            return { common_chat_schema::TYPE_STRING };
382        case common_chat_schema::KIND_ARRAY:
383        case common_chat_schema::KIND_TUPLE:
384            return { common_chat_schema::TYPE_ARRAY };
385        case common_chat_schema::KIND_OBJECT:
386            return { common_chat_schema::TYPE_OBJECT };
387        case common_chat_schema::KIND_CONST:
388            return { json_type(static_cast<const common_chat_schema_const &>(s).value) };
389        case common_chat_schema::KIND_ENUM: {
390            common_chat_schema::type_set types;
391            for (const auto & value : static_cast<const common_chat_schema_enum &>(s).values) {
392                types.add(json_type(value));
393            }
394            return types;
395        }
396        case common_chat_schema::KIND_REF: {
397            const auto * target = static_cast<const common_chat_schema_ref &>(s).target;
398            if (!target || !visited.insert(target).second) {
399                // a cycle contributes no type, to be safe
400                return {};
401            }
402            auto types = value_types_impl(*target, visited);
403            visited.erase(target);
404            return types;
405        }
406        case common_chat_schema::KIND_ANY_OF: {
407            common_chat_schema::type_set types;
408            for (const auto & child : static_cast<const common_chat_schema_any_of &>(s).children) {
409                types |= value_types_impl(*child, visited);
410            }
411            return types;
412        }
413        case common_chat_schema::KIND_ALL_OF: {
414            auto types = common_chat_schema::type_set::all();
415            for (const auto & child : static_cast<const common_chat_schema_all_of &>(s).children) {
416                types &= value_types_impl(*child, visited);
417            }
418            return types;
419        }
420    }
421    return {};
422}
423
424common_chat_schema::type_set common_chat_schema::value_types() const {
425    std::unordered_set<const common_chat_schema *> visited;
426    return value_types_impl(*this, visited);
427}
428
429static bool may_be_string_impl(const common_chat_schema & s, std::unordered_set<const common_chat_schema *> & visited) {
430    switch (s.kind()) {
431        case common_chat_schema::KIND_STRING:
432            return true;
433        case common_chat_schema::KIND_CONST:
434            return static_cast<const common_chat_schema_const &>(s).value.is_string();
435        case common_chat_schema::KIND_ENUM:
436            for (const auto & v : static_cast<const common_chat_schema_enum &>(s).values) {
437                if (v.is_string()) {
438                    return true;
439                }
440            }
441            return false;
442        case common_chat_schema::KIND_REF: {
443            // a cycle is taken as not a string, to be safe
444            const auto * target = static_cast<const common_chat_schema_ref &>(s).target;
445            if (!target || !visited.insert(target).second) {
446                return false;
447            }
448            bool result = may_be_string_impl(*target, visited);
449            visited.erase(target);
450            return result;
451        }
452        case common_chat_schema::KIND_ANY_OF:
453            for (const auto & child : static_cast<const common_chat_schema_any_of &>(s).children) {
454                if (may_be_string_impl(*child, visited)) {
455                    return true;
456                }
457            }
458            return false;
459        case common_chat_schema::KIND_ALL_OF: {
460            // every child must allow a string, an any child constrains nothing
461            bool any_string = false;
462            for (const auto & child : static_cast<const common_chat_schema_all_of &>(s).children) {
463                if (child->kind() == common_chat_schema::KIND_ANY) {
464                    continue;
465                }
466                if (!may_be_string_impl(*child, visited)) {
467                    return false;
468                }
469                any_string = true;
470            }
471            return any_string;
472        }
473        default:
474            return false;
475    }
476}
477
478bool common_chat_schema::may_be_string() const {
479    std::unordered_set<const common_chat_schema *> visited;
480    return may_be_string_impl(*this, visited);
481}
482
483const char * common_chat_schema::kind_name(node_kind kind) {
484    switch (kind) {
485        case KIND_ANY:     return "any";
486        case KIND_REF:     return "ref";
487        case KIND_ANY_OF:  return "anyOf";
488        case KIND_ALL_OF:  return "allOf";
489        case KIND_CONST:   return "const";
490        case KIND_ENUM:    return "enum";
491        case KIND_NULL:    return "null";
492        case KIND_BOOLEAN: return "boolean";
493        case KIND_NUMBER:  return "number";
494        case KIND_INTEGER: return "integer";
495        case KIND_STRING:  return "string";
496        case KIND_ARRAY:   return "array";
497        case KIND_TUPLE:   return "tuple";
498        case KIND_OBJECT:  return "object";
499    }
500    return "?";
501}
502
503const char * common_chat_schema::type_name(value_type type) {
504    switch (type) {
505        case TYPE_NULL:    return "null";
506        case TYPE_BOOLEAN: return "boolean";
507        case TYPE_NUMBER:  return "number";
508        case TYPE_INTEGER: return "integer";
509        case TYPE_STRING:  return "string";
510        case TYPE_ARRAY:   return "array";
511        case TYPE_OBJECT:  return "object";
512    }
513    return "?";
514}