rs-chunks 0.1.0

Fast, high-fidelity document chunking for RAG — a pure-Rust engine covering 36 file formats (Office, OpenDocument, PDF, email, ebooks, notebooks, and more).
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

//! EPUB → chunks/markdown by walking the spine (reading order) and reusing the
//! mature HTML chunker on each content document.

use serde_json::json;

use super::package::EpubPackage;
use crate::formats::html::common::ChunkRecordInput;
use crate::formats::html::page_aware::build_page_aware_chunks;
use crate::formats::html::section::build_section_chunks;
use crate::formats::html::semantic::build_semantic_chunks;
use crate::formats::html::sentence::build_sentence_chunks;
use crate::formats::html::sliding_window::build_sliding_window_chunks;
use crate::formats::html::structural::build_chunks_from_html_bytes;
use crate::formats::html::to_markdown::html_to_markdown_str;

fn build_for_mode(
    bytes: &[u8],
    mode: &str,
    window_size: usize,
    overlap: usize,
    sentences_per_chunk: usize,
    paragraphs_per_page: usize,
) -> Result<Vec<ChunkRecordInput>, String> {
    match mode {
        "default" | "structural" => build_chunks_from_html_bytes(bytes),
        "section" => build_section_chunks(bytes),
        "semantic" => build_semantic_chunks(bytes),
        "sentence" => build_sentence_chunks(bytes, sentences_per_chunk),
        "page_aware" => build_page_aware_chunks(bytes, paragraphs_per_page),
        "sliding_window" => build_sliding_window_chunks(bytes, window_size, overlap),
        other => Err(format!("Unknown EPUB mode: {other}")),
    }
}

fn doc_metadata(pkg: &EpubPackage) -> serde_json::Value {
    json!({
        "source_type": "epub",
        "title": pkg.title,
        "language": pkg.language,
        "creator": pkg.creator,
        "identifier": pkg.identifier,
        "epub_version": pkg.version,
        "spine_count": pkg.spine.len(),
    })
}

/// Chunk the whole book: each spine XHTML through the HTML chunker, merged in
/// reading order, with book-level and per-chapter provenance metadata injected.
pub fn chunk_package(
    pkg: &EpubPackage,
    mode: &str,
    window_size: usize,
    overlap: usize,
    sentences_per_chunk: usize,
    paragraphs_per_page: usize,
) -> Result<Vec<ChunkRecordInput>, String> {
    let meta = doc_metadata(pkg);
    let mut out: Vec<ChunkRecordInput> = Vec::new();

    for (spine_index, doc) in pkg.spine.iter().enumerate() {
        let mut records = build_for_mode(
            &doc.bytes,
            mode,
            window_size,
            overlap,
            sentences_per_chunk,
            paragraphs_per_page,
        )
        // An individual empty/nav-only XHTML shouldn't abort the whole book.
        .unwrap_or_default();

        for rec in &mut records {
            if let serde_json::Value::Object(map) = &mut rec.metadata {
                map.insert("document_metadata".to_string(), meta.clone());
                map.insert("spine_index".to_string(), json!(spine_index));
                map.insert("href".to_string(), json!(doc.href));
            }
        }
        out.append(&mut records);
    }
    Ok(out)
}

/// Convert the whole book to a single markdown string (reading order).
pub fn package_to_markdown(pkg: &EpubPackage) -> String {
    let mut parts: Vec<String> = Vec::new();
    if let Some(title) = &pkg.title {
        parts.push(format!("# {title}"));
    }
    for doc in &pkg.spine {
        let html = String::from_utf8_lossy(&doc.bytes);
        let md = html_to_markdown_str(&html);
        if !md.trim().is_empty() {
            parts.push(md.trim().to_string());
        }
    }
    parts.join("\n\n---\n\n").trim().to_string()
}