kernel/capabilities/
think_splitter.rs

1//! Separating a model's "thinking" spans (delimited by tags like
2//! `<think>…</think>`) from its visible text as the stream arrives.
3
4use crate::capabilities::held_suffix_len;
5
6/// A run of separated output: either visible text or hidden thinking.
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub enum Piece {
9    /// Visible text.
10    Text(String),
11    /// Thinking to hide or fold away.
12    Thinking(String),
13}
14
15/// A matching pair of open/close thinking delimiters.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct TagPair {
18    /// The opening delimiter.
19    pub open: String,
20    /// The closing delimiter.
21    pub close: String,
22}
23
24/// The default thinking delimiters recognized.
25pub fn default_pairs() -> Vec<TagPair> {
26    vec![
27        TagPair {
28            open: "<think>".to_owned(),
29            close: "</think>".to_owned(),
30        },
31        TagPair {
32            open: "<|START_THINKING|>".to_owned(),
33            close: "<|END_THINKING|>".to_owned(),
34        },
35    ]
36}
37
38/// Whether `text` contains any default thinking delimiter.
39pub fn has_visible_tags(text: &str) -> bool {
40    default_pairs()
41        .iter()
42        .any(|pair| text.contains(&pair.open) || text.contains(&pair.close))
43}
44
45enum Mode {
46    Text,
47    Thinking(String),
48}
49
50/// A streaming splitter that separates thinking spans from visible text.
51pub struct ThinkSplitter {
52    pairs: Vec<TagPair>,
53    open_tags: Vec<String>,
54    mode: Mode,
55    buffer: String,
56}
57
58impl Default for ThinkSplitter {
59    fn default() -> Self {
60        Self::with_pairs(default_pairs())
61    }
62}
63
64impl ThinkSplitter {
65    /// A splitter using the default delimiter pairs.
66    pub fn new() -> Self {
67        Self::default()
68    }
69
70    /// A splitter using the given delimiter pairs. Pairs with an empty open or
71    /// close delimiter are dropped — an empty delimiter would match everywhere
72    /// and make no progress.
73    pub fn with_pairs(pairs: Vec<TagPair>) -> Self {
74        let pairs: Vec<TagPair> = pairs
75            .into_iter()
76            .filter(|pair| !pair.open.is_empty() && !pair.close.is_empty())
77            .collect();
78        let open_tags = pairs.iter().map(|pair| pair.open.clone()).collect();
79        Self {
80            pairs,
81            open_tags,
82            mode: Mode::Text,
83            buffer: String::new(),
84        }
85    }
86
87    /// Feed a chunk and return the pieces that can be classified now. A partial
88    /// delimiter at the end of the buffer is held back until more text arrives.
89    pub fn feed(&mut self, chunk: &str) -> Vec<Piece> {
90        self.buffer.push_str(chunk);
91        let mut output = Vec::new();
92        loop {
93            match &self.mode {
94                Mode::Text => {
95                    let opening = self
96                        .pairs
97                        .iter()
98                        .filter_map(|pair| {
99                            self.buffer
100                                .find(&pair.open)
101                                .map(|at| (at, pair.open.len(), &pair.close))
102                        })
103                        .min_by_key(|(at, _, _)| *at);
104                    match opening {
105                        Some((at, open_len, close)) => {
106                            let close = close.clone();
107                            let before = self.buffer[..at].to_owned();
108                            if !before.is_empty() {
109                                output.push(Piece::Text(before));
110                            }
111                            self.buffer.drain(..at + open_len);
112                            self.mode = Mode::Thinking(close);
113                        }
114                        None => {
115                            Self::drain_prefix(
116                                &mut self.buffer,
117                                &self.open_tags,
118                                false,
119                                &mut output,
120                            );
121                            break;
122                        }
123                    }
124                }
125                Mode::Thinking(close) => {
126                    let close = close.clone();
127                    match self.buffer.find(&close) {
128                        Some(at) => {
129                            let before = self.buffer[..at].to_owned();
130                            if !before.is_empty() {
131                                output.push(Piece::Thinking(before));
132                            }
133                            self.buffer.drain(..at + close.len());
134                            self.mode = Mode::Text;
135                        }
136                        None => {
137                            Self::drain_prefix(&mut self.buffer, &[close], true, &mut output);
138                            break;
139                        }
140                    }
141                }
142            }
143        }
144        output
145    }
146
147    /// Emit any buffered text, classified by the current mode.
148    pub fn flush(&mut self) -> Vec<Piece> {
149        if self.buffer.is_empty() {
150            return Vec::new();
151        }
152        let buffer = std::mem::take(&mut self.buffer);
153        match &self.mode {
154            Mode::Thinking(_) => vec![Piece::Thinking(buffer)],
155            Mode::Text => vec![Piece::Text(buffer)],
156        }
157    }
158
159    fn drain_prefix(buffer: &mut String, tags: &[String], thinking: bool, output: &mut Vec<Piece>) {
160        let emit_len = buffer.len() - held_suffix_len(buffer, tags);
161        if emit_len > 0 {
162            let emit = buffer[..emit_len].to_owned();
163            output.push(if thinking {
164                Piece::Thinking(emit)
165            } else {
166                Piece::Text(emit)
167            });
168            buffer.drain(..emit_len);
169        }
170    }
171}
kernel/capabilities/think_splitter.rs

kernel/capabilities/
think_splitter.rs
