use crate::document_loaders::{DocumentLoader, LoadOptions};
use crate::retriever_engine::Document;
use crate::types::Layer3Result;
use async_trait::async_trait;
use std::path::PathBuf;
#[allow(dead_code)]
pub struct JsonLoader {
#[allow(dead_code)]
options: LoadOptions,
query: Option<String>,
}
impl JsonLoader {
pub fn new() -> Self {
Self {
options: LoadOptions::default(),
query: None,
}
}
pub fn with_query(query: impl Into<String>) -> Self {
Self {
options: LoadOptions::default(),
query: Some(query.into()),
}
}
}
impl Default for JsonLoader {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl DocumentLoader for JsonLoader {
async fn load(&self, path: PathBuf) -> Layer3Result<Document> {
let content = tokio::fs::read_to_string(&path).await?;
let _: serde_json::Value = serde_json::from_str(&content)?;
Ok(Document::new(content).with_source(path.to_string_lossy().to_string()))
}
async fn load_and_split(&self, path: PathBuf) -> Layer3Result<Vec<Document>> {
let content = tokio::fs::read_to_string(&path).await?;
let json: serde_json::Value = serde_json::from_str(&content)?;
if let serde_json::Value::Array(arr) = json {
return Ok(arr
.into_iter()
.enumerate()
.filter_map(|(i, v)| {
if let Ok(s) = serde_json::to_string(&v) {
Some(Document::new(s).with_source(format!(
"{}[{}]",
path.to_string_lossy(),
i
)))
} else {
None
}
})
.collect());
}
Ok(vec![self.load(path).await?])
}
fn supports(&self, path: &std::path::Path) -> bool {
path.extension()
.and_then(|e| e.to_str())
.map(|e| e == "json")
.unwrap_or(false)
}
fn extensions(&self) -> &[&str] {
&["json"]
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_json_loader_extensions() {
let loader = JsonLoader::new();
assert!(loader.extensions().contains(&"json"));
}
}