diff --git a/common/chat-auto-parser-generator.cpp b/common/chat-auto-parser-generator.cpp index d7e117e4d98b..dd89f4ac5e34 100644 --- a/common/chat-auto-parser-generator.cpp +++ b/common/chat-auto-parser-generator.cpp @@ -42,7 +42,7 @@ common_chat_params peg_generator::generate_parser(const common_chat_template & const autoparser & autoparser) { // Create the result structure common_chat_params data; - data.prompt = common_chat_template_direct_apply(tmpl, inputs); + data.prompt = common_chat_template_direct_apply(tmpl, inputs, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.preserved_tokens = autoparser.preserved_tokens; @@ -52,11 +52,27 @@ common_chat_params peg_generator::generate_parser(const common_chat_template & std::string parser_generation_prompt = data.generation_prompt; if (inputs.continue_final_message != COMMON_CHAT_CONTINUATION_NONE && !inputs.continue_msg.empty()) { - // Build up generation prompt manually + // Build up generation prompt manually, keeping the provenance of the + // template delimiters and the request-provided continuation content + // separate (the latter must not be parsed for special tokens). const auto & msg = inputs.continue_msg; + std::vector gen_parts; + if (!autoparser.reasoning.start.empty()) { - data.generation_prompt = data.generation_prompt.substr(0, data.generation_prompt.find(autoparser.reasoning.start)); + const size_t cut = data.generation_prompt.find(autoparser.reasoning.start); + + // Template-derived prefix of the generation prompt + gen_parts.push_back({false, data.generation_prompt.substr(0, cut)}); + // Reasoning markers are template text + gen_parts.push_back({false, autoparser.reasoning.start}); + // Request-provided reasoning content + gen_parts.push_back({true, msg.reasoning_content}); + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + gen_parts.push_back({false, autoparser.reasoning.end}); + } + + data.generation_prompt = data.generation_prompt.substr(0, cut); data.generation_prompt += autoparser.reasoning.start + msg.reasoning_content; if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += autoparser.reasoning.end; @@ -64,10 +80,19 @@ common_chat_params peg_generator::generate_parser(const common_chat_template & } if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + // Request-provided message content + gen_parts.push_back({true, msg.render_content()}); data.generation_prompt += msg.render_content(); } + if (gen_parts.empty()) { + // Reasoning-only continuation without a reasoning start marker: + // nothing request-provided was appended, the prompt is template text + gen_parts.push_back({false, data.generation_prompt}); + } + data.prompt += data.generation_prompt; + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto parser = autoparser.build_parser(inputs, parser_generation_prompt); diff --git a/common/chat.cpp b/common/chat.cpp index 743ecde0a77e..43e02c0a44ba 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -928,12 +928,92 @@ static void foreach_parameter(const json & } } +// The prompt parts mirror the rendered prompt string (their concatenation +// equals the prompt). These helpers edit the parts consistently with the +// prompt, even when the edited text spans two adjacent parts. + +// Erase [pos, pos + len) from the concatenation of parts. +// Returns false if the range is not fully covered by the parts. +static bool string_parts_erase_range(std::vector & parts, size_t pos, size_t len) { + size_t total = 0; + for (const auto & part : parts) { + total += part.val.size(); + } + if (pos + len > total) { + return false; + } + size_t off = 0; + size_t to_drop = len; + for (auto & part : parts) { + const size_t sz = part.val.size(); + if (off + sz > pos && to_drop > 0) { + const size_t skip = pos > off ? pos - off : 0; + const size_t drop = std::min(to_drop, sz - skip); + part.val.erase(skip, drop); + to_drop -= drop; + } + off += sz; + } + return to_drop == 0; +} + +// Replace the occurrence of "from" located at global position pos in the +// concatenation of parts with "to". The occurrence may span two adjacent +// parts; the replacement text is placed in the part where the occurrence +// starts, preserving its is_input provenance. +// Returns false if the parts do not contain "from" at pos (i.e. they no +// longer mirror the prompt they were derived from). +static bool string_parts_replace_at( + std::vector & parts, + size_t pos, + const std::string & from, + const std::string & to) { + std::string actual; + { + size_t off = 0; + for (const auto & part : parts) { + if (actual.size() >= from.size()) { + break; + } + const size_t sz = part.val.size(); + if (off <= pos && off + sz > pos) { + const size_t skip = pos - off; + actual += part.val.substr(skip, from.size() - actual.size()); + } + off += sz; + } + } + if (actual != from) { + return false; + } + if (!string_parts_erase_range(parts, pos, from.size())) { + return false; + } + size_t off = 0; + for (auto & part : parts) { + const size_t sz = part.val.size(); + if (off == pos) { + part.val.insert(0, to); + return true; + } + if (off < pos && pos < off + sz) { + part.val.insert(pos - off, to); + return true; + } + off += sz; + } + // pos is at the very end of the concatenation: append to the last part + parts.back().val += to; + return true; +} + static std::string common_chat_template_direct_apply_impl( const common_chat_template & tmpl, const autoparser::generation_params & inputs, const std::optional & messages_override = std::nullopt, const std::optional & tools_override = std::nullopt, - const std::optional & additional_context = std::nullopt) { + const std::optional & additional_context = std::nullopt, + std::vector * out_parts = nullptr) { jinja::context ctx(tmpl.source()); // messages_override is already built for this template, do not touch its content parts @@ -979,22 +1059,133 @@ static std::string common_chat_template_direct_apply_impl( const jinja::value results = runtime.execute(tmpl.prog); auto parts = jinja::runtime::gather_string_parts(results); + // Preserve the jinja::string parts (with is_input metadata) for the caller + if (out_parts) { + *out_parts = parts->as_string().parts; + } + std::string result = parts->as_string().str(); // TODO: improve this later if (inputs.add_bos && string_starts_with(result, tmpl.bos_token())) { result = result.substr(tmpl.bos_token().size()); + // Keep the parts in sync with the prompt (the token may span parts). + if (out_parts) { + string_parts_erase_range(*out_parts, 0, tmpl.bos_token().size()); + } } if (inputs.add_eos && string_ends_with(result, tmpl.eos_token())) { result = result.substr(0, result.size() - tmpl.eos_token().size()); + // Keep the parts in sync with the prompt (the token may span parts). + if (out_parts) { + size_t total = 0; + for (const auto & part : *out_parts) { + total += part.val.size(); + } + string_parts_erase_range(*out_parts, total - tmpl.eos_token().size(), tmpl.eos_token().size()); + } } return result; } std::string common_chat_template_direct_apply( const common_chat_template & tmpl, - const autoparser::generation_params & inputs) { - return common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt); + const autoparser::generation_params & inputs, + std::vector * out_parts) { + return common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, out_parts); +} + +bool common_chat_parts_have_special_input( + const struct llama_vocab * vocab, + const std::vector & parts) { + for (const auto & part : parts) { + if (!part.is_input || part.val.empty()) { + continue; + } + // Tokenize the input part in isolation with parse_special=true: if any + // of the resulting tokens is a control/unknown special token, the + // per-part parse_special handling changes the result (those tokens + // would be parsed as special in a whole-prompt parse_special=true pass, + // but are byte-fallbacked in the per-part parse_special=false pass). + // + // Note: this mirrors the tokenizer's own rule (see + // llama_vocab::impl::tokenizer_st_partition): with parse_special=false, + // special tokens with the control or unknown attribute are not parsed. + // + // The check is done per part, in isolation. A special token that spans + // a part boundary would require the template to emit a partial special + // token exactly at a user-content boundary, which well-formed chat + // templates do not do. + const auto tokens = common_tokenize(vocab, part.val, /*add_special=*/false, /*parse_special=*/true); + for (const auto tok : tokens) { + const auto attr = llama_vocab_get_attr(vocab, tok); + if (attr & (LLAMA_TOKEN_ATTR_CONTROL | LLAMA_TOKEN_ATTR_UNKNOWN)) { + return true; + } + } + } + return false; +} + +std::vector common_tokenize_parts( + const struct llama_vocab * vocab, + const std::vector & parts, + bool add_special) { + if (!common_chat_parts_have_special_input(vocab, parts)) { + // Fast path: no is_input part contains special-token text, so the + // per-part parse_special distinction is moot. Tokenize the + // concatenated prompt in a single pass so the token ids are identical + // to the legacy whole-prompt tokenization (this preserves normal + // tokenizer merges across part boundaries). + std::string full; + size_t total = 0; + for (const auto & part : parts) { + total += part.val.size(); + } + full.reserve(total); + for (const auto & part : parts) { + full += part.val; + } + return common_tokenize(vocab, full, add_special, /*parse_special=*/true); + } + + // Protection path: tokenize each part separately so that user-provided + // content (is_input) is never parsed for special tokens, while template + // parts keep parse_special=true so legitimate special tokens like + // <|im_start|>, <|im_end|>, etc. are properly recognized. + // + // Merge adjacent parts with the same provenance first: the rendered + // parts from the jinja runtime never contain adjacent parts of the same + // type (it merges them), but parts appended afterwards (e.g. the + // continuation generation prompt) can. Merging keeps normal tokenizer + // merges across same-type part boundaries. + std::vector merged; + merged.reserve(parts.size()); + for (const auto & part : parts) { + if (!merged.empty() && merged.back().is_input == part.is_input) { + merged.back().val += part.val; + } else { + merged.push_back(part); + } + } + + std::vector result; + + bool first = true; + for (const auto & part : merged) { + if (part.val.empty()) { + continue; + } + + // Only add special (BOS) on the very first non-empty part + const bool part_add_special = add_special && first; + first = false; + + const auto tokens = common_tokenize(vocab, part.val, part_add_special, /*parse_special=*/!part.is_input); + result.insert(result.end(), tokens.begin(), tokens.end()); + } + + return result; } static std::string common_chat_template_generation_prompt_impl( @@ -1076,7 +1267,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_tags = {"[/THINK]"}; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override = */ adjusted_messages, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override = */ adjusted_messages); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.preserved_tokens = { @@ -1088,13 +1279,22 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_ if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; data.generation_prompt = "[THINK]" + msg.reasoning_content; + gen_parts.push_back({false, "[THINK]"}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += "[/THINK]" + msg.render_content(); + gen_parts.push_back({false, "[/THINK]"}); + gen_parts.push_back({true, msg.render_content()}); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -1165,7 +1365,7 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_ const std::string GEN_PREFIX = "<|im_start|>assistant\n"; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; @@ -1204,17 +1404,27 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_ const auto & msg = inputs.continue_msg; data.generation_prompt = GEN_PREFIX; + std::vector gen_parts; + gen_parts.push_back({false, GEN_PREFIX}); if (supports_reasoning) { + // The delimiters are template text; the continued reasoning/content + // is request-provided and must not be parsed for special tokens. + gen_parts.push_back({false, "\n"}); + gen_parts.push_back({true, msg.reasoning_content}); data.generation_prompt += "\n" + msg.reasoning_content; if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + gen_parts.push_back({false, "\n\n\n"}); data.generation_prompt += "\n\n\n"; } } if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + gen_parts.push_back({true, msg.render_content()}); data.generation_prompt += msg.render_content(); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } std::vector tool_call_starts = { "" }; @@ -1355,15 +1565,23 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp adjusted_messages.push_back(msg); } - auto prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override= */ adjusted_messages); + auto prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override= */ adjusted_messages, std::nullopt, std::nullopt, &data.prompt_parts); // Check if we need to replace the return token with end token during // inference and without generation prompt. For more details see: // https://github.com/ggml-org/llama.cpp/issues/15417 + static constexpr std::string_view return_token = "<|return|>"; + static constexpr std::string_view end_token = "<|end|>"; if (inputs.is_inference && !inputs.add_generation_prompt) { - static constexpr std::string_view return_token = "<|return|>"; - static constexpr std::string_view end_token = "<|end|>"; if (size_t pos = prompt.rfind(return_token); pos != std::string::npos) { + // Apply the same replacement to the prompt parts so that the + // server-side tokenization of prompt_parts sees exactly the same + // text as data.prompt (in general, the token may span two + // adjacent parts). + if (!string_parts_replace_at(data.prompt_parts, pos, std::string(return_token), std::string(end_token))) { + LOG_WRN("%s: prompt parts do not mirror the prompt, the return->end replacement was not applied to them\n", + __func__); + } prompt.replace(pos, return_token.length(), end_token); } } @@ -1394,12 +1612,24 @@ static common_chat_params common_chat_params_init_gpt_oss(const common_chat_temp if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance of the continuation separate: the channel + // delimiters are template text, while the reasoning/content that + // comes from the request must not be parsed for special tokens. + std::vector gen_parts = { + {false, "<|start|>assistant<|channel|>analysis<|message|>"}, + {true, msg.reasoning_content}, + }; + data.generation_prompt = "<|start|>assistant<|channel|>analysis<|message|>" + msg.reasoning_content; if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + gen_parts.push_back({false, "<|end|><|start|>assistant<|channel|>final<|message|>"}); + gen_parts.push_back({true, msg.render_content()}); data.generation_prompt += "<|end|><|start|>assistant<|channel|>final<|message|>" + msg.render_content(); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); @@ -1509,7 +1739,7 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); if (inputs.add_generation_prompt && string_ends_with(data.prompt, "\n")) { @@ -1518,6 +1748,8 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ // from emitting its proper reasoning token sequence. data.generation_prompt = "<|turn>model\n"; data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, data.generation_prompt}); } data.message_delimiters = { @@ -1540,14 +1772,25 @@ static common_chat_params common_chat_params_init_gemma4(const common_chat_templ if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; - data.generation_prompt = string_ends_with(data.prompt, "\n") ? "<|turn>model\n" : ""; + const std::string turn_prefix = string_ends_with(data.prompt, "\n") ? "<|turn>model\n" : ""; + data.generation_prompt = turn_prefix; + gen_parts.push_back({false, turn_prefix}); data.generation_prompt += "<|channel>thought\n" + msg.reasoning_content; + gen_parts.push_back({false, "<|channel>thought\n"}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += "" + msg.render_content(); + gen_parts.push_back({false, ""}); + gen_parts.push_back({true, msg.render_content()}); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); @@ -1673,7 +1916,7 @@ static common_chat_params common_chat_params_init_functionary_v3_2(const common_ const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.preserved_tokens = { @@ -1685,8 +1928,14 @@ static common_chat_params common_chat_params_init_functionary_v3_2(const common_ if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + + // The header is template text; the continued content is + // request-provided and must not be parsed for special tokens. data.generation_prompt = "<|start_header_id|>assistant<|end_header_id|>\n\n>>>all\n" + msg.render_content(); data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, "<|start_header_id|>assistant<|end_header_id|>\n\n>>>all\n"}); + data.prompt_parts.push_back({true, msg.render_content()}); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -1774,7 +2023,7 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -1807,13 +2056,22 @@ static common_chat_params common_chat_params_init_kimi_k2(const common_chat_temp if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; data.generation_prompt = GEN_PROMPT + THINK_START + msg.reasoning_content; + gen_parts.push_back({false, GEN_PROMPT + THINK_START}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += THINK_END + msg.render_content(); + gen_parts.push_back({false, THINK_END}); + gen_parts.push_back({true, msg.render_content()}); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -1926,7 +2184,7 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat adjusted_messages.push_back(msg); } - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, adjusted_messages); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, adjusted_messages, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, adjusted_messages); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -1948,13 +2206,22 @@ static common_chat_params common_chat_params_init_lfm2(const common_chat_templat if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; data.generation_prompt = GEN_PROMPT + THINK_START + msg.reasoning_content; + gen_parts.push_back({false, GEN_PROMPT + THINK_START}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += THINK_END + msg.render_content(); + gen_parts.push_back({false, THINK_END}); + gen_parts.push_back({true, msg.render_content()}); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -2017,7 +2284,7 @@ static common_chat_params common_chat_params_init_gigachat_v3( common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = false; @@ -2028,8 +2295,14 @@ static common_chat_params common_chat_params_init_gigachat_v3( if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + + // The role header is template text; the continued content is + // request-provided and must not be parsed for special tokens. data.generation_prompt = "assistant<|role_sep|>\n" + msg.render_content(); data.prompt += data.generation_prompt; + + data.prompt_parts.push_back({false, "assistant<|role_sep|>\n"}); + data.prompt_parts.push_back({true, msg.render_content()}); } auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); @@ -2196,7 +2469,7 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha const std::string TC_SEPARATOR = "\n\n"; data.prompt = common_chat_template_direct_apply_impl( - tmpl, inputs, adjusted_messages, std::nullopt, additional_context); + tmpl, inputs, adjusted_messages, std::nullopt, additional_context, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl( tmpl, inputs, adjusted_messages, std::nullopt, additional_context); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; @@ -2211,20 +2484,31 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; if (is_v4 && msg.reasoning_content.empty()) { data.generation_prompt = GEN_PROMPT + THINK_END; + gen_parts.push_back({false, GEN_PROMPT + THINK_END}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += msg.render_content(); + gen_parts.push_back({true, msg.render_content()}); } } else { data.generation_prompt = GEN_PROMPT + THINK_START + msg.reasoning_content; + gen_parts.push_back({false, GEN_PROMPT + THINK_START}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += THINK_END + msg.render_content(); + gen_parts.push_back({false, THINK_END}); + gen_parts.push_back({true, msg.render_content()}); } } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } bool require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; @@ -2386,7 +2670,7 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -2434,13 +2718,22 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; data.generation_prompt = MSG_START + THINK_START + msg.reasoning_content; + gen_parts.push_back({false, MSG_START + THINK_START}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += THINK_END + RESP_START + msg.render_content(); + gen_parts.push_back({false, THINK_END + RESP_START}); + gen_parts.push_back({true, msg.render_content()}); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -2587,7 +2880,7 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t // Stable prefix of the generation prompt that precedes the (forced) <|START_THINKING|> marker. const std::string GEN_PREFIX = TURN_START + CHATBOT; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -2618,13 +2911,22 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; data.generation_prompt = GEN_PREFIX + THINK_START + msg.reasoning_content; + gen_parts.push_back({false, GEN_PREFIX + THINK_START}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += THINK_END + TEXT_START + msg.render_content(); + gen_parts.push_back({false, THINK_END + TEXT_START}); + gen_parts.push_back({true, msg.render_content()}); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -2703,7 +3005,7 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_MINIMAX_M3; data.supports_thinking = true; @@ -2746,13 +3048,22 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; data.generation_prompt = GEN_PROMPT + THINK_START + msg.reasoning_content; + gen_parts.push_back({false, GEN_PROMPT + THINK_START}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += THINK_END + msg.render_content(); + gen_parts.push_back({false, THINK_END}); + gen_parts.push_back({true, msg.render_content()}); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -3186,7 +3497,7 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -3216,13 +3527,22 @@ static common_chat_params common_chat_params_init_minicpm5(const common_chat_tem if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; data.generation_prompt = "<|im_start|>assistant\n\n" + msg.reasoning_content; + gen_parts.push_back({false, "<|im_start|>assistant\n\n"}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += "\n\n\n" + msg.render_content(); + gen_parts.push_back({false, "\n\n\n"}); + gen_parts.push_back({true, msg.render_content()}); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { @@ -3333,7 +3653,7 @@ static common_chat_params common_chat_params_init_muse_glimmer(const common_chat const autoparser::generation_params & inputs) { common_chat_params data; - data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = "<|start|>assistant"; data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; @@ -3354,13 +3674,22 @@ static common_chat_params common_chat_params_init_muse_glimmer(const common_chat if (inputs.has_continuation()) { const auto & msg = inputs.continue_msg; + // Keep the provenance separate: template delimiters vs request-provided + // continued content (must not be parsed for special tokens). + std::vector gen_parts; data.generation_prompt = "<|start|>assistant to=self<|message|>" + msg.reasoning_content; + gen_parts.push_back({false, "<|start|>assistant to=self<|message|>"}); + gen_parts.push_back({true, msg.reasoning_content}); if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { data.generation_prompt += "<|eom|><|start|>assistant to=user<|message|>" + msg.render_content(); + gen_parts.push_back({false, "<|eom|><|start|>assistant to=user<|message|>"}); + gen_parts.push_back({true, msg.render_content()}); } data.prompt += data.generation_prompt; + + data.prompt_parts.insert(data.prompt_parts.end(), gen_parts.begin(), gen_parts.end()); } auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; @@ -3700,7 +4029,7 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_ common_chat_params data; auto params_copy = params; params_copy.reasoning_format = COMMON_REASONING_FORMAT_NONE; - data.prompt = common_chat_template_direct_apply_impl(tmpl, params_copy); + data.prompt = common_chat_template_direct_apply_impl(tmpl, params_copy, std::nullopt, std::nullopt, std::nullopt, &data.prompt_parts); data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, params); data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; auto parser = build_chat_peg_parser([&data](common_chat_peg_builder &p) { diff --git a/common/chat.h b/common/chat.h index cb39e3458f44..522b8321aa87 100644 --- a/common/chat.h +++ b/common/chat.h @@ -7,6 +7,7 @@ #include "jinja/parser.h" #include "jinja/runtime.h" #include "jinja/caps.h" +#include "jinja/string.h" #include "json.h" @@ -269,6 +270,7 @@ struct common_chat_templates_inputs { struct common_chat_params { common_chat_format format = COMMON_CHAT_FORMAT_CONTENT_ONLY; std::string prompt; + std::vector prompt_parts; // preserves is_input metadata for safe tokenization std::string grammar; bool grammar_lazy = false; std::string generation_prompt; @@ -324,6 +326,30 @@ std::string common_chat_templates_source(const struct common_chat_templates * tm struct common_chat_params common_chat_templates_apply(const struct common_chat_templates * tmpls, const struct common_chat_templates_inputs & inputs); +// Returns true if at least one is_input part of parts contains text that the +// tokenizer would parse as special tokens (control/unknown tokens) when +// parse_special=true, i.e. tokenizing the parts with per-part parse_special +// would differ from a single parse_special=true pass over the concatenated +// text. +bool common_chat_parts_have_special_input( + const struct llama_vocab * vocab, + const std::vector & parts); + +// Tokenize prompt parts with input marking awareness. +// Parts marked as is_input (user/tool content) are tokenized with +// parse_special=false so that stray special tokens in request-provided +// content are not parsed as real special tokens (no special token injection). +// Parts not marked as is_input (template text) keep parse_special=true. +// +// If no is_input part contains special-token text, the parts are tokenized in +// a single pass over the concatenated text instead, so the resulting token +// ids are identical to the legacy whole-prompt tokenization (in particular, +// normal tokenizer merges across part boundaries are preserved). +std::vector common_tokenize_parts( + const struct llama_vocab * vocab, + const std::vector & parts, + bool add_special); + // Format single message, while taking into account the position of that message in chat history std::string common_chat_format_single(const struct common_chat_templates * tmpls, const std::vector & past_msg, @@ -365,7 +391,8 @@ std::map common_chat_templates_get_caps(const common_chat_tem std::string common_chat_template_direct_apply( const common_chat_template & tmpl, - const autoparser::generation_params & inputs); + const autoparser::generation_params & inputs, + std::vector * out_parts = nullptr); std::string common_chat_template_generation_prompt( const common_chat_template & tmpl, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index d2b01bee9957..7c790d859ae5 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -553,6 +553,19 @@ if (LLAMA_BUILD_SERVER) llama_build_and_test(test-server-prompt-checkpoint.cpp) target_include_directories(test-server-prompt-checkpoint PRIVATE ${PROJECT_SOURCE_DIR}/tools/server ${PROJECT_SOURCE_DIR}/tools/mtmd) target_link_libraries(test-server-prompt-checkpoint PRIVATE server-context) + + # input-marking-aware prompt tokenization (prompt_parts) regression tests; + # needs a vocab-only GGUF (the qwen2 vocab provides the control special + # tokens the tests rely on) and the GPT-OSS chat template + llama_build_and_test( + test-prompt-parts.cpp + WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} + ARGS + ${PROJECT_SOURCE_DIR}/models/ggml-vocab-qwen2.gguf + ${PROJECT_SOURCE_DIR}/models/templates/openai-gpt-oss-120b.jinja + ) + target_include_directories(test-prompt-parts PRIVATE ${PROJECT_SOURCE_DIR}/tools/server ${PROJECT_SOURCE_DIR}/tools/mtmd) + target_link_libraries(test-prompt-parts PRIVATE server-context) endif() # Upstream model-resolution coverage serves repositories through cpp-httplib. llama_build_and_test(test-model-resolution.cpp) diff --git a/tests/test-prompt-parts-mtmd.py b/tests/test-prompt-parts-mtmd.py new file mode 100644 index 000000000000..aeb1db6acc7b --- /dev/null +++ b/tests/test-prompt-parts-mtmd.py @@ -0,0 +1,221 @@ +"""Opt-in regression test for input-marking-aware tokenization in MTMD contexts. + +Runs llama-server with a multimodal model (chat GGUF + mmproj, i.e. a live +MTMD context) and verifies the OAI chat path: + 1. completion and token-count endpoints agree (prompt_parts shared path), + with and without a media attachment; + 2. the media file is counted/interleaved through the prompt_parts path; + 3. special-token text inside ordinary user content is NOT parsed as a + special token, even with an MTMD context present (the mctx branch must + not bypass the protection, with or without attached media). + +Like the other opt-in server tests, it does not download models: + + python tests/test-prompt-parts-mtmd.py --server build/bin/llama-server \ + --model model.gguf --mmproj mmproj.gguf --image image.png +""" +import argparse +import base64 +import json +import os +import socket +import subprocess +import sys +import time +import urllib.request + +CANDIDATE_SPECIAL_TOKENS = [ + "", "", "", + "<|im_start|>", "<|im_end|>", "<|user|>", "<|assistant|>", + "<|start|>", "<|end|>", "<|system|>", "<|model|>", + "", "<|turn>", "", "<|channel>", + "<|image|>", "<|eot_id|>", "<|eom_id|>", +] + +def log(msg): + print(msg, flush=True) + +def free_port(): + s = socket.socket() + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + +def http_json(url, payload=None, timeout=300): + data = None + headers = {"Content-Type": "application/json"} + if payload is not None: + data = json.dumps(payload).encode("utf-8") + req = urllib.request.Request(url, data=data, headers=headers) + with urllib.request.urlopen(req, timeout=timeout) as res: + return json.loads(res.read().decode("utf-8")) + +def wait_healthy(base, timeout_s): + deadline = time.time() + timeout_s + last_err = None + while time.time() < deadline: + try: + r = http_json(base + "/health", timeout=5) + if r.get("status") in ("ok", "loading model"): + if r.get("status") == "ok": + return + except Exception as e: + last_err = e + time.sleep(1) + raise RuntimeError(f"server did not become healthy in {timeout_s}s: {last_err}") + +def tokenize(base, content, parse_special): + r = http_json(base + "/tokenize", {"content": content, "add_special": False, + "parse_special": parse_special}) + return r["tokens"] + +def discover_special_token(base): + """Find a token the model parses as special (1 token with parse_special=true, + several with parse_special=false).""" + for cand in CANDIDATE_SPECIAL_TOKENS: + try: + with_ps = tokenize(base, cand, True) + without_ps = tokenize(base, cand, False) + except Exception: + continue + if len(with_ps) == 1 and len(without_ps) >= 2: + return cand, with_ps[0], without_ps + return None, None, None + +def prompt_tokens(base, body, endpoint): + r = http_json(base + endpoint, body, timeout=600) + if endpoint.endswith("input_tokens"): + return int(r["input_tokens"]) + return int(r["usage"]["prompt_tokens"]) + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--server", required=True, type=os.path.abspath) + ap.add_argument("--model", required=True, type=os.path.abspath) + ap.add_argument("--mmproj", required=True, type=os.path.abspath) + ap.add_argument("--image", required=True, type=os.path.abspath) + ap.add_argument("--port", type=int, default=0) + ap.add_argument("--ctx", type=int, default=4096) + ap.add_argument("--predict", type=int, default=1) + ap.add_argument("--load-timeout", type=int, default=600) + args = ap.parse_args() + + port = args.port or free_port() + base = f"http://127.0.0.1:{port}" + cmd = [args.server, "-m", args.model, "--mmproj", args.mmproj, + "-ngl", "0", "-c", str(args.ctx), "-np", "1", + "--host", "127.0.0.1", "--port", str(port), "--no-warmup", "-v"] + log("command: " + " ".join(cmd)) + # write the server log to a file: a PIPE would deadlock once its buffer + # fills up, since the server output is only read at shutdown + log_path = os.path.abspath(os.path.join(os.path.dirname(args.server), "..", "test-prompt-parts-mtmd.log")) + server_log = open(log_path, "w", encoding="utf-8") + proc = subprocess.Popen(cmd, stdout=server_log, stderr=subprocess.STDOUT) + try: + wait_healthy(base, args.load_timeout) + log("server healthy") + + special, special_id, special_fallback = discover_special_token(base) + if special is None: + log("FAIL: no candidate special token found in the model vocab") + return 1 + log(f"special token: {special!r} (id {special_id}, " + f"{len(special_fallback)} tokens without parse_special)") + + text_body = { + "messages": [{"role": "user", "content": "Describe the image."}], + "max_tokens": args.predict, + "stream": False, + } + n_base_c = prompt_tokens(base, text_body, "/v1/chat/completions") + n_base_t = prompt_tokens(base, text_body, "/v1/chat/completions/input_tokens") + log(f"baseline (no media): completion={n_base_c} count={n_base_t}") + if n_base_c != n_base_t: + log("FAIL: completion and count endpoints disagree (no media)") + return 1 + + # media attachment: the marker text comes from the template/OAI + # conversion, so the media path runs through prompt_parts + with open(args.image, "rb") as f: + b64 = base64.b64encode(f.read()).decode("ascii") + img_body = { + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "Describe the image."}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}, + ]}], + "max_tokens": args.predict, + "stream": False, + } + n_img_c = prompt_tokens(base, img_body, "/v1/chat/completions") + n_img_t = prompt_tokens(base, img_body, "/v1/chat/completions/input_tokens") + log(f"with media: completion={n_img_c} count={n_img_t}") + if n_img_c != n_img_t: + log("FAIL: completion and count endpoints disagree (with media)") + return 1 + if n_img_c <= n_base_c: + log(f"FAIL: media tokens not counted ({n_img_c} <= {n_base_c})") + return 1 + + # special-token injection in ordinary user content (no media): + # with the protection the injected text is tokenized as normal text + # (several tokens), not as a single special token. + inj_text = f"Describe the image. {special}injected{special} end" + inj_body = { + "messages": [{"role": "user", "content": inj_text}], + "max_tokens": args.predict, + "stream": False, + } + n_inj_c = prompt_tokens(base, inj_body, "/v1/chat/completions") + n_inj_t = prompt_tokens(base, inj_body, "/v1/chat/completions/input_tokens") + log(f"injected (no media): completion={n_inj_c} count={n_inj_t}") + if n_inj_c != n_inj_t: + log("FAIL: completion and count endpoints disagree (injected, no media)") + return 1 + delta = n_inj_t - n_base_t + log(f"injection token delta: {delta} (1 would mean the special token was parsed)") + if delta <= 1: + log("FAIL: injected special token was parsed as a special token") + return 1 + + # same, but with media attached (mtmd part interleaving + protection) + inj_img_body = { + "messages": [{"role": "user", "content": [ + {"type": "text", "text": inj_text}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}, + ]}], + "max_tokens": args.predict, + "stream": False, + } + n_inj_img_c = prompt_tokens(base, inj_img_body, "/v1/chat/completions") + n_inj_img_t = prompt_tokens(base, inj_img_body, "/v1/chat/completions/input_tokens") + log(f"injected (with media): completion={n_inj_img_c} count={n_inj_img_t}") + if n_inj_img_c != n_inj_img_t: + log("FAIL: completion and count endpoints disagree (injected, with media)") + return 1 + delta_img = n_inj_img_t - n_img_t + log(f"injection token delta with media: {delta_img}") + if delta_img <= 1: + log("FAIL: injected special token was parsed as a special token (with media)") + return 1 + + log("PASS: all MTMD prompt_parts checks succeeded") + return 0 + finally: + proc.terminate() + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + server_log.close() + try: + with open(log_path, "r", encoding="utf-8", errors="replace") as f: + tail = f.read().splitlines()[-25:] + log("--- server log tail ---") + log("\n".join(tail)) + except OSError: + pass + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test-prompt-parts.cpp b/tests/test-prompt-parts.cpp new file mode 100644 index 000000000000..b5c270e79957 --- /dev/null +++ b/tests/test-prompt-parts.cpp @@ -0,0 +1,639 @@ +// Regression tests for input-marking-aware prompt tokenization (prompt_parts). +// +// These tests cover: +// - special-token text inside ordinary user/tool input (no special token +// injection, template special tokens still parsed); +// - assistant continuation content provenance in the automatic parser path +// and the specialized parser paths (gpt-oss, qwen3-coder, gemma4: +// template delimiters vs request-provided reasoning/content); +// - the GPT-OSS <|return|> -> <|end|> replacement applied to prompt parts +// (specialized parser path, incl. continuation provenance); +// - consistency between the completion and token-count tokenization paths +// (both use the same shared tokenization function, through the OAI +// prompt_parts serialization); +// - unchanged token ids for normal prompts without injected special tokens. +// +// Usage: test-prompt-parts vocab-file [gpt-oss-template-file] +// +// vocab-file: a vocab-only GGUF, e.g. models/ggml-vocab-qwen2.gguf +// (must contain control special tokens like +// "<|im_start|>" / "<|im_end|>") +// gpt-oss-template-file: defaults to models/templates/openai-gpt-oss-120b.jinja + +#include "llama.h" +#include "common.h" +#include "chat.h" +#include "chat-auto-parser.h" +#include "jinja/string.h" + +#include "server-common.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +int n_checks = 0; +int n_fail = 0; + +#define CHECK_MSG(cond, ...) \ + do { \ + ++n_checks; \ + if (!(cond)) { \ + ++n_fail; \ + fprintf(stderr, "FAIL: %s (line %d): ", __func__, __LINE__); \ + fprintf(stderr, __VA_ARGS__); \ + fprintf(stderr, "\n"); \ + } \ + } while (0) + +#define CHECK(cond) CHECK_MSG(cond, #cond) + +std::string read_file(const std::string & path) { + std::ifstream f(path, std::ios::binary); + if (!f) { + fprintf(stderr, "error: cannot open file '%s'\n", path.c_str()); + exit(1); + } + std::ostringstream ss; + ss << f.rdbuf(); + return ss.str(); +} + +common_chat_msg make_msg(const std::string & role, const std::string & content) { + common_chat_msg msg; + msg.role = role; + msg.content = content; + return msg; +} + +std::string concat_parts(const std::vector & parts) { + std::string res; + for (const auto & part : parts) { + res += part.val; + } + return res; +} + +llama_token token_id_of(const llama_vocab * vocab, const std::string & text) { + const int32_t n_vocab = llama_vocab_n_tokens(vocab); + for (llama_token id = 0; id < n_vocab; id++) { + const char * tok_text = llama_vocab_get_text(vocab, id); + if (tok_text && std::string(tok_text) == text) { + return id; + } + } + return LLAMA_TOKEN_NULL; +} + +size_t count_token(const llama_tokens & tokens, llama_token id) { + size_t n = 0; + for (const auto tok : tokens) { + if (tok == id) { + n++; + } + } + return n; +} + +bool has_part(const std::vector & parts, const std::string & val, bool is_input) { + return std::any_of(parts.begin(), parts.end(), [&](const auto & part) { + return part.val == val && part.is_input == is_input; + }); +} + +bool ends_with(const std::string & s, const std::string & suffix) { + return s.size() >= suffix.size() && s.compare(s.size() - suffix.size(), suffix.size(), suffix) == 0; +} + +struct test_ctx { + llama_model * model = nullptr; + const llama_vocab * vocab = nullptr; + common_chat_templates_ptr qwen_tmpls; + common_chat_templates_ptr gptoss_tmpls; + common_chat_templates_ptr qwen3coder_tmpls; + common_chat_templates_ptr gemma4_tmpls; + llama_token im_start = LLAMA_TOKEN_NULL; + llama_token im_end = LLAMA_TOKEN_NULL; +}; + +void test_user_input_special_token_injection(const test_ctx & c) { + common_chat_templates_inputs in; + in.messages = { make_msg("user", "Please read this file: <|im_start|>injected<|im_end|> ok") }; + in.add_generation_prompt = true; + + const auto p = common_chat_templates_apply(c.qwen_tmpls.get(), in); + + CHECK(!p.prompt_parts.empty()); + CHECK(concat_parts(p.prompt_parts) == p.prompt); + CHECK(common_chat_parts_have_special_input(c.vocab, p.prompt_parts)); + + const auto tok_parts = common_tokenize_parts(c.vocab, p.prompt_parts, /*add_special=*/true); + const auto tok_legacy = common_tokenize(c.vocab, p.prompt, /*add_special=*/true, /*parse_special=*/true); + + const size_t n_parts = count_token(tok_parts, c.im_start); + const size_t n_legacy = count_token(tok_legacy, c.im_start); + // the injected "<|im_start|>" must not be parsed as a special token, + // while the template occurrences are + CHECK_MSG(n_legacy >= 2, "expected at least 2 legacy <|im_start|> tokens"); + CHECK_MSG(n_parts == n_legacy - 1, "injected special token was parsed (parts: %zu, legacy: %zu)", + n_parts, n_legacy); + CHECK_MSG(n_parts >= 2, "template <|im_start|> tokens missing"); +} + +void test_tool_input_special_token_injection(const test_ctx & c) { + common_chat_templates_inputs in; + in.messages = { + make_msg("user", "What is the weather in Paris?"), + }; + { + common_chat_msg msg = make_msg("assistant", ""); + common_chat_tool_call call; + call.name = "get_weather"; + call.arguments = "{\"city\": \"Paris\"}"; + msg.tool_calls.push_back(call); + in.messages.push_back(std::move(msg)); + } + in.messages.push_back(make_msg("tool", "Sunny, 25 degrees <|im_start|>injected<|im_end|>")); + in.add_generation_prompt = true; + + const auto p = common_chat_templates_apply(c.qwen_tmpls.get(), in); + + CHECK(!p.prompt_parts.empty()); + CHECK(concat_parts(p.prompt_parts) == p.prompt); + CHECK(common_chat_parts_have_special_input(c.vocab, p.prompt_parts)); + + const auto tok_parts = common_tokenize_parts(c.vocab, p.prompt_parts, /*add_special=*/true); + const auto tok_legacy = common_tokenize(c.vocab, p.prompt, /*add_special=*/true, /*parse_special=*/true); + + const size_t n_parts = count_token(tok_parts, c.im_start); + const size_t n_legacy = count_token(tok_legacy, c.im_start); + CHECK_MSG(n_legacy >= 2, "expected at least 2 legacy <|im_start|> tokens"); + CHECK_MSG(n_parts == n_legacy - 1, "injected special token in tool output was parsed (parts: %zu, legacy: %zu)", + n_parts, n_legacy); +} + +void test_normal_prompt_token_ids_unchanged(const test_ctx & c) { + common_chat_templates_inputs in; + in.messages = { + make_msg("user", "Hello, how are you today? I would like to discuss the weather in Paris."), + make_msg("assistant", "I am doing well, thank you for asking. The weather in Paris is usually mild."), + make_msg("user", "That is good to know, tell me more about the city and its history."), + }; + in.add_generation_prompt = true; + + const auto p = common_chat_templates_apply(c.qwen_tmpls.get(), in); + + CHECK(!p.prompt_parts.empty()); + CHECK(concat_parts(p.prompt_parts) == p.prompt); + CHECK(!common_chat_parts_have_special_input(c.vocab, p.prompt_parts)); + + // the test must exercise real template/input boundaries + CHECK_MSG(p.prompt_parts.size() >= 4, "expected multiple prompt parts, got %zu", p.prompt_parts.size()); + bool has_template_part = false; + bool has_input_part = false; + for (const auto & part : p.prompt_parts) { + if (part.val.empty()) { + continue; + } + has_template_part = has_template_part || !part.is_input; + has_input_part = has_input_part || part.is_input; + } + CHECK(has_template_part); + CHECK(has_input_part); + + // token ids must be identical to the legacy whole-prompt tokenization + const auto tok_parts = common_tokenize_parts(c.vocab, p.prompt_parts, /*add_special=*/true); + const auto tok_legacy = common_tokenize(c.vocab, p.prompt, /*add_special=*/true, /*parse_special=*/true); + if (tok_parts != tok_legacy) { + fprintf(stderr, " parts : %zu tokens\n legacy: %zu tokens\n", tok_parts.size(), tok_legacy.size()); + for (size_t i = 0; i < std::min(tok_parts.size(), tok_legacy.size()); i++) { + if (tok_parts[i] != tok_legacy[i]) { + fprintf(stderr, " first diff at %zu: %d vs %d\n", i, (int) tok_parts[i], (int) tok_legacy[i]); + break; + } + } + } + CHECK_MSG(tok_parts == tok_legacy, + "token ids changed for a normal prompt without injected special tokens"); + + // Synthetic boundary case: a BPE merge straddles a template/input part + // boundary ("Hello" = "Hel" + "lo"). Tokenizing the parts separately + // breaks the merge; the fast path must keep the legacy token ids. + const std::vector boundary_parts = { + {false, "Hel"}, + {true, "lo"}, + }; + const auto tok_boundary = common_tokenize_parts(c.vocab, boundary_parts, /*add_special=*/false); + const auto tok_whole = common_tokenize(c.vocab, "Hello", /*add_special=*/false, /*parse_special=*/true); + const auto tok_split = + common_tokenize(c.vocab, "Hel", /*add_special=*/false, /*parse_special=*/true); + const auto tok_split_2 = + common_tokenize(c.vocab, "lo", /*add_special=*/false, /*parse_special=*/false); + llama_tokens tok_naive; + tok_naive.insert(tok_naive.end(), tok_split.begin(), tok_split.end()); + tok_naive.insert(tok_naive.end(), tok_split_2.begin(), tok_split_2.end()); + CHECK_MSG(tok_naive != tok_whole, "test case does not straddle a tokenizer merge"); + CHECK_MSG(tok_boundary == tok_whole, + "part-boundary merge was broken for a normal prompt (legacy: %zu tokens, parts: %zu tokens)", + tok_whole.size(), tok_boundary.size()); + + // Protection path: adjacent template parts (e.g. pieces of the + // continuation generation prompt appended after rendering) must be + // merged before tokenizing, so normal merges across their boundary are + // preserved while the is_input part keeps the parse_special=false + // protection. + const std::vector protect_parts = { + {false, "Hel"}, + {false, "lo"}, + {true, " and <|im_start|>injected"}, + }; + const auto tok_protect = common_tokenize_parts(c.vocab, protect_parts, /*add_special=*/false); + const auto tok_ref_hello = common_tokenize(c.vocab, "Hello", /*add_special=*/false, /*parse_special=*/true); + const auto tok_ref_inj = common_tokenize(c.vocab, " and <|im_start|>injected", /*add_special=*/false, /*parse_special=*/false); + llama_tokens tok_ref; + tok_ref.insert(tok_ref.end(), tok_ref_hello.begin(), tok_ref_hello.end()); + tok_ref.insert(tok_ref.end(), tok_ref_inj.begin(), tok_ref_inj.end()); + CHECK_MSG(tok_protect == tok_ref, + "same-type part boundary merge was broken in the protection path (ref: %zu tokens, parts: %zu tokens)", + tok_ref.size(), tok_protect.size()); +} + +void test_continuation_provenance_autoparser(const test_ctx & c) { + // CONTENT continuation with injected special tokens in both the + // reasoning content and the message content + { + common_chat_msg last = make_msg("assistant", "The answer is 42 <|im_start|>injected"); + last.reasoning_content = "Let me think <|im_start|>injected <|im_end|> step by step."; + + common_chat_templates_inputs in; + in.messages = { make_msg("user", "What is 6 times 7?"), last }; + in.add_generation_prompt = true; + in.continue_final_message = COMMON_CHAT_CONTINUATION_CONTENT; + + const auto p = common_chat_templates_apply(c.qwen_tmpls.get(), in); + + CHECK(concat_parts(p.prompt_parts) == p.prompt); + // request-provided continuation content must be marked as input + CHECK_MSG(has_part(p.prompt_parts, last.reasoning_content, true), + "reasoning content not marked as input in prompt parts"); + CHECK_MSG(has_part(p.prompt_parts, last.content, true), + "message content not marked as input in prompt parts"); + + const auto tok_parts = common_tokenize_parts(c.vocab, p.prompt_parts, /*add_special=*/true); + const auto tok_legacy = common_tokenize(c.vocab, p.prompt, /*add_special=*/true, /*parse_special=*/true); + + const size_t n_parts = count_token(tok_parts, c.im_start); + const size_t n_legacy = count_token(tok_legacy, c.im_start); + CHECK_MSG(n_legacy >= 2, "expected legacy <|im_start|> tokens"); + CHECK_MSG(n_parts == n_legacy - 2, + "injected special tokens in continuation content were parsed (parts: %zu, legacy: %zu)", + n_parts, n_legacy); + } + // REASONING-only continuation + { + common_chat_msg last; + last.role = "assistant"; + last.reasoning_content = "Let me think <|im_start|>injected step by step."; + + common_chat_templates_inputs in; + in.messages = { make_msg("user", "What is 6 times 7?"), last }; + in.add_generation_prompt = true; + in.continue_final_message = COMMON_CHAT_CONTINUATION_REASONING; + + const auto p = common_chat_templates_apply(c.qwen_tmpls.get(), in); + + CHECK(concat_parts(p.prompt_parts) == p.prompt); + CHECK_MSG(has_part(p.prompt_parts, last.reasoning_content, true), + "reasoning content not marked as input in prompt parts"); + + const auto tok_parts = common_tokenize_parts(c.vocab, p.prompt_parts, /*add_special=*/true); + const auto tok_legacy = common_tokenize(c.vocab, p.prompt, /*add_special=*/true, /*parse_special=*/true); + + const size_t n_parts = count_token(tok_parts, c.im_start); + const size_t n_legacy = count_token(tok_legacy, c.im_start); + CHECK_MSG(n_parts == n_legacy - 1, + "injected special token in reasoning continuation was parsed (parts: %zu, legacy: %zu)", + n_parts, n_legacy); + } +} + +void test_gpt_oss_return_replacement(const test_ctx & c) { + common_chat_templates_inputs in; + in.messages = { + make_msg("user", "Hi"), + make_msg("assistant", "Hello there"), + }; + in.add_generation_prompt = false; // the template ends with <|return|> for the final assistant turn + + const auto p = common_chat_templates_apply(c.gptoss_tmpls.get(), in); + + CHECK(ends_with(p.prompt, "<|end|>")); + CHECK(p.prompt.find("<|return|>") == std::string::npos); + + // the parts must carry the same replacement, so server-side tokenization + // of prompt_parts sees the same text as data.prompt + CHECK(!p.prompt_parts.empty()); + CHECK(concat_parts(p.prompt_parts) == p.prompt); + CHECK(concat_parts(p.prompt_parts).find("<|return|>") == std::string::npos); + if (!p.prompt_parts.empty()) { + CHECK(ends_with(p.prompt_parts.back().val, "<|end|>")); + CHECK(!p.prompt_parts.back().is_input); + } +} + +void test_gpt_oss_continuation_provenance(const test_ctx & c) { + common_chat_msg last = make_msg("assistant", "final <|im_start|>injected"); + last.reasoning_content = "analysis <|im_start|>injected"; + + common_chat_templates_inputs in; + in.messages = { make_msg("user", "Hi"), last }; + in.add_generation_prompt = true; + in.continue_final_message = COMMON_CHAT_CONTINUATION_CONTENT; + + const auto p = common_chat_templates_apply(c.gptoss_tmpls.get(), in); + + CHECK(concat_parts(p.prompt_parts) == p.prompt); + + // specialized parser path: delimiters are template text, the + // request-provided reasoning/content is input + CHECK_MSG(has_part(p.prompt_parts, "<|start|>assistant<|channel|>analysis<|message|>", false), + "gpt-oss analysis delimiter not marked as template text"); + CHECK_MSG(has_part(p.prompt_parts, last.reasoning_content, true), + "gpt-oss reasoning content not marked as input"); + CHECK_MSG(has_part(p.prompt_parts, "<|end|><|start|>assistant<|channel|>final<|message|>", false), + "gpt-oss final delimiter not marked as template text"); + CHECK_MSG(has_part(p.prompt_parts, last.content, true), + "gpt-oss message content not marked as input"); + + const auto tok_parts = common_tokenize_parts(c.vocab, p.prompt_parts, /*add_special=*/true); + const auto tok_legacy = common_tokenize(c.vocab, p.prompt, /*add_special=*/true, /*parse_special=*/true); + + const size_t n_parts = count_token(tok_parts, c.im_start); + const size_t n_legacy = count_token(tok_legacy, c.im_start); + CHECK_MSG(n_parts == n_legacy - 2, + "injected special tokens in gpt-oss continuation were parsed (parts: %zu, legacy: %zu)", + n_parts, n_legacy); +} + +void test_qwen3_coder_continuation_provenance(const test_ctx & c) { + // specialized qwen3-coder parser path (non-reasoning template): the + // generation prompt prefix is template text, the continued content is + // request-provided + common_chat_msg last = make_msg("assistant", "final answer <|im_start|>injected"); + + common_chat_templates_inputs in; + in.messages = { make_msg("user", "Hi"), last }; + in.add_generation_prompt = true; + in.continue_final_message = COMMON_CHAT_CONTINUATION_CONTENT; + + const auto p = common_chat_templates_apply(c.qwen3coder_tmpls.get(), in); + + CHECK(concat_parts(p.prompt_parts) == p.prompt); + CHECK_MSG(has_part(p.prompt_parts, last.content, true), + "qwen3-coder continued content not marked as input"); + + // the generation prompt prefix (im_start/assistant) is the last template + // part right before the continued content + bool prefix_before_content = false; + for (size_t i = 1; i < p.prompt_parts.size(); i++) { + if (p.prompt_parts[i].val == last.content && p.prompt_parts[i].is_input && + p.prompt_parts[i - 1].val == "<|im_start|>assistant\n" && !p.prompt_parts[i - 1].is_input) { + prefix_before_content = true; + } + } + CHECK_MSG(prefix_before_content, "qwen3-coder generation prompt prefix not marked as template text"); + + const auto tok_parts = common_tokenize_parts(c.vocab, p.prompt_parts, /*add_special=*/true); + const auto tok_legacy = common_tokenize(c.vocab, p.prompt, /*add_special=*/true, /*parse_special=*/true); + const size_t n_parts = count_token(tok_parts, c.im_start); + const size_t n_legacy = count_token(tok_legacy, c.im_start); + CHECK_MSG(n_parts == n_legacy - 1, + "injected special token in qwen3-coder continuation was parsed (parts: %zu, legacy: %zu)", + n_parts, n_legacy); +} + +void test_gemma4_continuation_provenance(const test_ctx & c) { + // specialized gemma4 parser path: the turn prefix, the thought channel + // delimiter and the channel terminator are template text, the + // request-provided reasoning/content are input + common_chat_msg last = make_msg("assistant", "the answer <|im_start|>injected"); + last.reasoning_content = "thinking <|im_start|>injected"; + + common_chat_templates_inputs in; + in.messages = { make_msg("user", "Hi"), last }; + in.add_generation_prompt = true; + in.continue_final_message = COMMON_CHAT_CONTINUATION_CONTENT; + + const auto p = common_chat_templates_apply(c.gemma4_tmpls.get(), in); + + CHECK(concat_parts(p.prompt_parts) == p.prompt); + CHECK_MSG(has_part(p.prompt_parts, "<|turn>model\n", false), + "gemma4 turn prefix not marked as template text"); + CHECK_MSG(has_part(p.prompt_parts, "<|channel>thought\n", false), + "gemma4 thought delimiter not marked as template text"); + CHECK_MSG(has_part(p.prompt_parts, last.reasoning_content, true), + "gemma4 reasoning content not marked as input"); + CHECK_MSG(has_part(p.prompt_parts, "", false), + "gemma4 channel terminator not marked as template text"); + CHECK_MSG(has_part(p.prompt_parts, last.content, true), + "gemma4 message content not marked as input"); + + const auto tok_parts = common_tokenize_parts(c.vocab, p.prompt_parts, /*add_special=*/true); + const auto tok_legacy = common_tokenize(c.vocab, p.prompt, /*add_special=*/true, /*parse_special=*/true); + const size_t n_parts = count_token(tok_parts, c.im_start); + const size_t n_legacy = count_token(tok_legacy, c.im_start); + CHECK_MSG(n_parts == n_legacy - 2, + "injected special tokens in gemma4 continuation were parsed (parts: %zu, legacy: %zu)", + n_parts, n_legacy); +} + +void test_completion_and_count_consistency(const test_ctx & c) { + // Serialize the OAI chat request the same way the server does, then run + // the shared tokenization path the way both endpoints do: + // - completion route: server_tokenize_prompt_parts(..., is_placeholder=false) + // - count route: server_tokenize_prompt_parts(..., is_placeholder=true) + server_chat_params opt; + opt.use_jinja = true; + opt.prefill_assistant = false; + opt.reasoning_format = COMMON_REASONING_FORMAT_NONE; + opt.tmpls = common_chat_templates_init(nullptr, read_file("models/templates/Qwen-Qwen3-0.6B.jinja")); + opt.allow_image = false; + opt.allow_audio = false; + opt.allow_video = false; + opt.enable_thinking = true; + opt.reasoning_budget = -1; + opt.reasoning_budget_message = ""; + opt.media_path = ""; + opt.force_pure_content = false; + + json body = { + {"messages", json::array({ + {{"role", "user"}, {"content", "Hello <|im_start|>injected <|im_end|> world"}}, + })}, + {"max_tokens", 5}, + }; + std::vector files; + const json llama_params = oaicompat_chat_params_parse(body, opt, files); + + CHECK(llama_params.contains("prompt_parts")); + CHECK(llama_params.at("prompt_parts").is_array()); + CHECK(!llama_params.at("prompt_parts").empty()); + + // parse prompt_parts exactly like handle_completions_impl / handle_count_tokens + std::vector parts; + for (const auto & p : llama_params.at("prompt_parts")) { + parts.push_back({p.at("is_input").get(), p.at("text").get()}); + } + + const auto init_opt = mtmd_helper_init_opt_default(); + const auto completion_tokens = + server_tokenize_prompt_parts(c.vocab, nullptr, parts, files, init_opt, /*add_special=*/true, + /*is_placeholder=*/false).get_tokens(); + const auto count_tokens = + server_tokenize_prompt_parts(c.vocab, nullptr, parts, files, init_opt, /*add_special=*/true, + /*is_placeholder=*/true).get_tokens(); + + CHECK_MSG(count_tokens.size() == completion_tokens.size(), + "completion (%zu) and count (%zu) token paths disagree", + completion_tokens.size(), count_tokens.size()); + CHECK_MSG(count_tokens == completion_tokens, + "completion and count token paths produced different tokens"); + + // and the protection must hold through the shared path + const size_t n_tokens = count_token(completion_tokens, c.im_start); + const size_t n_legacy = count_token(common_tokenize(c.vocab, llama_params.at("prompt").get(), + /*add_special=*/true, /*parse_special=*/true), c.im_start); + CHECK_MSG(n_tokens == n_legacy - 1, + "shared tokenization path lost the special-token protection (parts: %zu, legacy: %zu)", + n_tokens, n_legacy); +} + +void test_mtmd_part_layout() { + const std::string marker = "<__media__>"; + + { + // ordinary layout: marker in its own part + const std::vector parts = { + {false, "describe "}, + {true, marker}, + {true, " and <|im_start|>hack "}, + {false, "the picture"}, + }; + std::vector segs; + CHECK(server_build_mtmd_part_layout(marker, parts, /*n_files=*/1, segs)); + CHECK_MSG(segs.size() == 4, "expected 4 segments, got %zu", segs.size()); + if (segs.size() == 4) { + CHECK_MSG(!segs[0].is_bitmap && segs[0].text == "describe " && !segs[0].is_input, "seg0 mismatch"); + CHECK_MSG(segs[1].is_bitmap, "seg1 must be the bitmap"); + CHECK_MSG(!segs[2].is_bitmap && segs[2].text == " and <|im_start|>hack " && segs[2].is_input, "seg2 mismatch"); + CHECK_MSG(!segs[3].is_bitmap && segs[3].text == "the picture" && !segs[3].is_input, "seg3 mismatch"); + } + + // marker count mismatch must be reported + CHECK(!server_build_mtmd_part_layout(marker, parts, /*n_files=*/2, segs)); + } + { + // marker spanning a part boundary (template part ends mid-marker, + // input part completes it) + const std::vector parts = { + {false, "A <__media"}, + {true, "__> B"}, + }; + std::vector segs; + CHECK(server_build_mtmd_part_layout(marker, parts, /*n_files=*/1, segs)); + CHECK_MSG(segs.size() == 3, "expected 3 segments, got %zu", segs.size()); + if (segs.size() == 3) { + CHECK_MSG(!segs[0].is_bitmap && segs[0].text == "A " && !segs[0].is_input, "seg0 mismatch"); + CHECK_MSG(segs[1].is_bitmap, "seg1 must be the bitmap"); + CHECK_MSG(!segs[2].is_bitmap && segs[2].text == " B" && segs[2].is_input, "seg2 mismatch"); + } + } + { + // multiple markers, interleaved with input/template text + const std::vector parts = { + {false, "one "}, + {true, marker}, + {true, " two "}, + {false, marker}, + {false, " three"}, + }; + std::vector segs; + CHECK(server_build_mtmd_part_layout(marker, parts, /*n_files=*/2, segs)); + CHECK_MSG(segs.size() == 5, "expected 5 segments, got %zu", segs.size()); + if (segs.size() == 5) { + CHECK_MSG(!segs[0].is_bitmap && segs[0].text == "one ", "seg0 mismatch"); + CHECK_MSG(segs[1].is_bitmap, "seg1 must be a bitmap"); + CHECK_MSG(!segs[2].is_bitmap && segs[2].text == " two " && segs[2].is_input, "seg2 mismatch"); + CHECK_MSG(segs[3].is_bitmap, "seg3 must be a bitmap"); + CHECK_MSG(!segs[4].is_bitmap && segs[4].text == " three" && !segs[4].is_input, "seg4 mismatch"); + } + } +} + +} // namespace + +int main(int argc, char ** argv) { + if (argc < 2) { + fprintf(stderr, "Usage: %s vocab-file [gpt-oss-template-file]\n", argv[0]); + return 1; + } + const std::string vocab_path = argv[1]; + const std::string gptoss_template_path = + argc > 2 ? argv[2] : "models/templates/openai-gpt-oss-120b.jinja"; + + llama_backend_init(); + + test_ctx c; + + { + auto mparams = llama_model_default_params(); + mparams.vocab_only = true; + c.model = llama_model_load_from_file(vocab_path.c_str(), mparams); + if (c.model == nullptr) { + fprintf(stderr, "error: failed to load vocab '%s'\n", vocab_path.c_str()); + return 1; + } + c.vocab = llama_model_get_vocab(c.model); + } + + c.im_start = token_id_of(c.vocab, "<|im_start|>"); + c.im_end = token_id_of(c.vocab, "<|im_end|>"); + if (c.im_start == LLAMA_TOKEN_NULL || c.im_end == LLAMA_TOKEN_NULL) { + fprintf(stderr, "error: vocab '%s' lacks <|im_start|>/<|im_end|>; use e.g. models/ggml-vocab-qwen2.gguf\n", + vocab_path.c_str()); + llama_model_free(c.model); + return 1; + } + if (!llama_vocab_is_control(c.vocab, c.im_start)) { + fprintf(stderr, "error: <|im_start|> is not a control token in '%s'; the test needs a vocab where it is special\n", + vocab_path.c_str()); + llama_model_free(c.model); + return 1; + } + + c.qwen_tmpls = common_chat_templates_init(nullptr, read_file("models/templates/Qwen-Qwen3-0.6B.jinja")); + c.gptoss_tmpls = common_chat_templates_init(nullptr, read_file(gptoss_template_path)); + c.qwen3coder_tmpls = common_chat_templates_init(nullptr, read_file("models/templates/Qwen3-Coder.jinja")); + c.gemma4_tmpls = common_chat_templates_init(nullptr, read_file("models/templates/google-gemma-4-31B-it.jinja")); + + test_user_input_special_token_injection(c); + test_tool_input_special_token_injection(c); + test_normal_prompt_token_ids_unchanged(c); + test_continuation_provenance_autoparser(c); + test_gpt_oss_return_replacement(c); + test_gpt_oss_continuation_provenance(c); + test_qwen3_coder_continuation_provenance(c); + test_gemma4_continuation_provenance(c); + test_completion_and_count_consistency(c); + test_mtmd_part_layout(); + + llama_model_free(c.model); + + fprintf(stderr, "%s: %d checks, %d failures\n", n_fail ? "FAIL" : "PASS", n_checks, n_fail); + return n_fail ? 1 : 0; +} diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 06450076777b..e86bb770766e 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -957,6 +957,156 @@ server_tokens process_mtmd_prompt( return result; } +bool server_build_mtmd_part_layout( + const std::string & marker, + const std::vector & parts, + size_t n_files, + std::vector & out_segs) { + std::string prompt; + // (start offset, is_input) ranges covering [0, prompt.size()) + std::vector> ranges; + size_t total = 0; + for (const auto & part : parts) { + if (!part.val.empty()) { + ranges.push_back({total, part.is_input}); + total += part.val.size(); + } + } + prompt.reserve(total); + for (const auto & part : parts) { + prompt += part.val; + } + + // locate all media markers (a marker may span two adjacent parts) + std::vector marker_starts; + for (size_t pos = prompt.find(marker); pos != std::string::npos; pos = prompt.find(marker, pos + marker.size())) { + marker_starts.push_back(pos); + } + if (marker_starts.size() != n_files) { + return false; + } + + out_segs.clear(); + + auto add_text_range = [&](size_t begin, size_t end) { + for (size_t r = 0; r < ranges.size(); r++) { + const size_t r_end = (r + 1 < ranges.size()) ? ranges[r + 1].first : prompt.size(); + const size_t s = std::max(begin, ranges[r].first); + const size_t e = std::min(end, r_end); + if (s < e) { + out_segs.push_back({false, prompt.substr(s, e - s), ranges[r].second}); + } + } + }; + + size_t prev = 0; + for (size_t i = 0; i < marker_starts.size(); i++) { + add_text_range(prev, marker_starts[i]); + out_segs.push_back({true, "", false}); + prev = marker_starts[i] + marker.size(); + } + add_text_range(prev, prompt.size()); + + return true; +} + +server_tokens server_tokenize_prompt_parts( + const llama_vocab * vocab, + mtmd_context * mctx, + const std::vector & parts, + const std::vector & files, + const mtmd_helper_init_opt & init_opt, + bool add_special, + bool is_placeholder) { + const bool protect = common_chat_parts_have_special_input(vocab, parts); + + // No media to interleave: plain input-marking-aware tokenization. + // Covers text-only models as well as multimodal-capable models without + // attachments, so the special-token protection applies in both cases. + if (mctx == nullptr || files.empty()) { + const auto tokens = common_tokenize_parts(vocab, parts, add_special); + return server_tokens(tokens, /*has_mtmd=*/mctx != nullptr); + } + + // Media present, but no is_input part contains special-token text: keep + // the legacy mtmd path (single-pass tokenization of the flattened prompt) + // so the token ids are identical to the pre-input-marking behavior. + if (!protect) { + std::string prompt; + size_t total = 0; + for (const auto & part : parts) { + total += part.val.size(); + } + prompt.reserve(total); + for (const auto & part : parts) { + prompt += part.val; + } + return process_mtmd_prompt(mctx, prompt, files, init_opt, is_placeholder); + } + + // Media present and special-token text detected in request-provided + // content: interleave the media bitmaps at the marker positions and + // tokenize every text segment with its own parse_special flag. + + const std::string marker = get_media_marker(); + + std::vector segs; + if (!server_build_mtmd_part_layout(marker, parts, files.size(), segs)) { + throw std::runtime_error(string_format( + "number of media markers in prompt does not match number of files (%zu)", + files.size())); + } + + // these will be freed upon going out of scope + mtmd::bitmaps bitmaps; + std::vector videos; + for (auto & file : files) { + auto out = mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder, init_opt); + if (!out.bitmap) { + throw std::runtime_error("Failed to load image or audio file"); + } + bitmaps.entries.emplace_back(out.bitmap); + if (out.video_ctx) { + videos.emplace_back(out.video_ctx); + } + } + auto bitmaps_c_ptr = bitmaps.c_ptr(); + + // Stable storage: segs is fully built here, so the text pointers it + // provides remain valid for the lifetime of mtmd_texts/mtmd_parts. + std::vector mtmd_texts; + std::vector mtmd_parts; + mtmd_texts.reserve(segs.size()); + mtmd_parts.reserve(segs.size()); + size_t i_bm = 0; + for (const auto & sg : segs) { + if (sg.is_bitmap) { + mtmd_parts.push_back({nullptr, bitmaps_c_ptr[i_bm++]}); + } else if (!sg.text.empty()) { + mtmd_texts.push_back({ + sg.text.data(), + sg.text.size(), + /* add_special */ false, // per-part add_special is ignored + /* parse_special */ !sg.is_input, + }); + mtmd_parts.push_back({&mtmd_texts.back(), nullptr}); + } + } + + mtmd::input_chunks chunks(mtmd_input_chunks_init()); + std::vector part_ptrs(mtmd_parts.size()); + for (size_t i = 0; i < mtmd_parts.size(); i++) { + part_ptrs[i] = &mtmd_parts[i]; + } + const int32_t tokenized = mtmd_tokenize_from_parts(mctx, chunks.ptr.get(), part_ptrs.data(), part_ptrs.size(), add_special); + if (tokenized != 0) { + throw std::runtime_error("Failed to tokenize prompt"); + } + + auto result = server_tokens(chunks, true); + return result; +} + /** * break the input "prompt" object into multiple prompt if needed, then tokenize them * use tokenize_input_prompts() if the input could be an array. @@ -1344,6 +1494,15 @@ json oaicompat_chat_params_parse( llama_params["chat_format"] = static_cast(chat_params.format); llama_params["prompt"] = chat_params.prompt; + // Store prompt parts with is_input metadata for safe tokenization + // (prevents special token injection from user content) + if (!chat_params.prompt_parts.empty()) { + json parts_arr = json::array(); + for (const auto & part : chat_params.prompt_parts) { + parts_arr.push_back({{"is_input", part.is_input}, {"text", part.val}}); + } + llama_params["prompt_parts"] = parts_arr; + } if (!chat_params.grammar.empty()) { llama_params["grammar"] = chat_params.grammar; llama_params["grammar_type"] = std::string("tool_calls"); diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 6a2395714a29..6dccc692f407 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -277,6 +277,53 @@ server_tokens process_mtmd_prompt( const mtmd_helper_init_opt & init_opt, bool is_placeholder = false); +// A segment of the mtmd part list built for mtmd_tokenize_from_parts(): +// either a text segment (keeping its is_input provenance) or a placeholder +// for one media file; bitmaps are interleaved in order of marker occurrence. +struct server_mtmd_seg { + bool is_bitmap = false; + std::string text; + bool is_input = false; +}; + +/** + * Build the text/bitmap segment layout for mtmd_tokenize_from_parts() from + * prompt parts and the media marker: the text is split at every marker + * occurrence (keeping the is_input metadata of each part) and one bitmap + * segment is inserted per media file, in order of marker occurrence. + * A marker may span two adjacent parts. + * Returns false if the number of markers in the prompt does not match n_files. + */ +bool server_build_mtmd_part_layout( + const std::string & marker, + const std::vector & parts, + size_t n_files, + std::vector & out_segs); + +/** + * Tokenize a provenance-tagged chat prompt (jinja string parts carrying the + * is_input metadata), optionally interleaving the media files at the media + * marker positions. + * + * This is the single input-marking-aware tokenization path shared by the + * inference route (handle_completions_impl) and the token counting route + * (handle_count_tokens), so both see identical tokenization: + * - request-provided parts (is_input) are never parsed for special tokens, + * so stray special tokens in user/tool content cannot be injected; + * - when no is_input part contains special-token text, the prompt is + * tokenized as a whole (single pass, or the legacy mtmd path when media + * files are present), so the token ids are identical to the legacy + * whole-prompt tokenization. + */ +server_tokens server_tokenize_prompt_parts( + const llama_vocab * vocab, + mtmd_context * mctx, + const std::vector & parts, + const std::vector & files, + const mtmd_helper_init_opt & init_opt, + bool add_special, + bool is_placeholder = false); + /** * break the input "prompt" object into multiple prompt if needed, then tokenize them * this supports these cases: diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index c789c203c41e..27c8b987e754 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4985,8 +4985,21 @@ std::unique_ptr server_routes::handle_completions_impl( // process prompt std::vector inputs; - if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) { - // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below. + if (data.contains("prompt_parts") && data.at("prompt_parts").is_array() && !data.at("prompt_parts").empty()) { + // OAI-compat chat path with input-marking metadata from the jinja + // template: use the shared input-marking-aware tokenization, which + // also handles media interleaving. This must run before the MTMD + // branch so that chat requests served by multimodal-capable models + // (mctx != nullptr) receive the special-token protection too, + // even when no media is attached. + std::vector parts; + for (const auto & p : data.at("prompt_parts")) { + parts.push_back({p.at("is_input").get(), p.at("text").get()}); + } + inputs.push_back(server_tokenize_prompt_parts(ctx_server.vocab, ctx_server.mctx, parts, files, ctx_server.init_opt, /*add_special=*/true)); + } else if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) { + // This is the case used by OAI compatible chat path with MTMD (no prompt_parts, e.g. legacy templates). + // TODO It can be moved to the path below. inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files, ctx_server.init_opt)); } else { // Everything else, including multimodal completions. @@ -6214,9 +6227,19 @@ std::unique_ptr server_routes::handle_count_tokens(const l json prompt = body_parsed.at("prompt"); // SRV_DBG("prompt = %s\n", prompt.dump().c_str()); - // TODO @ngxson : refactor this code block, move this to server-common and reuse it in other places + // Use the same input-marking-aware tokenization path as the inference + // route (handle_completions_impl) when the chat template application + // produced prompt parts, so the count matches the actual completion + // tokenization (including the special-token protection and media + // interleaving). size_t n_tokens; - if (mctx != nullptr) { + if (body_parsed.contains("prompt_parts") && body_parsed.at("prompt_parts").is_array() && !body_parsed.at("prompt_parts").empty()) { + std::vector parts; + for (const auto & p : body_parsed.at("prompt_parts")) { + parts.push_back({p.at("is_input").get(), p.at("text").get()}); + } + n_tokens = server_tokenize_prompt_parts(vocab, mctx, parts, files, init_opt, /*add_special=*/true, /*is_placeholder=*/true).size(); + } else if (mctx != nullptr) { if (!prompt.is_string()) { throw std::runtime_error("for mtmd, input prompt must be a string."); }