adk_graph/
graph.rs

1//! StateGraph builder for constructing graphs
2
3use crate::checkpoint::Checkpointer;
4use crate::edge::{Edge, EdgeTarget, RouterFn, END, START};
5use crate::error::{GraphError, Result};
6use crate::node::{FunctionNode, Node, NodeContext, NodeOutput};
7use crate::state::{State, StateSchema};
8use std::collections::{HashMap, HashSet};
9use std::future::Future;
10use std::sync::Arc;
11
12/// Builder for constructing graphs
13pub struct StateGraph {
14    /// State schema
15    pub schema: StateSchema,
16    /// Registered nodes
17    pub nodes: HashMap<String, Arc<dyn Node>>,
18    /// Registered edges
19    pub edges: Vec<Edge>,
20}
21
22impl StateGraph {
23    /// Create a new graph with the given state schema
24    pub fn new(schema: StateSchema) -> Self {
25        Self { schema, nodes: HashMap::new(), edges: vec![] }
26    }
27
28    /// Create with a simple schema (just channel names, all overwrite)
29    pub fn with_channels(channels: &[&str]) -> Self {
30        Self::new(StateSchema::simple(channels))
31    }
32
33    /// Add a node to the graph
34    pub fn add_node<N: Node + 'static>(mut self, node: N) -> Self {
35        self.nodes.insert(node.name().to_string(), Arc::new(node));
36        self
37    }
38
39    /// Add a function as a node
40    pub fn add_node_fn<F, Fut>(self, name: &str, func: F) -> Self
41    where
42        F: Fn(NodeContext) -> Fut + Send + Sync + 'static,
43        Fut: Future<Output = Result<NodeOutput>> + Send + 'static,
44    {
45        self.add_node(FunctionNode::new(name, func))
46    }
47
48    /// Add a direct edge from source to target
49    pub fn add_edge(mut self, source: &str, target: &str) -> Self {
50        let target = EdgeTarget::from(target);
51
52        if source == START {
53            // Find existing entry or create new one
54            let entry_idx = self.edges.iter().position(|e| matches!(e, Edge::Entry { .. }));
55
56            match entry_idx {
57                Some(idx) => {
58                    if let Edge::Entry { targets } = &mut self.edges[idx] {
59                        if let EdgeTarget::Node(node) = &target {
60                            if !targets.contains(node) {
61                                targets.push(node.clone());
62                            }
63                        }
64                    }
65                }
66                None => {
67                    if let EdgeTarget::Node(node) = target {
68                        self.edges.push(Edge::Entry { targets: vec![node] });
69                    }
70                }
71            }
72        } else {
73            self.edges.push(Edge::Direct { source: source.to_string(), target });
74        }
75
76        self
77    }
78
79    /// Add a conditional edge with a router function
80    pub fn add_conditional_edges<F, I>(mut self, source: &str, router: F, targets: I) -> Self
81    where
82        F: Fn(&State) -> String + Send + Sync + 'static,
83        I: IntoIterator<Item = (&'static str, &'static str)>,
84    {
85        let targets_map: HashMap<String, EdgeTarget> =
86            targets.into_iter().map(|(k, v)| (k.to_string(), EdgeTarget::from(v))).collect();
87
88        self.edges.push(Edge::Conditional {
89            source: source.to_string(),
90            router: Arc::new(router),
91            targets: targets_map,
92        });
93
94        self
95    }
96
97    /// Add a conditional edge with an Arc router (for pre-built routers)
98    pub fn add_conditional_edges_arc<I>(
99        mut self,
100        source: &str,
101        router: RouterFn,
102        targets: I,
103    ) -> Self
104    where
105        I: IntoIterator<Item = (&'static str, &'static str)>,
106    {
107        let targets_map: HashMap<String, EdgeTarget> =
108            targets.into_iter().map(|(k, v)| (k.to_string(), EdgeTarget::from(v))).collect();
109
110        self.edges.push(Edge::Conditional {
111            source: source.to_string(),
112            router,
113            targets: targets_map,
114        });
115
116        self
117    }
118
119    /// Compile the graph for execution
120    pub fn compile(self) -> Result<CompiledGraph> {
121        self.validate()?;
122
123        Ok(CompiledGraph {
124            schema: self.schema,
125            nodes: self.nodes,
126            edges: self.edges,
127            checkpointer: None,
128            interrupt_before: HashSet::new(),
129            interrupt_after: HashSet::new(),
130            recursion_limit: 50,
131        })
132    }
133
134    /// Validate the graph structure
135    fn validate(&self) -> Result<()> {
136        // Check for entry point
137        let has_entry = self.edges.iter().any(|e| matches!(e, Edge::Entry { .. }));
138        if !has_entry {
139            return Err(GraphError::NoEntryPoint);
140        }
141
142        // Check all node references exist
143        for edge in &self.edges {
144            match edge {
145                Edge::Direct { source, target } => {
146                    if source != START && !self.nodes.contains_key(source) {
147                        return Err(GraphError::NodeNotFound(source.clone()));
148                    }
149                    if let EdgeTarget::Node(name) = target {
150                        if !self.nodes.contains_key(name) {
151                            return Err(GraphError::EdgeTargetNotFound(name.clone()));
152                        }
153                    }
154                }
155                Edge::Conditional { source, targets, .. } => {
156                    if !self.nodes.contains_key(source) {
157                        return Err(GraphError::NodeNotFound(source.clone()));
158                    }
159                    for target in targets.values() {
160                        if let EdgeTarget::Node(name) = target {
161                            if !self.nodes.contains_key(name) {
162                                return Err(GraphError::EdgeTargetNotFound(name.clone()));
163                            }
164                        }
165                    }
166                }
167                Edge::Entry { targets } => {
168                    for target in targets {
169                        if !self.nodes.contains_key(target) {
170                            return Err(GraphError::EdgeTargetNotFound(target.clone()));
171                        }
172                    }
173                }
174            }
175        }
176
177        Ok(())
178    }
179}
180
181/// A compiled graph ready for execution
182pub struct CompiledGraph {
183    pub(crate) schema: StateSchema,
184    pub(crate) nodes: HashMap<String, Arc<dyn Node>>,
185    pub(crate) edges: Vec<Edge>,
186    pub(crate) checkpointer: Option<Arc<dyn Checkpointer>>,
187    pub(crate) interrupt_before: HashSet<String>,
188    pub(crate) interrupt_after: HashSet<String>,
189    pub(crate) recursion_limit: usize,
190}
191
192impl CompiledGraph {
193    /// Configure checkpointing
194    pub fn with_checkpointer<C: Checkpointer + 'static>(mut self, checkpointer: C) -> Self {
195        self.checkpointer = Some(Arc::new(checkpointer));
196        self
197    }
198
199    /// Configure checkpointing with Arc
200    pub fn with_checkpointer_arc(mut self, checkpointer: Arc<dyn Checkpointer>) -> Self {
201        self.checkpointer = Some(checkpointer);
202        self
203    }
204
205    /// Configure interrupt before specific nodes
206    pub fn with_interrupt_before(mut self, nodes: &[&str]) -> Self {
207        self.interrupt_before = nodes.iter().map(|s| s.to_string()).collect();
208        self
209    }
210
211    /// Configure interrupt after specific nodes
212    pub fn with_interrupt_after(mut self, nodes: &[&str]) -> Self {
213        self.interrupt_after = nodes.iter().map(|s| s.to_string()).collect();
214        self
215    }
216
217    /// Set recursion limit for cycles
218    pub fn with_recursion_limit(mut self, limit: usize) -> Self {
219        self.recursion_limit = limit;
220        self
221    }
222
223    /// Get entry nodes
224    pub fn get_entry_nodes(&self) -> Vec<String> {
225        for edge in &self.edges {
226            if let Edge::Entry { targets } = edge {
227                return targets.clone();
228            }
229        }
230        vec![]
231    }
232
233    /// Get next nodes after executing the given nodes
234    pub fn get_next_nodes(&self, executed: &[String], state: &State) -> Vec<String> {
235        let mut next = Vec::new();
236
237        for edge in &self.edges {
238            match edge {
239                Edge::Direct { source, target: EdgeTarget::Node(n) }
240                    if executed.contains(source) =>
241                {
242                    if !next.contains(n) {
243                        next.push(n.clone());
244                    }
245                }
246                Edge::Conditional { source, router, targets } if executed.contains(source) => {
247                    let route = router(state);
248                    if let Some(EdgeTarget::Node(n)) = targets.get(&route) {
249                        if !next.contains(n) {
250                            next.push(n.clone());
251                        }
252                    }
253                    // If route leads to END or not found in targets, next will be empty for this path
254                }
255                _ => {}
256            }
257        }
258
259        next
260    }
261
262    /// Check if any of the executed nodes lead to END
263    pub fn leads_to_end(&self, executed: &[String], state: &State) -> bool {
264        for edge in &self.edges {
265            match edge {
266                Edge::Direct { source, target } if executed.contains(source) => {
267                    if target.is_end() {
268                        return true;
269                    }
270                }
271                Edge::Conditional { source, router, targets } if executed.contains(source) => {
272                    let route = router(state);
273                    if route == END {
274                        return true;
275                    }
276                    if let Some(target) = targets.get(&route) {
277                        if target.is_end() {
278                            return true;
279                        }
280                    }
281                }
282                _ => {}
283            }
284        }
285        false
286    }
287
288    /// Get the state schema
289    pub fn schema(&self) -> &StateSchema {
290        &self.schema
291    }
292
293    /// Get the checkpointer if configured
294    pub fn checkpointer(&self) -> Option<&Arc<dyn Checkpointer>> {
295        self.checkpointer.as_ref()
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302    use serde_json::json;
303
304    #[test]
305    fn test_basic_graph_construction() {
306        let graph = StateGraph::with_channels(&["input", "output"])
307            .add_node_fn("process", |_ctx| async { Ok(NodeOutput::new()) })
308            .add_edge(START, "process")
309            .add_edge("process", END)
310            .compile();
311
312        assert!(graph.is_ok());
313    }
314
315    #[test]
316    fn test_graph_missing_entry() {
317        let graph = StateGraph::with_channels(&["input"])
318            .add_node_fn("process", |_ctx| async { Ok(NodeOutput::new()) })
319            .add_edge("process", END) // No START -> process edge
320            .compile();
321
322        assert!(matches!(graph, Err(GraphError::NoEntryPoint)));
323    }
324
325    #[test]
326    fn test_graph_missing_node() {
327        let graph = StateGraph::with_channels(&["input"]).add_edge(START, "nonexistent").compile();
328
329        assert!(matches!(graph, Err(GraphError::EdgeTargetNotFound(_))));
330    }
331
332    #[test]
333    fn test_conditional_edges() {
334        let graph = StateGraph::with_channels(&["next"])
335            .add_node_fn("router", |_ctx| async { Ok(NodeOutput::new()) })
336            .add_node_fn("path_a", |_ctx| async { Ok(NodeOutput::new()) })
337            .add_node_fn("path_b", |_ctx| async { Ok(NodeOutput::new()) })
338            .add_edge(START, "router")
339            .add_conditional_edges(
340                "router",
341                |state| state.get("next").and_then(|v| v.as_str()).unwrap_or(END).to_string(),
342                [("path_a", "path_a"), ("path_b", "path_b"), (END, END)],
343            )
344            .compile()
345            .unwrap();
346
347        // Test routing
348        let mut state = State::new();
349        state.insert("next".to_string(), json!("path_a"));
350        let next = graph.get_next_nodes(&["router".to_string()], &state);
351        assert_eq!(next, vec!["path_a".to_string()]);
352
353        state.insert("next".to_string(), json!("path_b"));
354        let next = graph.get_next_nodes(&["router".to_string()], &state);
355        assert_eq!(next, vec!["path_b".to_string()]);
356    }
357}
adk_graph/graph.rs

adk_graph/
graph.rs
