diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 4cf580a056c8..99688f53b87b 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -100,6 +100,8 @@ add_library(${TARGET} sampling.h speculative.cpp speculative.h + trie.cpp + trie.h unicode.cpp unicode.h jinja/lexer.cpp diff --git a/common/chat.cpp b/common/chat.cpp index 22d2ee4a2a11..9ecd7c7cd78a 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1022,7 +1022,7 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_ data.supports_thinking = true; data.thinking_start_tag = "[THINK]"; - data.thinking_end_tag = "[/THINK]"; + data.thinking_end_tags = {"[/THINK]"}; data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; @@ -1148,6 +1148,9 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; + data.thinking_start_tag = "<|channel|>analysis<|message|>"; + data.thinking_end_tags = {"<|end|>"}; + // These special tokens are required to parse properly, so we include them // even if parse_tool_calls is false. data.preserved_tokens = { @@ -1292,7 +1295,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ data.format = COMMON_CHAT_FORMAT_PEG_GEMMA4; data.supports_thinking = true; data.thinking_start_tag = "<|channel>thought"; - data.thinking_end_tag = ""; + data.thinking_end_tags = {""}; data.preserved_tokens = { "<|channel>", @@ -1567,7 +1570,7 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp const std::string GEN_PROMPT = "<|im_assistant|>assistant<|im_middle|>"; data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; + data.thinking_end_tags = {THINK_END}; if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; @@ -1701,7 +1704,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat } data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; + data.thinking_end_tags = {THINK_END}; auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); @@ -1864,7 +1867,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; data.thinking_start_tag = ""; - data.thinking_end_tag = ""; + data.thinking_end_tags = {""}; data.preserved_tokens = { "|DSML|", "", @@ -2035,6 +2038,233 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha return data; } +static common_chat_params common_chat_params_init_deepseek_v4(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + const bool has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + const bool has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); + + std::optional additional_context; + if (has_response_format) { + additional_context = json{ { "response_format", inputs.json_schema } }; + } + + data.prompt = common_chat_template_direct_apply_impl( + tmpl, inputs, std::nullopt, std::nullopt, additional_context); + data.generation_prompt = common_chat_template_generation_prompt_impl( + tmpl, inputs, std::nullopt, std::nullopt, additional_context); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + data.thinking_start_tag = ""; + data.preserved_tokens = { + "|DSML|", + "", + "", + }; + + const bool extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + const bool parse_tool_calls = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + const bool include_grammar = has_response_format || parse_tool_calls; + + const std::string DSML = "|DSML|"; + const std::string THINK_START = ""; + const std::string THINK_END = ""; + const std::string FC_START = "<" + DSML + "tool_calls>"; + const std::string FC_TRIGGER = "\n\n" + FC_START; + const std::string FC_END = ""; + const std::string INVOKE_START = "<" + DSML + "invoke"; + const std::string INVOKE_END = ""; + const std::string PARAM_START = "<" + DSML + "parameter"; + const std::string PARAM_END = ""; + const std::string GEN_PROMPT = "<|Assistant|>"; + + data.thinking_end_tags = {THINK_END}; + if (parse_tool_calls) { + data.thinking_end_tags.push_back(FC_TRIGGER); + } + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + if (inputs.enable_thinking) { + data.generation_prompt = GEN_PROMPT + THINK_START + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += THINK_END + msg.render_content(); + } + } else { + data.generation_prompt = GEN_PROMPT + THINK_END; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += msg.render_content(); + } + } + + data.prompt += data.generation_prompt; + } + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.literal(GEN_PROMPT); + auto end = p.end(); + + auto response_format = p.eps(); + if (has_response_format) { + response_format = p.rule( + "response-format", + p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema))); + } + + auto after_reasoning = p.eps(); + auto tool_calls = p.eps(); + if (!parse_tool_calls) { + after_reasoning = has_response_format ? response_format : p.content(p.rest()); + } else { + auto string_value = p.ac( + p.tool_arg_string_value(p.until(PARAM_END)) + + p.tool_arg_close(p.literal(PARAM_END)), + PARAM_END); + + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + auto params = function.contains("parameters") ? function.at("parameters") : json::object(); + const auto & props = params.contains("properties") ? params.at("properties") : json::object(); + + std::set required; + if (params.contains("required")) { + params.at("required").get_to(required); + } + + auto schema_info = common_schema_info(); + schema_info.resolve_refs(params); + + std::vector required_args; + std::vector arg_parsers; + for (const auto & [param_name, param_schema] : props.items()) { + const bool is_required = required.find(param_name) != required.end(); + const bool is_string = schema_info.resolves_to_string(param_schema); + + auto value = is_string + ? string_value + : p.tool_arg_json_value(p.schema( + p.json(), + "tool-" + name + "-arg-" + param_name + "-schema", + param_schema, + false)) + + p.tool_arg_close(p.literal(PARAM_END)); + + auto arg = p.tool_arg( + p.tool_arg_open( + p.literal(PARAM_START + " name=\"") + + p.tool_arg_name(p.literal(param_name)) + + p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) + + value); + + arg_parsers.push_back(p.rule("tool-" + name + "-arg-" + param_name, arg)); + required_args.push_back(is_required); + } + + auto args = p.eps(); + if (!arg_parsers.empty()) { + auto any_arg = p.choice(arg_parsers); + for (size_t i = 0; i < arg_parsers.size(); i++) { + if (!required_args[i]) { + continue; + } + args = args + p.peek( + p.zero_or_more(p.negate(arg_parsers[i]) + any_arg + p.space()) + arg_parsers[i]); + } + args = args + p.zero_or_more(any_arg + p.space()); + } + + auto func_parser = p.tool( + p.tool_open( + p.literal(INVOKE_START + " name=\"") + + p.tool_name(p.literal(name)) + + p.literal("\">\n")) + + args + p.space() + + p.tool_close(p.literal(INVOKE_END))); + + tool_choice |= p.rule("tool-" + name, func_parser); + }); + + auto tool_calls_body = p.literal(FC_TRIGGER) + p.space() + tool_choice; + if (inputs.parallel_tool_calls) { + tool_calls_body = tool_calls_body + p.zero_or_more(p.space() + tool_choice); + } + tool_calls = p.trigger_rule( + "tool-call", + tool_calls_body + p.space() + p.literal(FC_END) + + p.optional(p.space() + p.literal(THINK_END))); + + const bool require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; + auto content_before_tools = p.content(p.until(FC_TRIGGER)); + if (has_response_format) { + auto tool_call_response = content_before_tools + tool_calls; + after_reasoning = require_tools + ? tool_call_response + : p.choice({tool_call_response, response_format}); + } else { + after_reasoning = content_before_tools + (require_tools ? tool_calls : p.optional(tool_calls)); + } + } + + if (extract_reasoning && inputs.enable_thinking) { + auto reasoning_with_tool_calls = + p.literal(THINK_START) + + p.reasoning(p.until_one_of({THINK_END, FC_TRIGGER})) + + tool_calls; + auto closed_reasoning = + p.literal(THINK_START) + + p.reasoning(p.until(THINK_END)) + + p.literal(THINK_END) + + after_reasoning; + auto open_reasoning = p.literal(THINK_START) + p.reasoning(p.rest()); + + if (parse_tool_calls) { + return generation_prompt + + p.choice({reasoning_with_tool_calls, closed_reasoning, open_reasoning}) + end; + } + return generation_prompt + p.choice({closed_reasoning, open_reasoning}) + end; + } + + if (extract_reasoning) { + auto reasoning = p.optional(p.choice({ + p.literal(THINK_START) + p.until(THINK_END) + p.literal(THINK_END), + p.literal(THINK_END), + })); + return generation_prompt + reasoning + after_reasoning + end; + } + + return generation_prompt + after_reasoning + end; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = !(has_response_format || + (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + auto schema = function.contains("parameters") ? function.at("parameters") : json::object(); + builder.resolve_refs(schema); + }); + if (has_response_format) { + auto schema = inputs.json_schema; + builder.resolve_refs(schema); + } + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + {COMMON_GRAMMAR_TRIGGER_TYPE_WORD, FC_TRIGGER}, + }; + } + + return data; +} + // Cohere2 MoE (a.k.a. "North Code") parser. // // The assistant turn is fully marker-wrapped: @@ -2077,7 +2307,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; data.thinking_start_tag = THINK_START; - data.thinking_end_tag = THINK_END; + data.thinking_end_tags = {THINK_END}; data.preserved_tokens = { TURN_START, TURN_END, CHATBOT, USER, SYSTEM, THINK_START, THINK_END, @@ -2418,7 +2648,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem }; data.thinking_start_tag = ""; - data.thinking_end_tag = ""; + data.thinking_end_tags = {""}; data.message_delimiters = { { COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" }, @@ -2621,6 +2851,14 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_deepseek_v3_2(tmpl, params); } + // DeepSeek V4 uses the DSML invoke/parameter format with a tool_calls outer block. + if (src.find("dsml_token") != std::string::npos && + src.find("tool_calls>") != std::string::npos && + src.find("DSML") != std::string::npos) { + LOG_DBG("Using specialized template: DeepSeek V4\n"); + return common_chat_params_init_deepseek_v4(tmpl, params); + } + // Gemma4 format detection if (src.find("'<|tool_call>call:'") != std::string::npos) { if (src.find("{#- OpenAI Chat Completions:") == std::string::npos) { @@ -2772,7 +3010,10 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_ auto_params.supports_thinking = autoparser.reasoning.mode != autoparser::reasoning_mode::NONE; if (auto_params.supports_thinking) { auto_params.thinking_start_tag = trim_whitespace(autoparser.reasoning.start); - auto_params.thinking_end_tag = trim_whitespace(autoparser.reasoning.end); + auto end_tag = trim_whitespace(autoparser.reasoning.end); + if (!end_tag.empty()) { + auto_params.thinking_end_tags = {std::move(end_tag)}; + } } common_peg_arena arena; arena.load(auto_params.parser); diff --git a/common/chat.h b/common/chat.h index 7898f1623f54..dc4ae9d7068c 100644 --- a/common/chat.h +++ b/common/chat.h @@ -274,7 +274,7 @@ struct common_chat_params { std::string generation_prompt; bool supports_thinking = false; std::string thinking_start_tag; // e.g., "" - std::string thinking_end_tag; // e.g., "" + std::vector thinking_end_tags; // e.g., "" std::vector grammar_triggers; std::vector preserved_tokens; std::vector additional_stops; @@ -297,7 +297,7 @@ struct common_chat_parser_params { common_peg_arena parser = {}; common_chat_parser_params() = default; common_chat_parser_params(const common_chat_params & chat_params) { - format = chat_params.format; + format = chat_params.format; generation_prompt = chat_params.generation_prompt; } }; diff --git a/common/common.h b/common/common.h index bffc1767a7d1..1b883322cd55 100644 --- a/common/common.h +++ b/common/common.h @@ -283,12 +283,12 @@ struct common_params_sampling { // reasoning budget sampler parameters // these are populated by the server/CLI based on chat template params - int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget - std::vector reasoning_budget_start; // start tag token sequence - std::vector reasoning_budget_end; // end tag token sequence - std::vector reasoning_budget_forced; // forced sequence (message + end tag) - std::string reasoning_budget_message; // message injected before end tag when budget exhausted - bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime + int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget + std::vector reasoning_budget_start; // start tag token sequence + std::vector reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence + std::vector reasoning_budget_forced; // forced sequence (message + first end tag) + std::string reasoning_budget_message; // message injected before end tag when budget exhausted + bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime bool backend_sampling = false; diff --git a/common/jinja/caps.cpp b/common/jinja/caps.cpp index ae378ebd4fd0..ba13dc17b8aa 100644 --- a/common/jinja/caps.cpp +++ b/common/jinja/caps.cpp @@ -481,6 +481,7 @@ caps caps_get(jinja::program & prog) { }); }, [&](context & ctx) { + ctx.set_val("enable_thinking", mk_val(true)); caps_apply_preserve_reasoning(ctx, true); }, nullptr, // tools_fn diff --git a/common/peg-parser.cpp b/common/peg-parser.cpp index 807e952d902c..ef290ed7c057 100644 --- a/common/peg-parser.cpp +++ b/common/peg-parser.cpp @@ -3,10 +3,10 @@ #include "common.h" #include "json-schema-to-grammar.h" #include "log.h" +#include "trie.h" #include "unicode.h" #include -#include #include #include #include @@ -32,154 +32,6 @@ static bool is_hex_digit(const char c) { return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F'); } -// Trie for matching multiple literals. -// This is used in common_peg_until_parser and to build a GBNF exclusion grammar -struct trie { - struct node { - std::map children; // Use uint32_t to store Unicode codepoints - bool is_word; - }; - - std::vector nodes; - - trie(const std::vector & words) { - create_node(); // root node - for (const auto & w : words) { - insert(w); - } - } - - enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH }; - - // Check if a delimiter starts at the given position - match_result check_at(std::string_view sv, size_t start_pos) const { - size_t current = 0; // Start at root - size_t pos = start_pos; - - // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str()); - - while (pos < sv.size()) { - auto result = common_parse_utf8_codepoint(sv, pos); - if (result.status != utf8_parse_result::SUCCESS) { - break; - } - - auto it = nodes[current].children.find(result.codepoint); - if (it == nodes[current].children.end()) { - // Can't continue matching - return match_result{match_result::NO_MATCH}; - } - - current = it->second; - pos += result.bytes_consumed; - - // Check if we've matched a complete word - if (nodes[current].is_word) { - return match_result{match_result::COMPLETE_MATCH}; - } - } - - // Reached end of input while still in the trie (not at root) - if (current != 0) { - // We're in the middle of a potential match - return match_result{match_result::PARTIAL_MATCH}; - } - - // Reached end at root (no match) - return match_result{match_result::NO_MATCH}; - } - - private: - size_t create_node() { - size_t index = nodes.size(); - nodes.emplace_back(); - return index; - } - - void insert(const std::string & word) { - size_t current = 0; - size_t pos = 0; - while (pos < word.length()) { - auto result = common_parse_utf8_codepoint(word, pos); - if (result.status != utf8_parse_result::SUCCESS) { - break; - } - - uint32_t ch = result.codepoint; - pos += result.bytes_consumed; - - auto it = nodes[current].children.find(ch); - if (it == nodes[current].children.end()) { - size_t child = create_node(); - nodes[current].children[ch] = child; - current = child; - } else { - current = it->second; - } - } - nodes[current].is_word = true; - } -}; - -// Aho-Corasick automaton -struct aho_corasick { - trie t; - std::vector fail; // failure links - std::vector order; // states in BFS order - std::vector terminal; // match states (directly or via a suffix link) - std::set alphabet; // every character with a transition - - aho_corasick(const std::vector & strings) : t(strings) { - const auto & nodes = t.nodes; - const size_t n = nodes.size(); - - fail.assign(n, 0); - order.reserve(n); - - std::deque queue{ 0 }; - while (!queue.empty()) { - size_t u = queue.front(); - queue.pop_front(); - order.push_back(u); - for (const auto & [ch, v] : nodes[u].children) { - if (u != 0) { - size_t f = fail[u]; - while (f && nodes[f].children.find(ch) == nodes[f].children.end()) { - f = fail[f]; - } - auto it = nodes[f].children.find(ch); - fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0; - } - queue.push_back(v); - } - } - - terminal.assign(n, false); - for (size_t u : order) { - terminal[u] = nodes[u].is_word || (u != 0 && terminal[fail[u]]); - } - - for (const auto & node : nodes) { - for (const auto & [ch, v] : node.children) { - alphabet.insert(ch); - } - } - } - - size_t num_states() const { return t.nodes.size(); } - bool is_terminal(size_t s) const { return terminal[s]; } - - // follow failure links until a transition on `ch` exists. - size_t next(size_t state, uint32_t ch) const { - const auto & nodes = t.nodes; - while (state && nodes[state].children.find(ch) == nodes[state].children.end()) { - state = fail[state]; - } - auto it = nodes[state].children.find(ch); - return it != nodes[state].children.end() ? it->second : 0; - } -}; - static std::pair parse_hex_escape(const std::string & str, size_t pos, int hex_count) { if (pos + hex_count > str.length()) { return {0, 0}; @@ -797,7 +649,7 @@ struct parser_executor { } common_peg_parse_result operator()(const common_peg_until_parser & p) const { - trie matcher(p.delimiters); + common_trie matcher(p.delimiters); // Scan input and check for delimiters size_t pos = start_pos; @@ -824,12 +676,12 @@ struct parser_executor { // Check if a delimiter starts at this position auto match = matcher.check_at(ctx.input, pos); - if (match == trie::COMPLETE_MATCH) { + if (match == common_trie::COMPLETE_MATCH) { // Found a complete delimiter, return everything before it return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); } - if (match == trie::PARTIAL_MATCH) { + if (match == common_trie::PARTIAL_MATCH) { // Found a partial match extending to end of input, return everything before it return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos); } @@ -1559,7 +1411,7 @@ static std::string gbnf_ac_grammar( const std::map> &, const std::vector &, const std::function &)> & build_rule) { - aho_corasick ac(strings); + common_aho_corasick ac(strings); auto state_name = [&](size_t s) -> std::string { if (s == 0) { diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 7da0bb1c57ce..1fe242d062d1 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -1,39 +1,52 @@ #include "reasoning-budget.h" #include "common.h" +#include "trie.h" #include "unicode.h" #include "log.h" +#include #include #include #include #include struct token_matcher { - std::vector tokens; - size_t pos = 0; + std::vector seqs; + common_aho_corasick ac; + size_t state = 0; - bool advance(llama_token token) { - if (tokens.empty()) { - return false; - } + token_matcher(const std::vector & seqs) : seqs(collect(seqs)), ac(build_trie(this->seqs)) {} - if (token == tokens[pos]) { - pos++; - if (pos >= tokens.size()) { - pos = 0; - return true; - } - } else { - pos = 0; - if (token == tokens[0]) { - pos = 1; + static std::vector collect(const std::vector & seqs) { + std::vector res; + for (const auto & seq : seqs) { + if (!seq.empty() && std::find(res.begin(), res.end(), seq) == res.end()) { + res.push_back(seq); } } - return false; + return res; + } + + static common_trie build_trie(const std::vector & seqs) { + common_trie t; + for (const auto & seq : seqs) { + t.insert(std::vector(seq.begin(), seq.end())); + } + return t; } - void reset() { pos = 0; } + // returns the index into seqs of the longest sequence ending at this token, or -1 + int32_t advance(llama_token token) { + state = ac.next(state, (uint32_t) token); + const int32_t p = ac.match_pattern(state); + if (p >= 0) { + state = 0; + } + return p; + } + + void reset() { state = 0; } }; struct common_reasoning_budget_ctx { @@ -41,7 +54,7 @@ struct common_reasoning_budget_ctx { token_matcher start_matcher; token_matcher end_matcher; - std::vector forced_tokens; + llama_tokens forced_tokens; int32_t budget; // maximum tokens in reasoning block int32_t remaining; // tokens remaining in budget @@ -50,6 +63,8 @@ struct common_reasoning_budget_ctx { // for forcing size_t force_pos; // next position in forced_tokens to force + + int32_t end_match; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none }; static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) { @@ -62,7 +77,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to switch (ctx->state) { case REASONING_BUDGET_IDLE: { - if (ctx->start_matcher.advance(token)) { + if (ctx->start_matcher.advance(token) >= 0) { ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; COM_TRC("activated, budget=%d tokens\n", ctx->budget); @@ -78,8 +93,10 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to case REASONING_BUDGET_COUNTING: case REASONING_BUDGET_WAITING_UTF8: { - if (ctx->end_matcher.advance(token)) { + const int32_t match = ctx->end_matcher.advance(token); + if (match >= 0) { ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; COM_TRC("%s", "deactivated (natural end)\n"); break; } @@ -115,19 +132,25 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to break; } case REASONING_BUDGET_FORCING: + { + // track the end sequence within forced_tokens so it is also reported on DONE + const int32_t match = ctx->end_matcher.advance(token); ctx->force_pos++; if (ctx->force_pos >= ctx->forced_tokens.size()) { ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; COM_TRC("%s", "forced sequence complete, done\n"); } break; + } case REASONING_BUDGET_DONE: // Re-arm on a new start tag: some models emit multiple blocks // per response, and each should get a fresh budget window. - if (ctx->start_matcher.advance(token)) { + if (ctx->start_matcher.advance(token) >= 0) { ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; ctx->end_matcher.reset(); + ctx->end_match = -1; COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget); if (ctx->remaining <= 0) { @@ -169,11 +192,12 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) { ctx->start_matcher.reset(); ctx->end_matcher.reset(); ctx->force_pos = 0; + ctx->end_match = -1; } static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, const std::vector & start_tokens, - const std::vector & end_tokens, const std::vector & forced_tokens, + const struct llama_vocab * vocab, const std::vector & start_seqs, + const std::vector & end_seqs, const llama_tokens & forced_tokens, int32_t budget, common_reasoning_budget_state initial_state); static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl); @@ -205,12 +229,12 @@ static struct llama_sampler * common_reasoning_budget_clone(const struct llama_s } static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state) { // promote COUNTING with budget <= 0 to FORCING if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) { initial_state = REASONING_BUDGET_FORCING; @@ -220,25 +244,26 @@ static struct llama_sampler * common_reasoning_budget_init_state( /* .iface = */ &common_reasoning_budget_i, /* .ctx = */ new common_reasoning_budget_ctx { /* .vocab = */ vocab, - /* .start_matcher = */ { start_tokens, 0 }, - /* .end_matcher = */ { end_tokens, 0 }, + /* .start_matcher = */ token_matcher(start_seqs), + /* .end_matcher = */ token_matcher(end_seqs), /* .forced_tokens = */ forced_tokens, /* .budget = */ budget, /* .remaining = */ budget, /* .state = */ initial_state, /* .force_pos = */ 0, + /* .end_match = */ -1, } ); } struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { - return common_reasoning_budget_init_state(vocab, start_tokens, end_tokens, forced_tokens, budget, initial_state); + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state) { + return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state); } common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) { @@ -248,6 +273,19 @@ common_reasoning_budget_state common_reasoning_budget_get_state(const struct lla return ((const common_reasoning_budget_ctx *)smpl->ctx)->state; } +const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl) { + if (!smpl) { + return nullptr; + } + + const auto * ctx = (const common_reasoning_budget_ctx *) smpl->ctx; + if (ctx->end_match < 0) { + return nullptr; + } + + return &ctx->end_matcher.seqs[ctx->end_match]; +} + bool common_reasoning_budget_force(struct llama_sampler * smpl) { if (!smpl) { return false; diff --git a/common/reasoning-budget.h b/common/reasoning-budget.h index 0cf689a56637..1b89a04c42e8 100644 --- a/common/reasoning-budget.h +++ b/common/reasoning-budget.h @@ -2,6 +2,8 @@ #include "llama.h" +#include "common.h" + #include #include @@ -17,30 +19,34 @@ enum common_reasoning_budget_state { // reasoning block (e.g. between and ). // // State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE -// IDLE: passthrough, watching for start_tokens sequence -// COUNTING: counting down remaining tokens, watching for natural end_tokens +// IDLE: passthrough, watching for a start sequence +// COUNTING: counting down remaining tokens, watching for a natural end sequence // WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence // FORCING: forces forced_tokens token-by-token (all other logits -> -inf) // DONE: passthrough forever // // Parameters: // vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr) -// start_tokens - token sequence that activates counting -// end_tokens - token sequence for natural deactivation +// start_seqs - token sequences, any of which activates counting +// end_seqs - token sequences, any of which naturally deactivates // forced_tokens - token sequence forced when budget expires // budget - max tokens allowed in the reasoning block // initial_state - initial state // struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_tokens, - const std::vector & end_tokens, - const std::vector & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + int32_t budget, + common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl); +// The end sequence that transitioned the sampler to DONE, or nullptr if none +// was recorded. Cleared when a new start sequence re-arms the sampler. +const llama_tokens * common_reasoning_budget_get_end_match(const struct llama_sampler * smpl); + // Manually transition the reasoning budget sampler into the FORCING state. // Returns true if the transition occurred. bool common_reasoning_budget_force(struct llama_sampler * smpl); diff --git a/common/sampling.cpp b/common/sampling.cpp index 75a299e23ece..7b241e34f77f 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -299,7 +299,7 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, st if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0 || params.reasoning_control)) { rbudget = common_reasoning_budget_init( vocab, - params.reasoning_budget_start, + {params.reasoning_budget_start}, params.reasoning_budget_end, params.reasoning_budget_forced, params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens); @@ -453,6 +453,17 @@ void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, boo if (gsmpl->rbudget && is_generated) { llama_sampler_accept(gsmpl->rbudget, token); + + // if done, replay end sequence which may contain a grammar trigger + const bool is_done = common_reasoning_budget_get_state(gsmpl->rbudget) == REASONING_BUDGET_DONE; + if (gsmpl->grmr && !accept_grammar && is_done) { + const llama_tokens * end_seq = common_reasoning_budget_get_end_match(gsmpl->rbudget); + if (end_seq) { + for (const llama_token end_token : *end_seq) { + llama_sampler_accept(gsmpl->grmr, end_token); + } + } + } } if (gsmpl->grmr && accept_grammar) { diff --git a/common/trie.cpp b/common/trie.cpp new file mode 100644 index 000000000000..b5c9666ba2ee --- /dev/null +++ b/common/trie.cpp @@ -0,0 +1,123 @@ +#include "trie.h" + +#include "unicode.h" + +#include + +common_trie::match_result common_trie::check_at(std::string_view sv, size_t start_pos) const { + size_t current = 0; // Start at root + size_t pos = start_pos; + + // LOG_DBG("%s: checking at pos %zu, sv='%s'\n", __func__, start_pos, std::string(sv).c_str()); + + while (pos < sv.size()) { + auto result = common_parse_utf8_codepoint(sv, pos); + if (result.status != utf8_parse_result::SUCCESS) { + break; + } + + auto it = nodes[current].children.find(result.codepoint); + if (it == nodes[current].children.end()) { + // Can't continue matching + return match_result{match_result::NO_MATCH}; + } + + current = it->second; + pos += result.bytes_consumed; + + // Check if we've matched a complete word + if (nodes[current].pattern >= 0) { + return match_result{match_result::COMPLETE_MATCH}; + } + } + + // Reached end of input while still in the trie (not at root) + if (current != 0) { + // We're in the middle of a potential match + return match_result{match_result::PARTIAL_MATCH}; + } + + // Reached end at root (no match) + return match_result{match_result::NO_MATCH}; +} + +int32_t common_trie::insert(const std::string & word) { + std::vector symbols; + size_t pos = 0; + while (pos < word.length()) { + auto result = common_parse_utf8_codepoint(word, pos); + if (result.status != utf8_parse_result::SUCCESS) { + break; + } + + symbols.push_back(result.codepoint); + pos += result.bytes_consumed; + } + return insert(symbols); +} + +int32_t common_trie::insert(const std::vector & symbols) { + size_t current = 0; + for (uint32_t ch : symbols) { + auto it = nodes[current].children.find(ch); + if (it == nodes[current].children.end()) { + size_t child = create_node(); + nodes[current].children[ch] = child; + current = child; + } else { + current = it->second; + } + } + if (nodes[current].pattern < 0) { + nodes[current].pattern = n_patterns++; + } + return nodes[current].pattern; +} + +common_aho_corasick::common_aho_corasick(common_trie trie) : t(std::move(trie)) { + const auto & nodes = t.nodes; + const size_t n = nodes.size(); + + fail.assign(n, 0); + order.reserve(n); + + std::deque queue{ 0 }; + while (!queue.empty()) { + size_t u = queue.front(); + queue.pop_front(); + order.push_back(u); + for (const auto & [ch, v] : nodes[u].children) { + if (u != 0) { + size_t f = fail[u]; + while (f && nodes[f].children.find(ch) == nodes[f].children.end()) { + f = fail[f]; + } + auto it = nodes[f].children.find(ch); + fail[v] = (it != nodes[f].children.end() && it->second != v) ? it->second : 0; + } + queue.push_back(v); + } + } + + // fail[u] points to a strictly shorter suffix, so the first pattern found on + // the fail chain (including u itself) is the longest pattern ending at u + match.assign(n, -1); + for (size_t u : order) { + match[u] = nodes[u].pattern >= 0 ? nodes[u].pattern : (u != 0 ? match[fail[u]] : -1); + } + + for (const auto & node : nodes) { + for (const auto & [ch, v] : node.children) { + alphabet.insert(ch); + } + } +} + +size_t common_aho_corasick::next(size_t state, uint32_t ch) const { + const auto & nodes = t.nodes; + while (state && nodes[state].children.find(ch) == nodes[state].children.end()) { + state = fail[state]; + } + auto it = nodes[state].children.find(ch); + return it != nodes[state].children.end() ? it->second : 0; +} diff --git a/common/trie.h b/common/trie.h new file mode 100644 index 000000000000..0f7b16a36ad2 --- /dev/null +++ b/common/trie.h @@ -0,0 +1,73 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +// Trie for matching multiple literals. +// This is used in common_peg_until_parser and to build a GBNF exclusion grammar +struct common_trie { + struct node { + std::map children; // Use uint32_t to store Unicode codepoints + int32_t pattern = -1; // index of the pattern ending at this node, -1 if none + }; + + std::vector nodes; + + common_trie() { + create_node(); // root node + } + + common_trie(const std::vector & words) : common_trie() { + for (const auto & w : words) { + insert(w); + } + } + + enum match_result { NO_MATCH, PARTIAL_MATCH, COMPLETE_MATCH }; + + // Check if a delimiter starts at the given position + match_result check_at(std::string_view sv, size_t start_pos) const; + + // Insert a word as a sequence of Unicode codepoints, returns its pattern index + int32_t insert(const std::string & word); + + // Insert a raw symbol sequence, returns its pattern index (insertion order, + // duplicates keep the first index) + int32_t insert(const std::vector & symbols); + + private: + int32_t n_patterns = 0; + + size_t create_node() { + size_t index = nodes.size(); + nodes.emplace_back(); + return index; + } +}; + +// Aho-Corasick automaton +struct common_aho_corasick { + common_trie t; + std::vector fail; // failure links + std::vector order; // states in BFS order + std::vector match; // longest pattern ending at each state (directly or via a suffix link), -1 if none + std::set alphabet; // every character with a transition + + common_aho_corasick(common_trie trie); + + common_aho_corasick(const std::vector & strings) + : common_aho_corasick(common_trie(strings)) {} + + size_t num_states() const { return t.nodes.size(); } + bool is_terminal(size_t s) const { return match[s] >= 0; } + + // index of the longest pattern ending at this state, -1 if none + int32_t match_pattern(size_t s) const { return match[s]; } + + // follow failure links until a transition on `ch` exists. + size_t next(size_t state, uint32_t ch) const; +}; diff --git a/models/templates/deepseek-ai-DeepSeek-V4.jinja b/models/templates/deepseek-ai-DeepSeek-V4.jinja index f19f787b1b7e..bb2b429bf60d 100644 --- a/models/templates/deepseek-ai-DeepSeek-V4.jinja +++ b/models/templates/deepseek-ai-DeepSeek-V4.jinja @@ -11,6 +11,9 @@ {%- set dsml_token = '|DSML|' -%} {%- set thinking_start_token = '' -%} {%- set thinking_end_token = '' -%} +{%- set reasoning_effort_max = 'Reasoning Effort: Absolute maximum with no shortcuts permitted.\nYou MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\nExplicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n' -%} +{%- set response_format_template = '## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n' -%} +{%- set has_tools = false -%} {%- set tools_header = '## Tools\n\nYou have access to a set of tools to help answer the user\'s question. You can invoke tools by writing a "<' + dsml_token + 'tool_calls>" block like the following:\n\n<' + dsml_token + 'tool_calls>\n<' + dsml_token + 'invoke name="$TOOL_NAME">\n<' + dsml_token + 'parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE\n...\n\n<' + dsml_token + 'invoke name="$TOOL_NAME2">\n...\n\n\n\nString parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.\n\nIf thinking_mode is enabled (triggered by ' + thinking_start_token + '), you MUST output your complete reasoning inside ' + thinking_start_token + '...' + thinking_end_token + ' BEFORE any tool calls or final response.\n\nOtherwise, output directly after ' + thinking_end_token + ' with tool calls or final response.\n\n### Available Tool Schemas\n\n' -%} {%- set tools_footer = '\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n' -%} {%- set ns = namespace(system_prompt='', is_first_sp=true) -%} @@ -25,6 +28,7 @@ {%- endif -%} {%- endfor -%} {%- if tools is defined and tools -%} + {%- set has_tools = true -%} {%- set ts = namespace(schemas='') -%} {%- for tool in tools -%} {%- if tool['type'] == 'function' -%} @@ -37,7 +41,16 @@ {%- set ns.system_prompt = tools_header + ts.schemas + tools_footer -%} {%- endif -%} {%- endif -%} +{%- if response_format is defined -%} + {%- if ns.system_prompt -%} + {%- set ns.system_prompt = ns.system_prompt + '\n\n' -%} + {%- endif -%} + {%- set ns.system_prompt = ns.system_prompt + response_format_template + (response_format | tojson) -%} +{%- endif -%} {{- bos_token -}} +{%- if messages and thinking and reasoning_effort is defined and reasoning_effort == 'max' -%} + {{- reasoning_effort_max -}} +{%- endif -%} {{- ns.system_prompt -}} {%- set last_user_idx = namespace(value=-1) -%} {%- for message in messages -%} @@ -67,7 +80,8 @@ {%- set state.in_user = false -%} {{- '<|Assistant|>' -}} {%- set is_after_last_user = loop.index0 > last_user_idx.value -%} - {%- if is_after_last_user and thinking -%} + {%- set keep_reasoning = thinking and (has_tools or (preserve_thinking is defined and preserve_thinking) or is_after_last_user) -%} + {%- if keep_reasoning -%} {{- thinking_start_token -}} {%- if message['reasoning_content'] is defined and message['reasoning_content'] -%} {{- message['reasoning_content'] -}} @@ -95,6 +109,9 @@ {{- '<' + dsml_token + 'parameter name="' + key + '" string="false">' + (val | tojson) + '\n' -}} {%- endif -%} {%- endfor -%} + {%- if not args -%} + {{- '\n' -}} + {%- endif -%} {{- '\n' -}} {%- endfor -%} {{- '' -}} @@ -109,4 +126,4 @@ {%- else -%} {{- thinking_end_token -}} {%- endif -%} -{%- endif -%} \ No newline at end of file +{%- endif -%} diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 93685ec8ff42..0f11f97bc3c8 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -1135,7 +1135,7 @@ static void test_peg_parser(common_chat_templates * tmpls, // budget sampler inhibits grammar application while inside thinking blocks — // triggers inside ... are suppressed. bool use_reasoning_budget_path = false; - if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tag.empty()) { + if (parser.params_.grammar_lazy && !parser.params_.thinking_end_tags.empty()) { use_reasoning_budget_path = true; for (const auto & trigger : parser.params_.grammar_triggers) { if (trigger.type != COMMON_GRAMMAR_TRIGGER_TYPE_WORD) { @@ -1153,7 +1153,7 @@ static void test_peg_parser(common_chat_templates * tmpls, // Walk through full_input tracking thinking state; only match triggers // when outside thinking blocks. const auto & think_start = parser.params_.thinking_start_tag; - const auto & think_end = parser.params_.thinking_end_tag; + const auto & think_ends = parser.params_.thinking_end_tags; bool in_thinking = false; for (size_t i = 0; i < full_input.size(); ++i) { @@ -1163,12 +1163,14 @@ static void test_peg_parser(common_chat_templates * tmpls, i += think_start.size() - 1; continue; } - if (in_thinking && full_input.compare(i, think_end.size(), think_end) == 0) { - in_thinking = false; - i += think_end.size() - 1; - continue; - } if (in_thinking) { + for (const auto & think_end : think_ends) { + if (full_input.compare(i, think_end.size(), think_end) == 0) { + in_thinking = false; + i += think_end.size() - 1; + break; + } + } continue; } // Outside thinking — check if any trigger word starts here @@ -1192,7 +1194,7 @@ static void test_peg_parser(common_chat_templates * tmpls, } } - if (!use_reasoning_budget_path) { + if (!use_reasoning_budget_path && parser.params_.grammar_lazy) { // Legacy path: find triggers without thinking-awareness for (const auto & trigger : parser.params_.grammar_triggers) { size_t pos = std::string::npos; @@ -4016,6 +4018,281 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } + // DeepSeek V4 tests - format uses DSML markup: + // <|DSML|tool_calls> + // <|DSML|invoke name="foo"> + // <|DSML|parameter name="bar" string="true|false">value + // + // + { + auto tst = peg_tester("models/templates/deepseek-ai-DeepSeek-V4.jinja", detailed_debug); + + tst.test("Hello, world!\nWhat's up?") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .expect(message_assist) + .run(); + + tst.test("Let me thinkHello, world!\nWhat's up?") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .expect_reasoning("Let me think") + .expect_content("Hello, world!\nWhat's up?") + .run(); + + tst.test( + "Let me check the time\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"get_time\">\n" + "<|DSML|parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time")) + .run(); + + tst.test( + "Let me check the time\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"get_time\">\n" + "<|DSML|parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .expect_reasoning("Let me check the time") + .expect_tool_calls({ { "get_time", R"({"city": "Tokyo"})", {} } }) + .run(); + + tst.test( + "I need to output the invoice details in JSON" + R"({"amount": 123.45, "date": "2025-12-03"})") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .json_schema(invoice_schema) + .expect_reasoning("I need to output the invoice details in JSON") + .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})") + .run(); + + tst.test(R"({"amount": 123.45, "date": "2025-12-03"})") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .json_schema(invoice_schema) + .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})") + .run(); + + tst.test( + "Let me call a tool first\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"get_time\">\n" + "<|DSML|parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .json_schema(invoice_schema) + .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me call a tool first")) + .run(); + + const std::string inline_tool_markup = + "This remains content: <|DSML|tool_calls>\n" + "<|DSML|invoke name=\"get_time\">\n" + "<|DSML|parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + ""; + tst.test(inline_tool_markup) + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .expect_content(inline_tool_markup) + .run(); + + tst.test( + "Calling both\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"get_time\">\n" + "<|DSML|parameter name=\"city\" string=\"true\">Paris\n" + "\n" + "<|DSML|invoke name=\"get_weather\">\n" + "<|DSML|parameter name=\"city\" string=\"true\">Paris\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .parallel_tool_calls(true) + .tools({ get_time_tool, get_weather_tool }) + .expect(message_with_reasoning_content_and_multiple_tool_calls( + "Calling both", "", + { { "get_time", R"({"city": "Paris"})" }, { "get_weather", R"({"city": "Paris"})" } })) + .run(); + + tst.test( + "Optional first\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"magic_int\">\n" + "<|DSML|parameter name=\"name\" string=\"true\">foo bar\n" + "<|DSML|parameter name=\"ref\" string=\"false\">42\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ magic_int_tool }) + .expect_reasoning("Optional first") + .expect_tool_calls({ + { "magic_int", R"({"name": "foo bar", "ref": 42})", {} }, + }) + .run(); + + tst.test( + "Free-form string\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"get_time\">\n" + "<|DSML|parameter name=\"city\" string=\"true\">line one\n & line two\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .expect_reasoning("Free-form string") + .expect_tool_calls({ + { "get_time", "{\"city\":\"line one\\n & line two\"}", {} }, + }) + .run(); + + tst.test( + "Still thinking\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"get_time\">\n" + "<|DSML|parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .expect_reasoning("Still thinking") + .expect_tool_calls({ { "get_time", R"({"city": "Tokyo"})", {} } }) + .run(); + + tst.test("world!\nWhat's up?") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .enable_thinking(false) + .messages({ message_user, message_assist_prefill_content }) + .add_generation_prompt(false) + .continue_final_message(COMMON_CHAT_CONTINUATION_CONTENT) + .expect_content("Hello, world!\nWhat's up?") + .run(); + + { + auto tmpls = read_templates("models/templates/deepseek-ai-DeepSeek-V4.jinja"); + + common_chat_templates_inputs inputs; + inputs.messages = { message_user }; + inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; + inputs.enable_thinking = true; + inputs.tools = { magic_int_tool }; + + make_peg_parser parser(tmpls.get(), inputs, detailed_debug); + + auto expect_parse_error = [&](const std::string & output) { + bool got_error = false; + try { + parser.parse(output, false); + } catch (const std::runtime_error &) { + got_error = true; + } + + if (!got_error) { + throw std::runtime_error("Expected DeepSeek V4 parser to reject invalid parameters"); + } + }; + + expect_parse_error( + "Missing required\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"magic_int\">\n" + "<|DSML|parameter name=\"name\" string=\"true\">foo bar\n" + "\n" + ""); + + } + + tst.test( + "Duplicate required\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"magic_int\">\n" + "<|DSML|parameter name=\"ref\" string=\"false\">1\n" + "<|DSML|parameter name=\"ref\" string=\"false\">2\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ magic_int_tool }) + .expect_reasoning("Duplicate required") + .expect_tool_calls({ { "magic_int", R"({"ref": 2})", {} } }) + .run(); + + tst.test( + "Duplicate optional\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"magic_int\">\n" + "<|DSML|parameter name=\"ref\" string=\"false\">1\n" + "<|DSML|parameter name=\"name\" string=\"true\">first\n" + "<|DSML|parameter name=\"name\" string=\"true\">second\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ magic_int_tool }) + .expect_reasoning("Duplicate optional") + .expect_tool_calls({ { "magic_int", R"({"ref": 1, "name": "second"})", {} } }) + .run(); + + json many_required_schema = { + { "type", "object" }, + { "properties", json::object() }, + { "required", json::array() }, + }; + json many_required_args = json::object(); + std::string many_required_call = + "Many required\n\n" + "<|DSML|tool_calls>\n" + "<|DSML|invoke name=\"many_required\">\n"; + for (int i = 0; i < 16; i++) { + const auto param_name = "arg" + std::to_string(i); + many_required_schema["properties"][param_name] = { { "type", "integer" } }; + many_required_schema["required"].push_back(param_name); + } + for (int i = 15; i >= 0; i--) { + const auto param_name = "arg" + std::to_string(i); + many_required_args[param_name] = i; + many_required_call += + "<|DSML|parameter name=\"" + param_name + "\" string=\"false\">" + std::to_string(i) + + "\n"; + } + many_required_call += + "\n" + ""; + + common_chat_tool many_required_tool{ + /* .name = */ "many_required", + /* .description = */ "Tool with many required arguments", + /* .parameters = */ many_required_schema.dump(), + }; + tst.test(many_required_call) + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ many_required_tool }) + .expect_reasoning("Many required") + .expect_tool_calls({ { "many_required", many_required_args.dump(), {} } }) + .run(); + } + // GLM-4.6 tests - format: function_name\n...\n...\n { auto tst = peg_tester("models/templates/GLM-4.6.jinja", detailed_debug); @@ -5743,6 +6020,7 @@ static void test_template_generation_prompt() { std::vector messages; bool add_generation_prompt = true; common_chat_continuation continue_final_message = COMMON_CHAT_CONTINUATION_NONE; + bool enable_thinking = true; }; auto basic = [&]() { @@ -5774,6 +6052,7 @@ static void test_template_generation_prompt() { inputs.messages = opts.messages; inputs.add_generation_prompt = opts.add_generation_prompt; inputs.continue_final_message = opts.continue_final_message; + inputs.enable_thinking = opts.enable_thinking; auto params = common_chat_templates_apply(tmpls.get(), inputs); @@ -5872,6 +6151,140 @@ static void test_template_generation_prompt() { check(tmpls, continuation_reasoning(), "<|Assistant|>I'm"); } + { + auto tmpls = read_templates("models/templates/deepseek-ai-DeepSeek-V4.jinja"); + check(tmpls, basic(), "<|Assistant|>"); + check(tmpls, continuation_content(), "<|Assistant|>I'm thinkingHello, "); + check(tmpls, continuation_reasoning(), "<|Assistant|>I'm"); + + auto continuation_content_no_thinking = continuation_content(); + continuation_content_no_thinking.enable_thinking = false; + check(tmpls, continuation_content_no_thinking, "<|Assistant|>Hello, "); + + const std::string reasoning_effort_max = + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" + "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n" + "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"; + + common_chat_templates_inputs max_inputs; + max_inputs.messages = { system_msg, message_user }; + max_inputs.chat_template_kwargs["reasoning_effort"] = R"("max")"; + auto max_params = common_chat_templates_apply(tmpls.get(), max_inputs); + assert_contains(max_params.prompt, reasoning_effort_max + system_msg.content); + + auto high_inputs = max_inputs; + high_inputs.chat_template_kwargs["reasoning_effort"] = R"("high")"; + auto high_params = common_chat_templates_apply(tmpls.get(), high_inputs); + assert_equals(std::string::npos, high_params.prompt.find(reasoning_effort_max)); + + auto non_thinking_max_inputs = max_inputs; + non_thinking_max_inputs.enable_thinking = false; + auto non_thinking_max_params = common_chat_templates_apply(tmpls.get(), non_thinking_max_inputs); + assert_equals(std::string::npos, non_thinking_max_params.prompt.find(reasoning_effort_max)); + + common_chat_templates_inputs response_format_inputs; + response_format_inputs.messages = { system_msg, message_user }; + response_format_inputs.tools = { get_time_tool }; + response_format_inputs.json_schema = + R"({"type":"object","properties":{"answer":{"type":"string"}}})"; + auto response_format_params = common_chat_templates_apply(tmpls.get(), response_format_inputs); + if (response_format_params.grammar_lazy || + response_format_params.grammar_triggers.size() != 1 || + response_format_params.grammar_triggers[0].value != "\n\n<|DSML|tool_calls>" || + response_format_params.thinking_end_tags != std::vector{ + "", "\n\n<|DSML|tool_calls>" }) { + throw std::runtime_error("Unexpected DeepSeek V4 grammar or reasoning end configuration"); + } + const auto tools_pos = response_format_params.prompt.find("## Tools"); + const auto response_format_pos = response_format_params.prompt.find( + "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n"); + if (tools_pos == std::string::npos || response_format_pos == std::string::npos || tools_pos > response_format_pos) { + LOG_ERR("Expected response format after tools\nActual: %s\n", response_format_params.prompt.c_str()); + common_log_flush(common_log_main()); + throw std::runtime_error("Test failed"); + } + assert_contains(response_format_params.prompt, R"("answer": {"type": "string"})"); + + response_format_inputs.json_schema = "{}"; + auto json_object_params = common_chat_templates_apply(tmpls.get(), response_format_inputs); + assert_contains(json_object_params.prompt, + "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{}"); + + common_chat_msg assistant_history; + assistant_history.role = "assistant"; + assistant_history.content = "Previous answer"; + assistant_history.reasoning_content = "Previous reasoning"; + + common_chat_msg user_followup; + user_followup.role = "user"; + user_followup.content = "Follow up"; + + common_chat_templates_inputs default_history_inputs; + default_history_inputs.messages = { message_user, assistant_history, user_followup }; + auto default_history_params = common_chat_templates_apply(tmpls.get(), default_history_inputs); + assert_contains(default_history_params.prompt, "<|Assistant|>Previous answer"); + + auto preserved_history_inputs = default_history_inputs; + preserved_history_inputs.chat_template_kwargs["preserve_thinking"] = "true"; + auto preserved_history_params = common_chat_templates_apply(tmpls.get(), preserved_history_inputs); + assert_contains(preserved_history_params.prompt, "<|Assistant|>Previous reasoningPrevious answer"); + + auto preserve_reasoning_inputs = default_history_inputs; + preserve_reasoning_inputs.chat_template_kwargs["preserve_reasoning"] = "true"; + auto preserve_reasoning_params = common_chat_templates_apply(tmpls.get(), preserve_reasoning_inputs); + assert_contains(preserve_reasoning_params.prompt, "<|Assistant|>Previous reasoningPrevious answer"); + assert_equals(true, common_chat_templates_get_caps(tmpls.get()).at("supports_preserve_reasoning")); + + common_chat_msg empty_tool_call = simple_assist_msg("", "", "empty_args", "{}"); + common_chat_templates_inputs empty_tool_inputs; + empty_tool_inputs.messages = { message_user, empty_tool_call }; + empty_tool_inputs.tools = { empty_args_tool }; + auto empty_tool_params = common_chat_templates_apply(tmpls.get(), empty_tool_inputs); + assert_contains(empty_tool_params.prompt, + "<|DSML|invoke name=\"empty_args\">\n\n"); + + common_chat_msg user_start; + user_start.role = "user"; + user_start.content = "Check time and weather."; + + common_chat_msg assistant_tools; + assistant_tools.role = "assistant"; + assistant_tools.reasoning_content = "Need both"; + assistant_tools.tool_calls = { + { "get_time", R"({"city":"Paris"})", "call_time" }, + { "get_weather", R"({"city":"Paris"})", "call_weather" }, + }; + + common_chat_msg weather_result; + weather_result.role = "tool"; + weather_result.content = "weather result"; + weather_result.tool_call_id = "call_weather"; + + common_chat_msg time_result; + time_result.role = "tool"; + time_result.content = "time result"; + time_result.tool_call_id = "call_time"; + + common_chat_msg user_continue; + user_continue.role = "user"; + user_continue.content = "Continue."; + + common_chat_templates_inputs inputs; + inputs.messages = { user_start, assistant_tools, weather_result, time_result, user_continue }; + inputs.tools = { get_time_tool, get_weather_tool }; + + auto params = common_chat_templates_apply(tmpls.get(), inputs); + assert_contains(params.prompt, "<|Assistant|>Need both\n\n<|DSML|tool_calls>"); + + const auto weather_pos = params.prompt.find("weather result"); + const auto time_pos = params.prompt.find("time result"); + if (weather_pos == std::string::npos || time_pos == std::string::npos || weather_pos > time_pos) { + LOG_ERR("Expected tool results to preserve message order\nActual: %s\n", params.prompt.c_str()); + common_log_flush(common_log_main()); + throw std::runtime_error("Test failed"); + } + } + { auto tmpls = read_templates("models/templates/openbmb-MiniCPM5-1B.jinja"); check(tmpls, basic(), "<|im_start|>assistant\n\n"); diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index f54cff4f8a27..3bcc77e1733c 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -20,8 +20,8 @@ static void test_reasoning_budget( const char * test_name, const std::vector & sequence, - const std::vector & start_tokens, - const std::vector & end_tokens, + const std::vector & start_seqs, + const std::vector & end_seqs, const std::vector & forced_tokens, int32_t budget, common_reasoning_budget_state initial_state, @@ -31,8 +31,12 @@ static void test_reasoning_budget( // Find the maximum token ID to ensure our vocab covers all tokens llama_token max_token = 0; for (auto t : sequence) max_token = std::max(max_token, t); - for (auto t : start_tokens) max_token = std::max(max_token, t); - for (auto t : end_tokens) max_token = std::max(max_token, t); + for (const auto & seq : start_seqs) { + for (auto t : seq) max_token = std::max(max_token, t); + } + for (const auto & seq : end_seqs) { + for (auto t : seq) max_token = std::max(max_token, t); + } for (auto t : forced_tokens) max_token = std::max(max_token, t); // Create a minimal sampler with mock vocabulary @@ -40,8 +44,8 @@ static void test_reasoning_budget( // The UTF-8 boundary check will treat all tokens as complete (safe fallback) auto * sampler = common_reasoning_budget_init( nullptr, // vocab - not used for basic state machine tests - start_tokens, - end_tokens, + start_seqs, + end_seqs, forced_tokens, budget, initial_state @@ -152,7 +156,7 @@ static void test_reasoning_budget_clone_mid_counting() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 2, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 2, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=2 llama_sampler_accept(sampler, 50); // COUNTING, remaining=1 @@ -171,7 +175,7 @@ static void test_reasoning_budget_clone_mid_forcing() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token @@ -191,7 +195,7 @@ static void test_reasoning_budget_force_manual() { // if COUNTING, force() succeeds and begins forcing the end sequence from the start { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=5 llama_sampler_accept(sampler, 50); // COUNTING, remaining=4 @@ -212,7 +216,7 @@ static void test_reasoning_budget_force_manual() { // if IDLE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); GGML_ASSERT(!common_reasoning_budget_force(sampler) && "force() must not transition from IDLE"); GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_IDLE); @@ -222,7 +226,7 @@ static void test_reasoning_budget_force_manual() { // if DONE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING llama_sampler_accept(sampler, 101); // natural end -> DONE @@ -236,7 +240,7 @@ static void test_reasoning_budget_force_manual() { // if FORCING, force() is a no-op and must not rewind the force position { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token (force_pos=1) @@ -254,6 +258,81 @@ static void test_reasoning_budget_force_manual() { fprintf(stderr, " Test 'manual force transition' passed\n"); } +static void test_reasoning_budget_end_match() { + const std::vector start = {{100}}; + const std::vector end = {{101}, {103, 104}}; + + // natural end records the sequence that matched; re-arming clears it + { + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 101}, 5, REASONING_BUDGET_IDLE); + + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); + + llama_sampler_accept(sampler, 100); // COUNTING + llama_sampler_accept(sampler, 50); + llama_sampler_accept(sampler, 103); + llama_sampler_accept(sampler, 104); // end matched via {103, 104}, DONE + + const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler); + GGML_ASSERT(matched != nullptr); + GGML_ASSERT(*matched == llama_tokens({103, 104})); + + llama_sampler_accept(sampler, 100); // re-arm, COUNTING + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); + + llama_sampler_free(sampler); + } + + // overlapping end sequences: the longest one ending at the position wins + { + const std::vector end_overlap = {{104}, {103, 104}}; + + auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, {102, 104}, 5, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); // COUNTING + llama_sampler_accept(sampler, 103); + llama_sampler_accept(sampler, 104); // both {104} and {103, 104} end here + + const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler); + GGML_ASSERT(matched != nullptr); + GGML_ASSERT(*matched == llama_tokens({103, 104})); + + llama_sampler_free(sampler); + } + + // forcing records the end sequence terminating forced_tokens + { + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 103, 104}, 0, REASONING_BUDGET_FORCING); + + llama_sampler_accept(sampler, 102); + llama_sampler_accept(sampler, 103); + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); + llama_sampler_accept(sampler, 104); // forced sequence complete, DONE + + const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler); + GGML_ASSERT(matched != nullptr); + GGML_ASSERT(*matched == llama_tokens({103, 104})); + + llama_sampler_free(sampler); + } + + // forced_tokens not ending with a known end sequence records nothing + { + auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102}, 0, REASONING_BUDGET_FORCING); + + llama_sampler_accept(sampler, 102); // forced sequence complete, DONE + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); + + llama_sampler_free(sampler); + } + + // a null sampler is safely ignored + GGML_ASSERT(common_reasoning_budget_get_end_match(nullptr) == nullptr); + + fprintf(stderr, " Test 'matched end sequence' passed\n"); +} + // UTF-8 boundary detection unit test // Tests common_utf8_is_complete() from reasoning-budget.h static void test_utf8_boundary_detection() { @@ -290,7 +369,7 @@ int main(void) { const std::vector forced = {102}; // forced token (not used in this test) const std::vector sequence = {100, 50, 51, 101, 52}; // start, two tokens, end, one more - test_reasoning_budget("natural end before budget exhausted", sequence, start, end, forced, + test_reasoning_budget("natural end before budget exhausted", sequence, {start}, {end}, forced, 5, // budget of 5 tokens REASONING_BUDGET_IDLE, SIZE_MAX, SIZE_MAX); // no forcing expected (natural end) @@ -306,7 +385,7 @@ int main(void) { const std::vector forced = {102, 101}; // forced message + end const std::vector sequence = {100, 50, 51, 52, 53}; // start + 4 tokens (budget=2) - test_reasoning_budget("budget exhausted forcing", sequence, start, end, forced, + test_reasoning_budget("budget exhausted forcing", sequence, {start}, {end}, forced, 2, // budget of 2 tokens REASONING_BUDGET_IDLE, 3, // forcing starts at i=3 (accept at i=2 depletes budget, apply at i=3 forces) @@ -321,7 +400,7 @@ int main(void) { const std::vector forced = {102, 101}; const std::vector sequence = {100, 50, 51, 52}; // start token first, then 3 tokens - test_reasoning_budget("activate immediately budget=0", sequence, start, end, forced, + test_reasoning_budget("activate immediately budget=0", sequence, {start}, {end}, forced, 0, // budget of 0 tokens REASONING_BUDGET_COUNTING, // starts counting, promoted to FORCING since budget=0 0, // forcing starts at i=0 (initialized in FORCING, apply forces immediately) @@ -335,7 +414,7 @@ int main(void) { const std::vector forced = {102}; const std::vector sequence = {50, 51, 52, 53}; - test_reasoning_budget("no start/end configured", sequence, start, end, forced, + test_reasoning_budget("no start/end configured", sequence, {start}, {end}, forced, 2, // budget REASONING_BUDGET_IDLE, SIZE_MAX, SIZE_MAX); // no forcing (no start/end configured) @@ -350,7 +429,7 @@ int main(void) { const std::vector forced = {102, 101}; const std::vector sequence = {50, 51, 52, 53}; - test_reasoning_budget("activate immediately with budget", sequence, start, end, forced, + test_reasoning_budget("activate immediately with budget", sequence, {start}, {end}, forced, 2, // budget of 2 tokens REASONING_BUDGET_COUNTING, 2, // forcing starts at i=2 (after 2 accepts deplete budget, apply at i=2 forces) @@ -373,18 +452,50 @@ int main(void) { const std::vector forced = {102, 101}; const std::vector sequence = {100, 50, 101, 100, 60, 61, 62, 63}; - test_reasoning_budget("multi-block re-arms budget after DONE", sequence, start, end, forced, + test_reasoning_budget("multi-block re-arms budget after DONE", sequence, {start}, {end}, forced, 2, // budget of 2 tokens (per block) REASONING_BUDGET_IDLE, 6, // forcing starts at i=6 (after second block exhausts at i=5) 7); // forcing continues through i=7 } + // Test 7: Multiple start sequences - the second sequence activates counting + // Flow: i=0 accept(110), i=1 accept(111)->COUNTING rem=2; i=2 accept(50)->rem=1; + // i=3 accept(51)->rem=0->FORCING; i=4..5 apply() forces the end sequence + { + const std::vector start = {{100}, {110, 111}}; + const std::vector end = {{101}}; + const std::vector forced = {102, 101}; + const std::vector sequence = {110, 111, 50, 51, 52, 53}; + + test_reasoning_budget("multiple start sequences", sequence, start, end, forced, + 2, // budget of 2 tokens + REASONING_BUDGET_IDLE, + 4, // forcing starts at i=4 (accept at i=3 depletes budget) + 5); // forcing continues through i=5 + } + + // Test 8: Multiple end sequences - natural end via the second sequence + // Flow: i=0 accept(100)->COUNTING rem=5; i=1 accept(50)->rem=4; + // i=2 accept(103)->partial end, rem=3; i=3 accept(104)->end matched, DONE + { + const std::vector start = {{100}}; + const std::vector end = {{101}, {103, 104}}; + const std::vector forced = {102, 101}; + const std::vector sequence = {100, 50, 103, 104, 52}; + + test_reasoning_budget("multiple end sequences", sequence, start, end, forced, + 5, // budget of 5 tokens + REASONING_BUDGET_IDLE, + SIZE_MAX, SIZE_MAX); // no forcing expected (natural end) + } + test_reasoning_budget_clone_mid_counting(); test_reasoning_budget_clone_mid_forcing(); test_reasoning_budget_force_manual(); + test_reasoning_budget_end_match(); - printf("OK (9 tests passed)\n"); + printf("OK (12 tests passed)\n"); printf("Testing UTF-8 boundary detection... "); test_utf8_boundary_detection(); diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 78b5c6819b03..0ba22cb27391 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -1123,10 +1123,10 @@ json oaicompat_chat_params_parse( reasoning_budget = opt.reasoning_budget; } - if (!chat_params.thinking_end_tag.empty()) { + if (!chat_params.thinking_end_tags.empty()) { llama_params["reasoning_budget_tokens"] = reasoning_budget; llama_params["reasoning_budget_start_tag"] = chat_params.thinking_start_tag; - llama_params["reasoning_budget_end_tag"] = chat_params.thinking_end_tag; + llama_params["reasoning_budget_end_tags"] = chat_params.thinking_end_tags; llama_params["reasoning_budget_message"] = json_value(body, "reasoning_budget_message", opt.reasoning_budget_message); llama_params["reasoning_control"] = json_value(body, "reasoning_control", false); } diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 89026eb4e3f0..e880f4ca728d 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -390,21 +390,40 @@ std::vector> make_llama_cmpl_schema(const common_params & ctx.params.sampling.reasoning_budget_start = common_tokenize(ctx.vocab, data.at("reasoning_budget_start_tag").get(), false, true); })); - add((new field_str("reasoning_budget_end_tag")) - ->set_desc("Token string marking the end of the reasoning budget section") + add((new field_json("reasoning_budget_end_tags")) + ->add_alias("reasoning_budget_end_tag") + ->set_desc("Token strings marking the end of the reasoning budget section; the first is forced when the budget expires") ->set_handler([&](field_eval_context & ctx, const json & data) { GGML_ASSERT(ctx.vocab != nullptr); - std::string end_tag = data.at("reasoning_budget_end_tag").get(); - ctx.params.sampling.reasoning_budget_end = common_tokenize(ctx.vocab, end_tag, false, true); + ctx.params.sampling.reasoning_budget_end.clear(); + if (data.contains("reasoning_budget_end_tags")) { + for (const auto & t : data.at("reasoning_budget_end_tags")) { + std::string tag = t.get(); + if (!tag.empty()) { + ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true)); + } + } + } else if (data.contains("reasoning_budget_end_tag")) { + std::string tag = data.at("reasoning_budget_end_tag").get(); + if (!tag.empty()) { + ctx.params.sampling.reasoning_budget_end.push_back(common_tokenize(ctx.vocab, tag, false, true)); + } + } })); add((new field_str("reasoning_budget_message")) ->set_desc("Message to prepend to the reasoning budget end tag when forcing it") ->set_handler([&](field_eval_context & ctx, const json & data) { GGML_ASSERT(ctx.vocab != nullptr); - std::string end_tag = json_value(data, "reasoning_budget_end_tag", std::string()); - std::string message = data.at("reasoning_budget_message").get(); - ctx.params.sampling.reasoning_budget_forced = common_tokenize(ctx.vocab, message + end_tag, false, true); + if (!ctx.params.sampling.reasoning_budget_end.empty()) { + llama_tokens end_tag = ctx.params.sampling.reasoning_budget_end.front(); + std::string message = json_value(data, "reasoning_budget_message", std::string()); + if (!message.empty()) { + llama_tokens message_tokens = common_tokenize(ctx.vocab, message, false, true); + end_tag.insert(end_tag.begin(), message_tokens.begin(), message_tokens.end()); + } + ctx.params.sampling.reasoning_budget_forced = std::move(end_tag); + } })); add((new field_json("logit_bias")) @@ -546,7 +565,7 @@ task_params eval_llama_cmpl_schema( // debugging { auto budget = params.sampling.reasoning_budget_tokens; - SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu toks, forced=%zu toks\n", + SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu seqs, forced=%zu toks\n", budget, params.sampling.generation_prompt.c_str(), params.sampling.reasoning_budget_start.size(), params.sampling.reasoning_budget_end.size(),