synaptic_graph/
builder.rs1use 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
13pub 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 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 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 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 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 pub fn set_entry_point(mut self, name: impl Into<String>) -> Self {
84 self.entry_point = Some(name.into());
85 self
86 }
87
88 pub fn interrupt_before(mut self, nodes: Vec<String>) -> Self {
90 self.interrupt_before.extend(nodes);
91 self
92 }
93
94 pub fn interrupt_after(mut self, nodes: Vec<String>) -> Self {
96 self.interrupt_after.extend(nodes);
97 self
98 }
99
100 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 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 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}