pocketflow_batch/
lib.rs

1use pocketflow_core::*;
2use std::collections::HashMap;
3
4#[derive(Clone)]
5pub struct BatchNode {
6    name: String,
7    params: Params,
8    successors: HashMap<String, BatchNode>,
9    func: std::sync::Arc<
10        dyn Fn(&mut (dyn std::any::Any + Send), &Params) -> Result<Option<String>> + Send + Sync,
11    >,
12}
13
14impl BatchNode {
15    pub fn new<F>(name: impl Into<String>, func: F) -> Self
16    where
17        F: Fn(&mut (dyn std::any::Any + Send), &Params) -> Result<Option<String>>
18            + Send
19            + Sync
20            + 'static,
21    {
22        Self {
23            name: name.into(),
24            params: Params::new(),
25            successors: HashMap::new(),
26            func: std::sync::Arc::new(func),
27        }
28    }
29
30    pub fn add_successor(&mut self, action: impl Into<String>, node: BatchNode) -> &mut Self {
31        self.successors.insert(action.into(), node);
32        self
33    }
34
35    pub fn next(&mut self, node: BatchNode) -> &mut Self {
36        self.add_successor("default", node)
37    }
38
39    pub fn set_params(&mut self, params: Params) {
40        self.params = params;
41    }
42
43    pub fn get_params(&self) -> &Params {
44        &self.params
45    }
46
47    pub fn get_successor(&self, action: &str) -> Option<&BatchNode> {
48        self.successors.get(action)
49    }
50
51    pub fn has_successors(&self) -> bool {
52        !self.successors.is_empty()
53    }
54
55    pub fn run(&self, shared: &mut (dyn std::any::Any + Send)) -> Result<()> {
56        if self.has_successors() {
57            eprintln!("Warning: BatchNode won't run successors. Use BatchFlow.");
58        }
59        
60        (self.func)(shared, &self.params)?;
61        Ok(())
62    }
63
64    pub fn run_recursive(&self, shared: &mut (dyn std::any::Any + Send)) -> Result<()> {
65        let action = (self.func)(shared, &self.params)?;
66        
67        if let Some(next_node) = action
68            .as_ref()
69            .and_then(|a| self.successors.get(a))
70            .or_else(|| self.successors.get("default")) {
71            next_node.run_recursive(shared)?;
72        }
73        
74        Ok(())
75    }
76}
77
78#[derive(Clone)]
79pub struct BatchFlow {
80    start_node: Option<BatchNode>,
81    params: Params,
82}
83
84impl BatchFlow {
85    pub fn new() -> Self {
86        Self {
87            start_node: None,
88            params: Params::new(),
89        }
90    }
91
92    pub fn start(mut self, node: BatchNode) -> Self {
93        self.start_node = Some(node);
94        self
95    }
96
97    pub fn set_params(&mut self, params: Params) {
98        self.params = params;
99    }
100
101    pub fn run(&self, shared: &mut (dyn std::any::Any + Send)) -> Result<()> {
102        if let Some(ref node) = self.start_node {
103            let mut node = node.clone();
104            node.set_params(self.params.clone());
105            node.run_recursive(shared)?;
106        }
107        Ok(())
108    }
109
110    pub fn run_with_params(&self, shared: &mut (dyn std::any::Any + Send), params: Params) -> Result<()> {
111        if let Some(ref node) = self.start_node {
112            let mut node = node.clone();
113            let mut merged_params = self.params.clone();
114            merged_params.merge(&params);
115            node.set_params(merged_params);
116            node.run_recursive(shared)?;
117        }
118        Ok(())
119    }
120
121    pub fn run_batch(&self, shared: &mut (dyn std::any::Any + Send), batch_params: Vec<Params>) -> Result<Vec<()>> {
122        let mut results = Vec::with_capacity(batch_params.len());
123        for params in batch_params {
124            let flow = BatchFlow {
125                start_node: self.start_node.clone(),
126                params: self.params.clone(),
127            };
128            flow.run_with_params(shared, params)?;
129            results.push(());
130        }
131        Ok(results)
132    }
133}
134
135impl Default for BatchFlow {
136    fn default() -> Self {
137        Self::new()
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use std::sync::{Arc, Mutex};
145
146    #[derive(Default, Clone)]
147    struct TestShared {
148        pub counter: Arc<Mutex<i32>>,
149    }
150
151    #[test]
152    fn test_batch_flow() {
153        let mut shared = TestShared::default();
154        
155        let node = BatchNode::new("batch", |shared, params| {
156            if let Some(shared) = shared.downcast_mut::<TestShared>() {
157                let mut counter = shared.counter.lock().unwrap();
158                if let Some(value) = params.get("value") {
159                    if let Some(num) = value.as_i64() {
160                        *counter += num as i32;
161                    }
162                }
163            }
164            Ok(None)
165        });
166        
167        let flow = BatchFlow::new().start(node);
168        
169        let batch_params = vec![
170            {
171                let mut params = Params::new();
172                params.insert("value", 1);
173                params
174            },
175            {
176                let mut params = Params::new();
177                params.insert("value", 2);
178                params
179            },
180        ];
181        
182        flow.run_batch(&mut shared, batch_params).unwrap();
183        
184        let counter = shared.counter.lock().unwrap();
185        assert_eq!(*counter, 3);
186    }
187}
pocketflow_batch/lib.rs

pocketflow_batch/
lib.rs
