lc_langgraph/compiled/
graph.rs1use super::types::{DynamicTask, GraphExecution};
5use crate::checkpointer::Checkpointer;
6use crate::edge::{ConditionalEdge, GraphEdge};
7use crate::errors::{GraphError, GraphResult};
8use crate::node::GraphNode;
9use crate::state::{Reducer, StateSchema};
10use crate::subgraph::SubgraphNode;
11use std::collections::HashMap;
12use std::sync::Arc;
13use tokio::sync::Mutex;
14use tokio::sync::RwLock;
15
16#[allow(clippy::type_complexity)]
24pub struct CompiledGraph<S: StateSchema> {
25 pub(super) nodes: HashMap<String, Arc<dyn GraphNode<S>>>,
26 pub(super) edges: Vec<GraphEdge>,
27 pub(super) entry_point: String,
28 pub(super) default_reducer: Arc<dyn Reducer<S>>,
29 pub(super) conditional_routers: HashMap<String, Arc<dyn ConditionalEdge<S>>>,
30 pub(super) checkpointer: Option<Arc<Mutex<dyn Checkpointer<S> + Send>>>,
31 pub(super) recursion_limit: usize,
32 pub(super) interrupt_before: Vec<String>,
33 pub(super) interrupt_after: Vec<String>,
34
35 pub(super) runtime_nodes: Arc<RwLock<HashMap<String, Arc<dyn GraphNode<S>>>>>,
36 pub(super) runtime_edges: Arc<RwLock<Vec<GraphEdge>>>,
37 pub(super) runtime_conditional_routers:
38 Arc<RwLock<HashMap<String, Arc<dyn ConditionalEdge<S>>>>>,
39 pub(super) task_inbox: Arc<Mutex<Vec<DynamicTask>>>,
40}
41
42impl<S: StateSchema> std::fmt::Debug for CompiledGraph<S> {
43 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44 f.debug_struct("CompiledGraph")
45 .field("nodes", &self.nodes.keys().collect::<Vec<_>>())
46 .field("edges", &self.edges)
47 .field("entry_point", &self.entry_point)
48 .field("recursion_limit", &self.recursion_limit)
49 .field("interrupt_before", &self.interrupt_before)
50 .field("interrupt_after", &self.interrupt_after)
51 .finish()
52 }
53}
54
55impl<S: StateSchema> CompiledGraph<S> {
56 pub(crate) fn new(
57 nodes: HashMap<String, Arc<dyn GraphNode<S>>>,
58 edges: Vec<GraphEdge>,
59 entry_point: String,
60 default_reducer: Arc<dyn Reducer<S>>,
61 ) -> Self {
62 Self {
63 nodes,
64 edges,
65 entry_point,
66 default_reducer,
67 conditional_routers: HashMap::new(),
68 checkpointer: None,
69 recursion_limit: 25,
70 interrupt_before: Vec::new(),
71 interrupt_after: Vec::new(),
72 runtime_nodes: Arc::new(RwLock::new(HashMap::new())),
73 runtime_edges: Arc::new(RwLock::new(Vec::new())),
74 runtime_conditional_routers: Arc::new(RwLock::new(HashMap::new())),
75 task_inbox: Arc::new(Mutex::new(Vec::new())),
76 }
77 }
78
79 pub(crate) fn add_router(&mut self, name: String, router: Arc<dyn ConditionalEdge<S>>) {
80 self.conditional_routers.insert(name, router);
81 }
82
83 pub fn with_checkpointer<C: Checkpointer<S> + 'static>(mut self, checkpointer: C) -> Self {
84 self.checkpointer = Some(Arc::new(Mutex::new(checkpointer)));
85 self
86 }
87
88 pub fn with_recursion_limit(mut self, limit: usize) -> Self {
89 self.recursion_limit = limit;
90 self
91 }
92
93 pub fn with_interrupt_before(mut self, nodes: Vec<String>) -> Self {
94 self.interrupt_before = nodes;
95 self
96 }
97
98 pub fn with_interrupt_after(mut self, nodes: Vec<String>) -> Self {
99 self.interrupt_after = nodes;
100 self
101 }
102
103 pub fn node_names(&self) -> Vec<String> {
104 self.nodes.keys().cloned().collect()
105 }
106
107 pub fn get_edges(&self) -> &[GraphEdge] {
108 &self.edges
109 }
110
111 pub fn entry_point(&self) -> &str {
112 &self.entry_point
113 }
114
115 pub fn recursion_limit(&self) -> usize {
116 self.recursion_limit
117 }
118
119 pub fn interrupt_before(&self) -> &[String] {
120 &self.interrupt_before
121 }
122
123 pub fn interrupt_after(&self) -> &[String] {
124 &self.interrupt_after
125 }
126
127 pub async fn last_checkpoint_state(&self) -> Option<S> {
129 if let Some(ref cp) = self.checkpointer {
130 let guard = cp.lock().await;
131 let ids = guard.list().await.ok()?;
132 let last_id = ids.last()?.clone();
133 guard.load(&last_id).await.ok()
134 } else {
135 None
136 }
137 }
138
139 pub async fn create_resume_execution(
142 &self,
143 interrupted_node: &str,
144 ) -> Option<GraphExecution<S>> {
145 let state = self.last_checkpoint_state().await?;
146 let current = interrupted_node
147 .strip_prefix("after_")
148 .unwrap_or(interrupted_node);
149 Some(GraphExecution {
150 state,
151 current_node: current.to_string(),
152 steps: Vec::new(),
153 recursion_count: 0,
154 interrupted_at: interrupted_node.to_string(),
155 })
156 }
157
158 pub(super) async fn get_node(&self, name: &str) -> GraphResult<Arc<dyn GraphNode<S>>> {
159 if let Some(node) = self.nodes.get(name) {
160 return Ok(node.clone());
161 }
162 if let Some(node) = self.runtime_nodes.read().await.get(name) {
163 return Ok(node.clone());
164 }
165 Err(GraphError::ExecutionError(format!(
166 "Node '{}' not found",
167 name
168 )))
169 }
170
171 pub fn submit_task(&self, description: String) {
173 if let Ok(mut inbox) = self.task_inbox.try_lock() {
174 inbox.push(DynamicTask {
175 id: uuid::Uuid::new_v4().to_string(),
176 description,
177 });
178 }
179 }
180
181 pub async fn inject_node(&self, name: &str, node: Arc<dyn GraphNode<S>>) -> GraphResult<()> {
183 self.runtime_nodes
184 .write()
185 .await
186 .insert(name.to_string(), node);
187 Ok(())
188 }
189
190 pub async fn inject_edge(&self, source: &str, target: &str) -> GraphResult<()> {
192 self.runtime_edges
193 .write()
194 .await
195 .push(GraphEdge::fixed(source, target));
196 Ok(())
197 }
198
199 pub async fn inject_subgraph<SubS: StateSchema + 'static>(
201 &self,
202 name: &str,
203 subgraph: CompiledGraph<SubS>,
204 input_mapper: impl Fn(&S) -> SubS + Send + Sync + 'static,
205 output_mapper: impl Fn(&SubS, &mut S) + Send + Sync + 'static,
206 ) -> GraphResult<()> {
207 let node = SubgraphNode::new(name, subgraph, input_mapper, output_mapper);
208 self.runtime_nodes
209 .write()
210 .await
211 .insert(name.to_string(), Arc::new(node));
212 Ok(())
213 }
214}