lc_rag/loaders/
json.rs

1// src/retrieval/loaders/json.rs
2//! JSON 文档加载器实现
3//!
4//! 提供从 JSON 文件加载内容的功能，支持指定字段作为文档内容。
5
6use super::{Document, DocumentLoader, LoaderError};
7use async_trait::async_trait;
8use serde_json::Value;
9use std::path::PathBuf;
10
11/// JSON 文档加载器
12///
13/// 支持加载 JSON 文件，可以指定某个字段作为文档内容。
14/// - 对于 JSON 数组，每个元素生成一个文档
15/// - 对于 JSON 对象，整个对象作为一个文档（或指定字段）
16pub struct JSONLoader {
17    /// JSON 文件路径
18    pub path: PathBuf,
19
20    /// 作为文档内容的字段名（可选）
21    /// 如果指定，则提取该字段值作为 content
22    pub content_key: Option<String>,
23
24    /// 是否保留原始 JSON 作为元数据
25    pub preserve_raw: bool,
26}
27
28impl JSONLoader {
29    /// 创建新的 JSON 加载器
30    ///
31    /// # 参数
32    /// * `path` - JSON 文件路径
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    /// 创建带内容字段的 JSON 加载器
42    ///
43    /// # 参数
44    /// * `path` - JSON 文件路径
45    /// * `content_key` - 作为文档内容的字段名
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    /// 设置是否保留原始 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 文件不存在: {}",
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!(docs[0].metadata.get("title"), Some(&"A".to_string()));
219    }
220
221    #[tokio::test]
222    async fn test_json_loader_with_preserve_raw() {
223        let mut temp_file = NamedTempFile::new().unwrap();
224        write!(temp_file, "{{\"name\": \"test\", \"value\": 123}}").unwrap();
225
226        let loader = JSONLoader::new(temp_file.path()).with_preserve_raw(true);
227        let result = loader.load().await;
228
229        assert!(result.is_ok());
230        let docs = result.unwrap();
231        assert!(docs[0].metadata.contains_key("raw_json"));
232    }
233}
lc_rag/loaders/json.rs

lc_rag/loaders/
json.rs
