1use 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
16static 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
20pub struct DocxLoader {
25 path: String,
27}
28
29impl DocxLoader {
30 pub fn new(path: impl Into<String>) -> Self {
32 Self { path: path.into() }
33 }
34
35 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 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 Self::extract_text_from_xml(&xml_content)
67 }
68
69 fn extract_text_from_xml(xml: &str) -> Result<String, LoaderError> {
71 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 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 }
84 result.push_str(m.as_str());
85 last_end = m.end();
86 }
87
88 Ok(result)
89 }
90}
91
92#[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 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}