|
| 1 | +#include "parsers.h" |
| 2 | + |
| 3 | +// K2 Horizon format: |
| 4 | +// - Reasoning: <ifm|think>...</ifm|think>, or <ifm|think_fast>/<ifm|think_faster> for medium/low reasoning_effort |
| 5 | +// - Tool calls: <ifm|tool_calls><ifm|tool_call>...</ifm|tool_call>...</ifm|tool_calls>, one call per <ifm|tool_call>: |
| 6 | +// xml (default): name <ifm|arg_key>k</ifm|arg_key> [<ifm|arg_type>t</ifm|arg_type>] <ifm|arg_value>v</ifm|arg_value> ... |
| 7 | +// json: {"name": "...", "arguments": {...}} |
| 8 | +common_chat_params common_chat_params_init_k2_horizon(const common_chat_template & tmpl, |
| 9 | + const autoparser::generation_params & inputs) { |
| 10 | + common_chat_params data; |
| 11 | + |
| 12 | + // The template requires a thinking field on every assistant message |
| 13 | + auto messages = inputs.messages; |
| 14 | + for (auto & msg : messages) { |
| 15 | + if (msg.value("role", "") == "assistant" && !msg.contains("reasoning_content")) { |
| 16 | + msg["reasoning_content"] = ""; |
| 17 | + } |
| 18 | + } |
| 19 | + |
| 20 | + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, messages); |
| 21 | + data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, messages); |
| 22 | + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; |
| 23 | + data.supports_thinking = true; |
| 24 | + |
| 25 | + const std::string effort = inputs.extra_context.value("reasoning_effort", "high"); |
| 26 | + const std::string call_format = inputs.extra_context.value("tool_call_format", "xml"); |
| 27 | + |
| 28 | + // Templates that handle enable_thinking disable it with an empty <ifm|think></ifm|think> block for every effort |
| 29 | + const bool thinking_off = !inputs.enable_thinking && tmpl.source().find("enable_thinking") != std::string::npos; |
| 30 | + const std::string think = thinking_off ? "ifm|think" : |
| 31 | + effort == "medium" ? "ifm|think_fast" : |
| 32 | + effort == "low" ? "ifm|think_faster" : "ifm|think"; |
| 33 | + |
| 34 | + const std::string GEN_PREFIX = "<|ifm|im_start|>assistant\n"; |
| 35 | + const std::string THINK_START = "<" + think + ">"; |
| 36 | + const std::string THINK_END = "</" + think + ">"; |
| 37 | + const std::string SECTION_START = "<ifm|tool_calls>"; |
| 38 | + const std::string SECTION_END = "</ifm|tool_calls>"; |
| 39 | + const std::string CALL_START = "<ifm|tool_call>"; |
| 40 | + const std::string CALL_END = "</ifm|tool_call>"; |
| 41 | + const std::string ARG_KEY = "<ifm|arg_key>"; |
| 42 | + const std::string ARG_KEY_END = "</ifm|arg_key>"; |
| 43 | + const std::string ARG_TYPE = "<ifm|arg_type>"; |
| 44 | + const std::string ARG_TYPE_END = "</ifm|arg_type>"; |
| 45 | + const std::string ARG_VAL = "<ifm|arg_value>"; |
| 46 | + const std::string ARG_VAL_END = "</ifm|arg_value>"; |
| 47 | + |
| 48 | + data.thinking_start_tag = THINK_START; |
| 49 | + data.thinking_end_tags = { THINK_END }; |
| 50 | + |
| 51 | + data.preserved_tokens = data.thinking_end_tags; |
| 52 | + data.preserved_tokens.insert(data.preserved_tokens.end(), { |
| 53 | + THINK_START, SECTION_START, SECTION_END, CALL_START, CALL_END, |
| 54 | + ARG_KEY, ARG_KEY_END, ARG_TYPE, ARG_TYPE_END, ARG_VAL, ARG_VAL_END, |
| 55 | + }); |
| 56 | + |
| 57 | + data.message_delimiters = { |
| 58 | + { COMMON_CHAT_ROLE_ASSISTANT, "<|ifm|im_start|>assistant" }, |
| 59 | + { COMMON_CHAT_ROLE_USER, "<|ifm|im_start|>user" }, |
| 60 | + { COMMON_CHAT_ROLE_TOOL, "<|ifm|im_start|>tool" }, |
| 61 | + { COMMON_CHAT_ROLE_SYSTEM, "<|ifm|im_start|>system" }, |
| 62 | + }; |
| 63 | + |
| 64 | + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); |
| 65 | + auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty(); |
| 66 | + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; |
| 67 | + auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); |
| 68 | + |
| 69 | + if (inputs.has_continuation()) { |
| 70 | + const auto & msg = inputs.continue_msg; |
| 71 | + |
| 72 | + data.generation_prompt = GEN_PREFIX + THINK_START + "\n" + msg.reasoning_content; |
| 73 | + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { |
| 74 | + data.generation_prompt += THINK_END + msg.render_content(); |
| 75 | + } |
| 76 | + |
| 77 | + data.prompt += data.generation_prompt; |
| 78 | + } |
| 79 | + |
| 80 | + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { |
| 81 | + auto generation_prompt = p.literal(GEN_PREFIX); |
| 82 | + |
| 83 | + auto think_end = p.choice(); |
| 84 | + for (const auto & tag : data.thinking_end_tags) { |
| 85 | + think_end |= p.literal(tag); |
| 86 | + } |
| 87 | + auto think_body = p.until_one_of(data.thinking_end_tags); |
| 88 | + auto think_block = [&](const common_peg_parser & body) { |
| 89 | + return p.optional(THINK_START + p.space() + p.ac(body + think_end, data.thinking_end_tags)); |
| 90 | + }; |
| 91 | + auto reasoning = extract_reasoning ? think_block(p.reasoning(think_body)) : p.eps(); |
| 92 | + |
| 93 | + if (has_response_format) { |
| 94 | + // The answer must be bare JSON, so the think block is consumed even when it is not extracted |
| 95 | + auto thoughts = extract_reasoning ? reasoning : think_block(think_body); |
| 96 | + return generation_prompt + (thoughts << p.content(p.schema(p.json(), "response-format", inputs.json_schema))); |
| 97 | + } |
| 98 | + |
| 99 | + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { |
| 100 | + return generation_prompt + (reasoning << p.content(p.rest())); |
| 101 | + } |
| 102 | + |
| 103 | + auto tool_choice = p.choice(); |
| 104 | + if (call_format == "json") { |
| 105 | + tool_choice = p.standard_json_tools(CALL_START, CALL_END, inputs.tools, false, true); |
| 106 | + } else { |
| 107 | + auto arg_close = p.tool_arg_close(p.literal(ARG_VAL_END)); |
| 108 | + auto arg_string = p.rule("xml-arg-string", p.ac(p.tool_arg_string_value(p.until(ARG_VAL_END)) + arg_close, ARG_VAL_END)); |
| 109 | + |
| 110 | + // The models leave out <ifm|arg_type> even when asked for xml_typed |
| 111 | + auto arg_type = call_format == "xml_typed" ? p.optional(ARG_TYPE + p.until(ARG_TYPE_END) + ARG_TYPE_END + p.space()) : p.eps(); |
| 112 | + |
| 113 | + foreach_function(inputs.tools, [&](const json & tool) { |
| 114 | + const auto & function = tool.at("function"); |
| 115 | + std::string name = function.at("name"); |
| 116 | + |
| 117 | + std::vector<common_peg_parser> required_args; |
| 118 | + std::vector<common_peg_parser> optional_args; |
| 119 | + foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) { |
| 120 | + auto rule_name = "tool-" + name + "-arg-" + param.name; |
| 121 | + auto types = param.schema->value_types(); |
| 122 | + auto arg_value = arg_string; |
| 123 | + if (!types.has(common_chat_schema::TYPE_STRING)) { |
| 124 | + arg_value = p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", doc, *param.schema)) + arg_close; |
| 125 | + } |
| 126 | + if (types.has(common_chat_schema::TYPE_STRING) && !types.is_only(common_chat_schema::TYPE_STRING)) { |
| 127 | + // The string alternative accepts any text, so only the parser needs the JSON alternatives. |
| 128 | + auto json_value = p.choice(); |
| 129 | + if (types.has(common_chat_schema::TYPE_OBJECT)) { |
| 130 | + json_value |= p.json_object(); |
| 131 | + } |
| 132 | + if (types.has(common_chat_schema::TYPE_ARRAY)) { |
| 133 | + json_value |= p.json_array(); |
| 134 | + } |
| 135 | + if (types.has(common_chat_schema::TYPE_NUMBER) || types.has(common_chat_schema::TYPE_INTEGER)) { |
| 136 | + json_value |= p.json_number(); |
| 137 | + } |
| 138 | + if (types.has(common_chat_schema::TYPE_BOOLEAN)) { |
| 139 | + json_value |= p.json_bool(); |
| 140 | + } |
| 141 | + if (types.has(common_chat_schema::TYPE_NULL)) { |
| 142 | + json_value |= p.json_null(); |
| 143 | + } |
| 144 | + arg_value = p.gbnf(p.atomic(p.tool_arg_json_value(json_value) + arg_close) | arg_string, "xml-arg-string"); |
| 145 | + } |
| 146 | + |
| 147 | + auto arg = p.space() + p.tool_arg(p.tool_arg_open(ARG_KEY + p.tool_arg_name(p.literal(param.name)) + ARG_KEY_END) << |
| 148 | + arg_type + ARG_VAL + arg_value); |
| 149 | + (param.required ? required_args : optional_args).push_back(p.rule(rule_name, arg)); |
| 150 | + }); |
| 151 | + |
| 152 | + auto args = p.permute("tool-" + name + "-args", required_args); |
| 153 | + if (!optional_args.empty()) { |
| 154 | + args = args + p.zero_or_more(p.choice(optional_args)); |
| 155 | + } |
| 156 | + |
| 157 | + tool_choice |= p.rule("tool-" + name, p.tool( |
| 158 | + p.tool_open(CALL_START + p.tool_name(p.literal(name)) + "\n") + p.tool_args(args) << p.tool_close(p.literal(CALL_END)))); |
| 159 | + }); |
| 160 | + } |
| 161 | + |
| 162 | + auto required = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED; |
| 163 | + auto calls = inputs.parallel_tool_calls ? tool_choice + p.zero_or_more(p.space() + tool_choice) : tool_choice; |
| 164 | + auto tool_calls = p.trigger_rule("tool-calls", p.repeat(SECTION_START << calls << SECTION_END, required ? 1 : 0, 1)); |
| 165 | + |
| 166 | + // Keep thinking inline when required calls bypass the content parser. |
| 167 | + if (required && !extract_reasoning) { |
| 168 | + reasoning = p.content(think_block(think_body)); |
| 169 | + } |
| 170 | + |
| 171 | + // A required call follows the reasoning directly, the models otherwise keep writing content |
| 172 | + auto content = required ? p.eps() : p.content(p.until(SECTION_START)); |
| 173 | + |
| 174 | + return generation_prompt + (reasoning << content << tool_calls); |
| 175 | + }); |
| 176 | + |
| 177 | + data.parser = parser.save(); |
| 178 | + |
| 179 | + if (include_grammar) { |
| 180 | + data.grammar_lazy = !(has_response_format || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED); |
| 181 | + data.grammar = build_grammar([&](const common_grammar_builder & builder) { |
| 182 | + parser.build_grammar(builder, data.grammar_lazy); |
| 183 | + }); |
| 184 | + |
| 185 | + if (data.grammar_lazy) { |
| 186 | + data.grammar_triggers = { |
| 187 | + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, SECTION_START }, |
| 188 | + }; |
| 189 | + } |
| 190 | + } |
| 191 | + |
| 192 | + return data; |
| 193 | +} |
0 commit comments