agent_chain_core/load/
loader.rs

1//! Load LangChain objects from JSON strings or objects.
2//!
3//! This module provides functions for deserializing LangChain objects from JSON,
4//! mirroring `langchain_core.load.load`.
5//!
6//! # Warning
7//!
8//! The `load` and `loads` functions can instantiate arbitrary types based on the
9//! serialized data. Be careful when using with untrusted input.
10
11use serde_json::Value;
12use std::collections::HashMap;
13use std::env;
14
15use super::mapping::{DEFAULT_NAMESPACES, DISALLOW_LOAD_FROM_PATH, get_all_serializable_mappings};
16use super::serializable::LC_VERSION;
17use crate::error::{Error, Result};
18
19/// Configuration for the Reviver.
20#[derive(Debug, Clone)]
21pub struct ReviverConfig {
22    /// A map of secrets to load. If a secret is not found in the map,
23    /// it will be loaded from the environment if `secrets_from_env` is `true`.
24    pub secrets_map: HashMap<String, String>,
25    /// A list of additional namespaces (modules) to allow to be deserialized.
26    pub valid_namespaces: Vec<String>,
27    /// Whether to load secrets from the environment.
28    pub secrets_from_env: bool,
29    /// Additional import mappings to override or extend the default mappings.
30    pub additional_import_mappings: HashMap<Vec<String>, Vec<String>>,
31    /// Whether to ignore unserializable fields (return None instead of error).
32    pub ignore_unserializable_fields: bool,
33}
34
35impl Default for ReviverConfig {
36    fn default() -> Self {
37        Self {
38            secrets_map: HashMap::new(),
39            valid_namespaces: DEFAULT_NAMESPACES.iter().map(|s| s.to_string()).collect(),
40            secrets_from_env: true,
41            additional_import_mappings: HashMap::new(),
42            ignore_unserializable_fields: false,
43        }
44    }
45}
46
47impl ReviverConfig {
48    /// Create a new ReviverConfig with default settings.
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    /// Set the secrets map.
54    pub fn with_secrets_map(mut self, secrets_map: HashMap<String, String>) -> Self {
55        self.secrets_map = secrets_map;
56        self
57    }
58
59    /// Add additional valid namespaces.
60    pub fn with_valid_namespaces(mut self, namespaces: Vec<String>) -> Self {
61        self.valid_namespaces.extend(namespaces);
62        self
63    }
64
65    /// Set whether to load secrets from the environment.
66    pub fn with_secrets_from_env(mut self, secrets_from_env: bool) -> Self {
67        self.secrets_from_env = secrets_from_env;
68        self
69    }
70
71    /// Add additional import mappings.
72    pub fn with_additional_import_mappings(
73        mut self,
74        mappings: HashMap<Vec<String>, Vec<String>>,
75    ) -> Self {
76        self.additional_import_mappings = mappings;
77        self
78    }
79
80    /// Set whether to ignore unserializable fields.
81    pub fn with_ignore_unserializable_fields(mut self, ignore: bool) -> Self {
82        self.ignore_unserializable_fields = ignore;
83        self
84    }
85}
86
87/// Reviver for JSON objects.
88///
89/// The Reviver is responsible for transforming serialized LangChain objects
90/// back into their original form. It handles:
91/// - Secret resolution (from secrets_map or environment)
92/// - Namespace validation
93/// - Type mapping (for backwards compatibility)
94#[derive(Debug, Clone)]
95pub struct Reviver {
96    config: ReviverConfig,
97    import_mappings: HashMap<Vec<String>, Vec<String>>,
98}
99
100impl Reviver {
101    /// Create a new Reviver with the given configuration.
102    pub fn new(config: ReviverConfig) -> Self {
103        let mut import_mappings = get_all_serializable_mappings();
104        import_mappings.extend(config.additional_import_mappings.clone());
105
106        Self {
107            config,
108            import_mappings,
109        }
110    }
111
112    /// Create a new Reviver with default configuration.
113    pub fn with_defaults() -> Self {
114        Self::new(ReviverConfig::default())
115    }
116
117    /// Revive a JSON value.
118    ///
119    /// This method processes a JSON value and:
120    /// - Resolves secrets
121    /// - Validates namespaces
122    /// - Maps old namespaces to new ones
123    ///
124    /// # Arguments
125    ///
126    /// * `value` - The JSON value to revive.
127    ///
128    /// # Returns
129    ///
130    /// The revived value, or an error if revival fails.
131    pub fn revive(&self, value: &Value) -> Result<RevivedValue> {
132        let Some(obj) = value.as_object() else {
133            return Ok(RevivedValue::Value(value.clone()));
134        };
135
136        let lc = obj.get("lc").and_then(|v| v.as_i64());
137        let type_ = obj.get("type").and_then(|v| v.as_str());
138        let id = obj.get("id").and_then(|v| v.as_array());
139
140        // Check if this is a LangChain serialized object
141        if lc != Some(LC_VERSION as i64) || id.is_none() {
142            return Ok(RevivedValue::Value(value.clone()));
143        }
144
145        let id: Vec<String> = id
146            .unwrap()
147            .iter()
148            .filter_map(|v| v.as_str().map(|s| s.to_string()))
149            .collect();
150
151        match type_ {
152            Some("secret") => self.revive_secret(&id),
153            Some("not_implemented") => self.revive_not_implemented(value),
154            Some("constructor") => self.revive_constructor(&id, obj),
155            _ => Ok(RevivedValue::Value(value.clone())),
156        }
157    }
158
159    /// Revive a secret value.
160    fn revive_secret(&self, id: &[String]) -> Result<RevivedValue> {
161        if id.is_empty() {
162            return Ok(RevivedValue::None);
163        }
164
165        let key = &id[0];
166
167        // First check the secrets map
168        if let Some(secret) = self.config.secrets_map.get(key) {
169            return Ok(RevivedValue::String(secret.clone()));
170        }
171
172        // Then check the environment
173        if self.config.secrets_from_env
174            && let Ok(value) = env::var(key)
175            && !value.is_empty()
176        {
177            return Ok(RevivedValue::String(value));
178        }
179
180        Ok(RevivedValue::None)
181    }
182
183    /// Revive a not_implemented value.
184    fn revive_not_implemented(&self, value: &Value) -> Result<RevivedValue> {
185        if self.config.ignore_unserializable_fields {
186            return Ok(RevivedValue::None);
187        }
188
189        Err(Error::Other(format!(
190            "Trying to load an object that doesn't implement serialization: {:?}",
191            value
192        )))
193    }
194
195    /// Revive a constructor value.
196    fn revive_constructor(
197        &self,
198        id: &[String],
199        obj: &serde_json::Map<String, Value>,
200    ) -> Result<RevivedValue> {
201        if id.is_empty() {
202            return Err(Error::Other("Constructor id cannot be empty".to_string()));
203        }
204
205        let namespace: Vec<String> = id[..id.len() - 1].to_vec();
206        let name = id.last().unwrap().clone();
207
208        // Validate namespace
209        let root_namespace = namespace.first().map(|s| s.as_str()).unwrap_or("");
210
211        // The root namespace ["langchain"] alone is not valid
212        if namespace == vec!["langchain".to_string()] {
213            return Err(Error::Other(format!("Invalid namespace: {:?}", id)));
214        }
215
216        if !self
217            .config
218            .valid_namespaces
219            .iter()
220            .any(|ns| ns == root_namespace)
221        {
222            return Err(Error::Other(format!("Invalid namespace: {:?}", id)));
223        }
224
225        // Check for explicit import path
226        let mapping_key: Vec<String> = id.to_vec();
227        let resolved_path = if let Some(import_path) = self.import_mappings.get(&mapping_key) {
228            import_path.clone()
229        } else if DISALLOW_LOAD_FROM_PATH.contains(&root_namespace) {
230            return Err(Error::Other(format!(
231                "Trying to deserialize something that cannot be deserialized \
232                 in current version of langchain-core: {:?}",
233                mapping_key
234            )));
235        } else {
236            id.to_vec()
237        };
238
239        // Get kwargs
240        let kwargs = obj
241            .get("kwargs")
242            .and_then(|v| v.as_object())
243            .cloned()
244            .unwrap_or_default();
245
246        // Return the resolved constructor info
247        Ok(RevivedValue::Constructor(ConstructorInfo {
248            path: resolved_path,
249            name,
250            kwargs: Value::Object(kwargs),
251        }))
252    }
253}
254
255/// Result of reviving a JSON value.
256#[derive(Debug, Clone)]
257pub enum RevivedValue {
258    /// A simple value (unchanged from input).
259    Value(Value),
260    /// A string value (typically a resolved secret).
261    String(String),
262    /// A constructor to instantiate.
263    Constructor(ConstructorInfo),
264    /// No value (for ignored unserializable fields or missing secrets).
265    None,
266}
267
268impl RevivedValue {
269    /// Convert to a serde_json::Value.
270    pub fn to_value(&self) -> Value {
271        match self {
272            RevivedValue::Value(v) => v.clone(),
273            RevivedValue::String(s) => Value::String(s.clone()),
274            RevivedValue::Constructor(info) => {
275                serde_json::json!({
276                    "_type": "constructor",
277                    "path": info.path,
278                    "name": info.name,
279                    "kwargs": info.kwargs,
280                })
281            }
282            RevivedValue::None => Value::Null,
283        }
284    }
285
286    /// Check if this is a None value.
287    pub fn is_none(&self) -> bool {
288        matches!(self, RevivedValue::None)
289    }
290}
291
292/// Information about a constructor to instantiate.
293#[derive(Debug, Clone)]
294pub struct ConstructorInfo {
295    /// The full path to the type (namespace + name).
296    pub path: Vec<String>,
297    /// The name of the type.
298    pub name: String,
299    /// The constructor arguments.
300    pub kwargs: Value,
301}
302
303/// Revive a LangChain object from a JSON string.
304///
305/// # Warning
306///
307/// This function can instantiate arbitrary types based on the serialized data.
308/// Be careful when using with untrusted input.
309///
310/// # Arguments
311///
312/// * `text` - The JSON string to parse and revive.
313/// * `config` - Optional configuration for the reviver.
314///
315/// # Returns
316///
317/// The revived value.
318///
319/// # Errors
320///
321/// Returns an error if parsing or revival fails.
322pub fn loads(text: &str, config: Option<ReviverConfig>) -> Result<Value> {
323    let value: Value = serde_json::from_str(text)?;
324    load(value, config)
325}
326
327/// Revive a LangChain object from a JSON value.
328///
329/// Use this if you already have a parsed JSON object,
330/// e.g., from `serde_json::from_str` or similar.
331///
332/// # Warning
333///
334/// This function can instantiate arbitrary types based on the serialized data.
335/// Be careful when using with untrusted input.
336///
337/// # Arguments
338///
339/// * `obj` - The JSON value to revive.
340/// * `config` - Optional configuration for the reviver.
341///
342/// # Returns
343///
344/// The revived value.
345///
346/// # Errors
347///
348/// Returns an error if revival fails.
349pub fn load(obj: Value, config: Option<ReviverConfig>) -> Result<Value> {
350    let reviver = Reviver::new(config.unwrap_or_default());
351    load_recursive(&obj, &reviver)
352}
353
354/// Recursively load a JSON value.
355fn load_recursive(obj: &Value, reviver: &Reviver) -> Result<Value> {
356    match obj {
357        Value::Object(map) => {
358            // First recursively load all nested values
359            let mut loaded_obj = serde_json::Map::new();
360            for (k, v) in map {
361                loaded_obj.insert(k.clone(), load_recursive(v, reviver)?);
362            }
363
364            // Then revive this node
365            let loaded_value = Value::Object(loaded_obj);
366            let revived = reviver.revive(&loaded_value)?;
367            Ok(revived.to_value())
368        }
369        Value::Array(arr) => {
370            let loaded: Result<Vec<Value>> =
371                arr.iter().map(|v| load_recursive(v, reviver)).collect();
372            Ok(Value::Array(loaded?))
373        }
374        _ => Ok(obj.clone()),
375    }
376}
377
378/// Load with secrets from a map.
379///
380/// Convenience function that creates a ReviverConfig with the given secrets map.
381///
382/// # Arguments
383///
384/// * `text` - The JSON string to parse and revive.
385/// * `secrets_map` - A map of secret names to their values.
386///
387/// # Returns
388///
389/// The revived value.
390pub fn loads_with_secrets(text: &str, secrets_map: HashMap<String, String>) -> Result<Value> {
391    let config = ReviverConfig::new().with_secrets_map(secrets_map);
392    loads(text, Some(config))
393}
394
395/// Load with additional valid namespaces.
396///
397/// Convenience function that creates a ReviverConfig with additional namespaces.
398///
399/// # Arguments
400///
401/// * `text` - The JSON string to parse and revive.
402/// * `namespaces` - Additional namespaces to allow.
403///
404/// # Returns
405///
406/// The revived value.
407pub fn loads_with_namespaces(text: &str, namespaces: Vec<String>) -> Result<Value> {
408    let config = ReviverConfig::new().with_valid_namespaces(namespaces);
409    loads(text, Some(config))
410}
411
412#[cfg(test)]
413mod tests {
414    use super::*;
415
416    #[test]
417    fn test_reviver_default() {
418        let reviver = Reviver::with_defaults();
419        assert!(reviver.config.secrets_from_env);
420        assert!(!reviver.config.valid_namespaces.is_empty());
421    }
422
423    #[test]
424    fn test_revive_simple_value() {
425        let reviver = Reviver::with_defaults();
426        let value = serde_json::json!({"key": "value"});
427        let result = reviver.revive(&value).unwrap();
428
429        match result {
430            RevivedValue::Value(v) => {
431                assert_eq!(v.get("key").and_then(|v| v.as_str()), Some("value"));
432            }
433            _ => panic!("Expected Value"),
434        }
435    }
436
437    #[test]
438    fn test_revive_secret_from_map() {
439        let config = ReviverConfig::new().with_secrets_map(HashMap::from([(
440            "MY_SECRET".to_string(),
441            "secret_value".to_string(),
442        )]));
443        let reviver = Reviver::new(config);
444
445        let value = serde_json::json!({
446            "lc": 1,
447            "type": "secret",
448            "id": ["MY_SECRET"]
449        });
450
451        let result = reviver.revive(&value).unwrap();
452        match result {
453            RevivedValue::String(s) => assert_eq!(s, "secret_value"),
454            _ => panic!("Expected String"),
455        }
456    }
457
458    #[test]
459    fn test_revive_missing_secret() {
460        let config = ReviverConfig::new().with_secrets_from_env(false);
461        let reviver = Reviver::new(config);
462
463        let value = serde_json::json!({
464            "lc": 1,
465            "type": "secret",
466            "id": ["NONEXISTENT_SECRET"]
467        });
468
469        let result = reviver.revive(&value).unwrap();
470        assert!(result.is_none());
471    }
472
473    #[test]
474    fn test_revive_not_implemented_error() {
475        let reviver = Reviver::with_defaults();
476
477        let value = serde_json::json!({
478            "lc": 1,
479            "type": "not_implemented",
480            "id": ["some", "type"],
481            "repr": "SomeType(...)"
482        });
483
484        let result = reviver.revive(&value);
485        assert!(result.is_err());
486    }
487
488    #[test]
489    fn test_revive_not_implemented_ignored() {
490        let config = ReviverConfig::new().with_ignore_unserializable_fields(true);
491        let reviver = Reviver::new(config);
492
493        let value = serde_json::json!({
494            "lc": 1,
495            "type": "not_implemented",
496            "id": ["some", "type"],
497            "repr": "SomeType(...)"
498        });
499
500        let result = reviver.revive(&value).unwrap();
501        assert!(result.is_none());
502    }
503
504    #[test]
505    fn test_revive_constructor() {
506        let reviver = Reviver::with_defaults();
507
508        let value = serde_json::json!({
509            "lc": 1,
510            "type": "constructor",
511            "id": ["langchain_core", "messages", "ai", "AIMessage"],
512            "kwargs": {
513                "content": "Hello, world!"
514            }
515        });
516
517        let result = reviver.revive(&value).unwrap();
518        match result {
519            RevivedValue::Constructor(info) => {
520                assert_eq!(info.name, "AIMessage");
521                assert!(info.path.contains(&"langchain_core".to_string()));
522            }
523            _ => panic!("Expected Constructor"),
524        }
525    }
526
527    #[test]
528    fn test_revive_constructor_with_mapping() {
529        let reviver = Reviver::with_defaults();
530
531        // Old langchain.schema path should be mapped to langchain_core
532        let value = serde_json::json!({
533            "lc": 1,
534            "type": "constructor",
535            "id": ["langchain", "schema", "messages", "AIMessage"],
536            "kwargs": {
537                "content": "Hello!"
538            }
539        });
540
541        let result = reviver.revive(&value).unwrap();
542        match result {
543            RevivedValue::Constructor(info) => {
544                assert_eq!(info.name, "AIMessage");
545                // Should be remapped to langchain_core path
546                assert!(info.path.contains(&"langchain_core".to_string()));
547            }
548            _ => panic!("Expected Constructor"),
549        }
550    }
551
552    #[test]
553    fn test_revive_invalid_namespace() {
554        let reviver = Reviver::with_defaults();
555
556        let value = serde_json::json!({
557            "lc": 1,
558            "type": "constructor",
559            "id": ["invalid_namespace", "SomeClass"],
560            "kwargs": {}
561        });
562
563        let result = reviver.revive(&value);
564        assert!(result.is_err());
565    }
566
567    #[test]
568    fn test_loads_simple() {
569        let json = r#"{"key": "value"}"#;
570        let result = loads(json, None).unwrap();
571        assert_eq!(result.get("key").and_then(|v| v.as_str()), Some("value"));
572    }
573
574    #[test]
575    fn test_loads_nested() {
576        let json = r#"{
577            "outer": {
578                "lc": 1,
579                "type": "secret",
580                "id": ["TEST_KEY"]
581            }
582        }"#;
583
584        let config = ReviverConfig::new().with_secrets_map(HashMap::from([(
585            "TEST_KEY".to_string(),
586            "resolved".to_string(),
587        )]));
588
589        let result = loads(json, Some(config)).unwrap();
590        assert_eq!(
591            result.get("outer").and_then(|v| v.as_str()),
592            Some("resolved")
593        );
594    }
595
596    #[test]
597    fn test_load_array() {
598        let json = r#"[
599            {"lc": 1, "type": "secret", "id": ["KEY1"]},
600            {"lc": 1, "type": "secret", "id": ["KEY2"]}
601        ]"#;
602
603        let config = ReviverConfig::new().with_secrets_map(HashMap::from([
604            ("KEY1".to_string(), "value1".to_string()),
605            ("KEY2".to_string(), "value2".to_string()),
606        ]));
607
608        let result = loads(json, Some(config)).unwrap();
609        let arr = result.as_array().unwrap();
610        assert_eq!(arr[0].as_str(), Some("value1"));
611        assert_eq!(arr[1].as_str(), Some("value2"));
612    }
613}