parser updates

Author: aldehir

common/chat.cpp

@@ -1318,7 +1318,37 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp
     data.prompt = prompt;
     data.format = COMMON_CHAT_FORMAT_GPT_OSS;
 
-    // TODO: support tool calls in GPT-OSS?
+    if (inputs.tools.is_array() && !inputs.tools.empty()) {
+        data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;
+        data.grammar = build_grammar([&](const common_grammar_builder & builder) {
+            std::vector<std::string> tool_rules;
+            foreach_function(inputs.tools, [&](const json & tool) {
+                const auto & function = tool.at("function");
+                std::string name = function.at("name");
+                auto parameters = function.at("parameters");
+                builder.resolve_refs(parameters);
+
+                tool_rules.push_back(builder.add_rule(name + "-call",
+                    "\"" + name + "\"" + " space? \"<|constrain|>json<|message|>\" " + builder.add_schema(name + "-args", parameters)
+                ));
+            });
+
+            auto tool_call = builder.add_rule("tool_call", string_join(tool_rules, " | "));
+            builder.add_rule("root", "\"<|channel|>commentary\" space \"to=functions.\" " + tool_call);
+
+            data.grammar_triggers.push_back({
+                COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN,
+                "<\\|channel\\|>commentary\\s+to=functions\\."
+            });
+
+            data.preserved_tokens = {
+                "<|channel|>",
+                "<|constrain|>",
+                "<|message|>",
+                "<|start|>",
+            };
+        });
+    }
 
     return data;
 }

common/parsers/harmony.cpp

@@ -1,14 +1,25 @@
 #include "harmony.h"
 #include "regex-partial.h"
+#include "log.h"
 
 harmony_msg_parser::harmony_msg_parser(common_chat_msg_parser & builder)
     : builder(builder) {}
 
 void harmony_msg_parser::parse() {
-    // TODO @ngxson : this won't work with --special enabled, we should fix that
-    //builder.try_parse_reasoning("<|channel|>analysis<|message|>", "<|start|>assistant<|channel|>final<|message|>");
-    //builder.add_content(builder.consume_rest());
-    channel();
+    try {
+        channel();
+    } catch (const harmony_parse_error & e) {
+        LOG_ERR("Parse error: %s", e.what());
+    }
+
+    while (builder.try_find_literal("<|start|>")) {
+        try {
+            start();
+        } catch(const harmony_parse_error & e) {
+            LOG_ERR("Parse error: %s, skipping to next valid token", e.what());
+        }
+    }
+
     builder.add_content(builder.consume_rest());
 }
 
@@ -33,9 +44,6 @@ void harmony_msg_parser::analysis() {
         static const common_regex end("<\\|end\\|>");
         if (auto res = builder.try_find_regex(end, std::string::npos, false)) {
             builder.add_reasoning_content(res->prelude);
-            if (builder.try_consume_literal("<|start|>")) {
-                start();
-            }
         } else {
             builder.add_reasoning_content(builder.consume_rest());
         }
@@ -59,15 +67,21 @@ void harmony_msg_parser::commentary() {
     if (builder.try_consume_regex(to)) {
         user_function();
     }
+
+    if (builder.try_consume_literal("<|message|>")) {
+        static const common_regex end("<\\|end\\|>");
+        if (!builder.try_find_regex(end)) {
+            builder.add_content(builder.consume_rest());
+        }
+    }
 }
 
 void harmony_msg_parser::user_function() {
     static const common_regex tool_call_regex(
-        "functions\\.([a-zA-Z_][a-zA-Z0-9_]*)\\s*(<\\|constrain\\|>)?([a-z]+)<\\|message\\|>"
+        "functions\\.([a-zA-Z_][a-zA-Z0-9_]*)\\s?(?:<\\|constrain\\|>([a-zA-Z]+))?<\\|message\\|>"
     );
     if (auto res = builder.try_consume_regex(tool_call_regex)) {
         auto name = builder.str(res->groups[1]);
-        //auto constrain_type = builder.str(res->groups[3]);
         auto args = builder.consume_rest();
 
         builder.add_tool_call(name, "", args);

common/parsers/harmony.h

@@ -6,6 +6,11 @@
 #include <string>
 #include <vector>
 
+class harmony_parse_error: public std::runtime_error {
+  public:
+    harmony_parse_error(const std::string & message) : std::runtime_error(message) {}
+};
+
 class harmony_msg_parser {
     common_chat_msg_parser & builder;
 

tests/test-chat.cpp

@@ -1725,7 +1725,25 @@ static void test_template_output_parsers() {
                     /* .reasoning_format = */ COMMON_REASONING_FORMAT_AUTO,
                 }));
 
-
+        assert_msg_equals(simple_assist_msg("Hello, world!\nWhat's up?", "I'm\nthinking"),
+            common_chat_parse(
+                "<|channel|>analysis<|message|>I'm\nthinking<|end|>"
+                "<|start|>assistant<|channel|>commentary<|message|>Hello, world!\nWhat's up?",
+                /* is_partial= */ true,
+                {
+                    /* .format = */ COMMON_CHAT_FORMAT_GPT_OSS,
+                    /* .reasoning_format = */ COMMON_REASONING_FORMAT_AUTO,
+                }));
+        assert_msg_equals(simple_assist_msg("Hello, world!\nWhat's up?", "I'm\nthinking", "special_function", "{\"arg1\": 1}"),
+            common_chat_parse(
+                "<|channel|>analysis<|message|>I'm\nthinking<|end|>"
+                "<|start|>assistant<|channel|>commentary<|message|>Hello, world!\nWhat's up?<|end|>"
+                "<|start|>assistant<|channel|>commentary to=functions.special_function <|constrain|>json<|message|>{\"arg1\": 1}",
+                /* is_partial= */ true,
+                {
+                    /* .format = */ COMMON_CHAT_FORMAT_GPT_OSS,
+                    /* .reasoning_format = */ COMMON_REASONING_FORMAT_AUTO,
+                }));
     }
 }