llm_kernel/llm/
json_extract.rs

1//! Extract structured JSON from raw LLM text output.
2//!
3//! LLMs often wrap JSON in markdown code fences or add prose around it.
4//! This module handles the common cases:
5//!
6//! - ` ```json ... ``` ` fenced blocks
7//! - ` ``` ... ``` ` fences with JSON inside
8//! - Raw balanced-bracket JSON buried in text
9
10use serde::de::DeserializeOwned;
11
12/// Extract the first JSON object or array from raw LLM text.
13///
14/// Handles ```json fences, ``` fences, and raw JSON.
15pub fn extract_json(text: &str) -> Option<String> {
16    let text = text.trim();
17
18    // Try ```json ... ```
19    if let Some(extracted) = extract_fenced(text, "json") {
20        return Some(extracted);
21    }
22
23    // Try ``` ... ``` (generic fence)
24    if let Some(extracted) = extract_fenced(text, "")
25        && looks_like_json(&extracted)
26    {
27        return Some(extracted);
28    }
29
30    // Try raw balanced bracket matching
31    find_balanced_json(text)
32}
33
34/// Parse LLM output into a typed struct using `extract_json` + `serde_json`.
35///
36/// Returns a deserialization error if no JSON is found or parsing fails.
37pub fn parse_json<T: DeserializeOwned>(text: &str) -> Result<T, ParseJsonError> {
38    let json_str = extract_json(text).ok_or(ParseJsonError::NoJsonFound)?;
39    serde_json::from_str(&json_str).map_err(ParseJsonError::Deserialize)
40}
41
42/// Error from [`parse_json`].
43#[derive(Debug)]
44pub enum ParseJsonError {
45    /// No JSON found in the LLM output text.
46    NoJsonFound,
47    /// Found JSON but deserialization failed.
48    Deserialize(serde_json::Error),
49}
50
51impl std::fmt::Display for ParseJsonError {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        match self {
54            Self::NoJsonFound => write!(f, "no JSON found in LLM output"),
55            Self::Deserialize(e) => write!(f, "JSON parse error: {e}"),
56        }
57    }
58}
59
60impl std::error::Error for ParseJsonError {
61    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
62        match self {
63            Self::NoJsonFound => None,
64            Self::Deserialize(e) => Some(e),
65        }
66    }
67}
68
69/// Type-safe JSON extraction wrapper for use with LLM clients.
70pub struct JsonExtractor<T>(std::marker::PhantomData<T>);
71
72impl<T: DeserializeOwned> JsonExtractor<T> {
73    /// Parse a raw LLM response string into `T`.
74    pub fn parse(text: &str) -> Result<T, ParseJsonError> {
75        parse_json(text)
76    }
77}
78
79/// Extract the content of a named XML tag from raw LLM text.
80///
81/// Returns the text between `<tag>` and `</tag>`, trimmed of leading/trailing whitespace.
82/// Returns `None` if either the opening or closing tag is absent.
83///
84/// Useful for parsing Claude-style structured output that wraps results in XML tags:
85///
86/// ```
87/// # use llm_kernel::llm::json_extract::extract_xml_tag;
88/// let text = "Here is my answer:\n<result>\nhello world\n</result>";
89/// assert_eq!(extract_xml_tag(text, "result"), Some("hello world"));
90/// ```
91pub fn extract_xml_tag<'a>(text: &'a str, tag: &str) -> Option<&'a str> {
92    let open = format!("<{tag}>");
93    let close = format!("</{tag}>");
94    let start = text.find(open.as_str())? + open.len();
95    let end = text[start..].find(close.as_str())?;
96    Some(text[start..start + end].trim())
97}
98
99fn extract_fenced(text: &str, lang: &str) -> Option<String> {
100    let opener = if lang.is_empty() {
101        "```".to_string()
102    } else {
103        format!("```{}", lang)
104    };
105
106    let start = text.find(&opener)?;
107    let after_open = start + opener.len();
108
109    // Skip to end of opening line
110    let body_start = text[after_open..]
111        .find('\n')
112        .map_or(after_open, |i| after_open + i + 1);
113
114    // Find closing ```
115    let closer = text[body_start..].find("```")?;
116    let body = text[body_start..body_start + closer].trim();
117
118    if body.is_empty() {
119        return None;
120    }
121
122    Some(body.to_string())
123}
124
125fn looks_like_json(s: &str) -> bool {
126    let trimmed = s.trim();
127    trimmed.starts_with('{') || trimmed.starts_with('[')
128}
129
130fn find_balanced_json(text: &str) -> Option<String> {
131    for (i, b) in text.bytes().enumerate() {
132        if b == b'{' || b == b'[' {
133            let open = b;
134            let close = if open == b'{' { b'}' } else { b']' };
135            if let Some(end) = find_balanced_end(text.as_bytes(), i, open, close) {
136                return Some(text[i..=end].to_string());
137            }
138        }
139    }
140    None
141}
142
143fn find_balanced_end(bytes: &[u8], start: usize, open: u8, close: u8) -> Option<usize> {
144    let mut depth = 0i32;
145    let mut in_string = false;
146    let mut escape = false;
147
148    for (i, &b) in bytes[start..].iter().enumerate() {
149        if escape {
150            escape = false;
151            continue;
152        }
153        if b == b'\\' && in_string {
154            escape = true;
155            continue;
156        }
157        if b == b'"' {
158            in_string = !in_string;
159            continue;
160        }
161        if in_string {
162            continue;
163        }
164        if b == open {
165            depth += 1;
166        } else if b == close {
167            depth -= 1;
168            if depth == 0 {
169                return Some(start + i);
170            }
171        }
172    }
173    None
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179    use serde::Deserialize;
180
181    #[derive(Debug, Deserialize, PartialEq)]
182    struct TestOutput {
183        name: String,
184        value: i32,
185    }
186
187    #[test]
188    fn extract_from_json_fence() {
189        let input = "Here is the result:\n```json\n{\"name\":\"test\",\"value\":42}\n```\nDone.";
190        let json = extract_json(input).unwrap();
191        let parsed: TestOutput = serde_json::from_str(&json).unwrap();
192        assert_eq!(
193            parsed,
194            TestOutput {
195                name: "test".into(),
196                value: 42
197            }
198        );
199    }
200
201    #[test]
202    fn extract_from_generic_fence() {
203        let input = "```\n{\"name\":\"x\",\"value\":1}\n```";
204        let json = extract_json(input).unwrap();
205        assert!(json.contains("\"name\""));
206    }
207
208    #[test]
209    fn extract_raw_json_object() {
210        let input = "The answer is {\"name\":\"raw\",\"value\":7} as shown.";
211        let json = extract_json(input).unwrap();
212        let parsed: TestOutput = serde_json::from_str(&json).unwrap();
213        assert_eq!(
214            parsed,
215            TestOutput {
216                name: "raw".into(),
217                value: 7
218            }
219        );
220    }
221
222    #[test]
223    fn extract_raw_json_array() {
224        let input = "Items: [1, 2, 3]";
225        let json = extract_json(input).unwrap();
226        assert_eq!(json, "[1, 2, 3]");
227    }
228
229    #[test]
230    fn no_json_returns_none() {
231        assert!(extract_json("no json here").is_none());
232    }
233
234    #[test]
235    fn parse_json_works() {
236        let input = "```json\n{\"name\":\"ok\",\"value\":99}\n```";
237        let result: TestOutput = parse_json(input).unwrap();
238        assert_eq!(result.value, 99);
239    }
240
241    #[test]
242    fn parse_json_fails_gracefully() {
243        let result: Result<TestOutput, ParseJsonError> = parse_json("no json");
244        assert!(matches!(result, Err(ParseJsonError::NoJsonFound)));
245    }
246
247    #[test]
248    fn nested_json() {
249        let input = r#"{"outer": {"inner": [1,2]}}"#;
250        let json = extract_json(input).unwrap();
251        assert!(json.starts_with('{'));
252        assert!(json.ends_with('}'));
253    }
254
255    #[test]
256    fn escaped_quotes_in_string() {
257        let input = r#"{"name":"he said \"hello\"","value":0}"#;
258        let json = extract_json(input).unwrap();
259        assert!(json.contains("hello"));
260    }
261}
llm_kernel/llm/json_extract.rs

llm_kernel/llm/
json_extract.rs
