1use 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#[derive(Debug, Clone)]
21pub struct ReviverConfig {
22 pub secrets_map: HashMap<String, String>,
25 pub valid_namespaces: Vec<String>,
27 pub secrets_from_env: bool,
29 pub additional_import_mappings: HashMap<Vec<String>, Vec<String>>,
31 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 pub fn new() -> Self {
50 Self::default()
51 }
52
53 pub fn with_secrets_map(mut self, secrets_map: HashMap<String, String>) -> Self {
55 self.secrets_map = secrets_map;
56 self
57 }
58
59 pub fn with_valid_namespaces(mut self, namespaces: Vec<String>) -> Self {
61 self.valid_namespaces.extend(namespaces);
62 self
63 }
64
65 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 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 pub fn with_ignore_unserializable_fields(mut self, ignore: bool) -> Self {
82 self.ignore_unserializable_fields = ignore;
83 self
84 }
85}
86
87#[derive(Debug, Clone)]
95pub struct Reviver {
96 config: ReviverConfig,
97 import_mappings: HashMap<Vec<String>, Vec<String>>,
98}
99
100impl Reviver {
101 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 pub fn with_defaults() -> Self {
114 Self::new(ReviverConfig::default())
115 }
116
117 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 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 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 if let Some(secret) = self.config.secrets_map.get(key) {
169 return Ok(RevivedValue::String(secret.clone()));
170 }
171
172 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 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 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 let root_namespace = namespace.first().map(|s| s.as_str()).unwrap_or("");
210
211 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 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 let kwargs = obj
241 .get("kwargs")
242 .and_then(|v| v.as_object())
243 .cloned()
244 .unwrap_or_default();
245
246 Ok(RevivedValue::Constructor(ConstructorInfo {
248 path: resolved_path,
249 name,
250 kwargs: Value::Object(kwargs),
251 }))
252 }
253}
254
255#[derive(Debug, Clone)]
257pub enum RevivedValue {
258 Value(Value),
260 String(String),
262 Constructor(ConstructorInfo),
264 None,
266}
267
268impl RevivedValue {
269 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 pub fn is_none(&self) -> bool {
288 matches!(self, RevivedValue::None)
289 }
290}
291
292#[derive(Debug, Clone)]
294pub struct ConstructorInfo {
295 pub path: Vec<String>,
297 pub name: String,
299 pub kwargs: Value,
301}
302
303pub fn loads(text: &str, config: Option<ReviverConfig>) -> Result<Value> {
323 let value: Value = serde_json::from_str(text)?;
324 load(value, config)
325}
326
327pub 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
354fn load_recursive(obj: &Value, reviver: &Reviver) -> Result<Value> {
356 match obj {
357 Value::Object(map) => {
358 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 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
378pub 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
395pub 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 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 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}