lc_rag/loaders/
pdf.rs

1// src/retrieval/loaders/pdf.rs
2//! PDF 文档加载器实现
3//!
4//! 提供从 PDF 文件加载文本内容的功能。
5
6use super::{Document, DocumentLoader, LoaderError};
7use async_trait::async_trait;
8use std::path::PathBuf;
9
10/// PDF 文档加载器
11pub struct PDFLoader {
12    /// PDF 文件路径
13    pub path: PathBuf,
14}
15
16impl PDFLoader {
17    /// 创建新的 PDF 加载器
18    ///
19    /// # 参数
20    /// * `path` - PDF 文件路径
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        // 验证文件存在
30        if !self.path.exists() {
31            return Err(LoaderError::Other(format!(
32                "PDF 文件不存在: {}",
33                self.path.display()
34            )));
35        }
36
37        // 使用 pdf_extract 库提取文本
38        let text = pdf_extract::extract_text(&self.path)
39            .map_err(|e| LoaderError::PdfError(format!("PDF 解析失败: {}", e)))?;
40
41        // 创建文档对象，包含元数据
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("不存在")),
62            _ => panic!("Expected Other error"),
63        }
64    }
65
66    #[tokio::test]
67    #[ignore = "requires a sample PDF file"]
68    async fn test_pdf_loader() {
69        // 注意：这需要一个实际存在的 PDF 文件进行测试
70        let loader = PDFLoader::new("./sample.pdf");
71        let result = loader.load().await;
72
73        // 应该成功返回至少一个文档
74        if let Ok(docs) = result {
75            assert!(!docs.is_empty());
76            assert!(docs[0].content.contains("PDF"));
77        }
78    }
79}
lc_rag/loaders/pdf.rs

lc_rag/loaders/
pdf.rs
