Skip to main content

llama_cpp_bindings/tool_call_format/
bracketed_args.rs

1use llama_cpp_bindings_types::BracketedJsonShape;
2use llama_cpp_bindings_types::ParsedToolCall;
3use llama_cpp_bindings_types::ToolCallArguments;
4use llama_cpp_bindings_types::ToolCallMarkers;
5
6use crate::error::BracketedArgsFailure;
7
8enum ParseStep<'body> {
9    Done,
10    Call(ParsedToolCall, &'body str),
11}
12
13fn consume_optional_prefix<'body>(input: &'body str, literal: &str) -> &'body str {
14    input.strip_prefix(literal).unwrap_or(input)
15}
16
17fn split_at_separator<'body>(
18    input: &'body str,
19    separator: &str,
20) -> Option<(&'body str, &'body str)> {
21    let (name_raw, after_separator) = input.split_once(separator)?;
22    Some((name_raw, after_separator))
23}
24
25fn consume_one_json_value<'body>(
26    input: &'body str,
27    tool_name: &str,
28) -> Result<(serde_json::Value, &'body str), BracketedArgsFailure> {
29    let mut stream = serde_json::Deserializer::from_str(input).into_iter::<serde_json::Value>();
30    let value = stream
31        .next()
32        .ok_or_else(|| BracketedArgsFailure::UnterminatedArguments {
33            tool_name: tool_name.to_owned(),
34        })?
35        .map_err(|err| BracketedArgsFailure::InvalidJsonArguments {
36            tool_name: tool_name.to_owned(),
37            message: err.to_string(),
38        })?;
39    let consumed = stream.byte_offset();
40
41    Ok((value, &input[consumed..]))
42}
43
44fn parse_one_call<'body>(
45    input: &'body str,
46    markers: &ToolCallMarkers,
47    shape: &BracketedJsonShape,
48) -> Result<ParseStep<'body>, BracketedArgsFailure> {
49    if input.is_empty() {
50        return Ok(ParseStep::Done);
51    }
52
53    let after_open = consume_optional_prefix(input, markers.open.as_str());
54
55    let Some((name_raw, after_separator)) =
56        split_at_separator(after_open, shape.name_args_separator.as_str())
57    else {
58        return Ok(ParseStep::Done);
59    };
60
61    let name = name_raw.trim().to_owned();
62    if name.is_empty() {
63        return Ok(ParseStep::Done);
64    }
65
66    let (arguments_value, after_arguments) = consume_one_json_value(after_separator, &name)?;
67
68    let after_close = consume_optional_prefix(after_arguments, markers.close.as_str());
69
70    Ok(ParseStep::Call(
71        ParsedToolCall::new(
72            String::new(),
73            name,
74            ToolCallArguments::ValidJson(arguments_value),
75        ),
76        after_close,
77    ))
78}
79
80/// # Errors
81///
82/// Returns [`BracketedArgsFailure`] when the body looks like a bracketed-JSON
83/// tool-call block (matches the name/args separator) but contains a structural
84/// issue: invalid JSON arguments or a JSON value truncated mid-stream.
85pub fn parse(
86    body: &str,
87    markers: &ToolCallMarkers,
88    shape: &BracketedJsonShape,
89) -> Result<Vec<ParsedToolCall>, BracketedArgsFailure> {
90    if shape.name_args_separator.is_empty() {
91        return Ok(Vec::new());
92    }
93
94    let mut parsed = Vec::new();
95    let mut remaining = body.trim_start();
96
97    loop {
98        match parse_one_call(remaining, markers, shape)? {
99            ParseStep::Done => break,
100            ParseStep::Call(call, rest) => {
101                parsed.push(call);
102                remaining = rest.trim_start();
103            }
104        }
105    }
106
107    Ok(parsed)
108}
109
110#[cfg(test)]
111mod tests {
112    use llama_cpp_bindings_types::BracketedJsonShape;
113    use llama_cpp_bindings_types::ToolCallArgsShape;
114    use llama_cpp_bindings_types::ToolCallArguments;
115    use llama_cpp_bindings_types::ToolCallMarkers;
116    use serde_json::json;
117
118    use super::parse;
119    use crate::error::BracketedArgsFailure;
120
121    fn mistral3_markers() -> ToolCallMarkers {
122        ToolCallMarkers {
123            open: "[TOOL_CALLS]".to_owned(),
124            close: String::new(),
125            args_shape: ToolCallArgsShape::BracketedJson(BracketedJsonShape {
126                name_args_separator: "[ARGS]".to_owned(),
127            }),
128        }
129    }
130
131    fn mistral3_shape() -> BracketedJsonShape {
132        BracketedJsonShape {
133            name_args_separator: "[ARGS]".to_owned(),
134        }
135    }
136
137    #[test]
138    fn parses_single_tool_call_with_open_marker_present() {
139        let parsed = parse(
140            "[TOOL_CALLS]get_weather[ARGS]{\"location\":\"Paris\"}",
141            &mistral3_markers(),
142            &mistral3_shape(),
143        )
144        .expect("must parse");
145
146        assert_eq!(parsed.len(), 1);
147        assert_eq!(parsed[0].name, "get_weather");
148        assert_eq!(
149            parsed[0].arguments,
150            ToolCallArguments::ValidJson(json!({"location": "Paris"})),
151        );
152    }
153
154    #[test]
155    fn parses_single_tool_call_when_classifier_stripped_open_marker() {
156        let parsed = parse(
157            "get_weather[ARGS]{\"location\":\"Paris\"}",
158            &mistral3_markers(),
159            &mistral3_shape(),
160        )
161        .expect("must parse");
162
163        assert_eq!(parsed.len(), 1);
164        assert_eq!(parsed[0].name, "get_weather");
165        assert_eq!(
166            parsed[0].arguments,
167            ToolCallArguments::ValidJson(json!({"location": "Paris"})),
168        );
169    }
170
171    #[test]
172    fn parses_two_consecutive_tool_calls_with_repeated_open_marker() {
173        let parsed = parse(
174            "[TOOL_CALLS]a[ARGS]{\"x\":1}[TOOL_CALLS]b[ARGS]{\"y\":2}",
175            &mistral3_markers(),
176            &mistral3_shape(),
177        )
178        .expect("must parse");
179
180        assert_eq!(parsed.len(), 2);
181        assert_eq!(parsed[0].name, "a");
182        assert_eq!(
183            parsed[0].arguments,
184            ToolCallArguments::ValidJson(json!({"x": 1}))
185        );
186        assert_eq!(parsed[1].name, "b");
187        assert_eq!(
188            parsed[1].arguments,
189            ToolCallArguments::ValidJson(json!({"y": 2}))
190        );
191    }
192
193    #[test]
194    fn rejects_malformed_json_arguments_with_typed_failure() {
195        let result = parse(
196            "[TOOL_CALLS]get_weather[ARGS]{\"location\":}",
197            &mistral3_markers(),
198            &mistral3_shape(),
199        );
200
201        let failure = result.expect_err("malformed JSON must produce a typed failure");
202        match failure {
203            BracketedArgsFailure::InvalidJsonArguments { tool_name, .. } => {
204                assert_eq!(tool_name, "get_weather");
205            }
206            other @ BracketedArgsFailure::UnterminatedArguments { .. } => {
207                panic!("expected InvalidJsonArguments, got {other:?}")
208            }
209        }
210    }
211
212    #[test]
213    fn returns_empty_vec_for_empty_body() {
214        let parsed =
215            parse("", &mistral3_markers(), &mistral3_shape()).expect("empty body must parse");
216        assert!(parsed.is_empty());
217    }
218
219    #[test]
220    fn returns_empty_vec_when_body_lacks_separator() {
221        let parsed = parse(
222            "plain text without separator",
223            &mistral3_markers(),
224            &mistral3_shape(),
225        )
226        .expect("body without separator must parse");
227        assert!(parsed.is_empty());
228    }
229}