reasoning-budget: Implementing thinking-budget mechanism to control thought process

common/arg.cpp

@@ -3555,6 +3555,36 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING"));
+
+    // unescape \n, \t, \r, \\, \" in message strings passed via CLI/env
+    static auto unescape = [](const std::string & s) {
+        std::string r;
+        r.reserve(s.size());
+        for (size_t i = 0; i < s.size(); i++) {
+            if (s[i] == '\\' && i + 1 < s.size()) {
+                switch (s[i + 1]) {
+                    case 'n':  r += '\n'; i++; break;
+                    case 't':  r += '\t'; i++; break;
+                    case 'r':  r += '\r'; i++; break;
+                    case '\\': r += '\\'; i++; break;
+                    case '"':  r += '"';  i++; break;
+                    default:   r += s[i];      break;
+                }
+            } else {
+                r += s[i];
+            }
+        }
+        return r;
+    };
+
+    add_opt(common_arg(
+        {"--reasoning-budget-enable"},
+        {"--no-reasoning-budget-enable"},
+        "master switch for the reasoning budget mechanism (hard cutoff, soft warning, intro message, grace period, and the runtime reasoning-control endpoint); if disabled, none of the other --reasoning-budget-* settings take effect regardless of their values (default: disabled)",
+        [](common_params & params, bool value) {
+            params.sampling.reasoning_budget_enabled = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_ENABLE"));
     add_opt(common_arg(
         {"--reasoning-budget"}, "N",
         "token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)",
@@ -3565,11 +3595,40 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET"));
     add_opt(common_arg(
         {"--reasoning-budget-message"}, "MESSAGE",
-        "message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)",
+        "message forced when the reasoning budget is exhausted; should include the model's own closing tag (e.g. </think>) as it is not appended automatically, since the exact tag can differ between models/templates. If empty, falls back to forcing just the auto-detected closing tag alone so the block still always closes (default: none)",
         [](common_params & params, const std::string & value) {
-            params.sampling.reasoning_budget_message = value;
+            params.sampling.reasoning_budget_message = unescape(value);
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_MESSAGE"));
+    add_opt(common_arg(
+        {"--reasoning-budget-soft-ratio"}, "N",
+        "fraction of the reasoning budget consumed at which to inject a soft warning message before the hard cutoff: <= 0 disables, (0,1] enables (default: -1)",
+        [](common_params & params, const std::string & value) {
+            params.sampling.reasoning_budget_soft_ratio = std::stof(value);
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_SOFT_RATIO"));
+    add_opt(common_arg(
+        {"--reasoning-budget-soft-message"}, "MESSAGE",
+        "message injected at the soft reasoning budget threshold, before the hard cutoff (default: none)",
+        [](common_params & params, const std::string & value) {
+            params.sampling.reasoning_budget_soft_message = unescape(value);
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_SOFT_MESSAGE"));
+    add_opt(common_arg(
+        {"--reasoning-budget-intro-message"}, "MESSAGE",
+        string_format("message forced immediately when the reasoning block starts, announcing the token budget; use {budget} as a placeholder for the configured reasoning budget (default: '%s')",
+            params.sampling.reasoning_budget_intro_message.c_str()),
+        [](common_params & params, const std::string & value) {
+            params.sampling.reasoning_budget_intro_message = unescape(value);
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_INTRO_MESSAGE"));
+    add_opt(common_arg(
+        {"--reasoning-budget-grace-tokens"}, "N",
+        "once the reasoning budget is exhausted, wait up to N tokens for a paragraph break before forcing the cutoff, instead of forcing immediately (default: 0, disabled)",
+        [](common_params & params, int value) {
+            params.sampling.reasoning_budget_grace_tokens = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_GRACE_TOKENS"));
     add_opt(common_arg(
         {"--reasoning-preserve"},
         {"--no-reasoning-preserve"},

common/common.h

@@ -285,13 +285,24 @@ struct common_params_sampling {
 
     // reasoning budget sampler parameters
     // these are populated by the server/CLI based on chat template params
+    bool                      reasoning_budget_enabled  = false; // master switch: must be true for the budget/soft/intro/grace mechanism to run at all
     int32_t                   reasoning_budget_tokens   = -1;  // -1 = disabled, >= 0 = token budget
     std::vector<llama_token>  reasoning_budget_start;          // start tag token sequence
     std::vector<llama_tokens> reasoning_budget_end;            // end tag token sequences; the first tag is used as the forcing sequence
     std::vector<llama_token>  reasoning_budget_forced;         // forced sequence (message + first end tag)
     std::string               reasoning_budget_message;        // message injected before end tag when budget exhausted
     bool                      reasoning_control = false;       // create the budget sampler on demand so reasoning can be ended at runtime
 
+    float                     reasoning_budget_soft_ratio   = -1.0f; // <= 0 = disabled, (0,1] = fraction of budget at which to warn
+    std::vector<llama_token>  reasoning_budget_soft_forced;         // tokenized soft warning message (no end tag)
+    std::string               reasoning_budget_soft_message;        // soft warning message injected at the soft threshold
+
+    std::vector<llama_token>  reasoning_budget_intro_forced;        // tokenized intro message forced when the block starts (empty = disabled)
+    std::string               reasoning_budget_intro_message =      // intro message announcing the budget; supports a {budget} placeholder
+        "I'll keep this reasoning under {budget} tokens, so I'll stay focused and efficient. ";
+
+    int32_t                  reasoning_budget_grace_tokens = 0;    // <= 0 = force immediately, N>0 = wait up to N tokens for a paragraph break
+
     bool backend_sampling = false;
 
     // print the parameters into a string