Skip to main content

synaptic_graph/
builder.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::Arc;
3
4use synaptic_core::SynapseError;
5
6use crate::command::GraphContext;
7use crate::compiled::CompiledGraph;
8use crate::edge::{ConditionalEdge, Edge};
9use crate::node::Node;
10use crate::state::State;
11use crate::{END, START};
12
13/// Builder for constructing a state graph.
14pub struct StateGraph<S: State> {
15    nodes: HashMap<String, Box<dyn Node<S>>>,
16    edges: Vec<Edge>,
17    conditional_edges: Vec<ConditionalEdge<S>>,
18    entry_point: Option<String>,
19    interrupt_before: HashSet<String>,
20    interrupt_after: HashSet<String>,
21}
22
23impl<S: State> StateGraph<S> {
24    pub fn new() -> Self {
25        Self {
26            nodes: HashMap::new(),
27            edges: Vec::new(),
28            conditional_edges: Vec::new(),
29            entry_point: None,
30            interrupt_before: HashSet::new(),
31            interrupt_after: HashSet::new(),
32        }
33    }
34
35    /// Add a named node to the graph.
36    pub fn add_node(mut self, name: impl Into<String>, node: impl Node<S> + 'static) -> Self {
37        self.nodes.insert(name.into(), Box::new(node));
38        self
39    }
40
41    /// Add a fixed edge from source to target.
42    pub fn add_edge(mut self, source: impl Into<String>, target: impl Into<String>) -> Self {
43        self.edges.push(Edge {
44            source: source.into(),
45            target: target.into(),
46        });
47        self
48    }
49
50    /// Add a conditional edge with a routing function.
51    pub fn add_conditional_edges(
52        mut self,
53        source: impl Into<String>,
54        router: impl Fn(&S) -> String + Send + Sync + 'static,
55    ) -> Self {
56        self.conditional_edges.push(ConditionalEdge {
57            source: source.into(),
58            router: Arc::new(router),
59            path_map: None,
60        });
61        self
62    }
63
64    /// Add a conditional edge with a routing function and a path map for visualization.
65    ///
66    /// The `path_map` maps labels to target node names, enabling graph visualization
67    /// tools to show possible routing targets for conditional edges.
68    pub fn add_conditional_edges_with_path_map(
69        mut self,
70        source: impl Into<String>,
71        router: impl Fn(&S) -> String + Send + Sync + 'static,
72        path_map: HashMap<String, String>,
73    ) -> Self {
74        self.conditional_edges.push(ConditionalEdge {
75            source: source.into(),
76            router: Arc::new(router),
77            path_map: Some(path_map),
78        });
79        self
80    }
81
82    /// Set the entry point node for graph execution.
83    pub fn set_entry_point(mut self, name: impl Into<String>) -> Self {
84        self.entry_point = Some(name.into());
85        self
86    }
87
88    /// Mark nodes that should interrupt BEFORE execution (human-in-the-loop).
89    pub fn interrupt_before(mut self, nodes: Vec<String>) -> Self {
90        self.interrupt_before.extend(nodes);
91        self
92    }
93
94    /// Mark nodes that should interrupt AFTER execution (human-in-the-loop).
95    pub fn interrupt_after(mut self, nodes: Vec<String>) -> Self {
96        self.interrupt_after.extend(nodes);
97        self
98    }
99
100    /// Compile the graph into an executable CompiledGraph.
101    pub fn compile(self) -> Result<CompiledGraph<S>, SynapseError> {
102        let entry = self
103            .entry_point
104            .ok_or_else(|| SynapseError::Graph("no entry point set".to_string()))?;
105
106        if !self.nodes.contains_key(&entry) {
107            return Err(SynapseError::Graph(format!(
108                "entry point node '{entry}' not found"
109            )));
110        }
111
112        // Validate: every edge references existing nodes or END
113        for edge in &self.edges {
114            if edge.source != START && !self.nodes.contains_key(&edge.source) {
115                return Err(SynapseError::Graph(format!(
116                    "edge source '{}' not found",
117                    edge.source
118                )));
119            }
120            if edge.target != END && !self.nodes.contains_key(&edge.target) {
121                return Err(SynapseError::Graph(format!(
122                    "edge target '{}' not found",
123                    edge.target
124                )));
125            }
126        }
127
128        for ce in &self.conditional_edges {
129            if ce.source != START && !self.nodes.contains_key(&ce.source) {
130                return Err(SynapseError::Graph(format!(
131                    "conditional edge source '{}' not found",
132                    ce.source
133                )));
134            }
135            // Validate path_map targets reference existing nodes or END
136            if let Some(ref path_map) = ce.path_map {
137                for (label, target) in path_map {
138                    if target != END && !self.nodes.contains_key(target) {
139                        return Err(SynapseError::Graph(format!(
140                            "conditional edge path_map target '{target}' (label '{label}') not found"
141                        )));
142                    }
143                }
144            }
145        }
146
147        Ok(CompiledGraph {
148            nodes: self.nodes,
149            edges: self.edges,
150            conditional_edges: self.conditional_edges,
151            entry_point: entry,
152            interrupt_before: self.interrupt_before,
153            interrupt_after: self.interrupt_after,
154            checkpointer: None,
155            command_context: GraphContext::new(),
156        })
157    }
158}
159
160impl<S: State> Default for StateGraph<S> {
161    fn default() -> Self {
162        Self::new()
163    }
164}