Skip to main content

agent_line/tools/
parse.rs

1use crate::agent::StepError;
2
3/// Remove markdown code fences from LLM output.
4pub fn strip_code_fences(response: &str) -> String {
5    let trimmed = response.trim();
6    if trimmed.starts_with("```") {
7        let lines: Vec<&str> = trimmed.lines().collect();
8        // Skip first line (```rust) and last line (```)
9        lines[1..lines.len() - 1].join("\n")
10    } else {
11        trimmed.to_string()
12    }
13}
14
15/// Split LLM output into lines, stripping numbering, bullets, and whitespace.
16pub fn parse_lines(response: &str) -> Vec<String> {
17    response
18        .lines() // split into individual lines
19        .map(|line| line.trim().trim_start_matches(|c: char| c.is_ascii_digit()))
20        .map(|line| line.strip_prefix(".").unwrap_or(line))
21        .map(|line| line.strip_prefix("-").unwrap_or(line))
22        .map(|line| line.strip_prefix("*").unwrap_or(line))
23        .map(|line| line.trim())
24        .map(|line| line.to_string())
25        .filter(|line| !line.is_empty())
26        .collect()
27}
28
29/// Extract the first JSON object or array from text that may contain prose.
30pub fn extract_json(response: &str) -> Result<String, StepError> {
31    let no_fences = strip_code_fences(response);
32    // get opening { or  [
33    let trimmed = no_fences.trim();
34    let index = trimmed.find(['{', '[']);
35    if let Some(start) = index {
36        let slice = &trimmed[start..];
37        let parsed = serde_json::Deserializer::from_str(slice)
38            .into_iter::<serde_json::Value>()
39            .next();
40
41        if let Some(Ok(val)) = parsed {
42            return Ok(val.to_string());
43        }
44    }
45
46    Err(StepError::Invalid("invalid json".to_string()))
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    // --- parse_lines tests ---
54
55    #[test]
56    fn test_parse_lines_numbered() {
57        let input = "1. First item\n2. Second item\n3. Third item";
58        let result = parse_lines(input);
59        assert_eq!(result, vec!["First item", "Second item", "Third item"]);
60    }
61
62    #[test]
63    fn test_parse_lines_dashes() {
64        let input = "- Alpha\n- Beta\n- Gamma";
65        let result = parse_lines(input);
66        assert_eq!(result, vec!["Alpha", "Beta", "Gamma"]);
67    }
68
69    #[test]
70    fn test_parse_lines_asterisks() {
71        let input = "* One\n* Two";
72        let result = parse_lines(input);
73        assert_eq!(result, vec!["One", "Two"]);
74    }
75
76    #[test]
77    fn test_parse_lines_plain() {
78        let input = "First\nSecond\nThird";
79        let result = parse_lines(input);
80        assert_eq!(result, vec!["First", "Second", "Third"]);
81    }
82
83    #[test]
84    fn test_parse_lines_skips_empty_lines() {
85        let input = "One\n\nTwo\n\n";
86        let result = parse_lines(input);
87        assert_eq!(result, vec!["One", "Two"]);
88    }
89
90    #[test]
91    fn test_parse_lines_trims_whitespace() {
92        let input = "  1. Padded  \n  2. Also padded  ";
93        let result = parse_lines(input);
94        assert_eq!(result, vec!["Padded", "Also padded"]);
95    }
96
97    // --- extract_json tests ---
98
99    #[test]
100    fn test_extract_json_object_from_prose() {
101        let input = "Here is the result:\n{\"name\": \"test\", \"value\": 42}\nDone.";
102        let result = extract_json(input).unwrap();
103        let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
104        assert_eq!(parsed["name"], "test");
105        assert_eq!(parsed["value"], 42);
106    }
107
108    #[test]
109    fn test_extract_json_array() {
110        let input = "The topics are: [\"rust\", \"python\", \"go\"]";
111        let result = extract_json(input).unwrap();
112        let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
113        assert_eq!(parsed[0], "rust");
114    }
115
116    #[test]
117    fn test_extract_json_in_code_fence() {
118        let input = "```json\n{\"key\": \"value\"}\n```";
119        let result = extract_json(input).unwrap();
120        let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
121        assert_eq!(parsed["key"], "value");
122    }
123
124    #[test]
125    fn test_extract_json_no_json_returns_error() {
126        let input = "There is no JSON here at all.";
127        let result = extract_json(input);
128        assert!(result.is_err());
129    }
130
131    #[test]
132    fn test_extract_json_nested_object() {
133        let input = "Result: {\"outer\": {\"inner\": true}}";
134        let result = extract_json(input).unwrap();
135        let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
136        assert_eq!(parsed["outer"]["inner"], true);
137    }
138}