pocketflow 0.1.0

A minimalist LLM framework. Let Agents build Agents!
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use std::sync::Arc;

use dashmap::DashMap;
use eyre::Result;

use crate::{
    node::base::BaseNode,
    types::{ActionKey, DEFAULT_ACTION, Params, SharedState},
};

/// Flow orchestration structure
pub struct Flow {
    params: Params,
    pub(crate) start_node: Option<Arc<dyn BaseNode>>,
    node_map: DashMap<ActionKey, Arc<dyn BaseNode>>,
}

impl Default for Flow {
    fn default() -> Self {
        Self::new()
    }
}

impl Flow {
    pub fn new() -> Self {
        Self {
            params: Default::default(),
            start_node: None,
            node_map: Default::default(),
        }
    }

    pub fn set_params(&mut self, params: Params) -> &mut Self {
        self.params = params;
        self
    }

    pub fn start(&mut self, node: Arc<dyn BaseNode>) -> &mut Self {
        self.start_node = Some(node);
        self
    }

    pub fn get_node(&self, action: &str) -> Option<Arc<dyn BaseNode>> {
        self.node_map.get(action).map(|node| node.clone())
    }

    pub fn run(&self, shared: &SharedState) -> Result<String> {
        let start_node = self
            .start_node
            .as_ref()
            .ok_or_else(|| eyre::eyre!("Flow has no start node"))?;

        let mut current_node = start_node.clone();
        let _params = self.params.clone();
        // Initialize action_result directly when it's first needed

        // Start execution
        let mut prep_res = current_node.prep(shared)?;
        let mut exec_res = current_node.exec(&prep_res)?;
        let mut action_result = current_node.post(shared, &prep_res, &exec_res)?;

        loop {
            // Get successor from node's successors map
            let next_node_key = if action_result.is_empty() {
                DEFAULT_ACTION.to_string()
            } else {
                action_result.clone()
            };

            // Find the next node
            let next_node_opt = self.node_map.get(&next_node_key).map(|node| node.clone());

            // If there's no next node, break the loop
            if next_node_opt.is_none() {
                break;
            }

            // Update current node and continue execution
            current_node = next_node_opt.expect("Next node not found in flow");
            prep_res = current_node.prep(shared)?;
            exec_res = current_node.exec(&prep_res)?;
            action_result = current_node.post(shared, &prep_res, &exec_res)?;
        }

        Ok(action_result)
    }
}

/// Batch flow structure
pub struct BatchFlow {
    pub(crate) inner: Flow,
}

impl Default for BatchFlow {
    fn default() -> Self {
        Self::new()
    }
}

impl BatchFlow {
    pub fn new() -> Self {
        Self { inner: Flow::new() }
    }

    pub fn set_params(&mut self, params: Params) -> &mut Self {
        self.inner.set_params(params);
        self
    }

    pub fn start(&mut self, node: Arc<dyn BaseNode>) -> &mut Self {
        self.inner.start(node);
        self
    }

    pub fn run(&self, shared: &SharedState) -> Result<String> {
        let prep_res = self
            .inner
            .start_node
            .as_ref()
            .ok_or_else(|| eyre::eyre!("BatchFlow has no start node"))?
            .prep(shared)?;

        let batch_items = prep_res
            .as_array()
            .ok_or_else(|| eyre::eyre!("BatchFlow prep must return an array"))?;

        for item in batch_items {
            let mut params = self.inner.params.clone();
            if let Some(item_obj) = item.as_object() {
                for (k, v) in item_obj {
                    params.insert(k.clone(), v.clone());
                }
            }

            // Execute flow for each batch item
            self.inner.run(shared)?;
        }

        Ok(DEFAULT_ACTION.to_string())
    }
}