echodict/llama.cpp

1#include "json-partial.h"2 3#include "log.h"4 5#include <nlohmann/json.hpp>6 7#include <string>8#include <regex>9 10using json = nlohmann::ordered_json;11 12enum common_json_stack_element_type {13    COMMON_JSON_STACK_ELEMENT_OBJECT,14    COMMON_JSON_STACK_ELEMENT_KEY,15    COMMON_JSON_STACK_ELEMENT_ARRAY,16};17 18struct common_json_stack_element {19    common_json_stack_element_type type;20    std::string key;21};22 23bool common_json_parse(24    const std::string & input,25    const std::string & healing_marker,26    common_json & out)27{28    std::string::const_iterator it = input.begin();29    const auto end = input.end();30    return common_json_parse(it, end, healing_marker, out);31}32 33bool common_json_parse(34    std::string::const_iterator & it,35    const std::string::const_iterator & end,36    const std::string & healing_marker,37    common_json & out)38{39    // // https://json.nlohmann.me/features/parsing/sax_interface/40    struct json_error_locator : public nlohmann::json_sax<json> {41        std::size_t position;42        bool found_error;43        std::string last_token;44        std::string exception_message;45        std::vector<common_json_stack_element> stack;46 47        json_error_locator() : position(0), found_error(false) {}48 49        bool parse_error(std::size_t position, const std::string & last_token, const json::exception & ex) override { // NOLINT50            this->position = position - 1;51            this->found_error = true;52            this->last_token = last_token;53            this->exception_message = ex.what();54            return false;55        }56        void close_value() {57            if (!stack.empty() && (stack.back().type == COMMON_JSON_STACK_ELEMENT_KEY)) {58                stack.pop_back();59            }60        }61        bool null() override { // NOLINT62            close_value();63            return true;64        }65        bool boolean(bool) override { // NOLINT66            close_value();67            return true;68        }69        bool number_integer(number_integer_t) override { // NOLINT70            close_value();71            return true;72        }73        bool number_unsigned(number_unsigned_t) override { // NOLINT74            close_value();75            return true;76        }77        bool number_float(number_float_t, const string_t &) override { // NOLINT78            close_value();79            return true;80        }81        bool string(string_t &) override { // NOLINT82            close_value();83            return true;84        }85        bool binary(binary_t &) override { // NOLINT86            close_value();87            return true;88        }89        bool start_object(std::size_t) override { // NOLINT90            stack.push_back({COMMON_JSON_STACK_ELEMENT_OBJECT, ""});91            return true;92        }93        bool end_object() override {94            GGML_ASSERT(!stack.empty() && stack.back().type == COMMON_JSON_STACK_ELEMENT_OBJECT);95            stack.pop_back();96            close_value();97            return true;98        }99        bool key(string_t & key) override { // NOLINT100            stack.push_back({COMMON_JSON_STACK_ELEMENT_KEY, key});101            return true;102        }103        bool start_array(std::size_t) override { // NOLINT104            stack.push_back({COMMON_JSON_STACK_ELEMENT_ARRAY, ""});105            return true;106        }107        bool end_array() override {108            GGML_ASSERT(!stack.empty() && stack.back().type == COMMON_JSON_STACK_ELEMENT_ARRAY);109            stack.pop_back();110            close_value();111            return true;112        }113    };114    json_error_locator err_loc;115    auto start = it;116    json::sax_parse(it, end, &err_loc);117 118    if (err_loc.found_error) {119        it = start;120        auto temptative_end = it + err_loc.position;121        // LOG_DBG("Error at position %zu (is_end = %s): %s\n", err_loc.position, temptative_end == end ? "true" : "false", err_loc.exception_message.c_str());122 123        auto input = std::string(it, temptative_end);124        try {125            out.json = json::parse(input);126            // out.json = json::parse(it, temptative_end);127            it = temptative_end;128            return true;129        } catch (const std::exception & ex) {130            // No, needs healing.131            LOG_DBG("Failed to parse up to error: %s: <<<%s>>>\n", ex.what(), std::string(it, temptative_end).c_str());132        }133        auto can_parse = [](const std::string & str) {134            try {135                auto _ = json::parse(str); // NOLINT136                return true;137            } catch (const std::exception &) {138                return false;139            }140        };141        if (!healing_marker.empty() && !err_loc.stack.empty()) {142            std::string str(it, temptative_end);143            auto last_non_sp_pos = str.find_last_not_of(" \n\r\t");144            if (last_non_sp_pos == std::string::npos) {145                throw std::runtime_error("Cannot heal a truncated JSON that stopped in an unknown location");146            }147            auto last_non_sp_char = str[last_non_sp_pos];148            // Used to detect stops on a number, which may not be complete.149            auto was_maybe_number = [&]() {150                if (!str.empty() && std::isspace(str.back())) {151                    return false;152                }153                return std::isdigit(last_non_sp_char) ||154                    last_non_sp_char == '.' ||155                    last_non_sp_char == 'e' ||156                    last_non_sp_char == 'E' ||157                    last_non_sp_char == '-';158            };159 160            std::string closing;161            for (size_t i = err_loc.stack.size(); i > 0; i--) {162                auto & el = err_loc.stack[i - 1];163                if (el.type == COMMON_JSON_STACK_ELEMENT_OBJECT) {164                    closing += "}";165                } else if (el.type == COMMON_JSON_STACK_ELEMENT_ARRAY) {166                    closing += "]";167                } else if (el.type != COMMON_JSON_STACK_ELEMENT_KEY) {168                    throw std::runtime_error("Unexpected stack element type");169                }170            }171 172            // Matches a potentially partial unicode escape sequence, e.g. \u, \uX, \uXX, \uXXX, \uXXXX173            static const std::regex partial_unicode_regex(R"(\\u(?:[0-9a-fA-F](?:[0-9a-fA-F](?:[0-9a-fA-F](?:[0-9a-fA-F])?)?)?)?$)");174 175            auto is_high_surrogate = [&](const std::string & s) {176                // Check if a partial of a high surrogate (U+D800-U+DBFF)177                return s.length() >= 4 &&178                    s[0] == '\\' && s[1] == 'u' &&179                    std::tolower(s[2]) == 'd' &&180                    (s[3] == '8' || s[3] == '9' || std::tolower(s[3]) == 'a' || std::tolower(s[3]) == 'b');181            };182 183            // Initialize the unicode marker to a low surrogate to handle the edge case184            // where a high surrogate (U+D800-U+DBFF) is immediately followed by a185            // backslash (\)186            std::string unicode_marker_padding = "udc00";187            std::smatch last_unicode_seq;188 189            if (std::regex_search(str, last_unicode_seq, partial_unicode_regex)) {190                std::smatch second_last_seq;191                std::string prelude = str.substr(0, last_unicode_seq.position());192 193                // Pad the escape sequence with 0s until it forms a complete sequence of 6 characters194                unicode_marker_padding = std::string(6 - last_unicode_seq.length(), '0');195 196                if (is_high_surrogate(last_unicode_seq.str())) {197                    // If the sequence is a partial match for a high surrogate, add a low surrogate (U+DC00-U+UDFF)198                    unicode_marker_padding += "\\udc00";199                } else if (std::regex_search(prelude, second_last_seq, partial_unicode_regex)) {200                    if (is_high_surrogate(second_last_seq.str())) {201                        // If this follows a high surrogate, pad it to be a low surrogate202                        if (last_unicode_seq.length() == 2) {203                            unicode_marker_padding = "dc00";204                        } else if (last_unicode_seq.length() == 3) {205                            unicode_marker_padding = "c00";206                        } else {207                            // The original unicode_marker_padding is already padded with 0s208                        }209                    }210                }211            }212 213            const auto & magic_seed = out.healing_marker.marker = healing_marker;//"$llama.cpp.json$";214 215            if (err_loc.stack.back().type == COMMON_JSON_STACK_ELEMENT_KEY) {216                // We're inside an object value217                if (last_non_sp_char == ':' && can_parse(str + "1" + closing)) {218                    // Was about to create an object value219                    str += (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing;220                } else if (can_parse(str + ": 1" + closing)) {221                    str += (out.healing_marker.json_dump_marker = ":\"" + magic_seed) + "\"" + closing;222                } else if (last_non_sp_char == '{' && can_parse(str + closing)) {223                    // Was about to create an object224                    str += (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\": 1" + closing;225                } else if (can_parse(str + "\"" + closing)) {226                    // Was inside an object value string227                    str += (out.healing_marker.json_dump_marker = magic_seed) + "\"" + closing;228                } else if (str[str.length() - 1] == '\\' && can_parse(str + "\\\"" + closing)) {229                    // Was inside an object value string after an escape230                    str += (out.healing_marker.json_dump_marker = "\\" + magic_seed) + "\"" + closing;231                } else if (can_parse(str + unicode_marker_padding + "\"" + closing)) {232                    // Was inside an object value string after a partial unicode escape233                    str += (out.healing_marker.json_dump_marker = unicode_marker_padding + magic_seed) + "\"" + closing;234                } else {235                    // find last :236                    auto last_pos = str.find_last_of(':');237                    if (last_pos == std::string::npos) {238                        throw std::runtime_error("Cannot heal a truncated JSON that stopped in an unknown location");239                    }240                    // Cutting back to opening : for object value241                    str = str.substr(0, last_pos + 1) + (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing;242                }243            } else if (err_loc.stack.back().type == COMMON_JSON_STACK_ELEMENT_ARRAY) {244                if ((last_non_sp_char == ',' || last_non_sp_char == '[') && can_parse(str + "1" + closing)) {245                    // Was about to create an array value246                    str += (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing;247                } else if (can_parse(str + "\"" + closing)) {248                    // Was inside an array value string249                    str += (out.healing_marker.json_dump_marker = magic_seed) + "\"" + closing;250                } else if (str[str.length() - 1] == '\\' && can_parse(str + "\\\"" + closing)) {251                    // Was inside an array value string after an escape252                    str += (out.healing_marker.json_dump_marker = "\\" + magic_seed) + "\"" + closing;253                } else if (can_parse(str + unicode_marker_padding + "\"" + closing)) {254                    // Was inside an array value string after a partial unicode escape255                    str += (out.healing_marker.json_dump_marker = unicode_marker_padding + magic_seed) + "\"" + closing;256                } else if (!was_maybe_number() && can_parse(str + ", 1" + closing)) {257                    // Had just finished a value258                    str += (out.healing_marker.json_dump_marker = ",\"" + magic_seed) + "\"" + closing;259                } else {260                    auto last_pos = str.find_last_of("[,");261                    if (last_pos == std::string::npos) {262                        throw std::runtime_error("Cannot heal a truncated JSON array stopped in an unknown location");263                    }264                    // Cutting back to last [ or , for array value265                    str = str.substr(0, last_pos + 1) + (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing;266                }267            } else if (err_loc.stack.back().type == COMMON_JSON_STACK_ELEMENT_OBJECT) {268                if ((last_non_sp_char == '{' && can_parse(str + closing)) ||269                        (last_non_sp_char == ',' && can_parse(str + "\"\": 1" + closing))) {270                    // Was about to create an object key+value271                    str += (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\": 1" + closing;272                } else if (!was_maybe_number() && can_parse(str + ",\"\": 1" + closing)) {273                    // Was about to create an object key+value274                    str += (out.healing_marker.json_dump_marker = ",\"" + magic_seed) + "\": 1" + closing;275                } else if (can_parse(str + "\": 1" + closing)) {276                    // Was inside an object key string277                    str += (out.healing_marker.json_dump_marker = magic_seed) + "\": 1" + closing;278                } else if (str[str.length() - 1] == '\\' && can_parse(str + "\\\": 1" + closing)) {279                    // Was inside an object key string after an escape280                    str += (out.healing_marker.json_dump_marker = "\\" + magic_seed) + "\": 1" + closing;281                } else if (can_parse(str + unicode_marker_padding + "\": 1" + closing)) {282                    // Was inside an object key string after a partial unicode escape283                    str += (out.healing_marker.json_dump_marker = unicode_marker_padding + magic_seed) + "\": 1" + closing;284                } else {285                    auto last_pos = str.find_last_of(':');286                    if (last_pos == std::string::npos) {287                        throw std::runtime_error("Cannot heal a truncated JSON object stopped in an unknown location");288                    }289                    // fprintf(stderr, "Cutting back to last : for object key+value\n");290                    str = str.substr(0, last_pos + 1) + (out.healing_marker.json_dump_marker = "\"" + magic_seed) + "\"" + closing;291                }292            } else {293                throw std::runtime_error("Cannot heal a truncated JSON object stopped in an unknown location");294            }295            // fprintf(stderr, "HEALED:\nSTRING <<<\n%s\n>>>\n\nmagic_cut: <<<\n%s\n>>>\n\n", str.c_str(), out.healing_marker.json_dump_marker.c_str());296            out.json = json::parse(str);297            it = temptative_end;298            return true;299        }300        // handle unclosed top-level primitive301        if (err_loc.position != 0 && !healing_marker.empty() && err_loc.stack.empty()) {302            std::string str(it, temptative_end);303            const auto & magic_seed = out.healing_marker.marker = healing_marker;304            if (can_parse(str + "\"")) {305                // Was inside an string306                str += (out.healing_marker.json_dump_marker = magic_seed) + "\"";307            } else if (str[str.length() - 1] == '\\' && can_parse(str + "\\\"")) {308                // Was inside an string after an escape309                str += (out.healing_marker.json_dump_marker = "\\" + magic_seed) + "\"";310            } else {311                // TODO: handle more unclosed top-level primitive if the stack was empty but we got an error (e.g. "tru", "\"", etc...)312                // fprintf(stderr, "Closing: TODO\n");313                return false;314            }315            out.json = json::parse(str);316            it = temptative_end;317            return true;318        }319        return false;320    }321    out.json = json::parse(it, end);322    it = end;323    return true;324}325