stream_think_filter.h

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include <cstddef>
#include <string>
 
namespace entropic {
 
using TokenCallback = void (*)(const char*, size_t, void*);
 
class StreamThinkFilter {
public:
    StreamThinkFilter(TokenCallback cb, void* ud);
 
    StreamThinkFilter(TokenCallback cb, void* ud,
                      std::string open_marker, std::string close_marker);
 
    void set_raw_callback(TokenCallback cb, void* ud);
 
    void on_token(const char* chunk, size_t len);
 
    void flush();
 
private:
    TokenCallback cb_;      
    void* ud_;              
    TokenCallback raw_cb_ = nullptr; 
    void* raw_ud_ = nullptr;         
    bool match_tag(const std::string& buf, bool& is_open) const;
 
    bool in_think_ = false; 
 
    std::string open_marker_  = "<think>";   
    std::string close_marker_ = "</think>";  
    std::size_t max_marker_len_ = 8;
    std::string tag_buf_;   
    std::string utf8_buf_;  
 
    void emit_utf8_safe(const char* data, size_t len);
 
    void process_byte(char c);
};
 
} // namespace entropic