llama-cpp-sys-4 0.9.0

#include "parsers.h"

// K2 Horizon format:
// - Reasoning: <ifm|think>...</ifm|think>, or <ifm|think_fast>/<ifm|think_faster> for medium/low reasoning_effort
// - Tool calls: <ifm|tool_calls><ifm|tool_call>...</ifm|tool_call>...</ifm|tool_calls>, one call per <ifm|tool_call>:
//   xml (default): name <ifm|arg_key>k</ifm|arg_key> [<ifm|arg_type>t</ifm|arg_type>] <ifm|arg_value>v</ifm|arg_value> ...
//   json:          {"name": "...", "arguments": {...}}
common_chat_params common_chat_params_init_k2_horizon(const common_chat_template &          tmpl,
                                                      const autoparser::generation_params & inputs) {
    common_chat_params data;

    // The template requires a thinking field on every assistant message
    auto messages = inputs.messages;
    for (auto & msg : messages) {
        if (msg.value("role", "") == "assistant" && !msg.contains("reasoning_content")) {
            msg["reasoning_content"] = "";
        }
    }

    data.prompt            = common_chat_template_direct_apply_impl(tmpl, inputs, messages);
    data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, messages);
    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
    data.supports_thinking = true;

    const std::string effort      = inputs.extra_context.value("reasoning_effort", "high");
    const std::string call_format = inputs.extra_context.value("tool_call_format", "xml");

    // Templates that handle enable_thinking disable it with an empty <ifm|think></ifm|think> block for every effort
    const bool thinking_off = !inputs.enable_thinking && tmpl.source().find("enable_thinking") != std::string::npos;
    const std::string think = thinking_off       ? "ifm|think"        :
                              effort == "medium" ? "ifm|think_fast"   :
                              effort == "low"    ? "ifm|think_faster" : "ifm|think";

    const std::string GEN_PREFIX    = "<|ifm|im_start|>assistant\n";
    const std::string THINK_START   = "<" + think + ">";
    const std::string THINK_END     = "</" + think + ">";
    const std::string SECTION_START = "<ifm|tool_calls>";
    const std::string SECTION_END   = "</ifm|tool_calls>";
    const std::string CALL_START    = "<ifm|tool_call>";
    const std::string CALL_END      = "</ifm|tool_call>";
    const std::string ARG_KEY       = "<ifm|arg_key>";
    const std::string ARG_KEY_END   = "</ifm|arg_key>";
    const std::string ARG_TYPE      = "<ifm|arg_type>";
    const std::string ARG_TYPE_END  = "</ifm|arg_type>";
    const std::string ARG_VAL       = "<ifm|arg_value>";
    const std::string ARG_VAL_END   = "</ifm|arg_value>";

    data.thinking_start_tag = THINK_START;
    data.thinking_end_tags  = { THINK_END };

    data.preserved_tokens = data.thinking_end_tags;
    data.preserved_tokens.insert(data.preserved_tokens.end(), {
        THINK_START, SECTION_START, SECTION_END, CALL_START, CALL_END,
        ARG_KEY, ARG_KEY_END, ARG_TYPE, ARG_TYPE_END, ARG_VAL, ARG_VAL_END,
    });

    data.message_delimiters = {
        { COMMON_CHAT_ROLE_ASSISTANT, "<|ifm|im_start|>assistant" },
        { COMMON_CHAT_ROLE_USER,      "<|ifm|im_start|>user"      },
        { COMMON_CHAT_ROLE_TOOL,      "<|ifm|im_start|>tool"      },
        { COMMON_CHAT_ROLE_SYSTEM,    "<|ifm|im_start|>system"    },
    };

    auto has_tools           = inputs.tools.is_array() && !inputs.tools.empty();
    auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty();
    auto extract_reasoning   = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
    auto include_grammar     = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE);

    if (inputs.has_continuation()) {
        const auto & msg = inputs.continue_msg;

        data.generation_prompt = GEN_PREFIX + THINK_START + "\n" + msg.reasoning_content;
        if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
            data.generation_prompt += THINK_END + msg.render_content();
        }

        data.prompt.append(data.generation_prompt);
    }

    data.parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
        auto generation_prompt = p.literal(GEN_PREFIX);

        auto think_end = p.choice();
        for (const auto & tag : data.thinking_end_tags) {
            think_end |= p.literal(tag);
        }
        auto think_body  = p.until_one_of(data.thinking_end_tags);
        auto think_block = [&](const common_peg_parser & body) {
            return p.optional(THINK_START + p.space() + p.ac(body + think_end, data.thinking_end_tags));
        };
        auto reasoning = extract_reasoning ? think_block(p.reasoning(think_body)) : p.eps();

        if (has_response_format) {
            // The answer must be bare JSON, so the think block is consumed even when it is not extracted
            auto thoughts = extract_reasoning ? reasoning : think_block(think_body);
            return generation_prompt + (thoughts << p.content(p.schema(p.json(), "response-format", inputs.json_schema)));
        }

        if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) {
            return generation_prompt + (reasoning << p.content(p.rest()));
        }

        auto tool_choice = p.choice();
        if (call_format == "json") {
            tool_choice = p.standard_json_tools(CALL_START, CALL_END, inputs.tools, false, true);
        } else {
            auto arg_close  = p.tool_arg_close(p.literal(ARG_VAL_END));
            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));

            // The models leave out <ifm|arg_type> even when asked for xml_typed
            auto arg_type = call_format == "xml_typed" ? p.optional(ARG_TYPE + p.until(ARG_TYPE_END) + ARG_TYPE_END + p.space()) : p.eps();

            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
                const auto & function = tool.at("function");
                std::string  name     = function.at("name");

                std::vector<common_peg_parser> required_args;
                std::vector<common_peg_parser> optional_args;
                foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
                    auto rule_name = "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index);
                    auto types     = param.schema->value_types();
                    auto arg_value = arg_string;
                    if (!types.has(common_chat_schema::TYPE_STRING)) {
                        arg_value = p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", doc, *param.schema)) + arg_close;
                    }
                    if (types.has(common_chat_schema::TYPE_STRING) && !types.is_only(common_chat_schema::TYPE_STRING)) {
                        // The string alternative accepts any text, so only the parser needs the JSON alternatives.
                        auto json_value = p.choice();
                        if (types.has(common_chat_schema::TYPE_OBJECT)) {
                            json_value |= p.json_object();
                        }
                        if (types.has(common_chat_schema::TYPE_ARRAY)) {
                            json_value |= p.json_array();
                        }
                        if (types.has(common_chat_schema::TYPE_NUMBER) || types.has(common_chat_schema::TYPE_INTEGER)) {
                            json_value |= p.json_number();
                        }
                        if (types.has(common_chat_schema::TYPE_BOOLEAN)) {
                            json_value |= p.json_bool();
                        }
                        if (types.has(common_chat_schema::TYPE_NULL)) {
                            json_value |= p.json_null();
                        }
                        arg_value = p.gbnf(p.atomic(p.tool_arg_json_value(json_value) + arg_close) | arg_string, "xml-arg-string");
                    }

                    auto arg = p.space() + p.tool_arg(p.tool_arg_open(ARG_KEY + p.tool_arg_name(p.literal(param.name)) + ARG_KEY_END) <<
                                                      arg_type + ARG_VAL + arg_value);
                    (param.required ? required_args : optional_args).push_back(p.rule(rule_name, arg));
                });

                auto args = p.permute("tool-" + std::to_string(tool_index) + "-args", required_args);
                if (!optional_args.empty()) {
                    args = args + p.zero_or_more(p.choice(optional_args));
                }

                tool_choice |= p.rule("tool-" + std::to_string(tool_index), p.tool(
                    p.tool_open(CALL_START + p.tool_name(p.literal(name)) + "\n") + p.tool_args(args) << p.tool_close(p.literal(CALL_END))));
            });
        }

        auto required   = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;
        auto calls      = inputs.parallel_tool_calls ? tool_choice + p.zero_or_more(p.space() + tool_choice) : tool_choice;
        auto tool_calls = p.trigger_rule("tool-calls", p.repeat(SECTION_START << calls << SECTION_END, required ? 1 : 0, 1));

        // Keep thinking inline when required calls bypass the content parser.
        if (required && !extract_reasoning) {
            reasoning = p.content(think_block(think_body));
        }

        // A required call follows the reasoning directly, the models otherwise keep writing content
        auto content = required ? p.eps() : p.content(p.until(SECTION_START));

        return generation_prompt + (reasoning << content << tool_calls);
    });

    if (include_grammar) {
        data.grammar_lazy = !(has_response_format || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED);
        data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
            data.parser.build_grammar(builder, data.grammar_lazy);
        });

        if (data.grammar_lazy) {
            data.grammar_triggers = {
                { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, SECTION_START },
            };
        }
    }

    return data;
}