llama_cpp_bindings/
extract_tool_call_markers_from_haystack.rs

1use crate::tool_call_marker_pair::ToolCallMarkerPair;
2
3#[must_use]
4pub fn extract_tool_call_markers_from_haystack(haystack: &str) -> Option<ToolCallMarkerPair> {
5    if haystack.is_empty() {
6        return None;
7    }
8
9    let json_start = haystack.find('{')?;
10    let json_end = haystack.rfind('}')?;
11    if json_end < json_start {
12        return None;
13    }
14
15    let json_slice = &haystack[json_start..=json_end];
16    serde_json::from_str::<serde_json::Value>(json_slice).ok()?;
17
18    let open = haystack[..json_start].trim().to_owned();
19    let close = haystack[json_end + 1..].trim().to_owned();
20
21    if open.is_empty() || close.is_empty() {
22        return None;
23    }
24
25    Some(ToolCallMarkerPair { open, close })
26}
27
28#[cfg(test)]
29mod tests {
30    use super::ToolCallMarkerPair;
31    use super::extract_tool_call_markers_from_haystack;
32
33    #[test]
34    fn extracts_open_and_close_around_a_simple_json_payload() {
35        let pair = extract_tool_call_markers_from_haystack(
36            "<tool_call>{\"name\":\"x\",\"arguments\":{}}</tool_call>",
37        );
38
39        assert_eq!(
40            pair,
41            Some(ToolCallMarkerPair {
42                open: "<tool_call>".to_owned(),
43                close: "</tool_call>".to_owned(),
44            }),
45        );
46    }
47
48    #[test]
49    fn trims_surrounding_whitespace_from_each_marker() {
50        let pair = extract_tool_call_markers_from_haystack(
51            "  <tool_call>\n  {\"k\": 1}\n  </tool_call>  ",
52        );
53
54        assert_eq!(
55            pair,
56            Some(ToolCallMarkerPair {
57                open: "<tool_call>".to_owned(),
58                close: "</tool_call>".to_owned(),
59            }),
60        );
61    }
62
63    #[test]
64    fn returns_none_when_haystack_is_empty() {
65        assert_eq!(extract_tool_call_markers_from_haystack(""), None);
66    }
67
68    #[test]
69    fn returns_none_when_haystack_has_no_open_brace() {
70        assert_eq!(
71            extract_tool_call_markers_from_haystack("plain assistant text"),
72            None
73        );
74    }
75
76    #[test]
77    fn returns_none_when_haystack_has_open_brace_but_no_close() {
78        assert_eq!(
79            extract_tool_call_markers_from_haystack("<open>{ unclosed"),
80            None
81        );
82    }
83
84    #[test]
85    fn returns_none_when_close_brace_precedes_open_brace() {
86        assert_eq!(
87            extract_tool_call_markers_from_haystack("</close>}{<open>"),
88            None
89        );
90    }
91
92    #[test]
93    fn returns_none_when_brace_payload_is_not_valid_json() {
94        assert_eq!(
95            extract_tool_call_markers_from_haystack("<open>{not valid json}</close>"),
96            None
97        );
98    }
99
100    #[test]
101    fn returns_none_when_open_marker_resolves_to_empty_after_trim() {
102        assert_eq!(
103            extract_tool_call_markers_from_haystack("   {\"x\":1}</close>"),
104            None
105        );
106    }
107
108    #[test]
109    fn returns_none_when_close_marker_resolves_to_empty_after_trim() {
110        assert_eq!(
111            extract_tool_call_markers_from_haystack("<open>{\"x\":1}   "),
112            None
113        );
114    }
115
116    #[test]
117    fn extracts_around_an_object_that_contains_nested_braces() {
118        let pair = extract_tool_call_markers_from_haystack(
119            "<tool>{\"args\":{\"k\":[1,2,{\"deep\":true}]}}</tool>",
120        );
121
122        assert_eq!(
123            pair,
124            Some(ToolCallMarkerPair {
125                open: "<tool>".to_owned(),
126                close: "</tool>".to_owned(),
127            }),
128        );
129    }
130
131    #[test]
132    fn extracts_when_open_marker_contains_multibyte_utf8() {
133        let pair = extract_tool_call_markers_from_haystack("<|tool→call|>{\"k\":1}<|/tool→call|>");
134
135        assert_eq!(
136            pair,
137            Some(ToolCallMarkerPair {
138                open: "<|tool→call|>".to_owned(),
139                close: "<|/tool→call|>".to_owned(),
140            }),
141        );
142    }
143}
llama_cpp_bindings/extract_tool_call_markers_from_haystack.rs

llama_cpp_bindings/
extract_tool_call_markers_from_haystack.rs
