langchain_rust/output_parsers/
markdown_parser.rs

1use async_trait::async_trait;
2use regex::Regex;
3
4use super::{OutputParser, OutputParserError};
5
6pub struct MarkdownParser {
7    expresion: String,
8    trim: bool,
9}
10impl MarkdownParser {
11    pub fn new() -> Self {
12        Self {
13            expresion: r"```(?:\w+)?\s*([\s\S]+?)\s*```".to_string(),
14            trim: false,
15        }
16    }
17
18    pub fn with_custom_expresion(mut self, expresion: &str) -> Self {
19        self.expresion = expresion.to_string();
20        self
21    }
22
23    pub fn with_trim(mut self, trim: bool) -> Self {
24        self.trim = trim;
25        self
26    }
27}
28impl Default for MarkdownParser {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34#[async_trait]
35impl OutputParser for MarkdownParser {
36    async fn parse(&self, output: &str) -> Result<String, OutputParserError> {
37        let re = Regex::new(r"```(?:\w+)?\s*([\s\S]+?)\s*```")?;
38        if let Some(cap) = re.captures(output) {
39            let find = cap[1].to_string();
40            if self.trim {
41                Ok(find.trim().to_string())
42            } else {
43                Ok(find)
44            }
45        } else {
46            Err(OutputParserError::ParsingError(
47                "No code block found".into(),
48            ))
49        }
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[tokio::test]
58    async fn test_markdown_parser_finds_code_block() {
59        let parser = MarkdownParser::new();
60        let markdown_content = r#"
61```rust
62fn main() {
63    println!("Hello, world!");
64}
65```
66"#;
67        let result = parser.parse(markdown_content).await;
68        println!("{:?}", result);
69
70        let correct = r#"fn main() {
71    println!("Hello, world!");
72}"#;
73        assert!(result.is_ok());
74        assert_eq!(result.unwrap(), correct);
75    }
76}