Skip to main content

steeldb/
docs.rs

1//! Any-doc bridge — lower PDF / DOCX / PPTX / HTML / MD / TXT to plain text, natively in Rust, so the
2//! folder ingest can project every readable document into the hypergraph (not just tabular/text files).
3//! The extracted text feeds the same `TextEngine` (SPO tagger + SPLADE + gazetteer) as `.txt`.
4//!
5//! Text-layer extraction only: scanned/image PDFs (no text layer) yield little — OCR (the bundled
6//! PP-OCRv6 models) is the follow-up. Gated behind the `docs` feature (implies `onnx`).
7
8use std::io::Read;
9use std::path::Path;
10
11/// The document extensions this bridge can lower to text.
12pub const DOC_EXTS: &[&str] = &["pdf", "docx", "pptx", "html", "htm", "md", "markdown", "txt", "text"];
13
14pub fn is_doc_ext(ext: &str) -> bool {
15    DOC_EXTS.contains(&ext)
16}
17
18/// Extract plain text from a document by extension. `Ok(None)` = unsupported type; `Err` = read/parse
19/// failure the caller can report and skip.
20pub fn extract_text(path: &Path) -> std::io::Result<Option<String>> {
21    let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("").to_lowercase();
22    let text = match ext.as_str() {
23        "txt" | "text" | "md" | "markdown" => Some(std::fs::read_to_string(path)?),
24        "html" | "htm" => Some(strip_html(&std::fs::read_to_string(path)?)),
25        "docx" => Some(ooxml_text(path, "word/document.xml", &[])?),
26        "pptx" => Some(ooxml_text(path, "", &["ppt/slides/slide"])?),
27        "pdf" => extract_pdf(path),
28        _ => None,
29    };
30    Ok(text.map(|t| t.trim().to_string()).filter(|t| !t.is_empty()))
31}
32
33fn extract_pdf(path: &Path) -> Option<String> {
34    // Prefer poppler's `pdftotext` when available: it handles CID/CJK (Identity-H) fonts that the
35    // pure-Rust path chokes on. Falls back to pdf-extract for Latin text layers when poppler is absent.
36    if let Ok(out) = std::process::Command::new("pdftotext").arg("-q").arg(path).arg("-").output() {
37        if out.status.success() {
38            let t = String::from_utf8_lossy(&out.stdout).to_string();
39            if !t.trim().is_empty() {
40                return Some(t);
41            }
42        }
43    }
44    // pure-Rust fallback: pdf-extract panics on some encodings — contain AND silence it so one bad file
45    // neither crashes nor spams folder ingest.
46    let owned = path.to_path_buf();
47    let prev = std::panic::take_hook();
48    std::panic::set_hook(Box::new(|_| {}));
49    let r = std::panic::catch_unwind(move || pdf_extract::extract_text(&owned).ok()).ok().flatten();
50    std::panic::set_hook(prev);
51    if let Some(t) = &r {
52        if !t.trim().is_empty() {
53            return r;
54        }
55    }
56    // no text layer → scanned/image PDF: OCR the rasterized pages (PP-OCRv6).
57    #[cfg(feature = "ocr")]
58    {
59        return crate::ocr::ocr_pdf(path);
60    }
61    #[cfg(not(feature = "ocr"))]
62    r
63}
64
65/// Read text from an OOXML (zip) document. Either an exact `entry` (docx `word/document.xml`) or every
66/// entry whose name starts with one of `prefixes` (pptx `ppt/slides/slideN.xml`), in sorted order.
67fn ooxml_text(path: &Path, entry: &str, prefixes: &[&str]) -> std::io::Result<String> {
68    let file = std::fs::File::open(path)?;
69    let mut zip = zip::ZipArchive::new(file).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
70
71    let mut names: Vec<String> = (0..zip.len())
72        .filter_map(|i| zip.by_index(i).ok().map(|f| f.name().to_string()))
73        .filter(|n| (!entry.is_empty() && n == entry) || prefixes.iter().any(|p| n.starts_with(p) && n.ends_with(".xml")))
74        .collect();
75    names.sort();
76
77    let mut out = String::new();
78    for name in names {
79        let mut xml = String::new();
80        if let Ok(mut f) = zip.by_name(&name) {
81            if f.read_to_string(&mut xml).is_ok() {
82                out.push_str(&xml_text(&xml));
83                out.push('\n');
84            }
85        }
86    }
87    Ok(out)
88}
89
90/// Pull the visible text out of a WordprocessingML / PresentationML part: the `<w:t>` / `<a:t>` runs
91/// (local name `t`), with a newline at each paragraph end (`<w:p>` / `<a:p>`, local name `p`).
92fn xml_text(xml: &str) -> String {
93    use quick_xml::events::Event;
94    use quick_xml::reader::Reader;
95
96    let mut reader = Reader::from_str(xml);
97    let mut out = String::new();
98    let mut in_t = false;
99    let local = |name: &[u8]| -> Vec<u8> { name.rsplit(|&b| b == b':').next().unwrap_or(name).to_vec() };
100
101    loop {
102        match reader.read_event() {
103            Ok(Event::Start(e)) if local(e.name().as_ref()) == b"t" => in_t = true,
104            Ok(Event::End(e)) => match local(e.name().as_ref()).as_slice() {
105                b"t" => in_t = false,
106                b"p" => out.push('\n'),
107                _ => {}
108            },
109            Ok(Event::Text(t)) if in_t => {
110                if let Ok(s) = t.unescape() {
111                    out.push_str(&s);
112                }
113            }
114            Ok(Event::Eof) | Err(_) => break,
115            _ => {}
116        }
117    }
118    out
119}
120
121/// Minimal HTML → text: drop `<script>`/`<style>` blocks, strip tags, decode common entities, collapse
122/// whitespace. Enough for retrieval; not a full renderer.
123pub fn strip_html(html: &str) -> String {
124    let mut s = html.to_string();
125    for tag in ["script", "style"] {
126        loop {
127            let lower = s.to_lowercase();
128            let open = format!("<{tag}");
129            let close = format!("</{tag}>");
130            match (lower.find(&open), lower.find(&close)) {
131                (Some(a), Some(b)) if b > a => s.replace_range(a..b + close.len(), " "),
132                _ => break,
133            }
134        }
135    }
136    // strip tags
137    let mut out = String::with_capacity(s.len());
138    let mut in_tag = false;
139    for ch in s.chars() {
140        match ch {
141            '<' => in_tag = true,
142            '>' => in_tag = false,
143            _ if !in_tag => out.push(ch),
144            _ => {}
145        }
146    }
147    // decode a handful of entities + collapse whitespace
148    let out = out
149        .replace("&amp;", "&")
150        .replace("&lt;", "<")
151        .replace("&gt;", ">")
152        .replace("&quot;", "\"")
153        .replace("&#39;", "'")
154        .replace("&nbsp;", " ");
155    out.split_whitespace().collect::<Vec<_>>().join(" ")
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn html_stripped() {
164        let h = "<html><head><style>a{color:red}</style></head><body><h1>Title</h1><p>Hello &amp; welcome</p><script>x=1</script></body></html>";
165        let t = strip_html(h);
166        assert!(t.contains("Title"));
167        assert!(t.contains("Hello & welcome"));
168        assert!(!t.contains("color"));
169        assert!(!t.contains("x=1"));
170    }
171
172    #[test]
173    fn xml_runs_extracted() {
174        let x = r#"<w:document><w:body><w:p><w:r><w:t>Hello</w:t></w:r><w:r><w:t xml:space="preserve"> world</w:t></w:r></w:p><w:p><w:r><w:t>Second</w:t></w:r></w:p></w:body></w:document>"#;
175        let t = super::xml_text(x);
176        assert!(t.contains("Hello world"));
177        assert!(t.contains("Second"));
178    }
179}