agent_chain_core/load/
serializable.rs

1//! Serializable base trait and types.
2//!
3//! This module contains the core `Serializable` trait and related types,
4//! mirroring `langchain_core.load.serializable`.
5
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::any::Any;
9use std::collections::HashMap;
10
11#[cfg(feature = "specta")]
12use specta::Type;
13
14/// Serialization version.
15pub const LC_VERSION: i32 = 1;
16
17/// Base serialized structure containing common fields.
18#[cfg_attr(feature = "specta", derive(Type))]
19#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
20pub struct BaseSerialized {
21    /// The version of the serialization format.
22    pub lc: i32,
23    /// The unique identifier of the object (namespace path).
24    pub id: Vec<String>,
25    /// The name of the object (optional).
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub name: Option<String>,
28    /// The graph of the object (optional).
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub graph: Option<HashMap<String, Value>>,
31}
32
33impl Default for BaseSerialized {
34    fn default() -> Self {
35        Self {
36            lc: LC_VERSION,
37            id: Vec::new(),
38            name: None,
39            graph: None,
40        }
41    }
42}
43
44/// Serialized constructor representation.
45///
46/// Used when an object can be serialized and reconstructed from its constructor.
47#[cfg_attr(feature = "specta", derive(Type))]
48#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
49pub struct SerializedConstructor {
50    /// The version of the serialization format.
51    pub lc: i32,
52    /// The type of serialization. Always "constructor".
53    #[serde(rename = "type")]
54    pub type_: String,
55    /// The unique identifier of the object (namespace path).
56    pub id: Vec<String>,
57    /// The constructor arguments.
58    pub kwargs: HashMap<String, Value>,
59    /// The name of the object (optional).
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub name: Option<String>,
62    /// The graph of the object (optional).
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub graph: Option<HashMap<String, Value>>,
65}
66
67impl SerializedConstructor {
68    /// Create a new SerializedConstructor.
69    pub fn new(id: Vec<String>, kwargs: HashMap<String, Value>) -> Self {
70        Self {
71            lc: LC_VERSION,
72            type_: "constructor".to_string(),
73            id,
74            kwargs,
75            name: None,
76            graph: None,
77        }
78    }
79
80    /// Set the name.
81    pub fn with_name(mut self, name: impl Into<String>) -> Self {
82        self.name = Some(name.into());
83        self
84    }
85
86    /// Set the graph.
87    pub fn with_graph(mut self, graph: HashMap<String, Value>) -> Self {
88        self.graph = Some(graph);
89        self
90    }
91}
92
93/// Serialized secret representation.
94///
95/// Used to represent secret values that should not be serialized directly.
96#[cfg_attr(feature = "specta", derive(Type))]
97#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
98pub struct SerializedSecret {
99    /// The version of the serialization format.
100    pub lc: i32,
101    /// The type of serialization. Always "secret".
102    #[serde(rename = "type")]
103    pub type_: String,
104    /// The unique identifier of the secret (usually environment variable name).
105    pub id: Vec<String>,
106    /// The name of the object (optional).
107    #[serde(skip_serializing_if = "Option::is_none")]
108    pub name: Option<String>,
109    /// The graph of the object (optional).
110    #[serde(skip_serializing_if = "Option::is_none")]
111    pub graph: Option<HashMap<String, Value>>,
112}
113
114impl SerializedSecret {
115    /// Create a new SerializedSecret.
116    pub fn new(id: Vec<String>) -> Self {
117        Self {
118            lc: LC_VERSION,
119            type_: "secret".to_string(),
120            id,
121            name: None,
122            graph: None,
123        }
124    }
125
126    /// Create a secret from a single secret id (environment variable name).
127    pub fn from_secret_id(secret_id: impl Into<String>) -> Self {
128        Self::new(vec![secret_id.into()])
129    }
130}
131
132/// Serialized not implemented representation.
133///
134/// Used when an object cannot be serialized.
135#[cfg_attr(feature = "specta", derive(Type))]
136#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
137pub struct SerializedNotImplemented {
138    /// The version of the serialization format.
139    pub lc: i32,
140    /// The type of serialization. Always "not_implemented".
141    #[serde(rename = "type")]
142    pub type_: String,
143    /// The unique identifier of the object (namespace path).
144    pub id: Vec<String>,
145    /// The representation of the object (optional).
146    #[serde(skip_serializing_if = "Option::is_none")]
147    pub repr: Option<String>,
148    /// The name of the object (optional).
149    #[serde(skip_serializing_if = "Option::is_none")]
150    pub name: Option<String>,
151    /// The graph of the object (optional).
152    #[serde(skip_serializing_if = "Option::is_none")]
153    pub graph: Option<HashMap<String, Value>>,
154}
155
156impl SerializedNotImplemented {
157    /// Create a new SerializedNotImplemented.
158    pub fn new(id: Vec<String>) -> Self {
159        Self {
160            lc: LC_VERSION,
161            type_: "not_implemented".to_string(),
162            id,
163            repr: None,
164            name: None,
165            graph: None,
166        }
167    }
168
169    /// Set the repr.
170    pub fn with_repr(mut self, repr: impl Into<String>) -> Self {
171        self.repr = Some(repr.into());
172        self
173    }
174}
175
176/// Union type for all serialized representations.
177#[cfg_attr(feature = "specta", derive(Type))]
178#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
179#[serde(tag = "type")]
180pub enum Serialized {
181    /// A serialized constructor.
182    #[serde(rename = "constructor")]
183    Constructor(SerializedConstructorData),
184    /// A serialized secret.
185    #[serde(rename = "secret")]
186    Secret(SerializedSecretData),
187    /// A serialized not implemented.
188    #[serde(rename = "not_implemented")]
189    NotImplemented(SerializedNotImplementedData),
190}
191
192/// Data for SerializedConstructor without the type tag.
193#[cfg_attr(feature = "specta", derive(Type))]
194#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
195pub struct SerializedConstructorData {
196    /// The version of the serialization format.
197    pub lc: i32,
198    /// The unique identifier of the object (namespace path).
199    pub id: Vec<String>,
200    /// The constructor arguments.
201    pub kwargs: HashMap<String, Value>,
202    /// The name of the object (optional).
203    #[serde(skip_serializing_if = "Option::is_none")]
204    pub name: Option<String>,
205    /// The graph of the object (optional).
206    #[serde(skip_serializing_if = "Option::is_none")]
207    pub graph: Option<HashMap<String, Value>>,
208}
209
210/// Data for SerializedSecret without the type tag.
211#[cfg_attr(feature = "specta", derive(Type))]
212#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
213pub struct SerializedSecretData {
214    /// The version of the serialization format.
215    pub lc: i32,
216    /// The unique identifier of the secret (usually environment variable name).
217    pub id: Vec<String>,
218    /// The name of the object (optional).
219    #[serde(skip_serializing_if = "Option::is_none")]
220    pub name: Option<String>,
221    /// The graph of the object (optional).
222    #[serde(skip_serializing_if = "Option::is_none")]
223    pub graph: Option<HashMap<String, Value>>,
224}
225
226/// Data for SerializedNotImplemented without the type tag.
227#[cfg_attr(feature = "specta", derive(Type))]
228#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
229pub struct SerializedNotImplementedData {
230    /// The version of the serialization format.
231    pub lc: i32,
232    /// The unique identifier of the object (namespace path).
233    pub id: Vec<String>,
234    /// The representation of the object (optional).
235    #[serde(skip_serializing_if = "Option::is_none")]
236    pub repr: Option<String>,
237    /// The name of the object (optional).
238    #[serde(skip_serializing_if = "Option::is_none")]
239    pub name: Option<String>,
240    /// The graph of the object (optional).
241    #[serde(skip_serializing_if = "Option::is_none")]
242    pub graph: Option<HashMap<String, Value>>,
243}
244
245impl From<SerializedConstructor> for Serialized {
246    fn from(s: SerializedConstructor) -> Self {
247        Serialized::Constructor(SerializedConstructorData {
248            lc: s.lc,
249            id: s.id,
250            kwargs: s.kwargs,
251            name: s.name,
252            graph: s.graph,
253        })
254    }
255}
256
257impl From<SerializedSecret> for Serialized {
258    fn from(s: SerializedSecret) -> Self {
259        Serialized::Secret(SerializedSecretData {
260            lc: s.lc,
261            id: s.id,
262            name: s.name,
263            graph: s.graph,
264        })
265    }
266}
267
268impl From<SerializedNotImplemented> for Serialized {
269    fn from(s: SerializedNotImplemented) -> Self {
270        Serialized::NotImplemented(SerializedNotImplementedData {
271            lc: s.lc,
272            id: s.id,
273            repr: s.repr,
274            name: s.name,
275            graph: s.graph,
276        })
277    }
278}
279
280/// Trait for objects that can be serialized to JSON for LangChain compatibility.
281///
282/// This trait provides the core serialization functionality, allowing objects
283/// to be serialized and deserialized across different versions and implementations.
284///
285/// # Example
286///
287/// ```ignore
288/// use agent_chain_core::load::Serializable;
289///
290/// struct MyModel {
291///     name: String,
292///     value: i32,
293/// }
294///
295/// impl Serializable for MyModel {
296///     fn is_lc_serializable() -> bool {
297///         true
298///     }
299///
300///     fn get_lc_namespace() -> Vec<String> {
301///         vec!["my_package".to_string(), "models".to_string()]
302///     }
303/// }
304/// ```
305pub trait Serializable: Any + Send + Sync {
306    /// Is this class serializable?
307    ///
308    /// By design, even if a type implements `Serializable`, it is not serializable
309    /// by default. This is to prevent accidental serialization of objects that should
310    /// not be serialized.
311    ///
312    /// Returns `false` by default.
313    fn is_lc_serializable() -> bool
314    where
315        Self: Sized,
316    {
317        false
318    }
319
320    /// Get the namespace of the LangChain object.
321    ///
322    /// For example, if the class is `langchain.llms.openai.OpenAI`, then the
323    /// namespace is `["langchain", "llms", "openai"]`.
324    fn get_lc_namespace() -> Vec<String>
325    where
326        Self: Sized;
327
328    /// A map of constructor argument names to secret ids.
329    ///
330    /// For example, `{"openai_api_key": "OPENAI_API_KEY"}`.
331    fn lc_secrets(&self) -> HashMap<String, String> {
332        HashMap::new()
333    }
334
335    /// List of attribute names that should be included in the serialized kwargs.
336    ///
337    /// These attributes must be accepted by the constructor.
338    fn lc_attributes(&self) -> HashMap<String, Value> {
339        HashMap::new()
340    }
341
342    /// Return a unique identifier for this class for serialization purposes.
343    ///
344    /// The unique identifier is a list of strings that describes the path
345    /// to the object.
346    fn lc_id() -> Vec<String>
347    where
348        Self: Sized,
349    {
350        let mut id = Self::get_lc_namespace();
351        id.push(
352            std::any::type_name::<Self>()
353                .rsplit("::")
354                .next()
355                .unwrap_or("Unknown")
356                .to_string(),
357        );
358        id
359    }
360
361    /// Get the type name of this object.
362    fn lc_type_name(&self) -> &'static str {
363        std::any::type_name::<Self>()
364    }
365
366    /// Serialize this object to JSON.
367    ///
368    /// Returns a `SerializedConstructor` if the object is serializable,
369    /// or a `SerializedNotImplemented` if it is not.
370    fn to_json(&self) -> Serialized
371    where
372        Self: Sized + Serialize,
373    {
374        if !Self::is_lc_serializable() {
375            return self.to_json_not_implemented();
376        }
377
378        let kwargs = match serde_json::to_value(self) {
379            Ok(Value::Object(map)) => map.into_iter().collect(),
380            _ => HashMap::new(),
381        };
382
383        let secrets = self.lc_secrets();
384        let kwargs = if secrets.is_empty() {
385            kwargs
386        } else {
387            replace_secrets(kwargs, &secrets)
388        };
389
390        let mut final_kwargs = kwargs;
391        for (key, value) in self.lc_attributes() {
392            final_kwargs.insert(key, value);
393        }
394
395        SerializedConstructor::new(Self::lc_id(), final_kwargs).into()
396    }
397
398    /// Serialize a "not implemented" object.
399    fn to_json_not_implemented(&self) -> Serialized {
400        to_json_not_implemented_value(self.lc_type_name(), None)
401    }
402}
403
404/// Create a SerializedNotImplemented for an arbitrary object.
405///
406/// # Arguments
407///
408/// * `type_name` - The type name of the object.
409/// * `repr` - Optional string representation.
410pub fn to_json_not_implemented_value(type_name: &str, repr: Option<String>) -> Serialized {
411    let id: Vec<String> = type_name.split("::").map(|s| s.to_string()).collect();
412
413    let mut result = SerializedNotImplemented::new(id);
414    if let Some(r) = repr {
415        result = result.with_repr(r);
416    }
417
418    result.into()
419}
420
421/// Create a SerializedNotImplemented from a serde_json::Value.
422///
423/// Used as a fallback when normal serialization fails.
424pub fn to_json_not_implemented(value: &Value) -> Serialized {
425    let repr = serde_json::to_string_pretty(value).ok();
426    to_json_not_implemented_value("serde_json::Value", repr)
427}
428
429/// Replace secrets in kwargs with SerializedSecret placeholders.
430fn replace_secrets(
431    mut kwargs: HashMap<String, Value>,
432    secrets_map: &HashMap<String, String>,
433) -> HashMap<String, Value> {
434    for (path, secret_id) in secrets_map {
435        let parts: Vec<&str> = path.split('.').collect();
436
437        if parts.len() == 1 {
438            if kwargs.contains_key(path) {
439                kwargs.insert(
440                    path.clone(),
441                    serde_json::to_value(SerializedSecret::from_secret_id(secret_id)).unwrap(),
442                );
443            }
444        } else {
445            replace_nested_secret(&mut kwargs, &parts, secret_id);
446        }
447    }
448    kwargs
449}
450
451/// Replace a nested secret in kwargs.
452fn replace_nested_secret(current: &mut HashMap<String, Value>, parts: &[&str], secret_id: &str) {
453    if parts.is_empty() {
454        return;
455    }
456
457    let key = parts[0];
458
459    if parts.len() == 1 {
460        if current.contains_key(key) {
461            current.insert(
462                key.to_string(),
463                serde_json::to_value(SerializedSecret::from_secret_id(secret_id)).unwrap(),
464            );
465        }
466    } else if let Some(Value::Object(map)) = current.get_mut(key) {
467        let mut nested: HashMap<String, Value> =
468            map.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
469        replace_nested_secret(&mut nested, &parts[1..], secret_id);
470        *map = nested.into_iter().collect();
471    }
472}
473
474#[cfg(test)]
475mod tests {
476    use super::*;
477
478    #[test]
479    fn test_serialized_constructor() {
480        let mut kwargs = HashMap::new();
481        kwargs.insert("name".to_string(), Value::String("test".to_string()));
482
483        let constructor = SerializedConstructor::new(
484            vec![
485                "langchain".to_string(),
486                "llms".to_string(),
487                "OpenAI".to_string(),
488            ],
489            kwargs,
490        );
491
492        assert_eq!(constructor.lc, 1);
493        assert_eq!(constructor.type_, "constructor");
494        assert_eq!(constructor.id.len(), 3);
495    }
496
497    #[test]
498    fn test_serialized_secret() {
499        let secret = SerializedSecret::from_secret_id("OPENAI_API_KEY");
500
501        assert_eq!(secret.lc, 1);
502        assert_eq!(secret.type_, "secret");
503        assert_eq!(secret.id, vec!["OPENAI_API_KEY".to_string()]);
504    }
505
506    #[test]
507    fn test_serialized_not_implemented() {
508        let not_impl =
509            SerializedNotImplemented::new(vec!["my_module".to_string(), "MyClass".to_string()])
510                .with_repr("MyClass(...)".to_string());
511
512        assert_eq!(not_impl.lc, 1);
513        assert_eq!(not_impl.type_, "not_implemented");
514        assert_eq!(not_impl.repr, Some("MyClass(...)".to_string()));
515    }
516
517    #[test]
518    fn test_replace_secrets() {
519        let mut kwargs = HashMap::new();
520        kwargs.insert(
521            "api_key".to_string(),
522            Value::String("secret_value".to_string()),
523        );
524        kwargs.insert("model".to_string(), Value::String("gpt-4".to_string()));
525
526        let mut secrets = HashMap::new();
527        secrets.insert("api_key".to_string(), "OPENAI_API_KEY".to_string());
528
529        let result = replace_secrets(kwargs, &secrets);
530
531        assert!(result.get("api_key").unwrap().is_object());
532        assert_eq!(
533            result.get("model").unwrap(),
534            &Value::String("gpt-4".to_string())
535        );
536    }
537}