enable allow_toolcall_in_think for Kimi-K2

common/chat-parser-xml-toolcall.cpp

@@ -724,16 +724,10 @@ inline void parse_msg_with_xml_tool_calls(common_chat_msg_parser & builder, cons
         if (reasoning_unclosed) {
             if (auto pos = content.find(end_think); pos == std::string::npos && builder.pos() != builder.input().size()) {
                 unclosed_reasoning_content += content;
-                if (form.allow_toolcall_in_think) {
+                if (!(form.allow_toolcall_in_think && tc)) {
-                    builder.move_to(tc->groups[0].begin);
-                    if (!builder.try_consume_xml_tool_calls(form)) {
-                        unclosed_reasoning_content += tool_call_start;
-                        builder.move_to(tc->groups[0].end);
-                    }
-                } else {
                     unclosed_reasoning_content += tool_call_start;
+                    continue;
                 }
-                continue;
             } else {
                 reasoning_unclosed = false;
                 std::string reasoning_content;
@@ -781,8 +775,12 @@ inline void parse_msg_with_xml_tool_calls(common_chat_msg_parser & builder, cons
                 }
             } else {
                 // This <tool_call> start is in thinking block, skip this tool call
-                auto pos = think_start + start_think.size();
+                // This <tool_call> start is in thinking block
-                unclosed_reasoning_content = content.substr(pos) + tool_call_start;
+                if (form.allow_toolcall_in_think) {
+                    unclosed_reasoning_content = content.substr(think_start + start_think.size());
+                } else {
+                    unclosed_reasoning_content = content.substr(think_start + start_think.size()) + tool_call_start;
+                }
                 reasoning_unclosed = true;
                 content.resize(think_start);
                 toolcall_in_think = true;
@@ -805,14 +803,35 @@ inline void parse_msg_with_xml_tool_calls(common_chat_msg_parser & builder, cons
         }
 
         // remove potential partial suffix
-        if (content.size() > 0 && builder.pos() == builder.input().size() && unclosed_reasoning_content.empty()) {
+        if (builder.pos() == builder.input().size()) {
-            rstrip(content);
+            if (unclosed_reasoning_content.empty()) {
-            trim_potential_partial_word(content);
+                rstrip(content);
-            rstrip(content);
+                trim_potential_partial_word(content);
+                rstrip(content);
+            } else {
+                rstrip(unclosed_reasoning_content);
+                trim_potential_partial_word(unclosed_reasoning_content);
+                rstrip(unclosed_reasoning_content);
+            }
+        }
+
+        // consume unclosed_reasoning_content if allow_toolcall_in_think is set
+        if (form.allow_toolcall_in_think && !unclosed_reasoning_content.empty()) {
+            if (builder.syntax().reasoning_format != COMMON_REASONING_FORMAT_NONE && !builder.syntax().reasoning_in_content) {
+                builder.add_reasoning_content(unclosed_reasoning_content);
+            } else {
+                if (content.empty()) {
+                    content = start_think + unclosed_reasoning_content;
+                } else {
+                    content += "\n\n" + start_think;
+                    content += unclosed_reasoning_content;
+                }
+            }
+            unclosed_reasoning_content.clear();
         }
 
         // Add content
-        if (content.size() != 0) {
+        if (!content.empty()) {
             // If there are multiple content blocks
             if (builder.syntax().reasoning_format != COMMON_REASONING_FORMAT_NONE && !builder.syntax().reasoning_in_content && builder.result().content.size() != 0) {
                 builder.add_content("\n\n");
@@ -820,7 +839,7 @@ inline void parse_msg_with_xml_tool_calls(common_chat_msg_parser & builder, cons
             builder.add_content(content);
         }
 
-        // This <tool_call> start is in thinking block, skip this tool call
+        // This <tool_call> start is in thinking block and toolcall_in_think not set, skip this tool call
         if (toolcall_in_think && !form.allow_toolcall_in_think) {
             continue;
         }
@@ -829,7 +848,7 @@ inline void parse_msg_with_xml_tool_calls(common_chat_msg_parser & builder, cons
         if (!tc) {
             GGML_ASSERT(builder.pos() == builder.input().size());
             GGML_ASSERT(unclosed_reasoning_content.empty());
-            GGML_ASSERT(!reasoning_unclosed);
+            if (!form.allow_toolcall_in_think) GGML_ASSERT(!reasoning_unclosed);
             break;
         }
 

common/chat.cpp

@@ -1976,6 +1976,7 @@ static void common_chat_parse_kimi_k2(common_chat_msg_parser & builder) {
         form.scope_end   = "<|tool_calls_section_end|>";
         form.raw_argval  = false;
         form.last_val_end = "";
+        form.allow_toolcall_in_think = true;
         return form;
     })();
     builder.consume_reasoning_with_xml_tool_calls(form, "<think>", "</think>");

tests/test-chat.cpp

@@ -428,10 +428,38 @@ static void test_templates(const struct common_chat_templates * tmpls, const std
  */
 template <typename T>
 static void test_parser_with_streaming(const common_chat_msg & expected, const std::string & raw_message, T parse_msg) {
+    constexpr auto utf8_truncate_safe = [](const std::string_view s) -> size_t {
+        auto len = s.size();
+        if (len == 0) return 0;
+        auto i = len;
+        for (size_t back = 0; back < 4 && i > 0; ++back) {
+            --i;
+            unsigned char c = s[i];
+            if ((c & 0x80) == 0) {
+                return len;
+            } else if ((c & 0xC0) == 0xC0) {
+                size_t expected_len = 0;
+                if ((c & 0xE0) == 0xC0) expected_len = 2;
+                else if ((c & 0xF0) == 0xE0) expected_len = 3;
+                else if ((c & 0xF8) == 0xF0) expected_len = 4;
+                else return i;
+                if (len - i >= expected_len) {
+                    return len;
+                } else {
+                    return i;
+                }
+            }
+        }
+        return len - std::min(len, size_t(3));
+    };
+    constexpr auto utf8_truncate_safe_view = [utf8_truncate_safe](const std::string_view s) {
+        return s.substr(0, utf8_truncate_safe(s));
+    };
+
     auto merged = simple_assist_msg("");
     auto last_msg = parse_msg("");
     for (size_t i = 1; i <= raw_message.size(); ++i) {
-        auto curr_msg = parse_msg(raw_message.substr(0, i));
+        auto curr_msg = parse_msg(std::string(utf8_truncate_safe_view(std::string_view(raw_message).substr(0, i))));
         if (curr_msg == simple_assist_msg("")) continue;
         LOG_INF("Streaming msg: %s\n", common_chat_msgs_to_json_oaicompat<json>({curr_msg}).dump().c_str());
         for (auto diff: common_chat_msg_diff::compute_diffs(last_msg, curr_msg)) {
@@ -2752,6 +2780,93 @@ Hey there!<|im_end|>
                 "{\"args\": [{\"path\": \"src/providers/ThemeProvider.tsx\"}, {\"path\": \"src/components/Header.tsx\"}, {\"path\": \"src/components/ThemeToggle.tsx\"}, {\"path\": \"src/app/globals.css\"}, {\"path\": \"src/app/layout.tsx\"}]}"
                 "<|tool_call_end|><|tool_calls_section_end|>",
             [&](const std::string &msg) { return common_chat_parse(msg, /* is_partial= */ true, {COMMON_CHAT_FORMAT_KIMI_K2}); });
+        test_parser_with_streaming(
+                simple_assist_msg(
+                        "Let me start by examining the relevant files to understand the current implementation.", "",
+                        "read_file",
+                        "{\"files\": [{\"path\": \"src/app/Partners.tsx\", \"line_ranges\": [\"1-100\"]}]}"),
+                "Let me start by examining the relevant files to understand the current implementation."
+                "<|tool_calls_section_begin|><|tool_call_begin|>functions.read_file:0<|tool_call_argument_begin|>"
+                "{\"files\":[{\"path\":\"src/app/Partners.tsx\",\"line_ranges\":[\"1-100\"]}]}"
+                "<|tool_call_end|><|tool_calls_section_end|>",
+            [&](const std::string &msg) { return common_chat_parse(msg, /* is_partial= */ true, {COMMON_CHAT_FORMAT_KIMI_K2}); });
+        auto multi_tool_msg = simple_assist_msg("Let me call multiple tools.", "I'm thinking.");
+        multi_tool_msg.tool_calls.push_back({ "read_file", "{\"files\": [{\"path\": \"src/app/Partners.tsx\", \"line_ranges\": [\"1-100\"]}]}", "" });
+        multi_tool_msg.tool_calls.push_back({ "web_search", "{\"query\":\"\\\"From Zero\\\" Linkin Park album tracklist complete songs\",\"limit\":3,\"type\":\"text\"}", "" });
+        multi_tool_msg.tool_calls.push_back({ "complex_function", "{\"name\": \"John Doe\", \"age\": 30, \"active\": true, \"score\": 95.5}", "" });
+        multi_tool_msg.tool_calls.push_back({ "emoji_function", "{\"message\":\"Hello! 👋 🌟 🚀 Testing emojis: 😀😃😄😁 and symbols: ∑∏∆∇\"}", "" });
+        test_parser_with_streaming(multi_tool_msg,
+                "<think>I'm thinking.</think>Let me call multiple tools."
+                "<|tool_calls_section_begin|>"
+                "<|tool_call_begin|>functions.read_file:0<|tool_call_argument_begin|>"
+                "{\"files\":[{\"path\":\"src/app/Partners.tsx\",\"line_ranges\":[\"1-100\"]}]}"
+                "<|tool_call_end|>"
+                "<|tool_call_begin|>functions.web_search:1<|tool_call_argument_begin|>"
+                "{\"query\":\"\\\"From Zero\\\" Linkin Park album tracklist complete songs\",\"limit\":3,\"type\":\"text\"}"
+                "<|tool_call_end|>"
+                "<|tool_call_begin|>functions.complex_function:2<|tool_call_argument_begin|>"
+                "{\"name\": \"John Doe\", \"age\": 30, \"active\": true, \"score\": 95.5}"
+                "<|tool_call_end|>"
+                "<|tool_call_begin|>functions.emoji_function:3<|tool_call_argument_begin|>"
+                "{\"message\":\"Hello! 👋 🌟 🚀 Testing emojis: 😀😃😄😁 and symbols: ∑∏∆∇\"}"
+                "<|tool_call_end|>"
+                "<|tool_calls_section_end|>",
+            [&](const std::string &msg) { return common_chat_parse(msg, /* is_partial= */ true, {
+                    COMMON_CHAT_FORMAT_KIMI_K2,
+                    COMMON_REASONING_FORMAT_DEEPSEEK
+            }); });
+        test_parser_with_streaming(
+                simple_assist_msg("", "I'm thinking", "complex_function_in_think", "{\"name\":\"John Doe\",\"age\":30,\"active\":true,\"score\":95.5}"),
+                "<think>I'm thinking<|tool_calls_section_begin|><|tool_call_begin|>functions.complex_function_in_think:0<|tool_call_argument_begin|>"
+                "{\"name\": \"John Doe\", \"age\": 30, \"active\": true, \"score\": 95.5}"
+                "<|tool_call_end|><|tool_calls_section_end|>",
+            [&](const std::string &msg) { return common_chat_parse(msg, /* is_partial= */ true, {
+                    COMMON_CHAT_FORMAT_KIMI_K2,
+                    COMMON_REASONING_FORMAT_DEEPSEEK
+            }); });
+        test_parser_with_streaming(
+                simple_assist_msg("Hello", "I'm thinkingI'm still thinking", "complex_function_in_think", "{\"name\":\"John Doe\",\"age\":30,\"active\":true,\"score\":95.5}"),
+                "<think>I'm thinking<|tool_calls_section_begin|><|tool_call_begin|>functions.complex_function_in_think:0<|tool_call_argument_begin|>"
+                "{\"name\": \"John Doe\", \"age\": 30, \"active\": true, \"score\": 95.5}"
+                "<|tool_call_end|><|tool_calls_section_end|>I'm still thinking</think>Hello",
+            [&](const std::string &msg) { return common_chat_parse(msg, /* is_partial= */ true, {
+                    COMMON_CHAT_FORMAT_KIMI_K2,
+                    COMMON_REASONING_FORMAT_DEEPSEEK
+            }); });
+
+        // Test template rendering
+        common_chat_templates_inputs conversation_with_tools = inputs_tools;
+        conversation_with_tools.messages.push_back(simple_assist_msg("Let's do it", "Think first", "complex_function", "{\"name\":\"John Doe\",\"age\":30,\"active\":true,\"score\":95.5}"));
+        conversation_with_tools.messages.push_back({
+            "tool",
+            "Tool response 1",
+            /* .content_parts = */ {},
+            /* .tool_calls = */ {},
+            /* .reasoning_content = */ "",
+            /* .tool_name = */ "complex_function",
+            /* .tool_call_id = */ "",
+        });
+        conversation_with_tools.messages.push_back(simple_assist_msg("Continue", "Think next", "web_search", "{\"query\":\"\\\"From Zero\\\" Linkin Park album tracklist complete songs\",\"limit\":3,\"type\":\"text\"}"));
+        conversation_with_tools.messages.push_back({
+            "tool",
+            "Tool response 2",
+            /* .content_parts = */ {},
+            /* .tool_calls = */ {},
+            /* .reasoning_content = */ "",
+            /* .tool_name = */ "web_search",
+            /* .tool_call_id = */ "",
+        });
+        conversation_with_tools.messages.push_back(simple_assist_msg("CC", "Think last", "read_file", "{\"args\": [{\"path\": \"src/providers/ThemeProvider.tsx\"}, {\"path\": \"src/components/Header.tsx\"}, {\"path\": \"src/components/ThemeToggle.tsx\"}, {\"path\": \"src/app/globals.css\"}, {\"path\": \"src/app/layout.tsx\"}]}"));
+        conversation_with_tools.messages.push_back({
+            "tool",
+            "Tool response 3",
+            /* .content_parts = */ {},
+            /* .tool_calls = */ {},
+            /* .reasoning_content = */ "",
+            /* .tool_name = */ "read_file",
+            /* .tool_call_id = */ "",
+        });
+        assert_equals(common_chat_templates_apply(tmpls.get(), conversation_with_tools).prompt, std::string("<|im_system|>tool_declare<|im_middle|>[{\"type\": \"function\", \"function\": {\"name\": \"special_function\", \"description\": \"I'm special\", \"parameters\": {\"type\": \"object\", \"properties\": {\"arg1\": {\"type\": \"integer\", \"description\": \"The arg.\"}}, \"required\": [\"arg1\"]}}}]<|im_end|><|im_system|>system<|im_middle|>You are Kimi, an AI assistant created by Moonshot AI.<|im_end|><|im_user|>user<|im_middle|>Hey there!<|im_end|><|im_assistant|>assistant<|im_middle|><think>Think first</think>Let's do it<|tool_calls_section_begin|><|tool_call_begin|>functions.complex_function:0<|tool_call_argument_begin|>{\"name\":\"John Doe\",\"age\":30,\"active\":true,\"score\":95.5}<|tool_call_end|><|tool_calls_section_end|><|im_end|><|im_system|>complex_function<|im_middle|>## Return of functions.complex_function:0\nTool response 1<|im_end|><|im_assistant|>assistant<|im_middle|><think>Think next</think>Continue<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:1<|tool_call_argument_begin|>{\"query\":\"\\\"From Zero\\\" Linkin Park album tracklist complete songs\",\"limit\":3,\"type\":\"text\"}<|tool_call_end|><|tool_calls_section_end|><|im_end|><|im_system|>web_search<|im_middle|>## Return of functions.web_search:1\nTool response 2<|im_end|><|im_assistant|>assistant<|im_middle|><think>Think last</think>CC<|tool_calls_section_begin|><|tool_call_begin|>functions.read_file:2<|tool_call_argument_begin|>{\"args\": [{\"path\": \"src/providers/ThemeProvider.tsx\"}, {\"path\": \"src/components/Header.tsx\"}, {\"path\": \"src/components/ThemeToggle.tsx\"}, {\"path\": \"src/app/globals.css\"}, {\"path\": \"src/app/layout.tsx\"}]}<|tool_call_end|><|tool_calls_section_end|><|im_end|><|im_system|>read_file<|im_middle|>## Return of functions.read_file:2\nTool response 3<|im_end|><|im_assistant|>assistant<|im_middle|>"));
 
         // Test template generation for regular content
         test_templates(tmpls.get(), end_tokens, message_assist, tools,