lc_langgraph/
graph.rs

1// crates/lc-langgraph/src/graph.rs
2//! StateGraph - Main graph class for LangGraph
3
4use crate::compiled::CompiledGraph;
5use crate::edge::{ConditionalEdge, GraphEdge};
6use crate::errors::{GraphError, GraphResult};
7use crate::node::{AsyncFn, AsyncNode, GraphNode, SyncNode};
8use crate::state::{Reducer, ReplaceReducer, StateSchema, StateUpdate};
9use crate::subgraph::SubgraphNode;
10use std::collections::HashMap;
11use std::sync::Arc;
12
13/// START sentinel node identifier
14pub const START: &str = "__start__";
15/// END sentinel node identifier
16pub const END: &str = "__end__";
17
18/// StateGraph - Main graph builder
19///
20/// Manages nodes, edges, and state schema for graph execution.
21/// After building, compile to get an executable CompiledGraph.
22pub struct StateGraph<S: StateSchema> {
23    nodes: HashMap<String, Arc<dyn GraphNode<S>>>,
24    edges: Vec<GraphEdge>,
25    entry_point: Option<String>,
26    reducers: HashMap<String, Arc<dyn Reducer<S>>>,
27    default_reducer: Arc<dyn Reducer<S>>,
28    conditional_routers: HashMap<String, Arc<dyn ConditionalEdge<S>>>,
29}
30
31impl<S: StateSchema + 'static> StateGraph<S> {
32    pub fn new() -> Self {
33        Self {
34            nodes: HashMap::new(),
35            edges: Vec::new(),
36            entry_point: None,
37            reducers: HashMap::new(),
38            default_reducer: Arc::new(ReplaceReducer),
39            conditional_routers: HashMap::new(),
40        }
41    }
42
43    pub fn add_node<N: GraphNode<S> + 'static>(&mut self, node: N) -> &mut Self {
44        let name = node.name().to_string();
45        self.nodes.insert(name, Arc::new(node));
46        self
47    }
48
49    pub fn add_node_fn<F>(&mut self, name: impl Into<String>, func: F) -> &mut Self
50    where
51        F: Fn(&S) -> Result<StateUpdate<S>, GraphError> + Send + Sync + 'static,
52    {
53        let node = SyncNode::new(name, func);
54        let node_name = node.name().to_string();
55        self.nodes.insert(node_name, Arc::new(node));
56        self
57    }
58
59    pub fn add_async_node<F>(&mut self, name: impl Into<String>, func: F) -> &mut Self
60    where
61        F: AsyncFn<S> + 'static,
62    {
63        let node = AsyncNode::new(name, func);
64        let node_name = node.name().to_string();
65        self.nodes.insert(node_name, Arc::new(node));
66        self
67    }
68
69    pub fn add_subgraph<SubS: StateSchema + 'static>(
70        &mut self,
71        name: impl Into<String>,
72        subgraph: CompiledGraph<SubS>,
73        input_mapper: impl Fn(&S) -> SubS + Send + Sync + 'static,
74        output_mapper: impl Fn(&SubS, &mut S) + Send + Sync + 'static,
75    ) -> &mut Self {
76        let node = SubgraphNode::new(name, subgraph, input_mapper, output_mapper);
77        let node_name = node.name().to_string();
78        self.nodes.insert(node_name, Arc::new(node));
79        self
80    }
81
82    pub fn add_subgraph_same_state(
83        &mut self,
84        name: impl Into<String>,
85        subgraph: CompiledGraph<S>,
86    ) -> &mut Self {
87        let node: SubgraphNode<S, S> = SubgraphNode::same_state(name, subgraph);
88        let node_name = node.name().to_string();
89        self.nodes.insert(node_name, Arc::new(node));
90        self
91    }
92
93    pub fn add_edge(&mut self, source: impl Into<String>, target: impl Into<String>) -> &mut Self {
94        let edge = GraphEdge::fixed(source, target);
95        self.edges.push(edge);
96        self
97    }
98
99    pub fn add_conditional_edges(
100        &mut self,
101        source: impl Into<String>,
102        router_name: impl Into<String>,
103        targets: HashMap<String, String>,
104        default: Option<String>,
105    ) -> &mut Self {
106        let edge = GraphEdge::conditional(source, router_name, targets, default);
107        self.edges.push(edge);
108        self
109    }
110
111    pub fn add_fan_out(&mut self, source: impl Into<String>, targets: Vec<String>) -> &mut Self {
112        let edge = GraphEdge::fan_out(source, targets);
113        self.edges.push(edge);
114        self
115    }
116
117    pub fn add_fan_in(&mut self, sources: Vec<String>, target: impl Into<String>) -> &mut Self {
118        let edge = GraphEdge::fan_in(sources, target);
119        self.edges.push(edge);
120        self
121    }
122
123    pub fn set_conditional_router<R: ConditionalEdge<S> + 'static>(
124        &mut self,
125        name: impl Into<String>,
126        router: R,
127    ) -> &mut Self {
128        self.conditional_routers
129            .insert(name.into(), Arc::new(router));
130        self
131    }
132
133    pub fn set_entry_point(&mut self, node: impl Into<String>) -> &mut Self {
134        self.entry_point = Some(node.into());
135        self
136    }
137
138    pub fn set_reducer(
139        &mut self,
140        field: impl Into<String>,
141        reducer: Arc<dyn Reducer<S>>,
142    ) -> &mut Self {
143        self.reducers.insert(field.into(), reducer);
144        self
145    }
146
147    pub fn compile(&self) -> GraphResult<CompiledGraph<S>> {
148        if self.nodes.is_empty() {
149            return Err(GraphError::ValidationError(
150                "Graph has no nodes".to_string(),
151            ));
152        }
153
154        let entry = self
155            .entry_point
156            .clone()
157            .or_else(|| self.find_first_node_after_start())
158            .ok_or_else(|| GraphError::ValidationError("No entry point defined".to_string()))?;
159
160        if !self.nodes.contains_key(&entry) && entry != START {
161            return Err(GraphError::ValidationError(format!(
162                "Entry point '{}' not found",
163                entry
164            )));
165        }
166
167        let mut compiled = CompiledGraph::new(
168            self.nodes.clone(),
169            self.edges.clone(),
170            entry,
171            self.default_reducer.clone(),
172        );
173
174        for (name, router) in &self.conditional_routers {
175            compiled.add_router(name.clone(), router.clone());
176        }
177
178        compiled.validate()?;
179        Ok(compiled)
180    }
181
182    fn find_first_node_after_start(&self) -> Option<String> {
183        for edge in &self.edges {
184            if edge.source() == START {
185                if let Some(target) = edge.fixed_target() {
186                    return Some(target.to_string());
187                }
188            }
189        }
190        None
191    }
192}
193
194impl<S: StateSchema + 'static> Default for StateGraph<S> {
195    fn default() -> Self {
196        Self::new()
197    }
198}
199
200/// GraphBuilder - Fluent builder pattern for StateGraph
201pub struct GraphBuilder<S: StateSchema> {
202    graph: StateGraph<S>,
203}
204
205impl<S: StateSchema + 'static> GraphBuilder<S> {
206    pub fn new() -> Self {
207        Self {
208            graph: StateGraph::new(),
209        }
210    }
211
212    pub fn add_node<N: GraphNode<S> + 'static>(mut self, node: N) -> Self {
213        self.graph.add_node(node);
214        self
215    }
216
217    pub fn add_node_fn<F>(mut self, name: impl Into<String>, func: F) -> Self
218    where
219        F: Fn(&S) -> Result<StateUpdate<S>, GraphError> + Send + Sync + 'static,
220    {
221        self.graph.add_node_fn(name, func);
222        self
223    }
224
225    pub fn add_async_node<F>(mut self, name: impl Into<String>, func: F) -> Self
226    where
227        F: AsyncFn<S> + 'static,
228    {
229        self.graph.add_async_node(name, func);
230        self
231    }
232
233    pub fn add_subgraph<SubS: StateSchema + 'static>(
234        mut self,
235        name: impl Into<String>,
236        subgraph: CompiledGraph<SubS>,
237        input_mapper: impl Fn(&S) -> SubS + Send + Sync + 'static,
238        output_mapper: impl Fn(&SubS, &mut S) + Send + Sync + 'static,
239    ) -> Self {
240        self.graph
241            .add_subgraph(name, subgraph, input_mapper, output_mapper);
242        self
243    }
244
245    pub fn add_subgraph_same_state(
246        mut self,
247        name: impl Into<String>,
248        subgraph: CompiledGraph<S>,
249    ) -> Self {
250        self.graph.add_subgraph_same_state(name, subgraph);
251        self
252    }
253
254    pub fn add_edge(mut self, source: impl Into<String>, target: impl Into<String>) -> Self {
255        self.graph.add_edge(source, target);
256        self
257    }
258
259    pub fn add_conditional_edges(
260        mut self,
261        source: impl Into<String>,
262        router_name: impl Into<String>,
263        targets: HashMap<String, String>,
264        default: Option<String>,
265    ) -> Self {
266        self.graph
267            .add_conditional_edges(source, router_name, targets, default);
268        self
269    }
270
271    pub fn add_fan_out(mut self, source: impl Into<String>, targets: Vec<String>) -> Self {
272        self.graph.add_fan_out(source, targets);
273        self
274    }
275
276    pub fn add_fan_in(mut self, sources: Vec<String>, target: impl Into<String>) -> Self {
277        self.graph.add_fan_in(sources, target);
278        self
279    }
280
281    pub fn set_entry_point(mut self, node: impl Into<String>) -> Self {
282        self.graph.set_entry_point(node);
283        self
284    }
285
286    pub fn compile(self) -> GraphResult<CompiledGraph<S>> {
287        self.graph.compile()
288    }
289
290    pub fn build(self) -> StateGraph<S> {
291        self.graph
292    }
293}
294
295impl<S: StateSchema + 'static> Default for GraphBuilder<S> {
296    fn default() -> Self {
297        Self::new()
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304    use crate::state::AgentState;
305
306    #[test]
307    fn test_graph_creation() {
308        let graph: StateGraph<AgentState> = StateGraph::new();
309        assert!(graph.nodes.is_empty());
310        assert!(graph.edges.is_empty());
311    }
312
313    #[test]
314    fn test_add_node_fn() {
315        let mut graph: StateGraph<AgentState> = StateGraph::new();
316        graph.add_node_fn("test_node", |state: &AgentState| {
317            Ok(StateUpdate::full(state.clone()))
318        });
319        assert_eq!(graph.nodes.len(), 1);
320    }
321
322    #[test]
323    fn test_add_edge() {
324        let mut graph: StateGraph<AgentState> = StateGraph::new();
325        graph.add_edge(START, "node1");
326        graph.add_edge("node1", END);
327        assert_eq!(graph.edges.len(), 2);
328    }
329
330    #[test]
331    fn test_compile_empty_graph() {
332        let graph: StateGraph<AgentState> = StateGraph::new();
333        let result = graph.compile();
334        assert!(result.is_err());
335    }
336
337    #[test]
338    fn test_builder_pattern() {
339        let compiled = GraphBuilder::<AgentState>::new()
340            .add_node_fn("process", |state| Ok(StateUpdate::full(state.clone())))
341            .add_edge(START, "process")
342            .add_edge("process", END)
343            .compile();
344        assert!(compiled.is_ok());
345    }
346}
lc_langgraph/graph.rs

lc_langgraph/
graph.rs
