text-splitter 0.30.1

Split text into semantic chunks, up to a desired chunk size. Supports calculating length by characters and tokens, and is callable from Rust and Python.
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

//! Test for `CodeSplitter` behavior.
use std::fs;

use fake::{Fake, Faker};
use itertools::Itertools;
use more_asserts::assert_le;
#[cfg(feature = "code")]
use text_splitter::{ChunkConfig, CodeSplitter};

#[cfg(feature = "code")]
#[test]
fn random_chunk_size() {
    let text = fs::read_to_string("tests/inputs/code/hashbrown_set_rs.txt").unwrap();

    for _ in 0..10 {
        let max_characters = Faker.fake();
        let splitter = CodeSplitter::new(
            tree_sitter_rust::LANGUAGE,
            ChunkConfig::new(max_characters).with_trim(false),
        )
        .unwrap();
        let chunks = splitter.chunks(&text).collect::<Vec<_>>();

        assert_eq!(chunks.join(""), text);
        for chunk in chunks {
            assert_le!(chunk.chars().count(), max_characters);
        }
    }
}

#[cfg(feature = "code")]
#[test]
fn random_chunk_indices_increase() {
    let text = fs::read_to_string("tests/inputs/code/hashbrown_set_rs.txt").unwrap();

    for _ in 0..10 {
        let max_characters = Faker.fake::<usize>();
        let splitter =
            CodeSplitter::new(tree_sitter_rust::LANGUAGE, ChunkConfig::new(max_characters))
                .unwrap();
        let indices = splitter.chunk_indices(&text).map(|(i, _)| i);

        assert!(indices.tuple_windows().all(|(a, b)| a < b));
    }
}

#[cfg(feature = "code")]
#[test]
fn can_handle_invalid_code() {
    let text = "No code here";

    let splitter = CodeSplitter::new(
        tree_sitter_rust::LANGUAGE,
        ChunkConfig::new(5).with_trim(false),
    )
    .unwrap();
    let chunks = splitter.chunks(text).collect::<Vec<_>>();

    assert_eq!(chunks.join(""), text);
}

#[cfg(feature = "code")]
#[test]
fn groups_functions() {
    let text = "
fn fn1() {}
fn fn2() {}
fn fn3() {}
fn fn4() {}";

    let splitter = CodeSplitter::new(tree_sitter_rust::LANGUAGE, ChunkConfig::new(24)).unwrap();
    let chunks = splitter.chunks(text).collect::<Vec<_>>();

    assert_eq!(
        chunks,
        ["fn fn1() {}\nfn fn2() {}", "fn fn3() {}\nfn fn4() {}"]
    );
}

#[cfg(feature = "code")]
#[test]
fn groups_functions_with_children() {
    let text = "
fn fn1() {}
fn fn2() {
    let x = 4;
}
fn fn3() {}
fn fn4() {}";

    let splitter = CodeSplitter::new(tree_sitter_rust::LANGUAGE, ChunkConfig::new(30)).unwrap();
    let chunks = splitter.chunks(text).collect::<Vec<_>>();

    assert_eq!(
        chunks,
        [
            "fn fn1() {}",
            "fn fn2() {\n    let x = 4;\n}",
            "fn fn3() {}\nfn fn4() {}"
        ]
    );
}

#[cfg(feature = "code")]
#[test]
fn functions_overlap() {
    let text = "
fn fn1() {}
fn fn2() {}
fn fn3() {}
fn fn4() {}";

    let splitter = CodeSplitter::new(
        tree_sitter_rust::LANGUAGE,
        ChunkConfig::new(24).with_overlap(12).unwrap(),
    )
    .unwrap();
    let chunks = splitter.chunks(text).collect::<Vec<_>>();

    assert_eq!(
        chunks,
        [
            "fn fn1() {}\nfn fn2() {}",
            "fn fn2() {}\nfn fn3() {}",
            "fn fn3() {}\nfn fn4() {}"
        ]
    );
}