//! Document ingestion pipeline use std::path::Path; use walkdir::WalkDir; use crate::chunking::ChunkingEngine; use crate::config::IngestionConfig; use crate::embeddings::{EmbeddedDocument, EmbeddingEngine}; use crate::error::Result; /// Document ingester for RAG system pub struct DocumentIngester { config: IngestionConfig, embedding_engine: EmbeddingEngine, chunking_engine: ChunkingEngine, } impl DocumentIngester { /// Create a new document ingester pub fn new(config: IngestionConfig, embedding_engine: EmbeddingEngine) -> Self { let chunking_engine = ChunkingEngine::new(config.chunk_size, config.chunk_overlap); Self { config, embedding_engine, chunking_engine, } } /// Ingest a single document file pub async fn ingest_file(&self, file_path: &Path) -> Result> { let path_str = file_path.to_string_lossy().to_string(); // Read file content let content = std::fs::read_to_string(file_path).map_err(|e| { tracing::error!("Failed to read file {}: {}", path_str, e); crate::error::RagError::Io(e) })?; // Chunk the document let chunks = self.chunking_engine.chunk_file(&content, &path_str)?; tracing::debug!("Chunked {} into {} chunks", path_str, chunks.len()); // Embed the chunks let embedded = self.embedding_engine.embed_batch(&chunks).await?; tracing::info!( "Ingested {} with {} embedded chunks", path_str, embedded.len() ); Ok(embedded) } /// Ingest markdown documentation pub async fn ingest_markdown(&self, doc_path: &Path) -> Result> { let content = std::fs::read_to_string(doc_path)?; let path_str = doc_path.to_string_lossy().to_string(); let chunks = self.chunking_engine.chunk_markdown(&content, &path_str)?; let embedded = self.embedding_engine.embed_batch(&chunks).await?; tracing::info!( "Ingested markdown document: {} ({} chunks)", path_str, embedded.len() ); Ok(embedded) } /// Ingest KCL schema files pub async fn ingest_kcl(&self, schema_path: &Path) -> Result> { let content = std::fs::read_to_string(schema_path)?; let path_str = schema_path.to_string_lossy().to_string(); let chunks = self.chunking_engine.chunk_kcl(&content, &path_str)?; let embedded = self.embedding_engine.embed_batch(&chunks).await?; tracing::info!( "Ingested KCL schema: {} ({} chunks)", path_str, embedded.len() ); Ok(embedded) } /// Ingest Nushell scripts pub async fn ingest_nushell(&self, script_path: &Path) -> Result> { let content = std::fs::read_to_string(script_path)?; let path_str = script_path.to_string_lossy().to_string(); let chunks = self.chunking_engine.chunk_nushell(&content, &path_str)?; let embedded = self.embedding_engine.embed_batch(&chunks).await?; tracing::info!( "Ingested Nushell script: {} ({} chunks)", path_str, embedded.len() ); Ok(embedded) } /// Scan directory recursively for documents to ingest pub async fn scan_and_ingest_directory( &self, root_path: &Path, ) -> Result> { let mut all_documents = Vec::new(); for entry in WalkDir::new(root_path) .into_iter() .filter_map(|e| e.ok()) .filter(|e| e.path().is_file()) { let path = entry.path(); let ext = path.extension().and_then(|s| s.to_str()).unwrap_or(""); // Check if this file type should be ingested let should_ingest = self.config.doc_types.iter().any(|t| { (t == "markdown" && ext == "md") || (t == "kcl" && ext == "k") || (t == "nushell" && ext == "nu") }); if should_ingest { match self.ingest_file(path).await { Ok(docs) => { all_documents.extend(docs); } Err(e) => { tracing::warn!("Failed to ingest {}: {}", path.display(), e); // Continue with next file } } } } tracing::info!( "Ingested {} total documents from {}", all_documents.len(), root_path.display() ); Ok(all_documents) } /// Scan directory for documents without ingesting pub async fn scan_documents(&self, root_path: &Path) -> Result> { let mut documents = Vec::new(); for entry in WalkDir::new(root_path) .into_iter() .filter_map(|e| e.ok()) .filter(|e| e.path().is_file()) { let path = entry.path(); let ext = path.extension().and_then(|s| s.to_str()).unwrap_or(""); if self.config.doc_types.iter().any(|t| { (t == "markdown" && ext == "md") || (t == "kcl" && ext == "k") || (t == "nushell" && ext == "nu") }) { documents.push(path.to_string_lossy().to_string()); } } tracing::debug!("Found {} documents to ingest", documents.len()); Ok(documents) } /// Get reference to embedding engine pub fn embedding_engine(&self) -> &EmbeddingEngine { &self.embedding_engine } /// Get reference to chunking engine pub fn chunking_engine(&self) -> &ChunkingEngine { &self.chunking_engine } } #[cfg(test)] mod tests { use super::*; use crate::config::EmbeddingConfig; #[test] fn test_document_ingester_creation() { let embedding_config = EmbeddingConfig { provider: "local".to_string(), ..Default::default() }; let engine = EmbeddingEngine::new(embedding_config).unwrap(); let ingest_config = IngestionConfig::default(); let _ingester = DocumentIngester::new(ingest_config, engine); // Basic smoke test - object created successfully } #[tokio::test] async fn test_scan_documents() { let embedding_config = EmbeddingConfig { provider: "local".to_string(), ..Default::default() }; let engine = EmbeddingEngine::new(embedding_config).unwrap(); let ingest_config = IngestionConfig::default(); let ingester = DocumentIngester::new(ingest_config, engine); // Try to scan a directory that doesn't exist // This should return empty list, not error let result = ingester.scan_documents(Path::new("/nonexistent")).await; // Should either succeed with empty list or fail gracefully match result { Ok(docs) => { assert!(docs.is_empty()); } Err(_) => { // Also acceptable - directory doesn't exist } } } }