Skip to main content

synaptic_loaders/
json_loader.rs

1use crate::Document;
2use async_trait::async_trait;
3use serde_json::Value;
4use synaptic_core::SynapticError;
5
6use crate::Loader;
7
8/// Loads documents from a JSON string.
9///
10/// - If the JSON is an array of objects, each object becomes a Document.
11///   The `content_key` field is used to extract the document content (default: "content").
12///   The `id_key` field is used to extract the document id (default: "id").
13/// - If the JSON is a single object, it becomes one Document.
14pub struct JsonLoader {
15    json: String,
16    content_key: String,
17    id_key: String,
18}
19
20impl JsonLoader {
21    pub fn new(json: impl Into<String>) -> Self {
22        Self {
23            json: json.into(),
24            content_key: "content".to_string(),
25            id_key: "id".to_string(),
26        }
27    }
28
29    pub fn with_content_key(mut self, key: impl Into<String>) -> Self {
30        self.content_key = key.into();
31        self
32    }
33
34    pub fn with_id_key(mut self, key: impl Into<String>) -> Self {
35        self.id_key = key.into();
36        self
37    }
38}
39
40#[async_trait]
41impl Loader for JsonLoader {
42    async fn load(&self) -> Result<Vec<Document>, SynapticError> {
43        let value: Value = serde_json::from_str(&self.json)
44            .map_err(|e| SynapticError::Loader(format!("invalid JSON: {e}")))?;
45
46        match value {
47            Value::Array(arr) => {
48                let mut docs = Vec::with_capacity(arr.len());
49                for (i, item) in arr.iter().enumerate() {
50                    let id = item
51                        .get(&self.id_key)
52                        .and_then(|v| v.as_str())
53                        .map(|s| s.to_string())
54                        .unwrap_or_else(|| format!("doc-{i}"));
55                    let content = item
56                        .get(&self.content_key)
57                        .and_then(|v| v.as_str())
58                        .map(|s| s.to_string())
59                        .unwrap_or_else(|| item.to_string());
60                    docs.push(Document::new(id, content));
61                }
62                Ok(docs)
63            }
64            _ => {
65                let content = value
66                    .get(&self.content_key)
67                    .and_then(|v| v.as_str())
68                    .map(|s| s.to_string())
69                    .unwrap_or_else(|| value.to_string());
70                let id = value
71                    .get(&self.id_key)
72                    .and_then(|v| v.as_str())
73                    .map(|s| s.to_string())
74                    .unwrap_or_else(|| "doc-0".to_string());
75                Ok(vec![Document::new(id, content)])
76            }
77        }
78    }
79}