Skip to main content

llama_cpp_bindings/
streaming_json_probe.rs

1use serde_json::Value;
2use serde_json::error::Category;
3
4const NAME_FIELD: &str = "name";
5const ARGUMENTS_FIELD: &str = "arguments";
6
7#[derive(Copy, Clone, Debug, Eq, PartialEq)]
8pub enum JsonProbeOutcome {
9    StillPossiblyValid,
10    CompletedValid,
11    Failed,
12}
13
14impl JsonProbeOutcome {
15    #[must_use]
16    pub fn validate_prefix(buffer: &str) -> Self {
17        let trimmed = buffer.trim_start();
18        if trimmed.is_empty() {
19            return Self::StillPossiblyValid;
20        }
21        if !trimmed.starts_with('{') {
22            return Self::Failed;
23        }
24
25        let mut stream = serde_json::Deserializer::from_str(trimmed).into_iter::<Value>();
26        match stream.next() {
27            Some(Ok(value)) => evaluate_completed_value(&value, &trimmed[stream.byte_offset()..]),
28            Some(Err(parse_error)) => match parse_error.classify() {
29                Category::Eof => Self::StillPossiblyValid,
30                Category::Io | Category::Syntax | Category::Data => Self::Failed,
31            },
32            None => Self::StillPossiblyValid,
33        }
34    }
35}
36
37fn evaluate_completed_value(value: &Value, trailing: &str) -> JsonProbeOutcome {
38    let Value::Object(map) = value else {
39        return JsonProbeOutcome::Failed;
40    };
41
42    let Some(Value::String(name)) = map.get(NAME_FIELD) else {
43        return JsonProbeOutcome::Failed;
44    };
45    if name.is_empty() {
46        return JsonProbeOutcome::Failed;
47    }
48
49    if let Some(arguments) = map.get(ARGUMENTS_FIELD)
50        && !matches!(arguments, Value::Object(_))
51    {
52        return JsonProbeOutcome::Failed;
53    }
54
55    for key in map.keys() {
56        if key != NAME_FIELD && key != ARGUMENTS_FIELD {
57            return JsonProbeOutcome::Failed;
58        }
59    }
60
61    if trailing.trim().is_empty() {
62        JsonProbeOutcome::CompletedValid
63    } else {
64        JsonProbeOutcome::Failed
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::JsonProbeOutcome;
71
72    #[test]
73    fn empty_buffer_is_still_possibly_valid() {
74        assert_eq!(
75            JsonProbeOutcome::validate_prefix(""),
76            JsonProbeOutcome::StillPossiblyValid,
77        );
78    }
79
80    #[test]
81    fn whitespace_only_buffer_is_still_possibly_valid() {
82        assert_eq!(
83            JsonProbeOutcome::validate_prefix("   \n  "),
84            JsonProbeOutcome::StillPossiblyValid,
85        );
86    }
87
88    #[test]
89    fn single_open_brace_is_still_possibly_valid() {
90        assert_eq!(
91            JsonProbeOutcome::validate_prefix("{"),
92            JsonProbeOutcome::StillPossiblyValid,
93        );
94    }
95
96    #[test]
97    fn open_brace_with_trailing_space_is_still_possibly_valid() {
98        assert_eq!(
99            JsonProbeOutcome::validate_prefix("{ "),
100            JsonProbeOutcome::StillPossiblyValid,
101        );
102    }
103
104    #[test]
105    fn open_brace_with_quote_starting_key_is_still_possibly_valid() {
106        assert_eq!(
107            JsonProbeOutcome::validate_prefix(r#"{ ""#),
108            JsonProbeOutcome::StillPossiblyValid,
109        );
110    }
111
112    #[test]
113    fn partial_name_key_is_still_possibly_valid() {
114        assert_eq!(
115            JsonProbeOutcome::validate_prefix(r#"{ "name""#),
116            JsonProbeOutcome::StillPossiblyValid,
117        );
118    }
119
120    #[test]
121    fn partial_name_value_quote_is_still_possibly_valid() {
122        assert_eq!(
123            JsonProbeOutcome::validate_prefix(r#"{ "name": ""#),
124            JsonProbeOutcome::StillPossiblyValid,
125        );
126    }
127
128    #[test]
129    fn partial_name_value_letters_is_still_possibly_valid() {
130        assert_eq!(
131            JsonProbeOutcome::validate_prefix(r#"{ "name": "ge"#),
132            JsonProbeOutcome::StillPossiblyValid,
133        );
134    }
135
136    #[test]
137    fn complete_name_string_no_comma_is_still_possibly_valid() {
138        assert_eq!(
139            JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather""#),
140            JsonProbeOutcome::StillPossiblyValid,
141        );
142    }
143
144    #[test]
145    fn name_then_comma_is_still_possibly_valid() {
146        assert_eq!(
147            JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather","#),
148            JsonProbeOutcome::StillPossiblyValid,
149        );
150    }
151
152    #[test]
153    fn name_then_partial_arguments_key_is_still_possibly_valid() {
154        assert_eq!(
155            JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather", "argum"#),
156            JsonProbeOutcome::StillPossiblyValid,
157        );
158    }
159
160    #[test]
161    fn name_then_arguments_key_is_still_possibly_valid() {
162        assert_eq!(
163            JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather", "arguments""#),
164            JsonProbeOutcome::StillPossiblyValid,
165        );
166    }
167
168    #[test]
169    fn name_then_arguments_open_brace_is_still_possibly_valid() {
170        assert_eq!(
171            JsonProbeOutcome::validate_prefix(r#"{ "name": "get_weather", "arguments": {"#),
172            JsonProbeOutcome::StillPossiblyValid,
173        );
174    }
175
176    #[test]
177    fn arguments_with_partial_inner_key_value_is_still_possibly_valid() {
178        assert_eq!(
179            JsonProbeOutcome::validate_prefix(
180                r#"{ "name": "get_weather", "arguments": {"location":"#
181            ),
182            JsonProbeOutcome::StillPossiblyValid,
183        );
184    }
185
186    #[test]
187    fn arguments_with_partial_inner_string_value_is_still_possibly_valid() {
188        assert_eq!(
189            JsonProbeOutcome::validate_prefix(
190                r#"{ "name": "get_weather", "arguments": {"location": "Pa"#
191            ),
192            JsonProbeOutcome::StillPossiblyValid,
193        );
194    }
195
196    #[test]
197    fn complete_simple_tool_call_is_completed_valid() {
198        assert_eq!(
199            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{}}"#),
200            JsonProbeOutcome::CompletedValid,
201        );
202    }
203
204    #[test]
205    fn complete_tool_call_with_internal_whitespace_is_completed_valid() {
206        assert_eq!(
207            JsonProbeOutcome::validate_prefix(r#"{"name": "f", "arguments": {}}"#),
208            JsonProbeOutcome::CompletedValid,
209        );
210    }
211
212    #[test]
213    fn complete_tool_call_with_string_argument_is_completed_valid() {
214        assert_eq!(
215            JsonProbeOutcome::validate_prefix(
216                r#"{"name":"get_weather","arguments":{"location":"Paris"}}"#
217            ),
218            JsonProbeOutcome::CompletedValid,
219        );
220    }
221
222    #[test]
223    fn complete_tool_call_with_multiple_arguments_is_completed_valid() {
224        assert_eq!(
225            JsonProbeOutcome::validate_prefix(
226                r#"{"name":"book_flight","arguments":{"from":"NYC","to":"PAR","passengers":2}}"#
227            ),
228            JsonProbeOutcome::CompletedValid,
229        );
230    }
231
232    #[test]
233    fn complete_tool_call_with_nested_arguments_is_completed_valid() {
234        assert_eq!(
235            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{"a":{"b":[1,2,3]}}}"#),
236            JsonProbeOutcome::CompletedValid,
237        );
238    }
239
240    #[test]
241    fn complete_tool_call_with_close_brace_inside_string_is_completed_valid() {
242        assert_eq!(
243            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{"q":"a } b"}}"#),
244            JsonProbeOutcome::CompletedValid,
245        );
246    }
247
248    #[test]
249    fn complete_tool_call_with_escaped_quotes_in_string_is_completed_valid() {
250        assert_eq!(
251            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{"q":"he said \"hi\""}}"#),
252            JsonProbeOutcome::CompletedValid,
253        );
254    }
255
256    #[test]
257    fn complete_tool_call_with_unicode_strings_is_completed_valid() {
258        assert_eq!(
259            JsonProbeOutcome::validate_prefix(r#"{"name":"日本語","arguments":{"city":"パリ"}}"#),
260            JsonProbeOutcome::CompletedValid,
261        );
262    }
263
264    #[test]
265    fn complete_tool_call_with_trailing_whitespace_is_completed_valid() {
266        assert_eq!(
267            JsonProbeOutcome::validate_prefix("{\"name\":\"f\",\"arguments\":{}}\n"),
268            JsonProbeOutcome::CompletedValid,
269        );
270    }
271
272    #[test]
273    fn complete_tool_call_with_array_inside_arguments_is_completed_valid() {
274        assert_eq!(
275            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{"items":[1,2,3]}}"#),
276            JsonProbeOutcome::CompletedValid,
277        );
278    }
279
280    #[test]
281    fn complete_tool_call_without_arguments_field_is_completed_valid() {
282        assert_eq!(
283            JsonProbeOutcome::validate_prefix(r#"{"name":"ping"}"#),
284            JsonProbeOutcome::CompletedValid,
285        );
286    }
287
288    #[test]
289    fn top_level_array_is_failed() {
290        assert_eq!(
291            JsonProbeOutcome::validate_prefix("["),
292            JsonProbeOutcome::Failed
293        );
294    }
295
296    #[test]
297    fn top_level_scalar_number_is_failed() {
298        assert_eq!(
299            JsonProbeOutcome::validate_prefix("123"),
300            JsonProbeOutcome::Failed
301        );
302    }
303
304    #[test]
305    fn top_level_string_is_failed() {
306        assert_eq!(
307            JsonProbeOutcome::validate_prefix(r#""hi""#),
308            JsonProbeOutcome::Failed
309        );
310    }
311
312    #[test]
313    fn complete_object_with_wrong_first_key_is_failed() {
314        assert_eq!(
315            JsonProbeOutcome::validate_prefix(r#"{"foo":"bar"}"#),
316            JsonProbeOutcome::Failed,
317        );
318    }
319
320    #[test]
321    fn complete_object_with_non_string_name_is_failed() {
322        assert_eq!(
323            JsonProbeOutcome::validate_prefix(r#"{"name":123,"arguments":{}}"#),
324            JsonProbeOutcome::Failed,
325        );
326    }
327
328    #[test]
329    fn complete_object_with_null_name_is_failed() {
330        assert_eq!(
331            JsonProbeOutcome::validate_prefix(r#"{"name":null,"arguments":{}}"#),
332            JsonProbeOutcome::Failed,
333        );
334    }
335
336    #[test]
337    fn complete_object_with_arguments_as_array_is_failed() {
338        assert_eq!(
339            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":[]}"#),
340            JsonProbeOutcome::Failed,
341        );
342    }
343
344    #[test]
345    fn complete_object_with_arguments_as_string_is_failed() {
346        assert_eq!(
347            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":"hi"}"#),
348            JsonProbeOutcome::Failed,
349        );
350    }
351
352    #[test]
353    fn complete_object_with_third_top_level_key_is_failed() {
354        assert_eq!(
355            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{},"extra":1}"#),
356            JsonProbeOutcome::Failed,
357        );
358    }
359
360    #[test]
361    fn complete_object_with_empty_name_is_failed() {
362        assert_eq!(
363            JsonProbeOutcome::validate_prefix(r#"{"name":"","arguments":{}}"#),
364            JsonProbeOutcome::Failed,
365        );
366    }
367
368    #[test]
369    fn complete_object_with_trailing_garbage_is_failed() {
370        assert_eq!(
371            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":{}}garbage"#),
372            JsonProbeOutcome::Failed,
373        );
374    }
375
376    #[test]
377    fn empty_object_is_failed_due_to_missing_required_name() {
378        assert_eq!(
379            JsonProbeOutcome::validate_prefix("{}"),
380            JsonProbeOutcome::Failed
381        );
382    }
383
384    #[test]
385    fn complete_object_with_arguments_only_no_name_is_failed() {
386        assert_eq!(
387            JsonProbeOutcome::validate_prefix(r#"{"arguments":{}}"#),
388            JsonProbeOutcome::Failed,
389        );
390    }
391
392    #[test]
393    fn leading_whitespace_then_open_brace_is_still_possibly_valid() {
394        assert_eq!(
395            JsonProbeOutcome::validate_prefix("\n  \n{"),
396            JsonProbeOutcome::StillPossiblyValid,
397        );
398    }
399
400    #[test]
401    fn leading_whitespace_then_complete_tool_call_is_completed_valid() {
402        assert_eq!(
403            JsonProbeOutcome::validate_prefix("\n  {\"name\":\"f\",\"arguments\":{}}"),
404            JsonProbeOutcome::CompletedValid,
405        );
406    }
407
408    #[test]
409    fn complete_tool_call_followed_by_second_object_is_failed() {
410        assert_eq!(
411            JsonProbeOutcome::validate_prefix(
412                r#"{"name":"a","arguments":{}}{"name":"b","arguments":{}}"#
413            ),
414            JsonProbeOutcome::Failed,
415        );
416    }
417
418    #[test]
419    fn buffer_with_only_open_quote_is_still_possibly_valid() {
420        assert_eq!(
421            JsonProbeOutcome::validate_prefix(r#"{ "n"#),
422            JsonProbeOutcome::StillPossiblyValid,
423        );
424    }
425
426    #[test]
427    fn buffer_with_complete_first_field_unknown_second_key_is_failed() {
428        assert_eq!(
429            JsonProbeOutcome::validate_prefix(r#"{ "name": "f", "foo": 1}"#),
430            JsonProbeOutcome::Failed,
431        );
432    }
433
434    #[test]
435    fn unicode_letter_inside_name_value_completes_validly() {
436        assert_eq!(
437            JsonProbeOutcome::validate_prefix(r#"{"name":"éclair","arguments":{}}"#),
438            JsonProbeOutcome::CompletedValid,
439        );
440    }
441
442    #[test]
443    fn arguments_field_with_explicit_null_is_failed() {
444        assert_eq!(
445            JsonProbeOutcome::validate_prefix(r#"{"name":"f","arguments":null}"#),
446            JsonProbeOutcome::Failed,
447        );
448    }
449
450    #[test]
451    fn syntactically_malformed_object_is_failed() {
452        assert_eq!(
453            JsonProbeOutcome::validate_prefix("{,}"),
454            JsonProbeOutcome::Failed,
455        );
456    }
457}