lc_chains/
sequential_chain.rs

1// lc-chains/src/sequential_chain.rs
2//! Sequential Chain
3//!
4//! Execute multiple chains sequentially.
5
6use async_trait::async_trait;
7use serde_json::Value;
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use crate::base::{BaseChain, ChainError, ChainResult};
12
13/// Sequential Chain
14///
15/// Executes multiple chains sequentially, where the output of one chain
16/// can be used as the input of the next.
17pub struct SequentialChain {
18    /// Chain list.
19    chains: Vec<ChainStep>,
20
21    /// Chain name.
22    name: String,
23}
24
25/// Chain step.
26struct ChainStep {
27    /// Chain instance.
28    chain: Arc<dyn BaseChain>,
29
30    /// Input mapping (from global input or previous output).
31    input_mapping: HashMap<String, String>,
32
33    /// Output mapping (to global result).
34    output_mapping: HashMap<String, String>,
35}
36
37impl SequentialChain {
38    /// Create an empty SequentialChain.
39    pub fn new() -> Self {
40        Self {
41            chains: Vec::new(),
42            name: "sequential_chain".to_string(),
43        }
44    }
45
46    /// Set name.
47    pub fn with_name(mut self, name: impl Into<String>) -> Self {
48        self.name = name.into();
49        self
50    }
51
52    /// Add a Chain.
53    ///
54    /// # Arguments
55    /// * `chain` - Chain to add
56    /// * `input_keys` - Input keys (from global input)
57    /// * `output_keys` - Output keys (to global result)
58    pub fn add_chain(
59        mut self,
60        chain: Arc<dyn BaseChain>,
61        input_keys: Vec<&str>,
62        output_keys: Vec<&str>,
63    ) -> Self {
64        let input_mapping = input_keys
65            .into_iter()
66            .map(|k| (k.to_string(), k.to_string()))
67            .collect();
68
69        let output_mapping = output_keys
70            .into_iter()
71            .map(|k| (k.to_string(), k.to_string()))
72            .collect();
73
74        self.chains.push(ChainStep {
75            chain,
76            input_mapping,
77            output_mapping,
78        });
79
80        self
81    }
82
83    /// Add a Chain with mapping.
84    ///
85    /// # Arguments
86    /// * `chain` - Chain to add
87    /// * `input_mapping` - Input mapping {chain_input_key: global_key}
88    /// * `output_mapping` - Output mapping {chain_output_key: global_key}
89    pub fn add_chain_with_mapping(
90        mut self,
91        chain: Arc<dyn BaseChain>,
92        input_mapping: HashMap<String, String>,
93        output_mapping: HashMap<String, String>,
94    ) -> Self {
95        self.chains.push(ChainStep {
96            chain,
97            input_mapping,
98            output_mapping,
99        });
100
101        self
102    }
103}
104
105impl Default for SequentialChain {
106    fn default() -> Self {
107        Self::new()
108    }
109}
110
111#[async_trait]
112impl BaseChain for SequentialChain {
113    fn input_keys(&self) -> Vec<&str> {
114        if let Some(first) = self.chains.first() {
115            first.input_mapping.values().map(|s| s.as_str()).collect()
116        } else {
117            vec![]
118        }
119    }
120
121    fn output_keys(&self) -> Vec<&str> {
122        if let Some(last) = self.chains.last() {
123            last.output_mapping.values().map(|s| s.as_str()).collect()
124        } else {
125            vec![]
126        }
127    }
128
129    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
130        let mut current_state = inputs.clone();
131        let mut final_output = HashMap::new();
132
133        for (step_index, step) in self.chains.iter().enumerate() {
134            let mut chain_inputs = HashMap::new();
135            for (chain_key, global_key) in &step.input_mapping {
136                if let Some(value) = current_state.get(global_key) {
137                    chain_inputs.insert(chain_key.clone(), value.clone());
138                } else {
139                    return Err(ChainError::MissingInput(format!(
140                        "Step {}: missing input '{}' (mapped from '{}')",
141                        step_index, chain_key, global_key
142                    )));
143                }
144            }
145
146            let chain_output = step.chain.invoke(chain_inputs).await.map_err(|e| {
147                ChainError::ExecutionError(format!(
148                    "Step {} ({}) execution failed: {}",
149                    step_index,
150                    step.chain.name(),
151                    e
152                ))
153            })?;
154
155            for (chain_key, global_key) in &step.output_mapping {
156                if let Some(value) = chain_output.get(chain_key) {
157                    current_state.insert(global_key.clone(), value.clone());
158                    final_output.insert(global_key.clone(), value.clone());
159                } else {
160                    return Err(ChainError::OutputError(format!(
161                        "Step {} ({}) did not produce expected output key '{}' (mapped to '{}')",
162                        step_index,
163                        step.chain.name(),
164                        chain_key,
165                        global_key
166                    )));
167                }
168            }
169        }
170
171        Ok(final_output)
172    }
173
174    fn name(&self) -> &str {
175        &self.name
176    }
177}
178
179impl std::fmt::Debug for SequentialChain {
180    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
181        f.debug_struct("SequentialChain")
182            .field("steps", &self.chains.len())
183            .field("name", &self.name)
184            .finish()
185    }
186}
lc_chains/sequential_chain.rs

lc_chains/
sequential_chain.rs
