langchain_rust/chain/sequential/
builder.rs

1use std::collections::HashSet;
2
3use crate::chain::Chain;
4
5use super::SequentialChain;
6
7pub struct SequentialChainBuilder {
8    chains: Vec<Box<dyn Chain>>,
9}
10
11impl SequentialChainBuilder {
12    pub fn new() -> Self {
13        Self { chains: Vec::new() }
14    }
15
16    pub fn add_chain<C: Chain + 'static>(mut self, chain: C) -> Self {
17        self.chains.push(Box::new(chain));
18        self
19    }
20
21    pub fn build(self) -> SequentialChain {
22        let outputs: HashSet<String> = self
23            .chains
24            .iter()
25            .flat_map(|c| c.get_output_keys())
26            .collect();
27
28        let input_keys: HashSet<String> = self
29            .chains
30            .iter()
31            .flat_map(|c| c.get_input_keys())
32            .collect();
33
34        SequentialChain {
35            chains: self.chains,
36            input_keys,
37            outputs,
38        }
39    }
40}
41
42#[macro_export]
43macro_rules! sequential_chain {
44    ( $( $chain:expr ),* $(,)? ) => {
45        {
46            let mut builder = $crate::chain::SequentialChainBuilder::new();
47            $(
48                builder = builder.add_chain($chain);
49            )*
50            builder.build()
51        }
52    };
53}