langchain_rust/output_parsers/
markdown_parser.rs1use 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}