ri_agent_graph/
builder.rs

1#![allow(deprecated)] // Constructs GraphEvent with legacy trace_id/attempt fields during migration
2
3use crate::checkpoint_store::CheckpointStore;
4use crate::checkpointer::CheckpointSaver;
5use crate::edge::EdgeType;
6use crate::error::{AgentGraphError, Result};
7use crate::event_sink::EventSink;
8use crate::executor::Executor;
9use crate::graph::AgentGraph;
10use crate::graph::{END, START};
11use crate::interrupt::InterruptConfig;
12use crate::node::Node;
13use crate::reducer::Reducer;
14use crate::retry::RetryPolicy;
15use crate::router::RoutingFunction;
16use std::collections::HashMap;
17use std::sync::Arc;
18
19/// Builder for AgentGraph.
20pub struct AgentGraphBuilder {
21    pub(crate) nodes: HashMap<String, Arc<dyn Node>>,
22    pub(crate) edges: HashMap<String, Vec<EdgeType>>,
23    pub(crate) max_iterations: usize,
24    pub(crate) enable_cycle_detection: bool,
25    pub(crate) retry_policies: HashMap<String, RetryPolicy>,
26    pub(crate) interrupt_config: Option<InterruptConfig>,
27    pub(crate) checkpointer: Option<Arc<dyn CheckpointSaver>>,
28    pub(crate) checkpoint_store: Option<Arc<dyn CheckpointStore>>,
29    pub(crate) reducers: Vec<(String, Arc<dyn Reducer>)>,
30    pub(crate) graph_name: Option<String>,
31    pub(crate) event_sink: Option<Arc<dyn EventSink>>,
32    pub(crate) executor: Option<Arc<dyn Executor>>,
33}
34
35impl AgentGraphBuilder {
36    pub(crate) fn new() -> Self {
37        Self {
38            nodes: HashMap::new(),
39            edges: HashMap::new(),
40            max_iterations: 100,
41            enable_cycle_detection: true,
42            retry_policies: HashMap::new(),
43            interrupt_config: None,
44            checkpointer: None,
45            checkpoint_store: None,
46            reducers: Vec::new(),
47            graph_name: None,
48            event_sink: None,
49            executor: None,
50        }
51    }
52
53    /// Set the graph name (used in streaming events and debugging).
54    pub fn with_name(mut self, name: impl Into<String>) -> Self {
55        self.graph_name = Some(name.into());
56        self
57    }
58
59    /// Add a node to the graph.
60    pub fn add_node(mut self, name: impl Into<String>, node: Box<dyn Node>) -> Self {
61        self.nodes.insert(name.into(), Arc::from(node));
62        self
63    }
64
65    /// Add a node with a retry policy.
66    pub fn add_node_with_retry(
67        mut self,
68        name: impl Into<String>,
69        node: Box<dyn Node>,
70        retry: RetryPolicy,
71    ) -> Self {
72        let name = name.into();
73        self.nodes.insert(name.clone(), Arc::from(node));
74        self.retry_policies.insert(name, retry);
75        self
76    }
77
78    /// Add a subgraph as a node.
79    pub fn add_subgraph(mut self, name: impl Into<String>, subgraph: AgentGraph) -> Self {
80        self.nodes.insert(name.into(), Arc::new(subgraph));
81        self
82    }
83
84    /// Add a normal edge (always goes to next node).
85    /// Multiple edges from the same node create fan-out (parallel execution).
86    pub fn add_edge(mut self, from: impl Into<String>, to: impl Into<String>) -> Self {
87        let from = from.into();
88        let to = to.into();
89        self.edges
90            .entry(from)
91            .or_default()
92            .push(EdgeType::Normal(to));
93        self
94    }
95
96    /// Add a conditional edge (uses router to determine next node).
97    pub fn add_conditional_edge(
98        mut self,
99        from: impl Into<String>,
100        router: Box<dyn RoutingFunction>,
101    ) -> Self {
102        let from = from.into();
103        self.edges
104            .entry(from)
105            .or_default()
106            .push(EdgeType::Conditional(router));
107        self
108    }
109
110    /// Set the entry point (sugar for add_edge(START, node)).
111    pub fn set_entry_point(self, node: impl Into<String>) -> Self {
112        self.add_edge(START, node)
113    }
114
115    /// Set the finish point (sugar for add_edge(node, END)).
116    pub fn set_finish_point(self, node: impl Into<String>) -> Self {
117        self.add_edge(node, END)
118    }
119
120    /// Set maximum iterations before stopping (prevents infinite loops).
121    pub fn with_max_iterations(mut self, max: usize) -> Self {
122        self.max_iterations = max;
123        self
124    }
125
126    /// Enable or disable cycle detection.
127    pub fn with_cycle_detection(mut self, enable: bool) -> Self {
128        self.enable_cycle_detection = enable;
129        self
130    }
131
132    /// Register a state reducer for a key.
133    pub fn with_reducer(mut self, key: impl Into<String>, reducer: impl Reducer + 'static) -> Self {
134        self.reducers.push((key.into(), Arc::new(reducer)));
135        self
136    }
137
138    /// Set interrupt-before configuration.
139    pub fn with_interrupt_before(mut self, nodes: Vec<String>) -> Self {
140        let cfg = self
141            .interrupt_config
142            .get_or_insert_with(InterruptConfig::new);
143        cfg.interrupt_before.extend(nodes);
144        self
145    }
146
147    /// Set interrupt-after configuration.
148    pub fn with_interrupt_after(mut self, nodes: Vec<String>) -> Self {
149        let cfg = self
150            .interrupt_config
151            .get_or_insert_with(InterruptConfig::new);
152        cfg.interrupt_after.extend(nodes);
153        self
154    }
155
156    /// Set the legacy checkpointer for persistence (superstep-level).
157    pub fn with_checkpointer(mut self, checkpointer: impl CheckpointSaver + 'static) -> Self {
158        self.checkpointer = Some(Arc::new(checkpointer));
159        self
160    }
161
162    /// Set the granular checkpoint store (per-attempt recording).
163    pub fn with_checkpoint_store(mut self, store: Arc<dyn CheckpointStore>) -> Self {
164        self.checkpoint_store = Some(store);
165        self
166    }
167
168    /// Set a custom event sink for structured event handling.
169    pub fn with_event_sink(mut self, sink: Arc<dyn EventSink>) -> Self {
170        self.event_sink = Some(sink);
171        self
172    }
173
174    /// Set a custom executor for node execution.
175    pub fn with_executor(mut self, executor: Arc<dyn Executor>) -> Self {
176        self.executor = Some(executor);
177        self
178    }
179
180    /// Build the graph.
181    pub fn build(self) -> Result<AgentGraph> {
182        for (from, edge_list) in &self.edges {
183            for edge in edge_list {
184                if let EdgeType::Normal(to) = edge {
185                    if to != END && !self.nodes.contains_key(to) && to != START {
186                        return Err(AgentGraphError::NodeNotFound(format!(
187                            "Edge from '{}' points to non-existent node '{}'",
188                            from, to
189                        )));
190                    }
191                }
192            }
193        }
194
195        Ok(AgentGraph {
196            nodes: self.nodes,
197            edges: self.edges,
198            max_iterations: self.max_iterations,
199            enable_cycle_detection: self.enable_cycle_detection,
200            retry_policies: self.retry_policies,
201            interrupt_config: self.interrupt_config,
202            checkpointer: self.checkpointer,
203            checkpoint_store: self.checkpoint_store,
204            reducers: self.reducers,
205            graph_name: self.graph_name,
206            event_sink: self.event_sink,
207            executor: self.executor,
208        })
209    }
210}
ri_agent_graph/builder.rs

ri_agent_graph/
builder.rs
