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