llama_cpp_bindings/
streaming_markers.rs

1use crate::token::LlamaToken;
2
3#[derive(Copy, Clone, Debug, Eq, PartialEq)]
4pub enum MarkerKind {
5    ReasoningOpen,
6    ReasoningClose,
7    ToolCallOpen,
8    ToolCallClose,
9}
10
11/// Tokenized marker sequences (token IDs, not strings).
12///
13/// Each marker is a `Vec<LlamaToken>` of length `>= 1`; absent markers are
14/// `None`. Sequence matching at every `ingest()` is by token-ID equality,
15/// never by substring scanning of decoded text.
16#[derive(Clone, Debug, Default, Eq, PartialEq)]
17pub struct StreamingMarkers {
18    pub reasoning_open: Option<Vec<LlamaToken>>,
19    pub reasoning_close: Option<Vec<LlamaToken>>,
20    pub tool_call_open: Option<Vec<LlamaToken>>,
21    pub tool_call_close: Option<Vec<LlamaToken>>,
22}
23
24impl StreamingMarkers {
25    #[must_use]
26    pub const fn has_any(&self) -> bool {
27        self.reasoning_open.is_some()
28            || self.reasoning_close.is_some()
29            || self.tool_call_open.is_some()
30            || self.tool_call_close.is_some()
31    }
32
33    #[must_use]
34    pub fn max_token_len(&self) -> usize {
35        [
36            self.reasoning_open.as_deref(),
37            self.reasoning_close.as_deref(),
38            self.tool_call_open.as_deref(),
39            self.tool_call_close.as_deref(),
40        ]
41        .into_iter()
42        .flatten()
43        .map(<[LlamaToken]>::len)
44        .max()
45        .unwrap_or(0)
46    }
47
48    #[must_use]
49    pub fn lookup(&self, kind: MarkerKind) -> Option<&[LlamaToken]> {
50        match kind {
51            MarkerKind::ReasoningOpen => self.reasoning_open.as_deref(),
52            MarkerKind::ReasoningClose => self.reasoning_close.as_deref(),
53            MarkerKind::ToolCallOpen => self.tool_call_open.as_deref(),
54            MarkerKind::ToolCallClose => self.tool_call_close.as_deref(),
55        }
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use super::StreamingMarkers;
62    use crate::token::LlamaToken;
63
64    fn token(id: i32) -> LlamaToken {
65        LlamaToken::new(id)
66    }
67
68    #[test]
69    fn streaming_markers_with_no_markers_reports_none() {
70        let markers = StreamingMarkers::default();
71        assert!(!markers.has_any());
72        assert_eq!(markers.max_token_len(), 0);
73    }
74
75    #[test]
76    fn streaming_markers_max_token_len_takes_longest() {
77        let markers = StreamingMarkers {
78            reasoning_open: Some(vec![token(1)]),
79            reasoning_close: Some(vec![token(2), token(3), token(4)]),
80            tool_call_open: Some(vec![token(5), token(6)]),
81            tool_call_close: None,
82        };
83        assert_eq!(markers.max_token_len(), 3);
84    }
85}
llama_cpp_bindings/streaming_markers.rs

llama_cpp_bindings/
streaming_markers.rs
