Skip to main content

lc_langgraph/compiled/
graph.rs

1// crates/lc-langgraph/src/compiled/graph.rs
2//! CompiledGraph struct definition, constructors, getters, and runtime injection
3
4use 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/// CompiledGraph - Ready-to-execute graph
17///
18/// Created from StateGraph.compile(). Handles:
19/// - State management and updates
20/// - Edge routing (fixed and conditional)
21/// - Execution with recursion limits
22/// - Checkpointing for persistence
23#[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    /// Get the last checkpoint state (for interrupt recovery)
128    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    /// Create a resume execution context (from the last checkpoint)
140    /// interrupted_node may be "node_name" or "after_node_name"
141    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    /// Submit a new task for dynamic planning mid-execution
172    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    /// Inject a runtime node directly
182    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    /// Inject a runtime edge directly
191    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    /// Inject a subgraph as a runtime node
200    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}