echodict/llama.cpp

1#pragma once2 3#include <nlohmann/json_fwd.hpp>4 5#include <memory>6#include <unordered_map>7#include <unordered_set>8#include <string>9#include <string_view>10#include <functional>11#include <vector>12#include <variant>13 14struct common_grammar_builder;15 16class common_peg_parser_builder;17 18using common_peg_parser_id = size_t;19constexpr common_peg_parser_id COMMON_PEG_INVALID_PARSER_ID = static_cast<common_peg_parser_id>(-1);20 21using common_peg_ast_id = size_t;22constexpr common_peg_ast_id COMMON_PEG_INVALID_AST_ID = static_cast<common_peg_ast_id>(-1);23 24// Lightweight wrapper around common_peg_parser_id for convenience25class common_peg_parser {26    common_peg_parser_id id_;27    common_peg_parser_builder & builder_;28 29  public:30    common_peg_parser(const common_peg_parser & other) : id_(other.id_), builder_(other.builder_) {}31    common_peg_parser(common_peg_parser_id id, common_peg_parser_builder & builder) : id_(id), builder_(builder) {}32 33    common_peg_parser & operator=(const common_peg_parser & other);34    common_peg_parser & operator+=(const common_peg_parser & other);35    common_peg_parser & operator|=(const common_peg_parser & other);36 37    operator common_peg_parser_id() const { return id_; }38    common_peg_parser_id id() const { return id_; }39 40    common_peg_parser_builder & builder() const { return builder_; }41 42    // Creates a sequence43    common_peg_parser operator+(const common_peg_parser & other) const;44 45    // Creates a sequence separated by spaces.46    common_peg_parser operator<<(const common_peg_parser & other) const;47 48    // Creates a choice49    common_peg_parser operator|(const common_peg_parser & other) const;50 51    common_peg_parser operator+(const char * str) const;52    common_peg_parser operator+(const std::string & str) const;53    common_peg_parser operator<<(const char * str) const;54    common_peg_parser operator<<(const std::string & str) const;55    common_peg_parser operator|(const char * str) const;56    common_peg_parser operator|(const std::string & str) const;57};58 59common_peg_parser operator+(const char * str, const common_peg_parser & p);60common_peg_parser operator+(const std::string & str, const common_peg_parser & p);61common_peg_parser operator<<(const char * str, const common_peg_parser & p);62common_peg_parser operator<<(const std::string & str, const common_peg_parser & p);63common_peg_parser operator|(const char * str, const common_peg_parser & p);64common_peg_parser operator|(const std::string & str, const common_peg_parser & p);65 66enum common_peg_parse_result_type {67    COMMON_PEG_PARSE_RESULT_FAIL            = 0,68    COMMON_PEG_PARSE_RESULT_SUCCESS         = 1,69    COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT = 2,70};71 72const char * common_peg_parse_result_type_name(common_peg_parse_result_type type);73 74struct common_peg_ast_node {75    common_peg_ast_id id;76    std::string rule;77    std::string tag;78    size_t start;79    size_t end;80    std::string_view text;81    std::vector<common_peg_ast_id> children;82 83    bool is_partial = false;84};85 86struct common_peg_parse_result;87 88using common_peg_ast_visitor = std::function<void(const common_peg_ast_node & node)>;89 90class common_peg_ast_arena {91    std::vector<common_peg_ast_node> nodes_;92  public:93    common_peg_ast_id add_node(94        const std::string & rule,95        const std::string & tag,96        size_t start,97        size_t end,98        std::string_view text,99        std::vector<common_peg_ast_id> children,100        bool is_partial = false101    ) {102        common_peg_ast_id id = nodes_.size();103        nodes_.push_back({id, rule, tag, start, end, text, std::move(children), is_partial});104        return id;105    }106 107    const common_peg_ast_node & get(common_peg_ast_id id) const { return nodes_.at(id); }108 109    common_peg_ast_id find_by_tag(const common_peg_ast_node & parent, const std::string & tag, int max_depth = 3) const;110    common_peg_ast_id find_by_rule(const common_peg_ast_node & parent, const std::string & tag, int max_depth = 3) const;111 112    size_t size() const { return nodes_.size(); }113 114    void clear() { nodes_.clear(); }115 116    void visit(common_peg_ast_id id, const common_peg_ast_visitor & visitor) const;117    void visit(const common_peg_parse_result & result, const common_peg_ast_visitor & visitor) const;118 119    std::string dump();120};121 122struct common_peg_parse_result {123    common_peg_parse_result_type type = COMMON_PEG_PARSE_RESULT_FAIL;124    size_t start = 0;125    size_t end = 0;126 127    std::vector<common_peg_ast_id> nodes;128 129    common_peg_parse_result() = default;130 131    common_peg_parse_result(common_peg_parse_result_type type, size_t start)132        : type(type), start(start), end(start) {}133 134    common_peg_parse_result(common_peg_parse_result_type type, size_t start, size_t end)135        : type(type), start(start), end(end) {}136 137    common_peg_parse_result(common_peg_parse_result_type type, size_t start, size_t end, std::vector<common_peg_ast_id> nodes)138        : type(type), start(start), end(end), nodes(std::move(nodes)) {}139 140    bool fail() const { return type == COMMON_PEG_PARSE_RESULT_FAIL; }141    bool need_more_input() const { return type == COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT; }142    bool success() const { return type == COMMON_PEG_PARSE_RESULT_SUCCESS; }143};144 145enum common_peg_parse_flags {146    COMMON_PEG_PARSE_FLAG_NONE    = 0,147    COMMON_PEG_PARSE_FLAG_LENIENT = 1 << 0,148    COMMON_PEG_PARSE_FLAG_DEBUG   = 1 << 1,149};150 151inline common_peg_parse_flags operator|(common_peg_parse_flags a, common_peg_parse_flags b) {152    return static_cast<common_peg_parse_flags>(int(a) | int(b));153}154 155inline common_peg_parse_flags & operator|=(common_peg_parse_flags & a, common_peg_parse_flags b) {156    return a = a | b;157}158 159inline common_peg_parse_flags operator&(common_peg_parse_flags a, common_peg_parse_flags b) {160    return static_cast<common_peg_parse_flags>(int(a) & int(b));161}162 163inline common_peg_parse_flags operator~(common_peg_parse_flags a) {164    return static_cast<common_peg_parse_flags>(~int(a));165}166 167struct common_peg_parse_context {168    std::string input;169    common_peg_parse_flags flags;170    common_peg_ast_arena ast;171 172    int parse_depth;173 174    common_peg_parse_context(common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE)175        : flags(flags), parse_depth(0) {}176 177    common_peg_parse_context(const std::string & input, common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE)178        : input(input), flags(flags), parse_depth(0) {}179 180    bool is_lenient() const { return flags & COMMON_PEG_PARSE_FLAG_LENIENT; }181    bool is_debug() const { return flags & COMMON_PEG_PARSE_FLAG_DEBUG; }182};183 184class common_peg_arena;185 186// Parser variants187struct common_peg_epsilon_parser {};188 189struct common_peg_start_parser {};190 191struct common_peg_end_parser {};192 193struct common_peg_literal_parser {194    std::string literal;195};196 197struct common_peg_sequence_parser {198    std::vector<common_peg_parser_id> children;199};200 201struct common_peg_choice_parser {202    std::vector<common_peg_parser_id> children;203};204 205struct common_peg_repetition_parser {206    common_peg_parser_id child;207    int min_count;208    int max_count;  // -1 for unbounded209};210 211struct common_peg_and_parser {212    common_peg_parser_id child;213};214 215struct common_peg_not_parser {216    common_peg_parser_id child;217};218 219struct common_peg_any_parser {};220 221struct common_peg_space_parser {};222 223struct common_peg_chars_parser {224    struct char_range {225        uint32_t start;226        uint32_t end;227        bool contains(uint32_t codepoint) const { return codepoint >= start && codepoint <= end; }228    };229 230    std::string pattern;231    std::vector<char_range> ranges;232    bool negated;233    int min_count;234    int max_count;  // -1 for unbounded235};236 237struct common_peg_string_parser {238    char delimiter;239};240 241struct common_peg_until_parser {242    std::vector<std::string> delimiters;243};244 245struct common_peg_schema_parser {246    common_peg_parser_id child;247    std::string name;248    std::shared_ptr<nlohmann::ordered_json> schema;249 250    // Indicates if the GBNF should accept a raw string that matches the schema.251    bool raw;252};253 254struct common_peg_rule_parser {255    std::string name;256    common_peg_parser_id child;257    bool trigger;258};259 260struct common_peg_ref_parser {261    std::string name;262};263 264struct common_peg_atomic_parser {265    common_peg_parser_id child;266};267 268struct common_peg_tag_parser {269    common_peg_parser_id child;270    std::string tag;271};272 273struct common_peg_gbnf_parser {274    common_peg_parser_id child;275    std::string grammar;276};277 278// Variant holding all parser types279using common_peg_parser_variant = std::variant<280    common_peg_epsilon_parser,281    common_peg_start_parser,282    common_peg_end_parser,283    common_peg_literal_parser,284    common_peg_sequence_parser,285    common_peg_choice_parser,286    common_peg_repetition_parser,287    common_peg_and_parser,288    common_peg_not_parser,289    common_peg_any_parser,290    common_peg_space_parser,291    common_peg_chars_parser,292    common_peg_string_parser,293    common_peg_until_parser,294    common_peg_schema_parser,295    common_peg_rule_parser,296    common_peg_ref_parser,297    common_peg_atomic_parser,298    common_peg_tag_parser,299    common_peg_gbnf_parser300>;301 302class common_peg_arena {303    std::vector<common_peg_parser_variant> parsers_;304    std::unordered_map<std::string, common_peg_parser_id> rules_;305    common_peg_parser_id root_ = COMMON_PEG_INVALID_PARSER_ID;306 307  public:308    const common_peg_parser_variant & get(common_peg_parser_id id) const { return parsers_.at(id); }309    common_peg_parser_variant & get(common_peg_parser_id id) { return parsers_.at(id); }310 311    size_t size() const { return parsers_.size(); }312    bool empty() const { return parsers_.empty(); }313 314    common_peg_parser_id get_rule(const std::string & name) const;315    bool has_rule(const std::string & name) const { return rules_.find(name) != rules_.end(); }316 317    common_peg_parser_id root() const { return root_; }318    void set_root(common_peg_parser_id id) { root_ = id; }319 320    common_peg_parse_result parse(common_peg_parse_context & ctx, size_t start = 0) const;321    common_peg_parse_result parse(common_peg_parser_id id, common_peg_parse_context & ctx, size_t start) const;322 323    void resolve_refs();324 325    void build_grammar(const common_grammar_builder & builder, bool lazy = false) const;326 327    std::string dump(common_peg_parser_id id) const;328 329    nlohmann::json to_json() const;330    static common_peg_arena from_json(const nlohmann::json & j);331 332    std::string save() const;333    void load(const std::string & data);334 335    friend class common_peg_parser_builder;336 337  private:338    std::string dump_impl(common_peg_parser_id id, std::unordered_set<common_peg_parser_id> & visited) const;339 340    common_peg_parser_id add_parser(common_peg_parser_variant parser);341    void add_rule(const std::string & name, common_peg_parser_id id);342 343    common_peg_parser_id resolve_ref(common_peg_parser_id id);344};345 346class common_peg_parser_builder {347    common_peg_arena arena_;348 349    common_peg_parser wrap(common_peg_parser_id id) { return common_peg_parser(id, *this); }350    common_peg_parser add(const common_peg_parser_variant & p) { return wrap(arena_.add_parser(p)); }351 352  public:353    common_peg_parser_builder();354 355    // Match nothing, always succeed.356    //   S -> ε357    common_peg_parser eps() { return add(common_peg_epsilon_parser{}); }358 359    // Matches the start of the input.360    //   S -> ^361    common_peg_parser start() { return add(common_peg_start_parser{}); }362 363    // Matches the end of the input.364    //   S -> $365    common_peg_parser end() { return add(common_peg_end_parser{}); }366 367    // Matches an exact literal string.368    //   S -> "hello"369    common_peg_parser literal(const std::string & literal) { return add(common_peg_literal_parser{literal}); }370 371    // Matches a sequence of parsers in order, all must succeed.372    //   S -> A B C373    common_peg_parser sequence() { return add(common_peg_sequence_parser{}); }374    common_peg_parser sequence(const std::vector<common_peg_parser_id> & parsers);375    common_peg_parser sequence(const std::vector<common_peg_parser> & parsers);376    common_peg_parser sequence(std::initializer_list<common_peg_parser> parsers);377 378    // Matches the first parser that succeeds from a list of alternatives.379    //   S -> A | B | C380    common_peg_parser choice() { return add(common_peg_choice_parser{}); }381    common_peg_parser choice(const std::vector<common_peg_parser_id> & parsers);382    common_peg_parser choice(const std::vector<common_peg_parser> & parsers);383    common_peg_parser choice(std::initializer_list<common_peg_parser> parsers);384 385    // Matches one or more repetitions of a parser.386    //   S -> A+387    common_peg_parser one_or_more(const common_peg_parser & p) { return repeat(p, 1, -1); }388 389    // Matches zero or more repetitions of a parser, always succeeds.390    //   S -> A*391    common_peg_parser zero_or_more(const common_peg_parser & p) { return repeat(p, 0, -1); }392 393    // Matches zero or one occurrence of a parser, always succeeds.394    //   S -> A?395    common_peg_parser optional(const common_peg_parser & p) { return repeat(p, 0, 1); }396 397    // Positive lookahead: succeeds if child parser succeeds, consumes no input.398    //   S -> &A399    common_peg_parser peek(const common_peg_parser & p) { return add(common_peg_and_parser{p}); }400 401    // Negative lookahead: succeeds if child parser fails, consumes no input.402    //   S -> !A403    common_peg_parser negate(const common_peg_parser & p) { return add(common_peg_not_parser{p}); }404 405    // Matches any single character.406    //   S -> .407    common_peg_parser any() { return add(common_peg_any_parser{}); }408 409    // Matches between min and max repetitions of characters from a character class.410    //   S -> [a-z]{m,n}411    //412    // Use -1 for max to represent unbounded repetition (equivalent to {m,})413    common_peg_parser chars(const std::string & classes, int min = 1, int max = -1);414 415    // Creates a lightweight reference to a named rule (resolved during build()).416    // Use this for forward references in recursive grammars.417    //   expr_ref -> expr418    common_peg_parser ref(const std::string & name) { return add(common_peg_ref_parser{name}); }419 420    // Matches zero or more whitespace characters (space, tab, newline).421    //   S -> [ \t\n]*422    common_peg_parser space() { return add(common_peg_space_parser{}); }423 424    // Matches all characters until a delimiter is found (delimiter not consumed).425    //   S -> (!delim .)*426    common_peg_parser until(const std::string & delimiter) { return add(common_peg_until_parser{{delimiter}}); }427 428    // Matches all characters until one of the delimiters in the list is found (delimiter not consumed).429    //   S -> (!delim .)*430    common_peg_parser until_one_of(const std::vector<std::string> & delimiters) { return add(common_peg_until_parser{delimiters}); }431 432    // Matches everything433    //   S -> .*434    common_peg_parser rest() { return until_one_of({}); }435 436    // Matches between min and max repetitions of a parser (inclusive).437    //   S -> A{m,n}438    // Use -1 for max to represent unbounded repetition (equivalent to {m,})439    common_peg_parser repeat(const common_peg_parser & p, int min, int max) { return add(common_peg_repetition_parser{p, min,max}); }440 441    // Matches exactly n repetitions of a parser.442    //   S -> A{n}443    common_peg_parser repeat(const common_peg_parser & p, int n) { return repeat(p, n, n); }444 445    // Matches a double-quoted string: '"' content '"' space446    common_peg_parser double_quoted_string();447 448    // Matches a single-quoted string: "'" content "'" space449    common_peg_parser single_quoted_string();450 451    // Matches a string that accepts both double-quoted and single-quoted styles.452    common_peg_parser quoted_string();453 454    // Matches string content without the surrounding delimiter.455    common_peg_parser string_content(char delimiter);456 457    // Creates a complete JSON parser supporting objects, arrays, strings, numbers, booleans, and null.458    //   value -> object | array | string | number | true | false | null459    common_peg_parser json();460    common_peg_parser json_object();461    common_peg_parser json_string();462    common_peg_parser json_array();463    common_peg_parser json_number();464    common_peg_parser json_bool();465    common_peg_parser json_null();466 467    // Matches a JSON object member with a key and associated parser as the468    // value.469    common_peg_parser json_member(const std::string & key, const common_peg_parser & p);470 471    // Creates a complete Python format parser supporting dicts, arrays, strings, numbers, booleans, and None.472    // Differs from JSON: uses True/False/None, accepts both single and double-quoted strings.473    //   value -> dict | array | string | number | True | False | None474    common_peg_parser python_value();475    common_peg_parser python_dict();476    common_peg_parser python_string();477    common_peg_parser python_array();478    common_peg_parser python_number();479    common_peg_parser python_bool();480    common_peg_parser python_null();481 482    // A marker, i.e. text delimited by a pair of <> or []483    common_peg_parser marker();484 485    // Wraps a parser with JSON schema metadata for grammar generation.486    // Used internally to convert JSON schemas to GBNF grammar rules.487    common_peg_parser schema(const common_peg_parser & p, const std::string & name, const nlohmann::ordered_json & schema, bool raw = false);488 489    // Creates a named rule, stores it in the grammar, and returns a ref.490    // If trigger=true, marks this rule as an entry point for lazy grammar generation.491    //   auto json = p.rule("json", json_obj | json_arr | ...)492    common_peg_parser rule(const std::string & name, const common_peg_parser & p, bool trigger = false);493 494    // Creates a named rule using a builder function, and returns a ref.495    // If trigger=true, marks this rule as an entry point for lazy grammar generation.496    //   auto json = p.rule("json", [&]() { return json_object() | json_array() | ... })497    common_peg_parser rule(const std::string & name, const std::function<common_peg_parser()> & builder, bool trigger = false);498 499    // Creates a trigger rule. When generating a lazy grammar from the parser,500    // only trigger rules and descendents are emitted.501    common_peg_parser trigger_rule(const std::string & name, const common_peg_parser & p) { return rule(name, p, true); }502    common_peg_parser trigger_rule(const std::string & name, const std::function<common_peg_parser()> & builder) { return rule(name, builder, true); }503 504    // Creates an atomic parser. Atomic parsers do not create an AST node if505    // the child results in a partial parse, i.e. NEEDS_MORE_INPUT. This is506    // intended for situations where partial output is undesirable.507    common_peg_parser atomic(const common_peg_parser & p) { return add(common_peg_atomic_parser{p}); }508 509    // Tags create nodes in the generated AST for semantic purposes.510    // Unlike rules, you can tag multiple nodes with the same tag.511    common_peg_parser tag(const std::string & tag, const common_peg_parser & p) { return add(common_peg_tag_parser{p.id(), tag}); }512 513    // Wraps a child parser but emits a custom GBNF grammar string instead of514    // the child's grammar. Parsing delegates entirely to the child.515    common_peg_parser gbnf(const common_peg_parser & p, const std::string & grammar) { return add(common_peg_gbnf_parser{p, grammar}); }516 517    void set_root(const common_peg_parser & p);518 519    common_peg_arena build();520};521 522// Helper function for building parsers523common_peg_arena build_peg_parser(const std::function<common_peg_parser(common_peg_parser_builder & builder)> & fn);524