Brunobkr/llama.cpp_AlgMor24_github

1#include "chat-peg-parser.h"2#include "chat.h"3#include "common.h"4#include "json-schema-to-grammar.h"5#include "peg-parser.h"6#include "testing.h"7#include "peg-parser/simple-tokenize.h"8 9#include <iostream>10#include <numeric>11#include <regex>12#include <string>13 14#include "nlohmann/json.hpp"15 16using json = nlohmann::ordered_json;17 18static json create_tools();19static void test_example_native(testing & t);20static void test_example_qwen3_coder(testing & t);21static void test_example_qwen3_non_coder(testing & t);22static void test_command7_parser_compare(testing & t);23static void test_prefix_tool_names(testing & t);24static void test_tagged_peg_parser(testing & t);25static void test_permute(testing & t);26 27int main(int argc, char * argv[]) {28    testing t(std::cout);29    if (argc >= 2) {30        t.set_filter(argv[1]);31    }32 33    const char * verbose = getenv("LLAMA_TEST_VERBOSE");34    if (verbose) {35        t.verbose = std::string(verbose) == "1";36    }37 38    t.test("native", test_example_native);39    t.test("qwen3 coder", test_example_qwen3_coder);40    t.test("qwen3 non-coder", test_example_qwen3_non_coder);41    t.test("comparison", test_command7_parser_compare);42    t.test("prefix tool names", test_prefix_tool_names);43    t.test("tagged peg parser", test_tagged_peg_parser);44    t.test("permute", test_permute);45 46    return t.summary();47}48 49static json create_tools() {50    json tools = json::array();51 52    json tool_weather = {53        { "type",     "function" },54        { "function",55         {56              { "name", "get_current_weather" },57              { "description", "Get the current weather in a given location" },58              { "parameters",59                {60                    { "type", "object" },61                    { "properties",62                      { { "location",63                          { { "type", "string" }, { "description", "The city and state, e.g. San Francisco, CA" } } },64                        { "unit",65                          { { "type", "string" },66                            { "enum", { "celsius", "fahrenheit" } },67                            { "description",68                              "The temperature unit to use. Infer this from the users location." } } } } },69                    { "required", { "location", "unit" } },70                } },71          }                      }72    };73    tools.push_back(tool_weather);74 75    json tool_forecast = {76        { "type",     "function" },77        { "function",78         {79              { "name", "get_forecast" },80              { "description", "Get the weather forecast for a given location" },81              { "parameters",82                {83                    { "type", "object" },84                    { "properties",85                      { { "location",86                          { { "type", "string" }, { "description", "The city and state, e.g. San Francisco, CA" } } },87                        { "unit",88                          { { "type", "string" },89                            { "enum", { "celsius", "fahrenheit" } },90                            { "description", "The temperature unit to use. Infer this from the users location." } } },91                        { "days",92                          { { "type", "integer" },93                            { "description", "Number of days to forecast (1-10)" },94                            { "minimum", 1 },95                            { "maximum", 10 } } } } },96                    { "required", { "location", "unit" } },97                } },98          }                      }99    };100    tools.push_back(tool_forecast);101 102    json tool_search = {103        { "type",     "function" },104        { "function",105         { { "name", "search_knowledge_base" },106            { "description", "Search the internal technical documentation knowledge base." },107            { "parameters",108              { { "type", "object" },109                { "properties",110                  { { "query", { { "type", "string" }, { "description", "The search query string." } } },111                    { "max_results",112                      { { "type", "integer" },113                        { "description", "The maximum number of results to return." },114                        { "default", 5 } } },115                    { "category",116                      { { "type", "string" },117                        { "enum", { "api", "troubleshooting", "billing", "general" } },118                        { "description", "Filter search by specific category." } } } } },119                { "required", { "query", "category" } },120                { "additionalProperties", false } } },121            { "strict", true } } }122    };123    tools.push_back(tool_search);124 125    return tools;126}127 128struct tool_argument {129    std::string name;130    std::string type;131    bool        is_required;132    json        schema;133};134 135struct tool_definition {136    std::string                name;137    std::vector<tool_argument> arguments;138    json                       schema;139};140 141// Test fictitious model output that emits arguments as JSON.142static void test_example_native(testing & t) {143    struct test_case {144        // Parameters145        std::string             name;146        json                    tools;147        common_chat_tool_choice tool_choice;148        common_reasoning_format reasoning_format;149        json                    json_schema;150        bool                    parallel_tool_calls;151        std::string             generation_prompt;152        std::string             input;153 154        // Expect155        std::string                        expect_reasoning;156        std::string                        expect_content;157        std::vector<common_chat_tool_call> expect_tool_calls;158    };159 160    auto build_parser = [](const test_case & tc) {161        return build_chat_peg_parser([&](common_chat_peg_builder & p) {162            auto reasoning_in_content = (tc.reasoning_format == COMMON_REASONING_FORMAT_NONE);163            // Always use optional TAG_BASED pattern; generation_prompt is prepended to input164            auto reasoning = p.optional("<think>" + p.reasoning(p.until("</think>")) + "</think>" + p.space());165 166            // tool calling parser167            if (tc.tools.is_array() && !tc.tools.empty()) {168                auto tool_call =169                    p.standard_json_tools("<tool_call>[", "]</tool_call>", tc.tools, tc.parallel_tool_calls,170                                          tc.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED);171 172                return p.sequence({ (reasoning_in_content ? p.eps() : reasoning), p.content(p.until("<tool_call>")),173                                    p.optional(p.space() + tool_call), p.space(), p.end() });174            }175 176            // response_format parser177            if (tc.json_schema.is_object() && !tc.json_schema.empty()) {178                return p.sequence({ (reasoning_in_content ? p.eps() : reasoning),179                                    p.content(p.schema(p.json(), "response-output", tc.json_schema)), p.space(),180                                    p.end() });181            }182 183            // Content-only parser184            return p.sequence({ (reasoning_in_content ? p.eps() : reasoning), p.content(p.rest()), p.end() });185        });186    };187 188    std::vector<test_case> test_cases = std::vector<test_case>{189        {190         /* .name =                 */ "content with reasoning (no generation_prompt)",191         /* .tools =                */ {},192         /* .tool_choice =          */ COMMON_CHAT_TOOL_CHOICE_NONE,193         /* .reasoning_format =     */ COMMON_REASONING_FORMAT_AUTO,194         /* .json_schema =          */ {},195         /* .parallel_tool_calls =  */ false,196         /* .generation_prompt =    */ "",197         /* .input =                */ ("<think>The user said hello, I must say hello back</think>\nHello"),198         /* .expect_reasoning =     */ "The user said hello, I must say hello back",199         /* .expect_content =       */ "Hello",200         /* .expect_tool_calls =    */ {},201         },202        {203         /* .name =                 */ "content without reasoning (no generation_prompt)",204         /* .tools =                */ {},205         /* .tool_choice =          */ COMMON_CHAT_TOOL_CHOICE_NONE,206         /* .reasoning_format =     */ COMMON_REASONING_FORMAT_AUTO,207         /* .json_schema =          */ {},208         /* .parallel_tool_calls =  */ false,209         /* .generation_prompt =    */ "",210         /* .input =                */ ("Hello"),211         /* .expect_reasoning =     */ "",212         /* .expect_content =       */ "Hello",213         /* .expect_tool_calls =    */ {},214         },215        {216         /* .name =                 */ "content with reasoning_format = none (tags appear in content)",217         /* .tools =                */ {},218         /* .tool_choice =          */ COMMON_CHAT_TOOL_CHOICE_NONE,219         /* .reasoning_format =     */ COMMON_REASONING_FORMAT_NONE,220         /* .json_schema =          */ {},221         /* .parallel_tool_calls =  */ false,222         /* .generation_prompt =    */ "",223         /* .input =                */ ("<think>The user said hello, I must say hello back</think>\nHello"),224         /* .expect_reasoning =     */ "",225         /* .expect_content =       */ "<think>The user said hello, I must say hello back</think>\nHello",226         /* .expect_tool_calls =    */ {},227         },228        {229         /* .name =                 */ "content with reasoning generation_prompt",230         /* .tools =                */ {},231         /* .tool_choice =          */ COMMON_CHAT_TOOL_CHOICE_NONE,232         /* .reasoning_format =     */ COMMON_REASONING_FORMAT_AUTO,233         /* .json_schema =          */ {},234         /* .parallel_tool_calls =  */ false,235         /* .generation_prompt =    */ "<think>",236         /* .input =                */ ("The user said hello, I must say hello back</think>\nHello"),237         /* .expect_reasoning =     */ "The user said hello, I must say hello back",238         /* .expect_content =       */ "Hello",239         /* .expect_tool_calls =    */ {},240         },241        {242         /* .name =                 */ "content with reasoning generation_prompt and reasoning_format = none",243         /* .tools =                */ {},244         /* .tool_choice =          */ COMMON_CHAT_TOOL_CHOICE_NONE,245         /* .reasoning_format =     */ COMMON_REASONING_FORMAT_NONE,246         /* .json_schema =          */ {},247         /* .parallel_tool_calls =  */ false,248         /* .generation_prompt =    */ "",249         /* .input =                */ ("The user said hello, I must say hello back</think>\nHello"),250         /* .expect_reasoning =     */ "",251         /* .expect_content =       */ "The user said hello, I must say hello back</think>\nHello",252         /* .expect_tool_calls =    */ {},253         },254        {255         /* .name =                 */ "content with closed reasoning generation_prompt (empty reasoning discarded)",256         /* .tools =                */ {},257         /* .tool_choice =          */ COMMON_CHAT_TOOL_CHOICE_NONE,258         /* .reasoning_format =     */ COMMON_REASONING_FORMAT_AUTO,259         /* .json_schema =          */ {},260         /* .parallel_tool_calls =  */ false,261         /* .generation_prompt =    */ "<think></think>",262         /* .input =                */ ("Hello"),263         /* .expect_reasoning =     */ "",264         /* .expect_content =       */ "Hello",265         /* .expect_tool_calls =    */ {},266         },267        {268         /* .name =                 */ "tools with reasoning generation_prompt",269         /* .tools =                */ create_tools(),270         /* .tool_choice =          */ COMMON_CHAT_TOOL_CHOICE_AUTO,271         /* .reasoning_format =     */ COMMON_REASONING_FORMAT_AUTO,272         /* .json_schema =          */ {},273         /* .parallel_tool_calls =  */ false,274         /* .generation_prompt =    */ "<think>",275         /* .input =                */276            ("I must get the weather in New York</think>\n"277             "<tool_call>["278             R"({"name": "get_current_weather", "arguments": {"location": "New York City, NY", "unit": "fahrenheit"}})"279             "]</tool_call>"),280         /* .expect_reasoning =     */ "I must get the weather in New York",281         /* .expect_content =       */ "",282         /* .expect_tool_calls =    */283            { {284                /* .name =      */ "get_current_weather",285                /* .arguments = */ R"({"location": "New York City, NY", "unit": "fahrenheit"})",286                /* .id =        */ "",287            } },288         },289        {290         /* .name =                 */ "parallel tools with reasoning generation_prompt",291         /* .tools =                */ create_tools(),292         /* .tool_choice =          */ COMMON_CHAT_TOOL_CHOICE_AUTO,293         /* .reasoning_format =     */ COMMON_REASONING_FORMAT_AUTO,294         /* .json_schema =          */ {},295         /* .parallel_tool_calls =  */ true,296         /* .generation_prompt =    */ "<think>",297         /* .input =                */298            ("I must get the weather in New York and San Francisco and a 3 day forecast of each.</think>\nLet me "299             "search that for you."300             "<tool_call>["301             R"({"name": "get_current_weather", "arguments": {"location": "New York City, NY", "unit": "fahrenheit"}})"302             ", "303             R"({"name": "get_current_weather", "arguments": {"location": "San Francisco, CA", "unit": "fahrenheit"}})"304             ", "305             R"({"name": "get_forecast", "arguments": {"location": "New York City, NY", "unit": "fahrenheit", "days": 3}})"306             ", "307             R"({"name": "get_forecast", "arguments": {"location": "San Francisco, CA", "unit": "fahrenheit", "days": 3}})"308             "]</tool_call>"),309         /* .expect_reasoning =     */310            "I must get the weather in New York and San Francisco and a 3 day forecast of each.",                                                                     /* .expect_content =       */ "Let me search that for you.",311         /* .expect_tool_calls =    */312            { {313                  /* .name =      */ "get_current_weather",314                  /* .arguments = */ R"({"location": "New York City, NY", "unit": "fahrenheit"})",315                  /* .id =        */ "",316              },317              {318                  /* .name =      */ "get_current_weather",319                  /* .arguments = */ R"({"location": "San Francisco, CA", "unit": "fahrenheit"})",320                  /* .id =        */ "",321              },322              {323                  /* .name =      */ "get_forecast",324                  /* .arguments = */ R"({"location": "New York City, NY", "unit": "fahrenheit", "days": 3})",325                  /* .id =        */ "",326              },327              {328                  /* .name =      */ "get_forecast",329                  /* .arguments = */ R"({"location": "San Francisco, CA", "unit": "fahrenheit", "days": 3})",330                  /* .id =        */ "",331              } },332         },333        {334         /* .name =                 */ "response_format with reasoning generation_prompt",335         /* .tools =                */ {},336         /* .tool_choice =          */ COMMON_CHAT_TOOL_CHOICE_NONE,337         /* .reasoning_format =     */ COMMON_REASONING_FORMAT_AUTO,338         /* .json_schema =          */339            { { "type", "object" },340              { "properties",341                { { "invoice_number", { { "type", "string" } } },342                  { "amount", { { "type", "number" } } },343                  { "due_date", { { "type", "string" } } } } },344              { "required", { "invoice_number", "amount", "due_date" } } },345         /* .parallel_tool_calls =  */ false,346         /* .generation_prompt =    */ "<think>",347         /* .input =                */348            ("I must produce the invoice in the requested format</think>\n"349             R"({"invoice_number": "INV-2025-001", "amount": 1250.50, "due_date": "2025-12-31"})"),350         /* .expect_reasoning =     */ "I must produce the invoice in the requested format",351         /* .expect_content =       */352            R"({"invoice_number": "INV-2025-001", "amount": 1250.50, "due_date": "2025-12-31"})", /* .expect_tool_calls =    */ {},353         },354    };355 356    for (const auto & tc : test_cases) {357        t.test(tc.name, [&](testing & t) {358            auto parser  = build_parser(tc);359            auto lazy    = !tc.tools.empty() && tc.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;360            auto grammar = build_grammar([&](const common_grammar_builder & builder) {361                for (const auto & def : tc.tools) {362                    auto function   = def.at("function");363                    auto parameters = function.at("parameters");364                    builder.resolve_refs(parameters);365                };366                parser.build_grammar(builder, lazy);367            });368 369            t.log("Grammar:");370            for (const auto & line : string_split(grammar, "\n")) {371                t.log(line);372            }373 374            std::string              effective_input = tc.generation_prompt + tc.input;375            common_peg_parse_context ctx(effective_input);376            auto                     result = parser.parse(ctx);377 378            t.assert_true("success", result.success());379 380            common_chat_msg msg;381            auto            mapper = common_chat_peg_mapper(msg);382            mapper.from_ast(ctx.ast, result);383 384            t.assert_equal("content equal", tc.expect_content, msg.content);385            t.assert_equal("reasoning equal", tc.expect_reasoning, msg.reasoning_content);386            t.assert_equal("number of tool calls", tc.expect_tool_calls.size(), msg.tool_calls.size());387            for (auto i = 0u; i < std::min(tc.expect_tool_calls.size(), msg.tool_calls.size()); i++) {388                t.assert_equal("tool name", tc.expect_tool_calls[i].name, msg.tool_calls[i].name);389                t.assert_equal("tool args", tc.expect_tool_calls[i].arguments, msg.tool_calls[i].arguments);390            }391        });392    }393}394 395static void test_example_qwen3_coder(testing & t) {396    auto tools  = create_tools();397    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {398        auto content = p.rule("content", p.content(p.until("<tool_call>")));399 400        std::vector<common_peg_parser> tool_parsers;401        for (const auto & def : tools) {402            auto        function   = def.at("function");403            std::string name       = function.at("name");404            auto        parameters = function.at("parameters");405            auto        properties = parameters.at("properties");406 407            std::set<std::string> required_properties;408            if (function.contains("required")) {409                function.at("required").get_to(required_properties);410            }411 412            std::vector<common_peg_parser> arg_parsers;413            for (const auto & [param_name, param_schema] : properties.items()) {414                bool is_required = required_properties.find(param_name) != required_properties.end();415                auto type        = param_schema.value("type", "object");416 417                auto arg = p.tool_arg(418                    p.sequence({ p.tool_arg_open("<parameter=" + p.tool_arg_name(p.literal(param_name)) + ">"),419                                 (type == "string" ?420                                      p.tool_arg_string_value(p.schema(421                                          p.until_one_of({ "</parameter>\n<parameter=", "</parameter>\n</function>" }),422                                          "tool-" + name + "-arg-" + param_name + "-schema", param_schema, true)) :423                                      p.tool_arg_json_value(p.schema(424                                          p.json(), "tool-" + name + "-arg-" + param_name + "-schema", param_schema))),425                                 p.tool_arg_close("</parameter>\n" +426                                                  p.peek(p.literal("<parameter=") | p.literal("</function>"))) }));427 428                arg_parsers.push_back(is_required ? p.rule("tool-" + name + "-arg-" + param_name, arg) :429                                                    p.optional(p.rule("tool-" + name + "-arg-" + param_name, arg)));430            }431 432            tool_parsers.push_back(p.rule("tool-" + name, p.tool_open("<function=" + p.tool_name(p.literal(name)) + ">")433                                                              << p.sequence(arg_parsers)434                                                              << p.tool_close(p.literal("</function>"))));435        };436 437        auto tool_call = p.trigger_rule("tool-call", "<tool_call>" << p.choice(tool_parsers) << "</tool_call>");438 439        return content + p.zero_or_more(p.space() + tool_call) + p.end();440    });441 442    auto grammar = build_grammar([&](const common_grammar_builder & builder) {443        for (const auto & def : tools) {444            auto function   = def.at("function");445            auto parameters = function.at("parameters");446            builder.resolve_refs(parameters);447        };448        parser.build_grammar(builder);449    });450 451    t.log("Grammar:");452    for (const auto & line : string_split(grammar, "\n")) {453        t.log(line);454    }455 456    t.test("incremental parsing", [&](testing & t) {457        std::string input =458            "Let me search the knowledge base for cat pictures."459            "<tool_call>\n"460            "<function=search_knowledge_base>\n"461            "<parameter=query>cat pictures</parameter>\n"462            "<parameter=category>general</parameter>\n"463            "</function>\n"464            "</tool_call>";465 466        std::vector<std::string> tokens = simple_tokenize(input);467 468        common_chat_msg prev;469        for (auto it = tokens.begin(); it != tokens.end(); it++) {470            std::string in = std::accumulate(tokens.begin(), it + 1, std::string());471 472            common_peg_parse_context ctx(in, (it + 1 < tokens.end()) ? COMMON_PEG_PARSE_FLAG_LENIENT : COMMON_PEG_PARSE_FLAG_NONE);473 474            auto result = parser.parse(ctx);475            if (!t.assert_equal("not fail", false, result.fail())) {476                t.log(in.substr(0, result.end) + "[failed->]" + in.substr(result.end));477            }478 479            common_chat_msg msg;480            auto            mapper = common_chat_peg_mapper(msg);481            mapper.from_ast(ctx.ast, result);482 483            //t.log("Input: " + input);484            t.log("===========================================");485            t.log("Iteration " + std::to_string(in.size()));486            t.log("Reasoning: " + msg.reasoning_content);487            t.log("Content  : " + msg.content);488            for (const auto & tc : msg.tool_calls) {489                t.log("Tool name: " + tc.name);490                t.log("Tool args: " + tc.arguments);491            }492 493            try {494                // This shouldn't emit any runtime errors495                auto diffs = common_chat_msg_diff::compute_diffs(prev, msg);496            } catch (const std::exception & e) {497                t.log(in.substr(0, result.end) + "[failed->]" + in.substr(result.end));498                t.assert_true(std::string("failed with ") + e.what(), false);499            }500 501            prev = msg;502        }503    });504}505 506static void test_example_qwen3_non_coder(testing & t) {507    auto tools  = create_tools();508    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {509        // tool calling parser using standard JSON format510        auto tool_call = p.standard_json_tools("<tool_call>", "</tool_call>", tools, true, false);511 512        return p.sequence({ p.content(p.until("<tool_call>")), p.optional(p.space() + tool_call), p.end() });513    });514 515    auto grammar = build_grammar([&](const common_grammar_builder & builder) {516        for (const auto & def : tools) {517            auto function   = def.at("function");518            auto parameters = function.at("parameters");519            builder.resolve_refs(parameters);520        };521        parser.build_grammar(builder);522    });523 524    t.log("Grammar:");525    for (const auto & line : string_split(grammar, "\n")) {526        t.log(line);527    }528 529    t.test("tool call parsing", [&](testing & t) {530        std::string input =531            "I need to get the weather.\n"532            "<tool_call>"533            "{\"name\": \"get_current_weather\", \"arguments\": {\"location\": \"New York City, NY\", \"unit\": "534            "\"fahrenheit\"}}"535            "</tool_call>";536 537        common_peg_parse_context ctx(input);538        auto                     result = parser.parse(ctx);539 540        t.assert_true("success", result.success());541 542        common_chat_msg msg;543        auto            mapper = common_chat_peg_mapper(msg);544        mapper.from_ast(ctx.ast, result);545 546        t.assert_equal("content", "I need to get the weather.\n", msg.content);547        t.assert_equal("reasoning", "", msg.reasoning_content);548        t.assert_equal("tool calls count", 1u, msg.tool_calls.size());549        if (!msg.tool_calls.empty()) {550            t.assert_equal("tool name", "get_current_weather", msg.tool_calls[0].name);551            t.assert_equal("tool args", "{\"location\": \"New York City, NY\", \"unit\": \"fahrenheit\"}",552                           msg.tool_calls[0].arguments);553        }554    });555 556    t.test("incremental parsing", [&](testing & t) {557        std::string input =558            "I need to get the weather.\n"559            "<tool_call>"560            "{\"name\": \"get_current_weather\", \"arguments\": {\"location\": \"New York City, NY\", \"unit\": "561            "\"fahrenheit\"}}"562            "</tool_call>";563 564        std::vector<std::string> tokens = simple_tokenize(input);565 566        common_chat_msg prev;567        for (auto it = tokens.begin(); it != tokens.end(); it++) {568            std::string in = std::accumulate(tokens.begin(), it + 1, std::string());569 570            common_peg_parse_context ctx(in, (it + 1 < tokens.end()) ? COMMON_PEG_PARSE_FLAG_LENIENT : COMMON_PEG_PARSE_FLAG_NONE);571 572            auto result = parser.parse(ctx);573            if (!t.assert_equal("not fail", false, result.fail())) {574                t.log(in.substr(0, result.end) + "[failed->]" + in.substr(result.end));575            }576 577            common_chat_msg msg;578            auto            mapper = common_chat_peg_mapper(msg);579            mapper.from_ast(ctx.ast, result);580 581            //t.log("Input: " + input);582            t.log("===========================================");583            t.log("Iteration " + std::to_string(in.size()));584            t.log("Reasoning: " + msg.reasoning_content);585            t.log("Content  : " + msg.content);586            for (const auto & tc : msg.tool_calls) {587                t.log("Tool name: " + tc.name);588                t.log("Tool args: " + tc.arguments);589            }590 591            try {592                // This shouldn't emit any runtime errors593                auto diffs = common_chat_msg_diff::compute_diffs(prev, msg);594            } catch (const std::exception & e) {595                t.log(in.substr(0, result.end) + "[failed->]" + in.substr(result.end));596                t.assert_true(std::string("failed with ") + e.what(), false);597            }598 599            prev = msg;600        }601    });602}603 604void test_command7_parser_compare(testing & t) {605    auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {606        auto thinking =607            p.reasoning_block("<|START_THINKING|>" << p.reasoning(p.until("<|END_THINKING|>")) << "<|END_THINKING|>");608 609        auto response = "<|START_RESPONSE|>" << p.content(p.until("<|END_RESPONSE|>")) << "<|END_RESPONSE|>";610 611        auto tool_call_id = p.atomic("\"tool_call_id\"" << (":" << ("\"" + p.tool_id(p.string_content('"')) + "\"")));612        auto tool_call_name =613            p.atomic("\"tool_name\"" << (":" << ("\"" + p.tool_name(p.string_content('"')) + "\"")));614        auto tool_call_args = "\"parameters\"" << (":" << p.tool_args(p.json()));615 616        auto tool_call_fields = p.rule("tool-call-fields", tool_call_id | tool_call_name | tool_call_args);617        auto tool_call =618            p.rule("tool-call", p.tool(p.tool_open(p.literal("{"))619                                       << tool_call_fields << p.zero_or_more(p.literal(",") << tool_call_fields)620                                       << p.tool_close(p.literal("}"))));621 622        auto tool_calls = p.rule(623            "tool-calls", "<|START_ACTION|>" << ("[" << tool_call << p.zero_or_more(p.literal(",") << tool_call) << "]")624                                             << "<|END_ACTION|>");625 626        return p.optional(thinking) << (tool_calls | response) + p.end();627    });628 629    auto test_current = [&](const common_peg_arena & p, const std::string & input, bool is_partial,630                            bool print_results) {631        common_peg_parse_context ctx(input, is_partial ? COMMON_PEG_PARSE_FLAG_LENIENT : COMMON_PEG_PARSE_FLAG_NONE);632        auto                     result = p.parse(ctx);633 634        common_chat_msg msg;635        auto            mapper = common_chat_peg_mapper(msg);636        mapper.from_ast(ctx.ast, result);637 638        if (print_results) {639            std::cout << "== Parsed (new) ==\n";640            std::cout << "=== Reasoning ===\n";641            std::cout << msg.reasoning_content << "\n";642            std::cout << "\n\n=== Content ===\n";643            std::cout << msg.content << "\n";644            std::cout << "\n\n=== Tool Calls ===\n";645            for (const auto & tc : msg.tool_calls) {646                std::cout << "id: " << tc.id << "\n";647                std::cout << "name: " << tc.name << "\n";648                std::cout << "args: " << tc.arguments << "\n";649            }650        }651    };652 653    std::string reasoning =654        "To plan an effective trip to Japan that includes both historical sites and modern attractions within a "655        "budget of $4000 for a two-week stay, we need to:\n\n"656        "1. Identify key historical sites and modern attractions in Japan.\n"657        "2. Find affordable accommodation options that provide a balance between comfort and cost.\n"658        "3. Determine the best modes of transportation for getting around Japan.\n"659        "4. Create a day-by-day itinerary that ensures the user gets to see a variety of attractions without "660        "overspending.\n"661        "5. Provide a detailed cost breakdown that includes accommodation, transportation, meals, and entry fees "662        "to attractions.";663 664    std::vector<std::tuple<std::string, std::string, nlohmann::json>> tool_calls = {665        { "call_0", "plan_trip", nlohmann::json::parse(R"({666            "destination": "Japan",667            "duration": 14,668            "budget": 4000,669            "interests": ["historical sites", "modern attractions"],670            "accommodation_preferences": "affordable",671            "transportation_preferences": "efficient",672            "meal_preferences": "local cuisine"673        })") }674    };675 676    std::vector<std::string> tokens;677 678    // Build tokens679    if (!reasoning.empty()) {680        auto tokenized = simple_tokenize(reasoning);681        tokens.emplace_back("<|START_THINKING|>");682        tokens.insert(tokens.end(), tokenized.begin(), tokenized.end());683        tokens.emplace_back("<|END_THINKING|>");684    }685 686    if (!tool_calls.empty()) {687        tokens.emplace_back("<|START_ACTION|>");688 689        auto json = nlohmann::json::array();690        for (const auto & tc : tool_calls) {691            auto tc_json            = nlohmann::json::object();692            tc_json["tool_call_id"] = std::get<0>(tc);693            tc_json["tool_name"]    = std::get<1>(tc);694            tc_json["parameters"]   = std::get<2>(tc);695            json.push_back(tc_json);696        }697 698        auto tokenized = simple_tokenize(json.dump(-1, ' ', true));699        tokens.insert(tokens.end(), tokenized.begin(), tokenized.end());700 701        tokens.emplace_back("<|END_ACTION|>");702    }703 704    std::string input = std::accumulate(tokens.begin(), tokens.end(), std::string());705 706    t.test("current_parse", [&](testing & /* t */) { test_current(parser, input, false, false); });707    t.bench("current_parse_benchmark complete", [&]() { test_current(parser, input, false, false); }, 100);708    t.bench(709        "current_parse_benchmark incremental",710        [&]() {711            std::string in;712            for (auto i = 0u; i < tokens.size(); i++) {713                in += tokens[i];714                test_current(parser, in, i + 1 < tokens.size(), false);715            }716        },717        20);718}719 720// Test that tool names that are proper prefixes of other tool names don't cause721// premature matching during incremental parsing.722// For example, "special_function" should not match when parsing "special_function_with_opt".723static void test_prefix_tool_names(testing & t) {724    // Create tools where one name is a proper prefix of another725    json tools = json::array();726 727    json tool_short = {728        { "type", "function" },729        { "function",730          {731              { "name", "special_function" },732              { "description", "A special function" },733              { "parameters",734                {735                    { "type", "object" },736                    { "properties",737                      {738                          { "arg1", { { "type", "integer" } } },739                      } },740                    { "required", { "arg1" } },741                } },742          } }743    };744    tools.push_back(tool_short);745 746    json tool_long = {747        { "type", "function" },748        { "function",749          {750              { "name", "special_function_with_opt" },751              { "description", "A special function with optional params" },752              { "parameters",753                {754                    { "type", "object" },755                    { "properties",756                      {757                          { "arg1", { { "type", "integer" } } },758                          { "arg2", { { "type", "integer" } } },759                      } },760                    { "required", { "arg1" } },761                } },762          } }763    };764    tools.push_back(tool_long);765 766    // Use standard_constructed_tools which had the prefix matching bug767    std::map<std::string, std::string> markers = {768        { "tool_call_start_marker", "<tool_call>" },769        { "tool_call_end_marker", "</tool_call>" },770        { "function_opener", "<function=" },771        { "function_closer", "</function>" },772        { "function_name_suffix", ">" },773        { "parameter_key_prefix", "<param=" },774        { "parameter_key_suffix", ">" },775        { "parameter_closer", "</param>" },776    };777 778    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {779        auto content   = p.rule("content", p.content(p.until("<tool_call>")));780        auto tool_call = p.standard_constructed_tools(markers, tools, false, false);781        return content + p.zero_or_more(p.space() + tool_call) + p.end();782    });783 784    // Test parsing the long tool name - this should NOT trigger the short tool name785    t.test("parse long tool name", [&](testing & t) {786        std::string input =787            "Let me call the function."788            "<tool_call>"789            "<function=special_function_with_opt>"790            "<param=arg1>42</param>"791            "</function>"792            "</tool_call>";793 794        common_peg_parse_context ctx(input);795        auto                     result = parser.parse(ctx);796 797        t.assert_true("success", result.success());798 799        common_chat_msg msg;800        auto            mapper = common_chat_peg_mapper(msg);801        mapper.from_ast(ctx.ast, result);802 803        t.assert_equal("content", "Let me call the function.", msg.content);804        t.assert_equal("tool calls count", 1u, msg.tool_calls.size());805        if (!msg.tool_calls.empty()) {806            t.assert_equal("tool name", "special_function_with_opt", msg.tool_calls[0].name);807        }808    });809 810    // Test incremental parsing - the key test case811    // This ensures that when incrementally parsing "special_function_with_opt",812    // we don't prematurely emit "special_function" as a tool call813    t.test("incremental parse long tool name", [&](testing & t) {814        std::string input =815            "Let me call the function."816            "<tool_call>"817            "<function=special_function_with_opt>"818            "<param=arg1>42</param>"819            "</function>"820            "</tool_call>";821 822        std::vector<std::string> tokens = simple_tokenize(input);823 824        common_chat_msg prev;825        for (auto it = tokens.begin(); it != tokens.end(); it++) {826            std::string in = std::accumulate(tokens.begin(), it + 1, std::string());827 828            common_peg_parse_context ctx(in, (it + 1 < tokens.end()) ? COMMON_PEG_PARSE_FLAG_LENIENT : COMMON_PEG_PARSE_FLAG_NONE);829            auto                     result = parser.parse(ctx);830 831            if (!t.assert_equal("not fail", false, result.fail())) {832                t.log(in.substr(0, result.end) + "[failed->]" + in.substr(result.end));833                return;834            }835 836            common_chat_msg msg;837            auto            mapper = common_chat_peg_mapper(msg);838            mapper.from_ast(ctx.ast, result);839 840            // The critical check: during incremental parsing, we should never841            // see "special_function" as the tool name when parsing "special_function_with_opt"842            for (const auto & tc : msg.tool_calls) {843                if (!t.assert_equal("tool name should not be short prefix", false,844                                    tc.name == "special_function")) {845                    t.log("Premature tool name match at input: " + in);846                    return;847                }848            }849 850            try {851                auto diffs = common_chat_msg_diff::compute_diffs(prev, msg);852            } catch (const std::exception & e) {853                t.log(in.substr(0, result.end) + "[failed->]" + in.substr(result.end));854                t.assert_true(std::string("diff failed with ") + e.what(), false);855                return;856            }857 858            prev = msg;859        }860 861        // Final check: the complete parse should have the correct tool name862        t.assert_equal("final tool calls count", 1u, prev.tool_calls.size());863        if (!prev.tool_calls.empty()) {864            t.assert_equal("final tool name", "special_function_with_opt", prev.tool_calls[0].name);865        }866    });867 868    // Test parsing the short tool name still works869    t.test("parse short tool name", [&](testing & t) {870        std::string input =871            "Let me call the function."872            "<tool_call>"873            "<function=special_function>"874            "<param=arg1>42</param>"875            "</function>"876            "</tool_call>";877 878        common_peg_parse_context ctx(input);879        auto                     result = parser.parse(ctx);880 881        t.assert_true("success", result.success());882 883        common_chat_msg msg;884        auto            mapper = common_chat_peg_mapper(msg);885        mapper.from_ast(ctx.ast, result);886 887        t.assert_equal("content", "Let me call the function.", msg.content);888        t.assert_equal("tool calls count", 1u, msg.tool_calls.size());889        if (!msg.tool_calls.empty()) {890            t.assert_equal("tool name", "special_function", msg.tool_calls[0].name);891        }892    });893}894 895static void test_tagged_peg_parser(testing & t) {896    t.test("basic tag extraction", [&](testing & t) {897        auto parser = build_tagged_peg_parser([](common_peg_parser_builder & p) {898            return p.tag("greeting", p.until(" ")) + " " + p.tag("name", p.rest()) + p.end();899        });900 901        auto result = parser.parse_and_extract("Hello World");902        t.assert_true("success", result.result.success());903        t.assert_equal("greeting tag", "Hello", result.tags.at("greeting"));904        t.assert_equal("name tag", "World", result.tags.at("name"));905    });906 907    t.test("duplicate tags overwrite", [&](testing & t) {908        auto parser = build_tagged_peg_parser([](common_peg_parser_builder & p) {909            return p.tag("item", p.until(",")) + "," + p.tag("item", p.rest()) + p.end();910        });911 912        auto result = parser.parse_and_extract("first,second");913        t.assert_true("success", result.result.success());914        t.assert_equal("item tag", "second", result.tags.at("item"));915    });916 917    t.test("no tags extracted", [&](testing & t) {918        auto parser = build_tagged_peg_parser([](common_peg_parser_builder & p) {919            return p.rest() + p.end();920        });921 922        auto result = parser.parse_and_extract("Hello");923        t.assert_true("success", result.result.success());924        t.assert_equal("empty tags", 0u, result.tags.size());925    });926 927    t.test("structured extraction", [&](testing & t) {928        auto parser = build_tagged_peg_parser([](common_peg_parser_builder & p) {929            auto header = p.tag("header", p.until("\n"));930            auto body = p.tag("body", p.rest());931            return header + "\n" + body + p.end();932        });933 934        auto result = parser.parse_and_extract("Title\nBody content here");935        t.assert_true("success", result.result.success());936        t.assert_equal("header", "Title", result.tags.at("header"));937        t.assert_equal("body", "Body content here", result.tags.at("body"));938    });939 940    t.test("partial parse", [&](testing & t) {941        auto parser = build_tagged_peg_parser([](common_peg_parser_builder & p) {942            return p.tag("prefix", p.until(":")) + ":" + p.tag("value", p.rest()) + p.end();943        });944 945        auto result = parser.parse_and_extract("key:val", COMMON_PEG_PARSE_FLAG_LENIENT);946        t.assert_true("not fail", !result.result.fail());947        t.assert_equal("prefix tag", "key", result.tags.at("prefix"));948        t.assert_equal("value tag", "val", result.tags.at("value"));949    });950 951    t.test("find in the middle", [&](testing & t) {952        auto parser = build_tagged_peg_parser([](common_peg_parser_builder & p) {953            return p.choice({ p.literal("{"), p.literal(":") }) + p.space() + p.literal("\"") + p.atomic(p.literal("fun_name"));954        });955 956        std::string tpl = "This is a very long jinja template string. We have tools. We will try to call them now: <tool_call>{ \"fun_name\" : { \"arg\" : 1 }</tool_call>";957        auto result = parser.parse_anywhere_and_extract(tpl);958        t.assert_true("success", result.result.success());959    });960 961    t.test("fail find in the middle", [&](testing & t) {962        auto parser = build_tagged_peg_parser([](common_peg_parser_builder & p) {963            return p.choice({ p.literal("{"), p.literal(":") }) + p.space() + p.literal("\"") + p.atomic(p.literal("fun_name"));964        });965 966        std::string tpl = "This is a very long jinja template string. We have tools. We will try to call them now: <tool_call><fun=fun_name><arg name=arg>1</arg></tool_call>";967        auto result = parser.parse_anywhere_and_extract(tpl);968        t.assert_true("failure", result.result.fail());969    });970 971    t.test("find function tag with name", [&](testing &t) {972        std::string haystack = "\n<tool_call>\n<function=foofoo>\n<parameter=first>\nXXXX\n</parameter>\n<parameter=second>\nYYYY\n</parameter>\n</function>\n</tool_call>\n";973        auto parser = build_tagged_peg_parser([](common_peg_parser_builder & p) {974            std::string needle = "foofoo";975            return p.tag("fun_marker", p.choice({976            p.tag("fun_pre", p.literal("<") + p.until_one_of({ ">", needle })) + p.literal(needle) +977                p.tag("fun_post", p.negate(p.space() + p.literal("<")) + p.until(">") + p.literal(">")) + p.space(),978            p.tag("fun_pre", p.literal("[") + p.until_one_of({ "]", needle })) + p.literal(needle) +979                p.tag("fun_post", p.negate(p.space() + p.literal("[") + p.until("]") + p.literal("]")) + p.space()) }));980        });981        auto result = parser.parse_anywhere_and_extract(haystack);982        t.assert_true("success", result.result.success());983        t.assert_equal("fun_pre should be '<function='", "<function=", result.tags["fun_pre"]);984        t.assert_equal("fun_post should be '>'", ">", result.tags["fun_post"]);985    });986}987 988static void test_permute(testing & t) {989    auto accepts = [](const common_peg_arena & parser, const std::string & input) {990        common_peg_parse_context ctx(input);991        return parser.parse(ctx).success();992    };993 994    auto gbnf_of = [](const common_peg_arena & parser) {995        return build_grammar([&](const common_grammar_builder & builder) { parser.build_grammar(builder); });996    };997 998    auto assert_gbnf_equal = [](testing & t, const std::string & expected, const std::string & actual) {999        static const std::regex leading_ws_re = std::regex(R"((^|\n)\s+)");1000        t.assert_equal("gbnf are equal", std::regex_replace(expected, leading_ws_re, "$1"), actual);1001    };1002 1003    auto count_rules = [](const std::string & gbnf, const std::string & prefix) {1004        size_t count = 0;1005        for (const auto & line : string_split<std::string>(gbnf, '\n')) {1006            if (line.rfind(prefix, 0) == 0) {1007                count++;1008            }1009        }1010        return count;1011    };1012 1013    t.test("accepts every ordering", [&](testing & t) {1014        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {1015            return p.permute("abc", { p.literal("a"), p.literal("b"), p.literal("c") }) + p.end();1016        });1017 1018        for (const std::string input : { "abc", "acb", "bac", "bca", "cab", "cba" }) {1019            t.assert_true("accepts " + input, accepts(parser, input));1020        }1021    });1022 1023    t.test("single element", [&](testing & t) {1024        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {1025            return p.permute("a", { p.literal("a") }) + p.end();1026        });1027 1028        t.assert_true("accepts a", accepts(parser, "a"));1029        t.assert_true("rejects aa", !accepts(parser, "aa"));1030    });1031 1032    t.test("grammar left-factorizes shared tails", [&](testing & t) {1033        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {1034            return p.permute("abc", { p.literal("a"), p.literal("b"), p.literal("c") }) + p.end();1035        });1036 1037        // Every rule is one remaining subset, keyed by bitmask: abc-3 is {a,b}, abc-7 is {a,b,c}.1038        // Each subset is emitted once and shared by every branch that leads into it.1039        assert_gbnf_equal(t, R"""(1040            abc-1 ::= "a"1041            abc-2 ::= "b"1042            abc-3 ::= "a" abc-2 | "b" abc-11043            abc-4 ::= "c"1044            abc-5 ::= "a" abc-4 | "c" abc-11045            abc-6 ::= "b" abc-4 | "c" abc-21046            abc-7 ::= "a" abc-6 | "b" abc-5 | "c" abc-31047            root ::= abc-71048            space ::= | " " | "\n"{1,2} [ \t]{0,20}1049        )""", gbnf_of(parser));1050    });1051 1052    t.test("grammar emits one rule per remaining subset", [&](testing & t) {1053        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {1054            return p.permute("abcd", { p.literal("a"), p.literal("b"), p.literal("c"), p.literal("d") }) + p.end();1055        });1056 1057        // 2^4 - 1 non-empty subsets, one rule each - not the 4! = 24 orderings.1058        t.assert_equal("permute rule count", 15u, count_rules(gbnf_of(parser), "abcd-"));1059    });1060 1061    t.test("grammar emits no rules for a single element", [&](testing & t) {1062        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {1063            return p.permute("a", { p.literal("a") }) + p.end();1064        });1065 1066        assert_gbnf_equal(t, R"""(1067            root ::= "a"1068            space ::= | " " | "\n"{1,2} [ \t]{0,20}1069        )""", gbnf_of(parser));1070    });1071 1072    t.test("grammar falls back to the given order when too large", [&](testing & t) {1073        auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) {1074            std::vector<common_peg_parser> parsers;1075            for (size_t i = 0; i <= COMMON_CHAT_MAX_PERMUTE; i++) {1076                parsers.push_back(p.literal(std::string(1, (char) ('a' + i))));1077            }1078            return p.permute("big", parsers) + p.end();1079        });1080 1081        assert_gbnf_equal(t, R"""(1082            root ::= "a" "b" "c" "d" "e" "f" "g"1083            space ::= | " " | "\n"{1,2} [ \t]{0,20}1084        )""", gbnf_of(parser));1085    });1086}1087