lc_langgraph/
persistence.rs

1// crates/lc-langgraph/src/persistence.rs
2//! Graph persistence for serialization and storage
3//!
4//! This module provides persistence capabilities for graph definitions,
5//! allowing graphs to be saved, loaded, and shared across sessions.
6
7use async_trait::async_trait;
8use chrono::{DateTime, Utc};
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11use std::path::PathBuf;
12use tokio::sync::Mutex;
13use uuid::Uuid;
14
15/// GraphPersistence trait for storing and loading graph definitions
16#[async_trait]
17pub trait GraphPersistence: Send + Sync {
18    /// Save a graph definition with the given ID
19    async fn save(&self, id: &str, definition: &GraphDefinition) -> Result<(), PersistenceError>;
20
21    /// Load a graph definition by ID
22    async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError>;
23    /// Delete a graph definition by ID
24    async fn delete(&self, id: &str) -> Result<(), PersistenceError>;
25
26    /// Check if a graph definition exists
27    async fn exists(&self, id: &str) -> Result<bool, PersistenceError>;
28
29    /// List all stored graph IDs
30    async fn list(&self) -> Result<Vec<String>, PersistenceError>;
31}
32
33/// Persistence error types
34#[derive(Debug, thiserror::Error)]
35pub enum PersistenceError {
36    #[error("Graph '{0}' not found")]
37    NotFound(String),
38
39    #[error("Serialization error: {0}")]
40    SerializationError(String),
41
42    #[error("Deserialization error: {0}")]
43    DeserializationError(String),
44
45    #[error("IO error: {0}")]
46    IoError(String),
47
48    #[error("Invalid graph definition: {0}")]
49    InvalidDefinition(String),
50
51    #[error("MongoDB error: {0}")]
52    MongoError(String),
53
54    #[error("Connection error: {0}")]
55    ConnectionError(String),
56}
57
58/// GraphDefinition - Serializable graph structure
59#[derive(Debug, Clone, Serialize, Deserialize)]
60pub struct GraphDefinition {
61    /// Unique identifier
62    pub id: String,
63
64    /// Human-readable name
65    pub name: Option<String>,
66
67    /// Entry point node name
68    pub entry_point: String,
69
70    /// Node definitions
71    pub nodes: Vec<NodeDefinition>,
72
73    /// Edge definitions
74    pub edges: Vec<EdgeDefinition>,
75
76    /// Router definitions
77    pub routers: Vec<RouterDefinition>,
78
79    /// Maximum recursion limit
80    pub recursion_limit: usize,
81
82    /// Creation timestamp
83    pub created_at: DateTime<Utc>,
84
85    /// Last update timestamp
86    pub updated_at: DateTime<Utc>,
87
88    /// Custom metadata
89    pub metadata: HashMap<String, serde_json::Value>,
90}
91
92impl GraphDefinition {
93    /// Create a new graph definition with the given entry point
94    pub fn new(entry_point: String) -> Self {
95        let now = Utc::now();
96        Self {
97            id: Uuid::new_v4().to_string(),
98            name: None,
99            entry_point,
100            nodes: Vec::new(),
101            edges: Vec::new(),
102            routers: Vec::new(),
103            recursion_limit: 25,
104            created_at: now,
105            updated_at: now,
106            metadata: HashMap::new(),
107        }
108    }
109
110    /// Set a custom ID
111    pub fn with_id(mut self, id: String) -> Self {
112        self.id = id;
113        self
114    }
115
116    /// Set a human-readable name
117    pub fn with_name(mut self, name: String) -> Self {
118        self.name = Some(name);
119        self
120    }
121
122    /// Set recursion limit
123    pub fn with_recursion_limit(mut self, limit: usize) -> Self {
124        self.recursion_limit = limit;
125        self
126    }
127
128    /// Add a node definition
129    pub fn add_node(&mut self, node: NodeDefinition) {
130        self.nodes.push(node);
131        self.updated_at = Utc::now();
132    }
133
134    /// Add an edge definition
135    pub fn add_edge(&mut self, edge: EdgeDefinition) {
136        self.edges.push(edge);
137        self.updated_at = Utc::now();
138    }
139
140    /// Add a router definition
141    pub fn add_router(&mut self, router: RouterDefinition) {
142        self.routers.push(router);
143        self.updated_at = Utc::now();
144    }
145}
146
147/// NodeDefinition - Serializable node structure
148#[derive(Debug, Clone, Serialize, Deserialize)]
149pub struct NodeDefinition {
150    /// Node name
151    pub name: String,
152
153    /// Node type
154    pub node_type: NodeType,
155
156    /// Custom configuration
157    pub config: serde_json::Value,
158}
159
160/// NodeType - Type of node execution
161#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
162pub enum NodeType {
163    Sync,
164    Async,
165    Subgraph,
166    Custom,
167}
168
169/// EdgeDefinition - Serializable edge structure
170#[derive(Debug, Clone, Serialize, Deserialize)]
171pub struct EdgeDefinition {
172    /// Edge type
173    pub edge_type: EdgeType,
174
175    /// Source node name
176    pub source: String,
177
178    /// Target node name (for fixed edges)
179    pub target: Option<String>,
180
181    /// Multiple targets (for fan-out edges)
182    pub targets: Option<Vec<String>>,
183
184    /// Router name (for conditional edges)
185    pub router_name: Option<String>,
186
187    /// Conditional targets mapping (route -> target)
188    pub conditional_targets: Option<HashMap<String, String>>,
189
190    /// Default target for conditional edges
191    pub default_target: Option<String>,
192
193    /// Source nodes for fan-in edges
194    pub sources: Option<Vec<String>>,
195}
196
197impl EdgeDefinition {
198    /// Create a fixed edge
199    pub fn fixed(source: String, target: String) -> Self {
200        Self {
201            edge_type: EdgeType::Fixed,
202            source,
203            target: Some(target),
204            targets: None,
205            router_name: None,
206            conditional_targets: None,
207            default_target: None,
208            sources: None,
209        }
210    }
211
212    /// Create a conditional edge
213    pub fn conditional(
214        source: String,
215        router_name: String,
216        targets: HashMap<String, String>,
217        default_target: Option<String>,
218    ) -> Self {
219        Self {
220            edge_type: EdgeType::Conditional,
221            source,
222            target: None,
223            targets: None,
224            router_name: Some(router_name),
225            conditional_targets: Some(targets),
226            default_target,
227            sources: None,
228        }
229    }
230
231    /// Create a fan-out edge (parallel execution)
232    pub fn fan_out(source: String, targets: Vec<String>) -> Self {
233        Self {
234            edge_type: EdgeType::FanOut,
235            source,
236            target: None,
237            targets: Some(targets),
238            router_name: None,
239            conditional_targets: None,
240            default_target: None,
241            sources: None,
242        }
243    }
244
245    /// Create a fan-in edge (merge from parallel branches)
246    pub fn fan_in(sources: Vec<String>, target: String) -> Self {
247        Self {
248            edge_type: EdgeType::FanIn,
249            source: "__fan_in__".to_string(),
250            target: Some(target),
251            targets: None,
252            router_name: None,
253            conditional_targets: None,
254            default_target: None,
255            sources: Some(sources),
256        }
257    }
258}
259
260/// EdgeType - Type of edge connection
261#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
262pub enum EdgeType {
263    /// Fixed transition
264    Fixed,
265
266    /// Conditional routing
267    Conditional,
268
269    /// Parallel fan-out
270    FanOut,
271
272    /// Merge fan-in
273    FanIn,
274}
275
276/// RouterDefinition - Serializable router structure
277#[derive(Debug, Clone, Serialize, Deserialize)]
278pub struct RouterDefinition {
279    /// Router name
280    pub name: String,
281
282    /// Router type (e.g., "function", "state_key")
283    pub router_type: String,
284
285    /// Possible routes
286    pub routes: Vec<String>,
287
288    /// Custom configuration
289    pub config: serde_json::Value,
290}
291
292/// MemoryPersistence - In-memory graph storage
293pub struct MemoryPersistence {
294    graphs: Mutex<HashMap<String, GraphDefinition>>,
295}
296
297impl MemoryPersistence {
298    /// Create a new memory persistence instance
299    pub fn new() -> Self {
300        Self {
301            graphs: Mutex::new(HashMap::new()),
302        }
303    }
304}
305
306impl Default for MemoryPersistence {
307    fn default() -> Self {
308        Self::new()
309    }
310}
311
312#[async_trait]
313impl GraphPersistence for MemoryPersistence {
314    async fn save(&self, id: &str, definition: &GraphDefinition) -> Result<(), PersistenceError> {
315        let mut graphs = self.graphs.lock().await;
316        graphs.insert(id.to_string(), definition.clone());
317        Ok(())
318    }
319
320    async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError> {
321        let graphs = self.graphs.lock().await;
322        graphs
323            .get(id)
324            .cloned()
325            .ok_or_else(|| PersistenceError::NotFound(id.to_string()))
326    }
327
328    async fn delete(&self, id: &str) -> Result<(), PersistenceError> {
329        let mut graphs = self.graphs.lock().await;
330        graphs
331            .remove(id)
332            .map(|_| ())
333            .ok_or_else(|| PersistenceError::NotFound(id.to_string()))
334    }
335
336    async fn exists(&self, id: &str) -> Result<bool, PersistenceError> {
337        let graphs = self.graphs.lock().await;
338        Ok(graphs.contains_key(id))
339    }
340
341    async fn list(&self) -> Result<Vec<String>, PersistenceError> {
342        let graphs = self.graphs.lock().await;
343        Ok(graphs.keys().cloned().collect())
344    }
345}
346
347/// FilePersistence - File-based graph storage
348pub struct FilePersistence {
349    directory: PathBuf,
350}
351
352impl FilePersistence {
353    /// Create a new file persistence instance
354    pub fn new(directory: impl Into<PathBuf>) -> Result<Self, PersistenceError> {
355        let dir = directory.into();
356        if !dir.exists() {
357            std::fs::create_dir_all(&dir).map_err(|e| {
358                PersistenceError::IoError(format!(
359                    "Failed to create directory '{}': {}",
360                    dir.display(),
361                    e
362                ))
363            })?;
364        }
365        Ok(Self { directory: dir })
366    }
367
368    fn graph_path(&self, id: &str) -> Result<PathBuf, PersistenceError> {
369        // Sanitize id to prevent path traversal: reject ".." and absolute paths
370        if id.contains("..") || id.contains('/') || id.contains('\\') {
371            return Err(PersistenceError::InvalidDefinition(format!(
372                "Invalid graph id '{}': path traversal detected",
373                id
374            )));
375        }
376        if std::path::Path::new(id).is_absolute() {
377            return Err(PersistenceError::InvalidDefinition(format!(
378                "Invalid graph id '{}': absolute path not allowed",
379                id
380            )));
381        }
382        Ok(self.directory.join(format!("{}.json", id)))
383    }
384}
385
386impl Default for FilePersistence {
387    fn default() -> Self {
388        Self::new(".graph_definitions")
389            .expect("Failed to create default graph definitions directory")
390    }
391}
392
393#[async_trait]
394impl GraphPersistence for FilePersistence {
395    async fn save(&self, id: &str, definition: &GraphDefinition) -> Result<(), PersistenceError> {
396        let path = self.graph_path(id)?;
397
398        let json = serde_json::to_string_pretty(definition)
399            .map_err(|e| PersistenceError::SerializationError(e.to_string()))?;
400
401        tokio::fs::write(&path, json)
402            .await
403            .map_err(|e| PersistenceError::IoError(e.to_string()))?;
404
405        Ok(())
406    }
407
408    async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError> {
409        let path = self.graph_path(id)?;
410
411        if !path.exists() {
412            return Err(PersistenceError::NotFound(id.to_string()));
413        }
414
415        let json = tokio::fs::read_to_string(&path)
416            .await
417            .map_err(|e| PersistenceError::IoError(e.to_string()))?;
418
419        let definition: GraphDefinition = serde_json::from_str(&json)
420            .map_err(|e| PersistenceError::DeserializationError(e.to_string()))?;
421
422        Ok(definition)
423    }
424
425    async fn delete(&self, id: &str) -> Result<(), PersistenceError> {
426        let path = self.graph_path(id)?;
427
428        if path.exists() {
429            tokio::fs::remove_file(&path)
430                .await
431                .map_err(|e| PersistenceError::IoError(e.to_string()))?;
432        }
433
434        Ok(())
435    }
436
437    async fn exists(&self, id: &str) -> Result<bool, PersistenceError> {
438        let path = self.graph_path(id)?;
439        Ok(path.exists())
440    }
441
442    async fn list(&self) -> Result<Vec<String>, PersistenceError> {
443        let mut ids = Vec::new();
444
445        let mut entries = tokio::fs::read_dir(&self.directory)
446            .await
447            .map_err(|e| PersistenceError::IoError(e.to_string()))?;
448
449        while let Some(entry) = entries
450            .next_entry()
451            .await
452            .map_err(|e| PersistenceError::IoError(e.to_string()))?
453        {
454            let path = entry.path();
455            if path.extension().is_some_and(|ext| ext == "json") {
456                if let Some(id) = path.file_stem().and_then(|s| s.to_str()) {
457                    ids.push(id.to_string());
458                }
459            }
460        }
461
462        Ok(ids)
463    }
464}
465
466// ============================================================================
467// MongoDB 持久化实现
468// ============================================================================
469
470#[cfg(feature = "mongodb-persistence")]
471mod mongo_impl {
472    use super::*;
473    use mongodb::{
474        bson::{doc, from_document, to_document, Document},
475        options::{ClientOptions, FindOptions},
476        Client, Collection,
477    };
478
479    /// MongoDB配置
480    pub struct MongoConfig {
481        /// MongoDB连接URI (例如: mongodb://localhost:27017 或 mongodb+srv://user:pass@cluster.mongodb.net)
482        pub uri: String,
483
484        /// 数据库名称
485        pub database: String,
486
487        /// 集合名称
488        pub collection: String,
489    }
490
491    impl MongoConfig {
492        /// 创建新的MongoDB配置
493        pub fn new(uri: String, database: String, collection: String) -> Self {
494            Self {
495                uri,
496                database,
497                collection,
498            }
499        }
500
501        /// 从环境变量创建配置
502        ///
503        /// 环境变量:
504        /// - MONGO_URI: MongoDB连接URI
505        /// - MONGO_DATABASE: 数据库名称 (默认: langgraph)
506        /// - MONGO_COLLECTION: 集合名称 (默认: graph_definitions)
507        pub fn from_env() -> Result<Self, PersistenceError> {
508            Ok(Self {
509                uri: std::env::var("MONGO_URI").map_err(|_| {
510                    PersistenceError::ConnectionError(
511                        "MONGO_URI environment variable not set".to_string(),
512                    )
513                })?,
514                database: std::env::var("MONGO_DATABASE")
515                    .unwrap_or_else(|_| "langgraph".to_string()),
516                collection: std::env::var("MONGO_COLLECTION")
517                    .unwrap_or_else(|_| "graph_definitions".to_string()),
518            })
519        }
520    }
521
522    /// MongoPersistence - MongoDB图存储实现
523    pub struct MongoPersistence {
524        client: Client,
525        collection: Collection<Document>,
526        database_name: String,
527        collection_name: String,
528    }
529
530    impl MongoPersistence {
531        /// 创建新的MongoDB持久化实例
532        ///
533        /// # 参数
534        /// - config: MongoDB配置
535        ///
536        /// # 示例
537        /// ```ignore
538        /// let config = MongoConfig::new(
539        ///     "mongodb://localhost:27017",
540        ///     "langgraph",
541        ///     "graph_definitions"
542        /// );
543        /// let persistence = MongoPersistence::new(config).await?;
544        /// ```
545        pub async fn new(config: MongoConfig) -> Result<Self, PersistenceError> {
546            let client_options = ClientOptions::parse(&config.uri)
547                .await
548                .map_err(|e| PersistenceError::ConnectionError(e.to_string()))?;
549
550            let client = Client::with_options(client_options)
551                .map_err(|e| PersistenceError::ConnectionError(e.to_string()))?;
552
553            let database = client.database(&config.database);
554            let collection = database.collection(&config.collection);
555
556            Ok(Self {
557                client,
558                collection,
559                database_name: config.database,
560                collection_name: config.collection,
561            })
562        }
563
564        /// 从环境变量创建实例
565        pub async fn from_env() -> Result<Self, PersistenceError> {
566            let config = MongoConfig::from_env()?;
567            Self::new(config).await
568        }
569
570        /// 获取MongoDB客户端
571        pub fn client(&self) -> &Client {
572            &self.client
573        }
574
575        /// 获取集合名称
576        pub fn collection_name(&self) -> &str {
577            &self.collection_name
578        }
579
580        /// 获取数据库名称
581        pub fn database_name(&self) -> &str {
582            &self.database_name
583        }
584    }
585
586    #[async_trait]
587    impl GraphPersistence for MongoPersistence {
588        async fn save(
589            &self,
590            id: &str,
591            definition: &GraphDefinition,
592        ) -> Result<(), PersistenceError> {
593            let doc = to_document(definition)
594                .map_err(|e| PersistenceError::SerializationError(e.to_string()))?;
595
596            // 使用 upsert 操作：如果存在则更新，不存在则插入
597            self.collection
598                .update_one(
599                    doc! { "_id": id },
600                    doc! { "$set": doc },
601                    mongodb::options::UpdateOptions::builder()
602                        .upsert(true)
603                        .build(),
604                )
605                .await
606                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
607
608            Ok(())
609        }
610
611        async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError> {
612            let filter = doc! { "_id": id };
613
614            let result = self
615                .collection
616                .find_one(filter, None)
617                .await
618                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
619
620            match result {
621                Some(doc) => {
622                    let definition: GraphDefinition = from_document(doc)
623                        .map_err(|e| PersistenceError::DeserializationError(e.to_string()))?;
624                    Ok(definition)
625                }
626                None => Err(PersistenceError::NotFound(id.to_string())),
627            }
628        }
629
630        async fn delete(&self, id: &str) -> Result<(), PersistenceError> {
631            let filter = doc! { "_id": id };
632
633            let result = self
634                .collection
635                .delete_one(filter, None)
636                .await
637                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
638
639            if result.deleted_count == 0 {
640                Err(PersistenceError::NotFound(id.to_string()))
641            } else {
642                Ok(())
643            }
644        }
645
646        async fn exists(&self, id: &str) -> Result<bool, PersistenceError> {
647            let filter = doc! { "_id": id };
648
649            let count = self
650                .collection
651                .count_documents(filter, None)
652                .await
653                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
654
655            Ok(count > 0)
656        }
657
658        async fn list(&self) -> Result<Vec<String>, PersistenceError> {
659            let filter = doc! {};
660            let options = FindOptions::builder().projection(doc! { "_id": 1 }).build();
661
662            let mut cursor = self
663                .collection
664                .find(filter, options)
665                .await
666                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
667
668            let mut ids = Vec::new();
669            while cursor
670                .advance()
671                .await
672                .map_err(|e| PersistenceError::MongoError(e.to_string()))?
673            {
674                let doc = cursor
675                    .deserialize_current()
676                    .map_err(|e| PersistenceError::DeserializationError(e.to_string()))?;
677
678                if let Ok(id) = doc.get_str("_id") {
679                    ids.push(id.to_string());
680                }
681            }
682
683            Ok(ids)
684        }
685    }
686}
687
688#[cfg(feature = "mongodb-persistence")]
689pub use mongo_impl::{MongoConfig, MongoPersistence};
690
691#[cfg(test)]
692mod tests {
693    use super::*;
694
695    #[test]
696    fn test_node_type_serialization() {
697        let types = vec![
698            NodeType::Sync,
699            NodeType::Async,
700            NodeType::Subgraph,
701            NodeType::Custom,
702        ];
703
704        for t in types {
705            let json = serde_json::to_string(&t).unwrap();
706            let parsed: NodeType = serde_json::from_str(&json).unwrap();
707            assert_eq!(parsed, t);
708        }
709    }
710
711    #[test]
712    fn test_edge_type_serialization() {
713        let types = vec![
714            EdgeType::Fixed,
715            EdgeType::Conditional,
716            EdgeType::FanOut,
717            EdgeType::FanIn,
718        ];
719
720        for t in types {
721            let json = serde_json::to_string(&t).unwrap();
722            let parsed: EdgeType = serde_json::from_str(&json).unwrap();
723            assert_eq!(parsed, t);
724        }
725    }
726
727    #[test]
728    fn test_graph_definition_builder() {
729        let def = GraphDefinition::new("entry".to_string())
730            .with_id("test-id".to_string())
731            .with_name("Test Graph".to_string())
732            .with_recursion_limit(50);
733
734        assert_eq!(def.id, "test-id");
735        assert_eq!(def.name, Some("Test Graph".to_string()));
736        assert_eq!(def.entry_point, "entry");
737        assert_eq!(def.recursion_limit, 50);
738    }
739
740    #[tokio::test]
741    async fn test_memory_persistence() {
742        let persistence = MemoryPersistence::new();
743        let def = GraphDefinition::new("entry".to_string()).with_id("test-001".to_string());
744
745        persistence.save("test-001", &def).await.unwrap();
746        assert!(persistence.exists("test-001").await.unwrap());
747
748        let loaded = persistence.load("test-001").await.unwrap();
749        assert_eq!(loaded.id, "test-001");
750
751        persistence.delete("test-001").await.unwrap();
752        assert!(!persistence.exists("test-001").await.unwrap());
753    }
754}
lc_langgraph/persistence.rs

lc_langgraph/
persistence.rs
