refactor + test chat parser (try_consume_json_with_dumped_args, literal based thinking tags parsing)

Author: ochafik

common/chat-parser.cpp

@@ -37,42 +37,30 @@ void common_chat_msg_parser::add_reasoning_content(const std::string &reasoning_
     result_.reasoning_content += reasoning_content;
 }
 
-bool common_chat_msg_parser::add_tool_call(const std::string & name, const std::string & id, const std::string & arguments, const common_healing_marker & healing_marker) {
+bool common_chat_msg_parser::add_tool_call(const std::string & name, const std::string & id, const std::string & arguments) {
     if (name.empty()) {
         return false;
     }
 
-    auto marker_idx = std::string::npos;
-    if (!arguments.empty() && !healing_marker.marker.empty()) {
-        marker_idx = arguments.find(healing_marker.json_dump_marker);
-        if (marker_idx == std::string::npos) {
-            marker_idx = arguments.find(healing_marker.marker);
-        }
-    }
-
     common_chat_tool_call tool_call;
     tool_call.name = name;
-    tool_call.arguments = marker_idx != std::string::npos ? arguments.substr(0, marker_idx) : arguments;
+    tool_call.arguments = arguments;
     tool_call.id = id;
 
-    if (tool_call.arguments == "\"") {
-        // This happens because of completing `:"$magic` after `"arguments"`
-        tool_call.arguments = "";
-    }
     LOG_DBG("Tool call arguments:\n\traw: %s\n\tresult: %s\n", arguments.c_str(), tool_call.arguments.c_str());
     result_.tool_calls.emplace_back(tool_call);
     return true;
 }
-bool common_chat_msg_parser::add_tool_call(const json & tool_call, const common_healing_marker & healing_marker) {
+bool common_chat_msg_parser::add_tool_call(const json & tool_call) {
     std::string name = tool_call.contains("name") ? tool_call.at("name") : "";
     std::string id = tool_call.contains("id") ? tool_call.at("id") : "";
-    std::string arguments = tool_call.contains("arguments") ? tool_call.at("arguments").dump() : "";
-    return add_tool_call(name, id, arguments, healing_marker);
+    std::string arguments = tool_call.contains("arguments") ? tool_call.at("arguments") : "";
+    return add_tool_call(name, id, arguments);
 }
 
-bool common_chat_msg_parser::add_tool_calls(const json & arr, const common_healing_marker & healing_marker) {
+bool common_chat_msg_parser::add_tool_calls(const json & arr) {
     for (const auto & item : arr) {
-        if (!add_tool_call(item, healing_marker)) {
+        if (!add_tool_call(item)) {
             return false;
         }
     }
@@ -121,30 +109,71 @@ bool common_chat_msg_parser::try_consume_literal(const std::string & literal) {
     return true;
 }
 
+std::optional<common_chat_msg_parser::find_regex_result>  common_chat_msg_parser::try_find_literal(const std::string & literal) {
+    auto idx = input_.find(literal, pos_);
+    if (idx != std::string::npos) {
+        find_regex_result res;
+        res.prelude = input_.substr(pos_, idx - pos_);
+        auto end = idx + literal.size();
+        res.groups.emplace_back(common_string_range{idx, end});
+        move_to(end);
+        return res;
+    }
+    if (is_partial_) {
+        idx = string_find_partial_stop(input_, literal);
+        if (idx != std::string::npos && idx >= pos_) {
+            find_regex_result res;
+            res.prelude = input_.substr(pos_, idx - pos_);
+            auto end = input_.size();
+            res.groups.emplace_back(common_string_range{idx, end});
+            move_to(end);
+            return res;
+        }
+    }
+    return std::nullopt;
+}
+
 void common_chat_msg_parser::consume_literal(const std::string & literal) {
     if (!try_consume_literal(literal)) {
         incomplete("Expected literal '" + literal + "' at position " + std::to_string(pos_));
     }
 }
 
-void common_chat_msg_parser::try_consume_think_tags(const common_regex & start_think_regex, const common_regex & end_think_regex) {
+bool common_chat_msg_parser::try_parse_reasoning(const std::string & start_think, const std::string & end_think) {
+    auto handle_reasoning = [&](const std::string & reasoning, bool closed) {
+        if (syntax_.reasoning_in_content) {
+            add_content(syntax_.reasoning_format == COMMON_REASONING_FORMAT_DEEPSEEK ? "<think>" : start_think);
+            add_content(reasoning);
+            if (closed) {
+                add_content(syntax_.reasoning_format == COMMON_REASONING_FORMAT_DEEPSEEK ? "</think>" : end_think);
+            }
+        } else {
+            add_reasoning_content(reasoning);
+        }
+    };
     if (syntax_.reasoning_format != COMMON_REASONING_FORMAT_NONE) {
-        if (syntax_.thinking_forced_open || try_consume_regex(start_think_regex)) {
-            if (auto res = try_find_regex(end_think_regex)) {
-                result_.reasoning_content = res->prelude;
+        if (syntax_.thinking_forced_open || try_consume_literal(start_think)) {
+            if (auto res = try_find_literal(end_think)) {
+                handle_reasoning(res->prelude, /* closed */ true);
                 consume_spaces();
-            } else {
-                result_.reasoning_content = consume_rest();
-                if (!syntax_.thinking_forced_open) {
-                    incomplete("Failed to find end of reasoning tag " + end_think_regex.str());
-                }
-                return;
+                return true;
+            }
+            auto rest = consume_rest();
+            if (!rest.empty()) {
+                handle_reasoning(consume_rest(), /* closed */ !is_partial());
             }
-        } else if (auto res = try_find_regex(end_think_regex)) {
-            result_.reasoning_content = res->prelude;
+            if (!syntax_.thinking_forced_open) {
+                incomplete("Failed to find end of reasoning tag " + end_think);
+            }
+            return true;
+        }
+        if (auto res = try_find_literal(end_think)) {
+            handle_reasoning(res->prelude, /* closed */ true);
             consume_spaces();
+            return true;
         }
     }
+    return false;
 }
 
 std::string common_chat_msg_parser::consume_rest() {
@@ -195,19 +224,7 @@ std::optional<common_chat_msg_parser::consume_regex_result> common_chat_msg_pars
     return consume_regex_result{m.groups};
 }
 
-// Calls the callback, *then* explodes w/ a partial match exception if it's partial
-common_json common_chat_msg_parser::consume_json(
-    const std::vector<std::vector<std::string>> & args_paths
-) {
-    if (auto result = try_consume_json(args_paths)) {
-        return *result;
-    }
-    incomplete("Failed to consume JSON");
-}
-
-std::optional<common_json> common_chat_msg_parser::try_consume_json(
-    const std::vector<std::vector<std::string>> & args_paths
-) {
+std::optional<common_json> common_chat_msg_parser::try_consume_json() {
     auto it = input_.cbegin() + pos_;
     const auto end = input_.cend();
     common_json result;
@@ -222,45 +239,84 @@ std::optional<common_json> common_chat_msg_parser::try_consume_json(
     if (!is_partial()) {
         incomplete("JSON is incomplete");
     }
+    return result;
+}
 
-    LOG_DBG("Parsed partial JSON: %s (json_healing_marker: %s)\n", result.json.dump().c_str(), result.healing_marker.json_dump_marker.c_str());
+common_json common_chat_msg_parser::consume_json() {
+    if (auto result = try_consume_json()) {
+        return *result;
+    }
+    incomplete("Failed to consume JSON");
+}
 
-    // Healing marker found, we need to visit the json and removed objects that we didn't want to heal
+nlohmann::ordered_json common_chat_msg_parser::consume_json_with_dumped_args(
+    const std::vector<std::vector<std::string>> & args_paths
+) {
+    if (auto result = try_consume_json_with_dumped_args(args_paths)) {
+        return *result;
+    }
+    incomplete("Failed to consume JSON");
+}
+
+std::optional<nlohmann::ordered_json> common_chat_msg_parser::try_consume_json_with_dumped_args(
+    const std::vector<std::vector<std::string>> & args_paths
+) {
+    auto partial = try_consume_json();
+    if (!partial) {
+        return std::nullopt;
+    }
     auto is_arguments_path = [&](const std::vector<std::string> & path) {
         return std::find(args_paths.begin(), args_paths.end(), path) != args_paths.end();
     };
 
+    if (partial->healing_marker.marker.empty()) {
+        if (args_paths.empty()) {
+            // No arguments to dump, and JSON was parsed fully.
+            return partial->json;
+        }
+        if (is_arguments_path({})) {
+            // Entire JSON is the arguments and was parsed fully.
+            return partial->json.dump();
+        }
+    }
+
+    LOG_DBG("Parsed partial JSON: %s (json_healing_marker: %s)\n", partial->json.dump().c_str(), partial->healing_marker.json_dump_marker.c_str());
+
     std::vector<std::string> path;
-    std::function<json(const json &)> remove_unsupported_healings = [&](const json & j) {
+    std::function<json(const json &)> remove_unsupported_healings_and_dump_args = [&](const json & j) -> json {
+        if (is_arguments_path(path)) {
+            auto arguments = j.dump();
+            if (is_partial() && !partial->healing_marker.marker.empty()) {
+                auto idx = arguments.find(partial->healing_marker.json_dump_marker);
+                if (idx != std::string::npos) {
+                    arguments.resize(idx);
+                }
+                if (arguments == "\"") {
+                    // This happens because of completing `:"$magic` after `"arguments"`
+                    arguments = "";
+                }
+            }
+            return arguments;
+        }
         if (j.is_object()) {
             auto obj = json::object();
             for (const auto & p : j.items()) {
                 const auto & key = p.key();
                 const auto & value = p.value();
                 const std::string key_str = key; // NOLINT
                 auto idx = key_str.find(healing_marker_);
-                if (idx != std::string::npos) {//} && idx != 0) {
-                    // Don't heal keys halfway, cut just after their opening quotes
-                    obj[result.healing_marker.marker] = 1;
-                    if (idx != 0) {
-                        result.healing_marker.json_dump_marker = result.healing_marker.marker;
-                    }
+                if (idx != std::string::npos) {
                     break;
                 }
                 path.push_back(key_str);
-                auto is_args = is_arguments_path(path);
-                if (is_args) {
-                    obj[key] = value;
-                } else if (value.is_string()) {
+                if (value.is_string()) {
                     const std::string value_str = value;
-                    if (value_str.find(healing_marker_) == std::string::npos) {
-                        obj[key] = value;
-                    } else {
-                        obj[result.healing_marker.marker] = 1;
-                        result.healing_marker.json_dump_marker = result.healing_marker.marker;
+                    if (value_str.find(healing_marker_) != std::string::npos) {
+                        break;
                     }
+                    obj[key] = value;
                 } else {
-                    obj[key] = remove_unsupported_healings(value);
+                    obj[key] = remove_unsupported_healings_and_dump_args(value);
                 }
                 path.pop_back();
             }
@@ -274,23 +330,19 @@ std::optional<common_json> common_chat_msg_parser::try_consume_json(
                     auto idx = str.find(healing_marker_);
                     if (idx != std::string::npos) {
                         // Don't heal array values that aren't in the arguments.
-                        arr.push_back(result.healing_marker.marker);
-                        result.healing_marker.json_dump_marker = result.healing_marker.marker;
+                        // arr.push_back(partial->healing_marker.marker);
+                        // partial->healing_marker.json_dump_marker = partial->healing_marker.marker;
                         break;
                     }
                 }
-                arr.push_back(remove_unsupported_healings(value));
+                arr.push_back(remove_unsupported_healings_and_dump_args(value));
             }
             return arr;
         }
         return j;
     };
 
-    if (!is_arguments_path({})) {
-        auto cleaned = remove_unsupported_healings(result.json);
-        LOG_DBG("Cleaned up JSON %s to %s (json_healing_marker : '%s')\n", result.json.dump().c_str(), cleaned.dump().c_str(), result.healing_marker.json_dump_marker.c_str());
-        result.json = cleaned;
-    }
-    LOG_DBG("Half-healed json: %s\n", result.json.dump().c_str());
-    return result;
+    auto cleaned = remove_unsupported_healings_and_dump_args(partial->json);
+    LOG_DBG("Cleaned up JSON %s to %s (json_healing_marker : '%s')\n", partial->json.dump().c_str(), cleaned.dump().c_str(), partial->healing_marker.json_dump_marker.c_str());
+    return cleaned;
 }

common/chat-parser.h

@@ -2,6 +2,7 @@
 
 #include "chat.h"
 #include "json-partial.h"
+#include "json.hpp"
 #include "regex-partial.h"
 
 #include <optional>
@@ -53,13 +54,13 @@ class common_chat_msg_parser {
     void add_reasoning_content(const std::string & reasoning_content);
 
     // Adds a tool call to the result. If the tool call is too incomplete (e.g. name empty), it won't add anything.
-    bool add_tool_call(const std::string & name, const std::string & id, const std::string & arguments, const common_healing_marker & healing_marker);
+    bool add_tool_call(const std::string & name, const std::string & id, const std::string & arguments);
 
     // Adds a tool call using the "name", "id" and "arguments" fields of the json object
-    bool add_tool_call(const nlohmann::ordered_json & tool_call, const common_healing_marker & healing_marker);
+    bool add_tool_call(const nlohmann::ordered_json & tool_call);
 
     // Adds an array of tool calls using their "name", "id" and "arguments" fields.
-    bool add_tool_calls(const nlohmann::ordered_json & arr, const common_healing_marker & healing_marker);
+    bool add_tool_calls(const nlohmann::ordered_json & arr);
 
     void finish();
 
@@ -68,11 +69,9 @@ class common_chat_msg_parser {
 
     bool consume_spaces();
 
-    bool try_consume_literal(const std::string & literal);
-
     void consume_literal(const std::string & literal);
 
-    void try_consume_think_tags(const common_regex & start_think_regex, const common_regex & end_think_regex);
+    bool try_parse_reasoning(const std::string & start_think, const std::string & end_think);
 
     std::string consume_rest();
 
@@ -83,18 +82,24 @@ class common_chat_msg_parser {
 
     std::optional<find_regex_result> try_find_regex(const common_regex & regex, size_t from = std::string::npos);
 
+    bool try_consume_literal(const std::string & literal);
+
+    std::optional<find_regex_result> try_find_literal(const std::string & literal);
+
     struct consume_regex_result {
         std::vector<common_string_range> groups;
     };
     consume_regex_result consume_regex(const common_regex & regex);
 
     std::optional<consume_regex_result> try_consume_regex(const common_regex & regex);
 
-    common_json consume_json(
+    std::optional<common_json> try_consume_json();
+    common_json consume_json();
+
+    nlohmann::ordered_json consume_json_with_dumped_args(
         const std::vector<std::vector<std::string>> & args_paths = {}
     );
-
-    std::optional<common_json> try_consume_json(
+    std::optional<nlohmann::ordered_json> try_consume_json_with_dumped_args(
         const std::vector<std::vector<std::string>> & args_paths = {}
     );
 };