Skip to main content

lc_rag/loaders/
docx.rs

1//! DOCX 文档加载器
2//!
3//! 从 .docx 文件加载文档内容(纯文本提取)。
4//! 不依赖外部 crate,使用 ZIP 解压 + XML 解析提取文本。
5
6use std::collections::HashMap;
7use std::io::Read;
8
9use async_trait::async_trait;
10use regex::Regex;
11use std::sync::LazyLock;
12
13use super::{DocumentLoader, LoaderError};
14use lc_vector_stores::Document;
15
16/// M60: compile regexes once using LazyLock instead of on every call
17static WT_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"<w:t[^>]*>(.*?)</w:t>").unwrap());
18static PARA_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"</w:p>").unwrap());
19
20/// DOCX 文档加载器
21///
22/// 从 .docx 文件路径加载文档,提取正文文本。
23/// DOCX 本质是 ZIP 包,正文在 `word/document.xml` 中。
24pub struct DocxLoader {
25    /// 文件路径
26    path: String,
27}
28
29impl DocxLoader {
30    /// 从文件路径创建加载器
31    pub fn new(path: impl Into<String>) -> Self {
32        Self { path: path.into() }
33    }
34
35    /// 从 DOCX 字节流提取文本
36    ///
37    /// DOCX 是 ZIP 格式,正文在 `word/document.xml` 中,
38    /// 文本在 `<w:t>` 标签内。
39    fn extract_text_from_bytes(data: &[u8]) -> Result<String, LoaderError> {
40        let reader = std::io::Cursor::new(data);
41        let mut archive = zip::ZipArchive::new(reader)
42            .map_err(|e| LoaderError::Other(format!("DOCX 不是有效 ZIP: {}", e)))?;
43
44        // 读取 word/document.xml
45        let mut xml_content = String::new();
46        let mut found = false;
47        for i in 0..archive.len() {
48            let mut file = archive
49                .by_index(i)
50                .map_err(|e| LoaderError::Other(format!("读取 ZIP 条目失败: {}", e)))?;
51            if file.name() == "word/document.xml" {
52                file.read_to_string(&mut xml_content)
53                    .map_err(|e| LoaderError::Other(format!("读取 document.xml 失败: {}", e)))?;
54                found = true;
55                break;
56            }
57        }
58
59        if !found {
60            return Err(LoaderError::Other(
61                "DOCX 中未找到 word/document.xml".to_string(),
62            ));
63        }
64
65        // 提取 <w:t> 标签内容
66        Self::extract_text_from_xml(&xml_content)
67    }
68
69    /// 从 document.xml 提取文本
70    fn extract_text_from_xml(xml: &str) -> Result<String, LoaderError> {
71        // M60: use pre-compiled regexes from LazyLock
72        let mut result = String::new();
73        let mut last_end = 0;
74
75        for cap in WT_REGEX.captures_iter(xml) {
76            let m = cap.get(1).unwrap();
77            // 检查这个 <w:t> 之前是否有 </w:p>(新段落)
78            let before = &xml[last_end..m.start()];
79            if PARA_REGEX.is_match(before) && !result.is_empty() {
80                result.push('\n');
81            } else if !result.is_empty() {
82                // 同一段落内的连续文本
83            }
84            result.push_str(m.as_str());
85            last_end = m.end();
86        }
87
88        Ok(result)
89    }
90}
91
92// zip crate 在 dev-dependencies 中,我们需要在 features 中处理
93// 但为简单起见,直接添加为依赖
94// 注:此处使用 std::io + zip 最小依赖
95
96#[async_trait]
97impl DocumentLoader for DocxLoader {
98    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
99        let data = tokio::task::spawn_blocking({
100            let path = self.path.clone();
101            move || std::fs::read(&path)
102        })
103        .await
104        .map_err(|e| LoaderError::Other(format!("读取文件失败: {}", e)))?
105        .map_err(LoaderError::IoError)?;
106
107        let text = Self::extract_text_from_bytes(&data)?;
108
109        let mut metadata = HashMap::new();
110        metadata.insert("format".to_string(), "docx".to_string());
111        metadata.insert("source".to_string(), self.path.clone());
112
113        Ok(vec![Document {
114            content: text,
115            metadata,
116            id: None,
117        }])
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn test_extract_text_from_xml() {
127        let xml = r#"<?xml version="1.0"?>
128        <w:document>
129            <w:body>
130                <w:p><w:r><w:t>Hello</w:t></w:r><w:r><w:t> World</w:t></w:r></w:p>
131                <w:p><w:r><w:t>Second paragraph</w:t></w:r></w:p>
132            </w:body>
133        </w:document>"#;
134        let text = DocxLoader::extract_text_from_xml(xml).unwrap();
135        assert!(text.contains("Hello World"));
136        assert!(text.contains("Second paragraph"));
137    }
138
139    #[test]
140    fn test_extract_text_from_xml_empty() {
141        let xml = r#"<?xml version="1.0"?><w:document><w:body></w:body></w:document>"#;
142        let text = DocxLoader::extract_text_from_xml(xml).unwrap();
143        assert!(text.is_empty());
144    }
145
146    #[test]
147    fn test_extract_text_from_xml_with_xml_space() {
148        // w:t 可能带 xml:space="preserve" 属性
149        let xml = r#"<w:p><w:r><w:t xml:space="preserve">  spaced  </w:t></w:r></w:p>"#;
150        let text = DocxLoader::extract_text_from_xml(xml).unwrap();
151        assert_eq!(text, "  spaced  ");
152    }
153
154    #[test]
155    fn test_new() {
156        let loader = DocxLoader::new("/path/to/file.docx");
157        assert_eq!(loader.path, "/path/to/file.docx");
158    }
159}