Skip to main content

synaptic_parsers/
fixing_parser.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use synaptic_core::{ChatModel, ChatRequest, Message, RunnableConfig, SynapticError};
5use synaptic_runnables::Runnable;
6
7/// A parser that uses an LLM to fix outputs that fail to parse.
8///
9/// Wraps an inner `Runnable<String, O>`. If the inner parser fails,
10/// sends the output + error to the LLM and retries parsing.
11pub struct OutputFixingParser<O: Send + Sync + 'static> {
12    inner: Box<dyn Runnable<String, O>>,
13    llm: Arc<dyn ChatModel>,
14    max_retries: usize,
15}
16
17impl<O: Send + Sync + 'static> OutputFixingParser<O> {
18    /// Create a new `OutputFixingParser` wrapping the given inner parser and LLM.
19    /// Defaults to 1 retry attempt.
20    pub fn new(inner: Box<dyn Runnable<String, O>>, llm: Arc<dyn ChatModel>) -> Self {
21        Self {
22            inner,
23            llm,
24            max_retries: 1,
25        }
26    }
27
28    /// Set the maximum number of retry attempts.
29    pub fn with_max_retries(mut self, n: usize) -> Self {
30        self.max_retries = n;
31        self
32    }
33}
34
35#[async_trait]
36impl<O: Send + Sync + 'static> Runnable<String, O> for OutputFixingParser<O> {
37    async fn invoke(&self, input: String, config: &RunnableConfig) -> Result<O, SynapticError> {
38        // First attempt with the original input.
39        match self.inner.invoke(input.clone(), config).await {
40            Ok(value) => return Ok(value),
41            Err(first_err) => {
42                let mut last_err = first_err;
43                let mut current_input = input;
44
45                for _ in 0..self.max_retries {
46                    let prompt = format!(
47                        "The following output failed to parse:\n\n{}\n\nError: {}\n\nPlease provide a corrected version of the output that will parse successfully.",
48                        current_input, last_err
49                    );
50
51                    let request = ChatRequest::new(vec![
52                        Message::system("You are a helpful assistant that fixes parsing errors."),
53                        Message::human(prompt),
54                    ]);
55
56                    let response = self.llm.chat(request).await?;
57                    let fixed = response.message.content().to_string();
58
59                    match self.inner.invoke(fixed.clone(), config).await {
60                        Ok(value) => return Ok(value),
61                        Err(e) => {
62                            last_err = e;
63                            current_input = fixed;
64                        }
65                    }
66                }
67
68                Err(last_err)
69            }
70        }
71    }
72}