support partial content streaming in Generic mode

Author: ochafik

common/chat-parser.cpp

@@ -242,16 +242,18 @@ common_json common_chat_msg_parser::consume_json() {
 }
 
 common_chat_msg_parser::consume_json_result common_chat_msg_parser::consume_json_with_dumped_args(
-    const std::vector<std::vector<std::string>> & args_paths
+    const std::vector<std::vector<std::string>> & args_paths,
+    const std::vector<std::vector<std::string>> & content_paths
 ) {
-    if (auto result = try_consume_json_with_dumped_args(args_paths)) {
+    if (auto result = try_consume_json_with_dumped_args(args_paths, content_paths)) {
         return *result;
     }
     throw common_chat_msg_partial_exception("JSON");
 }
 
 std::optional<common_chat_msg_parser::consume_json_result> common_chat_msg_parser::try_consume_json_with_dumped_args(
-    const std::vector<std::vector<std::string>> & args_paths
+    const std::vector<std::vector<std::string>> & args_paths,
+    const std::vector<std::vector<std::string>> & content_paths
 ) {
     auto partial = try_consume_json();
     if (!partial) {
@@ -260,6 +262,9 @@ std::optional<common_chat_msg_parser::consume_json_result> common_chat_msg_parse
     auto is_arguments_path = [&](const std::vector<std::string> & path) {
         return std::find(args_paths.begin(), args_paths.end(), path) != args_paths.end();
     };
+    auto is_content_path = [&](const std::vector<std::string> & path) {
+        return std::find(content_paths.begin(), content_paths.end(), path) != content_paths.end();
+    };
 
     if (partial->healing_marker.marker.empty()) {
         if (args_paths.empty()) {
@@ -298,6 +303,18 @@ std::optional<common_chat_msg_parser::consume_json_result> common_chat_msg_parse
             }
             return arguments;
         }
+        if (is_content_path(path)) {
+            if (!j.is_string()) {
+                throw std::runtime_error("Content path must be a string");
+            }
+            std::string str = j;
+            auto idx = str.find(partial->healing_marker.marker); // not using json_dump_marker as we're inside a string
+            if (idx != std::string::npos) {
+                str.resize(idx);
+                found_healing_marker = true;
+            }
+            return str;
+        }
         if (j.is_object()) {
             auto obj = json::object();
             for (const auto & p : j.items()) {
@@ -314,6 +331,12 @@ std::optional<common_chat_msg_parser::consume_json_result> common_chat_msg_parse
                     const std::string value_str = value;
                     if (value_str.find(healing_marker_) != std::string::npos) {
                         found_healing_marker = true;
+                        if (is_content_path(path)) {
+                            if (partial->healing_marker.marker == partial->healing_marker.json_dump_marker) {
+                                // The healing occurred inside the string: good. Otherwise we just ditch the entire key/value pair.
+                                obj[key] = remove_unsupported_healings_and_dump_args(value);
+                            }
+                        }
                         break;
                     }
                     obj[key] = value;

common/chat-parser.h

@@ -95,10 +95,22 @@ class common_chat_msg_parser {
         bool is_partial;
     };
 
+    /*
+        Consume (possibly partial) json and converts specific subtrees to (possibly truncated) JSON strings.
+
+        By default, object keys can't be truncated, nor can string values (their corresponding key is removed,
+        e.g. `{"foo": "bar", "baz": "b` -> `{"foo": "bar"}`
+
+        But one can allow subpaths to be kept truncated, and possibly json-dumped to truncated json strings
+        - with `content_paths={{"foo"}}` -> `{"foo": "b` -> {"foo": "b"}`
+        - with `args_paths={{"foo"}}` -> `{"foo": {"b` -> `{"foo": "{b"}`
+    */
     consume_json_result consume_json_with_dumped_args(
-        const std::vector<std::vector<std::string>> & args_paths = {}
+        const std::vector<std::vector<std::string>> & args_paths = {},
+        const std::vector<std::vector<std::string>> & content_paths = {}
     );
     std::optional<consume_json_result> try_consume_json_with_dumped_args(
-        const std::vector<std::vector<std::string>> & args_paths = {}
+        const std::vector<std::vector<std::string>> & args_paths = {},
+        const std::vector<std::vector<std::string>> & content_paths = {}
     );
 };

common/chat.cpp

@@ -810,11 +810,14 @@ static common_chat_params common_chat_params_init_generic(const common_chat_temp
     return data;
 }
 static void common_chat_parse_generic(common_chat_msg_parser & builder) {
+    static const std::vector<std::vector<std::string>> content_paths = {
+        {"response"},
+    };
     static const std::vector<std::vector<std::string>> args_paths = {
         {"tool_call", "arguments"},
         {"tool_calls", "arguments"},
     };
-    auto data = builder.consume_json_with_dumped_args(args_paths);
+    auto data = builder.consume_json_with_dumped_args(args_paths, content_paths);
     if (data.value.contains("tool_calls")) {
         if (!builder.add_tool_calls(data.value.at("tool_calls")) || data.is_partial) {
             throw common_chat_msg_partial_exception("incomplete tool calls");

tests/test-chat-parser.cpp

@@ -165,40 +165,46 @@ const std::vector<std::string> barely_healable_jsons = {
   "{\"name\":\"python",
 };
 
-static void test(const std::string & input, bool is_partial, const std::vector<std::vector<std::string>> & args_paths, const std::string & expected) {
+static void test(const std::string & input, bool is_partial, const std::vector<std::vector<std::string>> & args_paths, const std::vector<std::vector<std::string>> & content_paths, const std::string & expected) {
   common_chat_msg_parser builder(input, is_partial, {});
-  auto js = builder.try_consume_json_with_dumped_args(args_paths);
+  auto js = builder.try_consume_json_with_dumped_args(args_paths, content_paths);
   assert_equals(true, js.has_value());
   assert_equals(is_partial, js->is_partial);
   assert_equals(expected, args_paths.size() == 1 && args_paths[0].empty() ? js->value.get<std::string>() : js->value.dump());
 }
 static void test_with_args(const std::string & input, const std::string & expected, bool parse_as_partial = true, bool is_partial = true) {
   common_chat_msg_parser builder(input, parse_as_partial, {});
-  auto js = builder.try_consume_json_with_dumped_args({{"args"}});
+  auto js = builder.try_consume_json_with_dumped_args({{"args"}}, {});
   assert_equals(true, js.has_value());
   assert_equals(is_partial, js->is_partial);
   assert_equals(expected, js->value.dump());
 }
 
 static void test_json_with_dumped_args_no_args() {
   // Normal JSON, nothing to heal, nothing to dump
-  test("{\"name\": \"python\"}", false, {}, "{\"name\":\"python\"}");
+  test("{\"name\": \"python\"}", false, {}, {}, "{\"name\":\"python\"}");
   // Full json is args
-  test("{\"name\": \"python\"}", false, {{}}, "{\"name\":\"python\"}");
+  test("{\"name\": \"python\"}", false, {{}}, {}, "{\"name\":\"python\"}");
 
   // If the arguments are further down, don't heal partial content.
   for (const auto & src : barely_healable_jsons) {
-    test(src, true, {{"arguments"}}, "{}");
+    test(src, true, {{"arguments"}}, {}, "{}");
   }
   // But heal content that isn't partial.
-  test("{\"name\": \"python\"", true, {{"arguments"}}, "{\"name\":\"python\"}");
+  test("{\"name\": \"python\"", true, {{"arguments"}}, {}, "{\"name\":\"python\"}");
 }
 
 static void test_json_with_dumped_args() {
+
+  // Partial content.
+  test("{\"content\": \"t", true, {}, {{"content"}}, "{\"content\":\"t\"}");
+  test("{\"content\": \"", true, {}, {{"content"}}, "{\"content\":\"\"}");
+  test("{\"content\": ", true, {}, {{"content"}}, "{}");
+
   // If the entire JSON is the arguments, healing it them dumping it produces the same output as the input (just reformatted).
-  test("{\"name\": \"python", true, {{}}, "{\"name\":\"python");
+  test("{\"name\": \"python", true, {{}}, {}, "{\"name\":\"python");
   for (const auto & src : barely_healable_jsons) {
-    test(src, true, {{}}, src);
+    test(src, true, {{}}, {}, src);
   }
 
   // Full JSON w/ args