Felipe97/llama-cpp-compiled

1#include "reasoning-budget.h"2#include "unicode.h"3 4#include "llama.h"5#include "ggml.h"6 7#ifdef NDEBUG8#undef NDEBUG9#endif10 11#include <cmath>12#include <cstddef>13#include <cstdio>14#include <string>15#include <vector>16 17// Reasoning budget sampler test helper18// These tests use nullptr vocab which safely falls back to treating all tokens as complete19// (The UTF-8 boundary detection logic is tested separately in test_utf8_boundary_detection)20static void test_reasoning_budget(21    const char * test_name,22    const std::vector<llama_token> & sequence,23    const std::vector<llama_tokens> & start_seqs,24    const std::vector<llama_tokens> & end_seqs,25    const std::vector<llama_token> & forced_tokens,26    int32_t budget,27    common_reasoning_budget_state initial_state,28    size_t expected_force_start,   // token index where forcing should start (SIZE_MAX = never)29    size_t expected_force_end      // token index where forcing should end (after this, no more forcing)30) {31    // Find the maximum token ID to ensure our vocab covers all tokens32    llama_token max_token = 0;33    for (auto t : sequence) max_token = std::max(max_token, t);34    for (const auto & seq : start_seqs) {35        for (auto t : seq) max_token = std::max(max_token, t);36    }37    for (const auto & seq : end_seqs) {38        for (auto t : seq) max_token = std::max(max_token, t);39    }40    for (auto t : forced_tokens) max_token = std::max(max_token, t);41 42    // Create a minimal sampler with mock vocabulary43    // For this test, we use nullptr as vocab since we're testing state transitions44    // The UTF-8 boundary check will treat all tokens as complete (safe fallback)45    auto * sampler = common_reasoning_budget_init(46        nullptr,  // vocab - not used for basic state machine tests47        start_seqs,48        end_seqs,49        forced_tokens,50        budget,51        initial_state52    );53 54    // Create a test token data array for checking forcing behavior55    // Vocab size must be large enough to include all tokens (start, end, forced, sequence)56    std::vector<llama_token_data> cur;57    const size_t n_vocab = (size_t)max_token + 1;58    for (size_t i = 0; i < n_vocab; i++) {59        cur.emplace_back(llama_token_data{(llama_token)i, logf((float)(i+1)), 0.0f});60    }61    llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false };62 63    size_t actual_force_start = SIZE_MAX;64    size_t actual_force_end = SIZE_MAX;65 66    // Feed the sequence and track when forcing occurs67    for (size_t i = 0; i < sequence.size(); i++) {68        // Check if we're in forcing state by applying and seeing if logits are modified69        cur_p.selected = -1;70        for (size_t j = 0; j < cur.size(); j++) {71            cur[j].logit = logf((float)(j+1));  // reset logits72        }73 74        llama_sampler_apply(sampler, &cur_p);75 76        // Check if forcing is active (all logits except one should be -INFINITY)77        size_t finite_count = 0;78        llama_token finite_token = -1;79        for (size_t j = 0; j < cur.size(); j++) {80            if (std::isfinite(cur[j].logit)) {81                finite_count++;82                finite_token = cur[j].id;83            }84        }85 86        llama_sampler_accept(sampler, sequence[i]);87 88        fprintf(stderr, "    i=%zu: token=%d, finite_count=%zu, finite_token=%d\n", i, (int)sequence[i], finite_count, (int)finite_token);89 90        if (finite_count == 1) {91            if (actual_force_start == SIZE_MAX) {92                actual_force_start = i;93            }94            actual_force_end = i;95        } else if (actual_force_start != SIZE_MAX && actual_force_end != SIZE_MAX) {96            // Forcing stopped97            break;98        }99    }100 101    llama_sampler_free(sampler);102 103    // Verify forcing occurred at expected positions104    if (expected_force_start == SIZE_MAX) {105        if (actual_force_start != SIZE_MAX) {106            fprintf(stderr, "Test '%s' FAILED: Expected no forcing, but forcing occurred at %zu\n", test_name, actual_force_start);107            GGML_ASSERT(false && "Expected no forcing, but forcing occurred");108        }109    } else {110        if (actual_force_start == SIZE_MAX) {111            fprintf(stderr, "Test '%s' FAILED: Expected forcing but none occurred\n", test_name);112            GGML_ASSERT(false && "Expected forcing but none occurred");113        }114        if (actual_force_start != expected_force_start) {115            fprintf(stderr, "Test '%s' FAILED: Forcing started at %zu, expected %zu\n", test_name, actual_force_start, expected_force_start);116            GGML_ASSERT(false && "Forcing started at wrong position");117        }118    }119 120    if (expected_force_end != SIZE_MAX) {121        if (actual_force_end < expected_force_end) {122            fprintf(stderr, "Test '%s' FAILED: Forcing ended at %zu, expected >= %zu\n", test_name, actual_force_end, expected_force_end);123            GGML_ASSERT(false && "Forcing ended too early");124        }125    }126 127    fprintf(stderr, "  Test '%s' passed (force_start=%zu, force_end=%zu)\n", test_name, actual_force_start, actual_force_end);128    (void)sequence;129}130 131static llama_token get_forced_token(struct llama_sampler * sampler, llama_token max_token) {132    std::vector<llama_token_data> cur;133    const size_t n_vocab = (size_t) max_token + 1;134    for (size_t i = 0; i < n_vocab; i++) {135        cur.emplace_back(llama_token_data{(llama_token) i, logf((float) (i + 1)), 0.0f});136    }137 138    llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false };139    llama_sampler_apply(sampler, &cur_p);140 141    size_t finite_count = 0;142    llama_token finite_token = LLAMA_TOKEN_NULL;143    for (size_t i = 0; i < cur.size(); i++) {144        if (std::isfinite(cur[i].logit)) {145            finite_count++;146            finite_token = cur[i].id;147        }148    }149 150    GGML_ASSERT(finite_count == 1 && "sampler is not forcing exactly one token");151    return finite_token;152}153 154static void test_reasoning_budget_clone_mid_counting() {155    const std::vector<llama_token> start = {100};156    const std::vector<llama_token> end = {101};157    const std::vector<llama_token> forced = {102, 101};158 159    auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 2, REASONING_BUDGET_IDLE);160 161    llama_sampler_accept(sampler, 100); // COUNTING, remaining=2162    llama_sampler_accept(sampler, 50);  // COUNTING, remaining=1163 164    auto * clone = llama_sampler_clone(sampler);165    llama_sampler_accept(clone, 51); // should exhaust the cloned remaining budget166 167    GGML_ASSERT(get_forced_token(clone, 102) == 102 && "cloned counting state lost remaining budget");168 169    llama_sampler_free(clone);170    llama_sampler_free(sampler);171}172 173static void test_reasoning_budget_clone_mid_forcing() {174    const std::vector<llama_token> start = {100};175    const std::vector<llama_token> end = {101};176    const std::vector<llama_token> forced = {102, 101};177 178    auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING);179 180    GGML_ASSERT(get_forced_token(sampler, 102) == 102);181    llama_sampler_accept(sampler, 102); // advance to the second forced token182 183    auto * clone = llama_sampler_clone(sampler);184 185    GGML_ASSERT(get_forced_token(clone, 102) == 101 && "cloned forcing state lost force position");186 187    llama_sampler_free(clone);188    llama_sampler_free(sampler);189}190 191static void test_reasoning_budget_force_manual() {192    const std::vector<llama_token> start  = {100};193    const std::vector<llama_token> end    = {101};194    const std::vector<llama_token> forced = {102, 101};195 196    // if COUNTING, force() succeeds and begins forcing the end sequence from the start197    {198        auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE);199 200        llama_sampler_accept(sampler, 100); // COUNTING, remaining=5201        llama_sampler_accept(sampler, 50);  // COUNTING, remaining=4202        GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING);203 204        GGML_ASSERT(common_reasoning_budget_force(sampler) && "force() should succeed from COUNTING");205        GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING);206 207        // forces the configured sequence from force_pos=0, then transitions to DONE208        GGML_ASSERT(get_forced_token(sampler, 102) == 102);209        llama_sampler_accept(sampler, 102);210        GGML_ASSERT(get_forced_token(sampler, 102) == 101);211        llama_sampler_accept(sampler, 101);212        GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE);213 214        llama_sampler_free(sampler);215    }216 217    // if IDLE, force() is a no-op218    {219        auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE);220 221        GGML_ASSERT(!common_reasoning_budget_force(sampler) && "force() must not transition from IDLE");222        GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_IDLE);223 224        llama_sampler_free(sampler);225    }226 227    // if DONE, force() is a no-op228    {229        auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE);230 231        llama_sampler_accept(sampler, 100); // COUNTING232        llama_sampler_accept(sampler, 101); // natural end -> DONE233        GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE);234 235        GGML_ASSERT(!common_reasoning_budget_force(sampler) && "force() must not transition from DONE");236        GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE);237 238        llama_sampler_free(sampler);239    }240 241    // if FORCING, force() is a no-op and must not rewind the force position242    {243        auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING);244 245        GGML_ASSERT(get_forced_token(sampler, 102) == 102);246        llama_sampler_accept(sampler, 102); // advance to the second forced token (force_pos=1)247 248        GGML_ASSERT(!common_reasoning_budget_force(sampler) && "force() must not transition from FORCING");249        GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING);250        GGML_ASSERT(get_forced_token(sampler, 102) == 101 && "force() must not rewind the force position");251 252        llama_sampler_free(sampler);253    }254 255    // a null sampler is safely ignored256    GGML_ASSERT(!common_reasoning_budget_force(nullptr));257 258    fprintf(stderr, "  Test 'manual force transition' passed\n");259}260 261static void test_reasoning_budget_end_match() {262    const std::vector<llama_tokens> start = {{100}};263    const std::vector<llama_tokens> end   = {{101}, {103, 104}};264 265    // natural end records the sequence that matched; re-arming clears it266    {267        auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 101}, 5, REASONING_BUDGET_IDLE);268 269        GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);270 271        llama_sampler_accept(sampler, 100); // COUNTING272        llama_sampler_accept(sampler, 50);273        llama_sampler_accept(sampler, 103);274        llama_sampler_accept(sampler, 104); // end matched via {103, 104}, DONE275 276        const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);277        GGML_ASSERT(matched != nullptr);278        GGML_ASSERT(*matched == llama_tokens({103, 104}));279 280        llama_sampler_accept(sampler, 100); // re-arm, COUNTING281        GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);282 283        llama_sampler_free(sampler);284    }285 286    // overlapping end sequences: the longest one ending at the position wins287    {288        const std::vector<llama_tokens> end_overlap = {{104}, {103, 104}};289 290        auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, {102, 104}, 5, REASONING_BUDGET_IDLE);291 292        llama_sampler_accept(sampler, 100); // COUNTING293        llama_sampler_accept(sampler, 103);294        llama_sampler_accept(sampler, 104); // both {104} and {103, 104} end here295 296        const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);297        GGML_ASSERT(matched != nullptr);298        GGML_ASSERT(*matched == llama_tokens({103, 104}));299 300        llama_sampler_free(sampler);301    }302 303    // forcing records the end sequence terminating forced_tokens304    {305        auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 103, 104}, 0, REASONING_BUDGET_FORCING);306 307        llama_sampler_accept(sampler, 102);308        llama_sampler_accept(sampler, 103);309        GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);310        llama_sampler_accept(sampler, 104); // forced sequence complete, DONE311 312        const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);313        GGML_ASSERT(matched != nullptr);314        GGML_ASSERT(*matched == llama_tokens({103, 104}));315 316        llama_sampler_free(sampler);317    }318 319    // forced_tokens not ending with a known end sequence records nothing320    {321        auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102}, 0, REASONING_BUDGET_FORCING);322 323        llama_sampler_accept(sampler, 102); // forced sequence complete, DONE324        GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE);325        GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr);326 327        llama_sampler_free(sampler);328    }329 330    // a null sampler is safely ignored331    GGML_ASSERT(common_reasoning_budget_get_end_match(nullptr) == nullptr);332 333    fprintf(stderr, "  Test 'matched end sequence' passed\n");334}335 336// UTF-8 boundary detection unit test337// Tests common_utf8_is_complete() from reasoning-budget.h338static void test_utf8_boundary_detection() {339    // Complete sequences340    GGML_ASSERT(common_utf8_is_complete("hello"));341    GGML_ASSERT(common_utf8_is_complete(""));342    GGML_ASSERT(common_utf8_is_complete("\xC2\xA0"));            // complete 2-byte UTF-8 (U+00A0)343    GGML_ASSERT(common_utf8_is_complete("\xE2\x80\x9C"));        // complete 3-byte UTF-8 (left double quote)344    GGML_ASSERT(common_utf8_is_complete("\xF0\x9F\x98\x80"));    // complete 4-byte UTF-8 (emoji)345    GGML_ASSERT(common_utf8_is_complete("abc\xC3\xA9"));         // ASCII + complete 2-byte346 347    // Incomplete sequences348    GGML_ASSERT(!common_utf8_is_complete(std::string("\xC2", 1)));            // 2-byte start, missing continuation349    GGML_ASSERT(!common_utf8_is_complete(std::string("\xE2\x80", 2)));        // 3-byte start + 1 cont, missing 1350    GGML_ASSERT(!common_utf8_is_complete(std::string("\xE2", 1)));            // 3-byte start, missing 2351    GGML_ASSERT(!common_utf8_is_complete(std::string("\xF0\x9F\x98", 3)));    // 4-byte start + 2 cont, missing 1352    GGML_ASSERT(!common_utf8_is_complete(std::string("\xF0\x9F", 2)));        // 4-byte start + 1 cont, missing 2353    GGML_ASSERT(!common_utf8_is_complete(std::string("\xF0", 1)));            // 4-byte start, missing 3354    GGML_ASSERT(!common_utf8_is_complete(std::string("\x80", 1)));            // orphan continuation byte355 356    // Mixed: ASCII followed by start of multi-byte357    GGML_ASSERT(!common_utf8_is_complete(std::string("hello\xC3", 6)));       // ASCII + incomplete 2-byte358    GGML_ASSERT(common_utf8_is_complete(std::string("hello\xC3\xA9", 7)));    // ASCII + complete 2-byte359}360 361int main(void) {362    // Reasoning budget sampler tests363    printf("Testing reasoning budget sampler... ");364 365    // Test 1: Basic budget with start/end tokens - no forcing (natural end before budget exhausted)366    {367        const std::vector<llama_token> start = {100};  // start token368        const std::vector<llama_token> end = {101};    // end token369        const std::vector<llama_token> forced = {102}; // forced token (not used in this test)370        const std::vector<llama_token> sequence = {100, 50, 51, 101, 52}; // start, two tokens, end, one more371 372        test_reasoning_budget("natural end before budget exhausted", sequence, {start}, {end}, forced,373            5,      // budget of 5 tokens374            REASONING_BUDGET_IDLE,375            SIZE_MAX, SIZE_MAX); // no forcing expected (natural end)376    }377 378    // Test 2: Budget exhausted, forcing should occur379    // Flow: i=0 apply()->passthrough, accept(100)->COUNTING; i=1 accept(50)->remaining=1380    // i=2 accept(51)->remaining=0->FORCING; i=3 apply() forces token[0]; i=4 apply() forces token[1]381    // At i=4, accept() advances force_pos to 2 which equals forced_tokens.size(), so state becomes DONE382    {383        const std::vector<llama_token> start = {100};384        const std::vector<llama_token> end = {101};385        const std::vector<llama_token> forced = {102, 101}; // forced message + end386        const std::vector<llama_token> sequence = {100, 50, 51, 52, 53}; // start + 4 tokens (budget=2)387 388        test_reasoning_budget("budget exhausted forcing", sequence, {start}, {end}, forced,389            2,      // budget of 2 tokens390            REASONING_BUDGET_IDLE,391            3,      // forcing starts at i=3 (accept at i=2 depletes budget, apply at i=3 forces)392            4);     // forcing continues through i=4 (accept at i=4 transitions to DONE)393    }394 395    // Test 3: Activate immediately with budget=0, forcing should start right away396    // Flow: init promotes COUNTING+budget=0 to FORCING, so apply() sees FORCING at i=0397    {398        const std::vector<llama_token> start = {100};399        const std::vector<llama_token> end = {101};400        const std::vector<llama_token> forced = {102, 101};401        const std::vector<llama_token> sequence = {100, 50, 51, 52}; // start token first, then 3 tokens402 403        test_reasoning_budget("activate immediately budget=0", sequence, {start}, {end}, forced,404            0,      // budget of 0 tokens405            REASONING_BUDGET_COUNTING, // starts counting, promoted to FORCING since budget=0406            0,      // forcing starts at i=0 (initialized in FORCING, apply forces immediately)407            1);     // forcing continues through i=1 (accept at i=1 transitions to DONE)408    }409 410    // Test 4: No start/end tokens configured - passthrough (no forcing)411    {412        const std::vector<llama_token> start = {};413        const std::vector<llama_token> end = {};414        const std::vector<llama_token> forced = {102};415        const std::vector<llama_token> sequence = {50, 51, 52, 53};416 417        test_reasoning_budget("no start/end configured", sequence, {start}, {end}, forced,418            2,      // budget419            REASONING_BUDGET_IDLE,420            SIZE_MAX, SIZE_MAX); // no forcing (no start/end configured)421    }422 423    // Test 5: Activate immediately with budget > 0, count down then force424    // Flow: i=0 accept(50)->remaining=1, i=1 accept(51)->remaining=0->FORCING425    // Forcing starts at i=2 (apply sees FORCING after accept at i=1 transitioned)426    {427        const std::vector<llama_token> start = {100};428        const std::vector<llama_token> end = {101};429        const std::vector<llama_token> forced = {102, 101};430        const std::vector<llama_token> sequence = {50, 51, 52, 53};431 432        test_reasoning_budget("activate immediately with budget", sequence, {start}, {end}, forced,433            2,      // budget of 2 tokens434            REASONING_BUDGET_COUNTING,435            2,      // forcing starts at i=2 (after 2 accepts deplete budget, apply at i=2 forces)436            3);     // forcing continues through i=3437    }438 439    // Test 6: Multi-block thinking. First block ends naturally at i=2, second440    // start tag at i=3 re-arms the budget, which then exhausts at i=5.441    // Regression: before this fix, DONE absorbed all subsequent tokens and a442    // second <think> block ran unbudgeted.443    // Flow: i=0 accept(100)->COUNTING rem=2; i=1 accept(50)->rem=1;444    //       i=2 accept(101)->end_matcher matches, DONE;445    //       i=3 accept(100)->re-arm, COUNTING rem=2;446    //       i=4 accept(60)->rem=1; i=5 accept(61)->rem=0->FORCING;447    //       i=6 apply()->forces token[0]=102, accept(62)->force_pos=1, stay FORCING;448    //       i=7 apply()->forces token[1]=101, accept(63)->force_pos=2->DONE.449    {450        const std::vector<llama_token> start = {100};451        const std::vector<llama_token> end = {101};452        const std::vector<llama_token> forced = {102, 101};453        const std::vector<llama_token> sequence = {100, 50, 101, 100, 60, 61, 62, 63};454 455        test_reasoning_budget("multi-block re-arms budget after DONE", sequence, {start}, {end}, forced,456            2,      // budget of 2 tokens (per block)457            REASONING_BUDGET_IDLE,458            6,      // forcing starts at i=6 (after second block exhausts at i=5)459            7);     // forcing continues through i=7460    }461 462    // Test 7: Multiple start sequences - the second sequence activates counting463    // Flow: i=0 accept(110), i=1 accept(111)->COUNTING rem=2; i=2 accept(50)->rem=1;464    //       i=3 accept(51)->rem=0->FORCING; i=4..5 apply() forces the end sequence465    {466        const std::vector<llama_tokens> start = {{100}, {110, 111}};467        const std::vector<llama_tokens> end = {{101}};468        const std::vector<llama_token> forced = {102, 101};469        const std::vector<llama_token> sequence = {110, 111, 50, 51, 52, 53};470 471        test_reasoning_budget("multiple start sequences", sequence, start, end, forced,472            2,      // budget of 2 tokens473            REASONING_BUDGET_IDLE,474            4,      // forcing starts at i=4 (accept at i=3 depletes budget)475            5);     // forcing continues through i=5476    }477 478    // Test 8: Multiple end sequences - natural end via the second sequence479    // Flow: i=0 accept(100)->COUNTING rem=5; i=1 accept(50)->rem=4;480    //       i=2 accept(103)->partial end, rem=3; i=3 accept(104)->end matched, DONE481    {482        const std::vector<llama_tokens> start = {{100}};483        const std::vector<llama_tokens> end = {{101}, {103, 104}};484        const std::vector<llama_token> forced = {102, 101};485        const std::vector<llama_token> sequence = {100, 50, 103, 104, 52};486 487        test_reasoning_budget("multiple end sequences", sequence, start, end, forced,488            5,      // budget of 5 tokens489            REASONING_BUDGET_IDLE,490            SIZE_MAX, SIZE_MAX); // no forcing expected (natural end)491    }492 493    test_reasoning_budget_clone_mid_counting();494    test_reasoning_budget_clone_mid_forcing();495    test_reasoning_budget_force_manual();496    test_reasoning_budget_end_match();497 498    printf("OK (12 tests passed)\n");499 500    printf("Testing UTF-8 boundary detection... ");501    test_utf8_boundary_detection();502    printf("OK\n");503 504    return 0;505}506