synaptic_parsers/
fixing_parser.rs1use std::sync::Arc;
2
3use async_trait::async_trait;
4use synaptic_core::{ChatModel, ChatRequest, Message, RunnableConfig, SynapticError};
5use synaptic_runnables::Runnable;
6
7pub 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 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 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 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}