llama_cpp_bindings/
streaming_markers.rs

1use crate::marker_kind::MarkerKind;
2use crate::token::LlamaToken;
3
4#[derive(Clone, Debug, Default, Eq, PartialEq)]
5pub struct StreamingMarkers {
6    pub reasoning_open: Option<Vec<LlamaToken>>,
7    pub reasoning_close: Option<Vec<LlamaToken>>,
8    pub tool_call_open: Option<Vec<LlamaToken>>,
9    pub tool_call_close: Option<Vec<LlamaToken>>,
10}
11
12impl StreamingMarkers {
13    #[must_use]
14    pub const fn has_any(&self) -> bool {
15        self.reasoning_open.is_some()
16            || self.reasoning_close.is_some()
17            || self.tool_call_open.is_some()
18            || self.tool_call_close.is_some()
19    }
20
21    #[must_use]
22    pub fn max_token_len(&self) -> usize {
23        [
24            self.reasoning_open.as_deref(),
25            self.reasoning_close.as_deref(),
26            self.tool_call_open.as_deref(),
27            self.tool_call_close.as_deref(),
28        ]
29        .into_iter()
30        .flatten()
31        .map(<[LlamaToken]>::len)
32        .max()
33        .unwrap_or(0)
34    }
35
36    #[must_use]
37    pub fn lookup(&self, kind: MarkerKind) -> Option<&[LlamaToken]> {
38        match kind {
39            MarkerKind::ReasoningOpen => self.reasoning_open.as_deref(),
40            MarkerKind::ReasoningClose => self.reasoning_close.as_deref(),
41            MarkerKind::ToolCallOpen => self.tool_call_open.as_deref(),
42            MarkerKind::ToolCallClose => self.tool_call_close.as_deref(),
43        }
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::StreamingMarkers;
50    use crate::token::LlamaToken;
51
52    fn token(id: i32) -> LlamaToken {
53        LlamaToken::new(id)
54    }
55
56    #[test]
57    fn streaming_markers_with_no_markers_reports_none() {
58        let markers = StreamingMarkers::default();
59        assert!(!markers.has_any());
60        assert_eq!(markers.max_token_len(), 0);
61    }
62
63    #[test]
64    fn streaming_markers_max_token_len_takes_longest() {
65        let markers = StreamingMarkers {
66            reasoning_open: Some(vec![token(1)]),
67            reasoning_close: Some(vec![token(2), token(3), token(4)]),
68            tool_call_open: Some(vec![token(5), token(6)]),
69            tool_call_close: None,
70        };
71        assert_eq!(markers.max_token_len(), 3);
72    }
73}
llama_cpp_bindings/streaming_markers.rs

llama_cpp_bindings/
streaming_markers.rs
