rs-chunks 0.6.1

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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

//! Chunking modes and the per-format options bag.
//!
//! [`ChunkOptions`] is consumed by the per-format `chunk_with_options` entry
//! points (e.g. `formats::docx::chunk_with_options`); its field defaults mirror
//! the keyword-argument defaults of the Python `get_chunks()`. The
//! source-agnostic dispatch layer (`dispatch::get_chunks` and friends) does
//! **not** take a `ChunkOptions` — it takes positional arguments (`mode: &str`,
//! `window_size`, `overlap`, `sentences_per_chunk`, `paragraphs_per_page`) to
//! match the Python entry point one-for-one. There is no dispatch-level
//! streaming API; streaming is per-format (`formats::<fmt>::stream`).

/// The chunking strategies across formats. Not every mode applies to every
/// format; each format validates and maps `mode` onto the strategies it
/// supports (e.g. spreadsheets use `Row`/`Table`/`Sheet`).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ChunkMode {
    /// Format's natural default strategy.
    Default,
    Section,
    Semantic,
    Sentence,
    SlidingWindow,
    PageAware,
    Structural,
    // Spreadsheet / delimited strategies
    Row,
    Table,
    Sheet,
}

impl ChunkMode {
    pub fn as_str(self) -> &'static str {
        match self {
            ChunkMode::Default => "default",
            ChunkMode::Section => "section",
            ChunkMode::Semantic => "semantic",
            ChunkMode::Sentence => "sentence",
            ChunkMode::SlidingWindow => "sliding_window",
            ChunkMode::PageAware => "page_aware",
            ChunkMode::Structural => "structural",
            ChunkMode::Row => "row",
            ChunkMode::Table => "table",
            ChunkMode::Sheet => "sheet",
        }
    }

    /// Parse a mode string. Inherent convenience that delegates to the
    /// [`std::str::FromStr`] impl but returns `Option` (kept for
    /// backwards-compatibility with existing callers).
    #[allow(clippy::should_implement_trait)]
    pub fn from_str(s: &str) -> Option<ChunkMode> {
        <ChunkMode as std::str::FromStr>::from_str(s).ok()
    }
}

impl std::str::FromStr for ChunkMode {
    type Err = crate::error::ChunkError;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        Ok(match s {
            "default" => ChunkMode::Default,
            "section" => ChunkMode::Section,
            "semantic" => ChunkMode::Semantic,
            "sentence" => ChunkMode::Sentence,
            "sliding_window" => ChunkMode::SlidingWindow,
            "page_aware" => ChunkMode::PageAware,
            "structural" => ChunkMode::Structural,
            "row" => ChunkMode::Row,
            "table" => ChunkMode::Table,
            "sheet" => ChunkMode::Sheet,
            other => {
                return Err(crate::error::ChunkError::InvalidArg(format!(
                    "unknown chunk mode '{other}'"
                )))
            }
        })
    }
}

impl Default for ChunkMode {
    fn default() -> Self {
        ChunkMode::Default
    }
}

/// Unified options. Field defaults match the Python API defaults so ported
/// tests observe identical behaviour.
#[derive(Debug, Clone)]
pub struct ChunkOptions {
    pub mode: ChunkMode,
    pub window_size: usize,
    pub overlap: usize,
    pub sentences_per_chunk: usize,
    pub paragraphs_per_page: usize,
    // Delimited / spreadsheet knobs
    pub rows_per_chunk: usize,
    pub include_headers: bool,
    pub delimiter: Option<u8>,
    pub encoding: String,
    pub skip_empty_rows: bool,
}

impl Default for ChunkOptions {
    fn default() -> Self {
        Self {
            mode: ChunkMode::Default,
            window_size: 3,
            overlap: 1,
            sentences_per_chunk: 3,
            paragraphs_per_page: 15,
            rows_per_chunk: 10,
            include_headers: true,
            delimiter: None,
            encoding: "utf-8".to_string(),
            skip_empty_rows: true,
        }
    }
}

impl ChunkOptions {
    pub fn new(mode: ChunkMode) -> Self {
        Self {
            mode,
            ..Default::default()
        }
    }

    pub fn with_window(mut self, window_size: usize, overlap: usize) -> Self {
        self.window_size = window_size;
        self.overlap = overlap;
        self
    }
}