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

llama_cpp_bindings/
streaming_markers.rs
