lc_rag/loaders/
json.rs

1// src/retrieval/loaders/json.rs
2//! JSON document loader implementation
3//!
4//! Loads content from JSON files, supporting a designated field as the document content.
5
6use super::{Document, DocumentLoader, LoaderError};
7use async_trait::async_trait;
8use serde_json::Value;
9use std::path::PathBuf;
10
11/// JSON document loader
12///
13/// Supports loading JSON files, optionally using a specific field as the document content.
14/// - For a JSON array, each element becomes one document
15/// - For a JSON object, the whole object becomes one document (or a specified field)
16pub struct JSONLoader {
17    /// JSON file path
18    pub path: PathBuf,
19
20    /// The field name used as the document content (optional)
21    /// If specified, that field's value is extracted as the content
22    pub content_key: Option<String>,
23
24    /// Whether to keep the raw JSON as metadata
25    pub preserve_raw: bool,
26}
27
28impl JSONLoader {
29    /// Creates a new JSON loader
30    ///
31    /// # Arguments
32    /// * `path` - the JSON file path
33    pub fn new(path: impl Into<PathBuf>) -> Self {
34        Self {
35            path: path.into(),
36            content_key: None,
37            preserve_raw: false,
38        }
39    }
40
41    /// Creates a JSON loader with a content field
42    ///
43    /// # Arguments
44    /// * `path` - the JSON file path
45    /// * `content_key` - the field name used as the document content
46    pub fn new_with_content_key(path: impl Into<PathBuf>, content_key: impl Into<String>) -> Self {
47        Self {
48            path: path.into(),
49            content_key: Some(content_key.into()),
50            preserve_raw: false,
51        }
52    }
53
54    /// Sets whether to keep the raw JSON
55    pub fn with_preserve_raw(mut self, preserve: bool) -> Self {
56        self.preserve_raw = preserve;
57        self
58    }
59}
60
61#[async_trait]
62impl DocumentLoader for JSONLoader {
63    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
64        if !self.path.exists() {
65            return Err(LoaderError::Other(format!(
66                "JSON file does not exist: {}",
67                self.path.display()
68            )));
69        }
70
71        let content = std::fs::read_to_string(&self.path)?;
72        let json: Value =
73            serde_json::from_str(&content).map_err(|e| LoaderError::JsonError(e.to_string()))?;
74
75        let documents = match json {
76            Value::Array(arr) => arr
77                .iter()
78                .filter_map(|item| self.json_value_to_document(item))
79                .collect(),
80            Value::Object(_) => match self.json_value_to_document(&json) {
81                Some(doc) => vec![doc],
82                None => vec![],
83            },
84            _ => {
85                vec![Document::new(json.to_string())
86                    .with_metadata("source", self.path.display().to_string())
87                    .with_metadata("format", "json")]
88            }
89        };
90
91        Ok(documents)
92    }
93}
94
95impl JSONLoader {
96    fn json_value_to_document(&self, value: &Value) -> Option<Document> {
97        match value {
98            Value::Object(obj) => {
99                let content = if let Some(key) = &self.content_key {
100                    obj.get(key)
101                        .map(|v| self.extract_string_value(v))
102                        .unwrap_or_else(|| value.to_string())
103                } else {
104                    value.to_string()
105                };
106
107                if content.is_empty() || content == "null" {
108                    return None;
109                }
110
111                let mut doc = Document::new(content);
112                doc = doc.with_metadata("source", self.path.display().to_string());
113                doc = doc.with_metadata("format", "json".to_string());
114
115                if let Some(key) = &self.content_key {
116                    doc = doc.with_metadata("content_key", key.clone());
117                }
118
119                for (k, v) in obj {
120                    if self.content_key.as_ref() != Some(k) {
121                        doc = doc.with_metadata(k.clone(), self.extract_string_value(v));
122                    }
123                }
124
125                if self.preserve_raw {
126                    doc = doc.with_metadata("raw_json", value.to_string());
127                }
128
129                Some(doc)
130            }
131            Value::String(s) => {
132                if s.is_empty() {
133                    return None;
134                }
135                Some(
136                    Document::new(s.clone())
137                        .with_metadata("source", self.path.display().to_string())
138                        .with_metadata("format", "json"),
139                )
140            }
141            Value::Number(_) | Value::Bool(_) => Some(
142                Document::new(value.to_string())
143                    .with_metadata("source", self.path.display().to_string())
144                    .with_metadata("format", "json"),
145            ),
146            Value::Null => None,
147            Value::Array(_) => Some(
148                Document::new(value.to_string())
149                    .with_metadata("source", self.path.display().to_string())
150                    .with_metadata("format", "json"),
151            ),
152        }
153    }
154
155    fn extract_string_value(&self, value: &Value) -> String {
156        match value {
157            Value::String(s) => s.clone(),
158            _ => value.to_string(),
159        }
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166    use std::io::Write;
167    use tempfile::NamedTempFile;
168
169    #[tokio::test]
170    async fn test_json_loader_nonexistent() {
171        let loader = JSONLoader::new("./nonexistent.json");
172        let result = loader.load().await;
173
174        assert!(result.is_err());
175    }
176
177    #[tokio::test]
178    async fn test_json_loader_invalid_json() {
179        let mut temp_file = NamedTempFile::new().unwrap();
180        write!(temp_file, "{{ invalid json }}").unwrap();
181
182        let loader = JSONLoader::new(temp_file.path());
183        let result = loader.load().await;
184
185        assert!(result.is_err());
186        match result.unwrap_err() {
187            LoaderError::JsonError(_) => {}
188            _ => panic!("Expected JsonError"),
189        }
190    }
191
192    #[tokio::test]
193    async fn test_json_loader_single_object() {
194        let mut temp_file = NamedTempFile::new().unwrap();
195        write!(temp_file, "{{\"title\": \"Test\", \"content\": \"Hello\"}}").unwrap();
196
197        let loader = JSONLoader::new_with_content_key(temp_file.path(), "content");
198        let result = loader.load().await;
199
200        assert!(result.is_ok());
201        let docs = result.unwrap();
202        assert_eq!(docs.len(), 1);
203        assert!(docs[0].content.contains("Hello"));
204    }
205
206    #[tokio::test]
207    async fn test_json_loader_array() {
208        let mut temp_file = NamedTempFile::new().unwrap();
209        write!(temp_file, "[{{\"title\": \"A\", \"content\": \"Content A\"}}, {{\"title\": \"B\", \"content\": \"Content B\"}}]").unwrap();
210
211        let loader = JSONLoader::new_with_content_key(temp_file.path(), "content");
212        let result = loader.load().await;
213
214        assert!(result.is_ok());
215        let docs = result.unwrap();
216        assert_eq!(docs.len(), 2);
217        assert!(docs[0].content.contains("Content A"));
218        assert_eq!(
219            docs[0].metadata.get("title"),
220            Some(&serde_json::Value::String("A".to_string()))
221        );
222    }
223
224    #[tokio::test]
225    async fn test_json_loader_with_preserve_raw() {
226        let mut temp_file = NamedTempFile::new().unwrap();
227        write!(temp_file, "{{\"name\": \"test\", \"value\": 123}}").unwrap();
228
229        let loader = JSONLoader::new(temp_file.path()).with_preserve_raw(true);
230        let result = loader.load().await;
231
232        assert!(result.is_ok());
233        let docs = result.unwrap();
234        assert!(docs[0].metadata.contains_key("raw_json"));
235    }
236}
lc_rag/loaders/json.rs

lc_rag/loaders/
json.rs
