1use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::any::Any;
9use std::collections::HashMap;
10
11#[cfg(feature = "specta")]
12use specta::Type;
13
14pub const LC_VERSION: i32 = 1;
16
17#[cfg_attr(feature = "specta", derive(Type))]
19#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
20pub struct BaseSerialized {
21 pub lc: i32,
23 pub id: Vec<String>,
25 #[serde(skip_serializing_if = "Option::is_none")]
27 pub name: Option<String>,
28 #[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#[cfg_attr(feature = "specta", derive(Type))]
48#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
49pub struct SerializedConstructor {
50 pub lc: i32,
52 #[serde(rename = "type")]
54 pub type_: String,
55 pub id: Vec<String>,
57 pub kwargs: HashMap<String, Value>,
59 #[serde(skip_serializing_if = "Option::is_none")]
61 pub name: Option<String>,
62 #[serde(skip_serializing_if = "Option::is_none")]
64 pub graph: Option<HashMap<String, Value>>,
65}
66
67impl SerializedConstructor {
68 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 pub fn with_name(mut self, name: impl Into<String>) -> Self {
82 self.name = Some(name.into());
83 self
84 }
85
86 pub fn with_graph(mut self, graph: HashMap<String, Value>) -> Self {
88 self.graph = Some(graph);
89 self
90 }
91}
92
93#[cfg_attr(feature = "specta", derive(Type))]
97#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
98pub struct SerializedSecret {
99 pub lc: i32,
101 #[serde(rename = "type")]
103 pub type_: String,
104 pub id: Vec<String>,
106 #[serde(skip_serializing_if = "Option::is_none")]
108 pub name: Option<String>,
109 #[serde(skip_serializing_if = "Option::is_none")]
111 pub graph: Option<HashMap<String, Value>>,
112}
113
114impl SerializedSecret {
115 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 pub fn from_secret_id(secret_id: impl Into<String>) -> Self {
128 Self::new(vec![secret_id.into()])
129 }
130}
131
132#[cfg_attr(feature = "specta", derive(Type))]
136#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
137pub struct SerializedNotImplemented {
138 pub lc: i32,
140 #[serde(rename = "type")]
142 pub type_: String,
143 pub id: Vec<String>,
145 #[serde(skip_serializing_if = "Option::is_none")]
147 pub repr: Option<String>,
148 #[serde(skip_serializing_if = "Option::is_none")]
150 pub name: Option<String>,
151 #[serde(skip_serializing_if = "Option::is_none")]
153 pub graph: Option<HashMap<String, Value>>,
154}
155
156impl SerializedNotImplemented {
157 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 pub fn with_repr(mut self, repr: impl Into<String>) -> Self {
171 self.repr = Some(repr.into());
172 self
173 }
174}
175
176#[cfg_attr(feature = "specta", derive(Type))]
178#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
179#[serde(tag = "type")]
180pub enum Serialized {
181 #[serde(rename = "constructor")]
183 Constructor(SerializedConstructorData),
184 #[serde(rename = "secret")]
186 Secret(SerializedSecretData),
187 #[serde(rename = "not_implemented")]
189 NotImplemented(SerializedNotImplementedData),
190}
191
192#[cfg_attr(feature = "specta", derive(Type))]
194#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
195pub struct SerializedConstructorData {
196 pub lc: i32,
198 pub id: Vec<String>,
200 pub kwargs: HashMap<String, Value>,
202 #[serde(skip_serializing_if = "Option::is_none")]
204 pub name: Option<String>,
205 #[serde(skip_serializing_if = "Option::is_none")]
207 pub graph: Option<HashMap<String, Value>>,
208}
209
210#[cfg_attr(feature = "specta", derive(Type))]
212#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
213pub struct SerializedSecretData {
214 pub lc: i32,
216 pub id: Vec<String>,
218 #[serde(skip_serializing_if = "Option::is_none")]
220 pub name: Option<String>,
221 #[serde(skip_serializing_if = "Option::is_none")]
223 pub graph: Option<HashMap<String, Value>>,
224}
225
226#[cfg_attr(feature = "specta", derive(Type))]
228#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
229pub struct SerializedNotImplementedData {
230 pub lc: i32,
232 pub id: Vec<String>,
234 #[serde(skip_serializing_if = "Option::is_none")]
236 pub repr: Option<String>,
237 #[serde(skip_serializing_if = "Option::is_none")]
239 pub name: Option<String>,
240 #[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
280pub trait Serializable: Any + Send + Sync {
306 fn is_lc_serializable() -> bool
314 where
315 Self: Sized,
316 {
317 false
318 }
319
320 fn get_lc_namespace() -> Vec<String>
325 where
326 Self: Sized;
327
328 fn lc_secrets(&self) -> HashMap<String, String> {
332 HashMap::new()
333 }
334
335 fn lc_attributes(&self) -> HashMap<String, Value> {
339 HashMap::new()
340 }
341
342 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 fn lc_type_name(&self) -> &'static str {
363 std::any::type_name::<Self>()
364 }
365
366 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 fn to_json_not_implemented(&self) -> Serialized {
400 to_json_not_implemented_value(self.lc_type_name(), None)
401 }
402}
403
404pub 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
421pub 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
429fn 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
451fn 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}