diff --git a/src/graph/structured.rs b/src/graph/structured.rs index 3b7a342..c1a4f15 100644 --- a/src/graph/structured.rs +++ b/src/graph/structured.rs @@ -88,11 +88,21 @@ async fn run_one_shot(prompt: &str, ctx: &mut RequestContext) -> Result } fn try_parse_json(raw: &str) -> Option { - let cleaned = strip_code_fences(raw.trim()); - + let cleaned = strip_code_fences(strip_thinking_blocks(raw.trim())); serde_json::from_str(cleaned).ok() } +fn strip_thinking_blocks(s: &str) -> &str { + let mut s = s.trim_start(); + while s.starts_with("") { + match s.find("") { + Some(end) => s = s[end + "".len()..].trim_start(), + None => break, + } + } + s +} + fn strip_code_fences(s: &str) -> &str { let after_open = s .strip_prefix("```json") @@ -148,6 +158,38 @@ mod tests { assert_eq!(v, json!({"x": true})); } + #[test] + fn try_parse_json_strips_thinking_blocks() { + let raw = "\nsome reasoning\n\n{\"a\": 1}"; + + let v = try_parse_json(raw).unwrap(); + + assert_eq!(v, json!({"a": 1})); + } + + #[test] + fn try_parse_json_strips_empty_thinking_block() { + let raw = "\n\n\n{\"a\": 1}"; + + let v = try_parse_json(raw).unwrap(); + + assert_eq!(v, json!({"a": 1})); + } + + #[test] + fn try_parse_json_strips_multiple_thinking_blocks() { + let raw = "first\nsecond\n{\"a\": 1}"; + + let v = try_parse_json(raw).unwrap(); + + assert_eq!(v, json!({"a": 1})); + } + + #[test] + fn try_parse_json_unclosed_think_tag_returns_none() { + assert!(try_parse_json("unclosed {\"a\": 1}").is_none()); + } + #[test] fn try_parse_json_returns_none_on_prose() { assert!(try_parse_json("Here is the result: it's good").is_none());