Skip to main content

lc_rag/loaders/
pdf.rs

1// src/retrieval/loaders/pdf.rs
2//! PDF document loader implementation
3//!
4//! Provides text-content loading from PDF files.
5
6use super::{Document, DocumentLoader, LoaderError};
7use async_trait::async_trait;
8use std::path::PathBuf;
9
10/// PDF document loader
11pub struct PDFLoader {
12    /// PDF file path
13    pub path: PathBuf,
14}
15
16impl PDFLoader {
17    /// Creates a new PDF loader
18    ///
19    /// # Arguments
20    /// * `path` - the PDF file path
21    pub fn new(path: impl Into<PathBuf>) -> Self {
22        Self { path: path.into() }
23    }
24}
25
26#[async_trait]
27impl DocumentLoader for PDFLoader {
28    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
29        // Verify the file exists
30        if !self.path.exists() {
31            return Err(LoaderError::Other(format!(
32                "PDF file does not exist: {}",
33                self.path.display()
34            )));
35        }
36
37        // Extract text using the pdf_extract library
38        let text = pdf_extract::extract_text(&self.path)
39            .map_err(|e| LoaderError::PdfError(format!("PDF parse failed: {}", e)))?;
40
41        // Create the document object, including metadata
42        let mut document = Document::new(text);
43        document = document.with_metadata("source".to_string(), self.path.display().to_string());
44        document = document.with_metadata("format".to_string(), "pdf".to_string());
45
46        Ok(vec![document])
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53
54    #[tokio::test]
55    async fn test_pdf_loader_nonexistent() {
56        let loader = PDFLoader::new("./nonexistent.pdf");
57        let result = loader.load().await;
58
59        assert!(result.is_err());
60        match result.unwrap_err() {
61            LoaderError::Other(msg) => assert!(msg.contains("does not exist")),
62            _ => panic!("Expected Other error"),
63        }
64    }
65}