From 8f5ab832ca7d8a7b4f23687693fb8b0ecbc227e7 Mon Sep 17 00:00:00 2001
From: Matt Thompson <111157855+boondocklabs@users.noreply.github.com>
Date: Fri, 24 Jul 2026 03:54:47 -0700
Subject: [PATCH] cohere2 moe template parser: enforce JSON schema for text
 responses if a response schema is provided (#26018)

---
 common/chat.cpp     | 19 ++++++++++++++-----
 tests/test-chat.cpp | 11 +++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/common/chat.cpp b/common/chat.cpp
index c441021b9..2b0b40ccb 100644
--- a/common/chat.cpp
+++ b/common/chat.cpp
@@ -2179,9 +2179,10 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
         { COMMON_CHAT_ROLE_SYSTEM,    TURN_START + SYSTEM },
     };
 
-    auto has_tools         = inputs.tools.is_array() && !inputs.tools.empty();
-    auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
-    auto include_grammar   = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;
+    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;
@@ -2212,7 +2213,11 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
                                              p.optional(p.literal(THINK_END))));
         }
 
-        auto text_content = p.literal(TEXT_START) + p.content(p.until(TEXT_END)) + p.optional(p.literal(TEXT_END));
+        auto text_content = has_response_format
+            ? p.literal(TEXT_START) +
+                p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) +
+                p.optional(p.literal(TEXT_END))
+            : p.literal(TEXT_START) + p.content(p.until(TEXT_END)) + p.optional(p.literal(TEXT_END));
 
         if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) {
             return generation_prompt + reasoning + text_content + p.optional(p.literal(TURN_END)) + end;
@@ -2240,13 +2245,17 @@ static common_chat_params common_chat_params_init_cohere2moe(const common_chat_t
     data.parser = parser.save();
 
     if (include_grammar) {
-        data.grammar_lazy = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
+        data.grammar_lazy = !has_response_format && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
         data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
             foreach_function(inputs.tools, [&](const json & tool) {
                 const auto & function = tool.at("function");
                 auto         schema   = function.at("parameters");
                 builder.resolve_refs(schema);
             });
+            if (has_response_format) {
+                auto schema = inputs.json_schema;
+                builder.resolve_refs(schema);
+            }
             parser.build_grammar(builder, data.grammar_lazy);
         });
 
diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp
index 28b28ca0f..371a305fc 100644
--- a/tests/test-chat.cpp
+++ b/tests/test-chat.cpp
@@ -2849,6 +2849,17 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .expect(message_assist)
             .run();
 
+        // JSON output schema
+        tst.test(
+                "I need to output the invoice details in JSON<|END_THINKING|>"
+                "<|START_TEXT|>{\"amount\": 123.45, \"date\": \"2025-12-03\"}<|END_TEXT|>")
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .json_schema(invoice_schema)
+            .tools({ special_function_tool })
+            .expect_reasoning("I need to output the invoice details in JSON")
+            .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})")
+                .run();
+
         // Single tool call with reasoning.
         tst.test(
                "I'm\nthinking<|END_THINKING|>"