stream_think_filter.cpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#include <entropic/core/stream_think_filter.h>
 
namespace entropic {
 
StreamThinkFilter::StreamThinkFilter(TokenCallback cb, void* ud)
    : cb_(cb), ud_(ud) {}
 
void StreamThinkFilter::set_raw_callback(TokenCallback cb, void* ud) {
    raw_cb_ = cb;
    raw_ud_ = ud;
}
 
static bool match_tag(const std::string& buf, bool& is_open) {
    if (buf == "<think>") { is_open = true; return true; }
    if (buf == "</think>") { is_open = false; return true; }
    return false;
}
 
static int utf8_char_len(unsigned char byte) {
    // Lookup: ASCII | continuation | 2-byte | 3-byte | 4-byte lead.
    if (byte < 0x80) { return 1; }
    int result = 0;  // continuation byte default
    if (byte >= 0xF0) { result = 4; }
    else if (byte >= 0xE0) { result = 3; }
    else if (byte >= 0xC0) { result = 2; }
    return result;
}
 
void StreamThinkFilter::emit_utf8_safe(const char* data, size_t len) {
    utf8_buf_.append(data, len);
 
    // Find the last complete codepoint boundary
    size_t safe = utf8_buf_.size();
    if (safe == 0) { return; }
 
    // Walk backward from end to find any incomplete trailing sequence
    auto* buf = reinterpret_cast<const unsigned char*>(utf8_buf_.data());
    for (size_t i = 1; i <= 4 && i <= safe; ++i) {
        unsigned char c = buf[safe - i];
        int expected = utf8_char_len(c);
        if (expected > 0) {
            // Found a lead byte at position (safe - i)
            size_t available = i; // bytes from lead to end
            if (available < static_cast<size_t>(expected)) {
                // Incomplete sequence — emit up to lead byte
                safe -= i;
            }
            break;
        }
    }
 
    if (safe > 0) {
        cb_(utf8_buf_.data(), safe, ud_);
    }
    utf8_buf_.erase(0, safe);
}
 
void StreamThinkFilter::process_byte(char c) {
    if (tag_buf_.empty() && c != '<') {
        if (!in_think_) { emit_utf8_safe(&c, 1); }
        return;
    }
    tag_buf_ += c;
    bool is_open = false;
    if (match_tag(tag_buf_, is_open)) {
        in_think_ = is_open;
        tag_buf_.clear();
        return;
    }
    if (tag_buf_.size() > 8) {
        if (!in_think_) {
            emit_utf8_safe(tag_buf_.data(), tag_buf_.size());
        }
        tag_buf_.clear();
    }
}
 
void StreamThinkFilter::on_token(const char* chunk, size_t len) {
    // Raw callback always gets everything (unfiltered, no UTF-8 alignment)
    if (raw_cb_) { raw_cb_(chunk, len, raw_ud_); }
 
    for (size_t i = 0; i < len; ++i) {
        process_byte(chunk[i]);
    }
}
 
void StreamThinkFilter::flush() {
    if (!tag_buf_.empty() && !in_think_) {
        emit_utf8_safe(tag_buf_.data(), tag_buf_.size());
    }
    tag_buf_.clear();
    // Flush any remaining UTF-8 buffer (may be incomplete at stream end)
    if (!utf8_buf_.empty()) {
        cb_(utf8_buf_.data(), utf8_buf_.size(), ud_);
        utf8_buf_.clear();
    }
}
 
} // namespace entropic