llm_kernel/tokens/
mod.rs

1//! Token estimation for LLM context budgeting.
2//!
3//! Provides a zero-dependency Unicode-script-based heuristic for estimating
4//! token counts, useful for budget management without pulling in tiktoken.
5//!
6//! ```
7//! use llm_kernel::tokens::estimate_tokens;
8//!
9//! let count = estimate_tokens("Hello, world! こんにちは世界");
10//! assert!(count > 0);
11//! ```
12
13/// Thread-safe token budget tracker.
14pub mod budget;
15
16/// Document chunking by sentence boundary and token budget.
17pub mod chunk;
18
19pub use chunk::{ChunkOptions, chunk_text};
20
21/// Characters-per-token ratio lookup using match on Unicode code point ranges.
22/// Compiles to a jump table — O(1) per character instead of linear scan.
23fn char_cpt(ch: char) -> f32 {
24    let cp = ch as u32;
25    match cp {
26        // Emoji emoticons, Misc symbols, Transport, Misc symbols
27        0x1F600..=0x1F64F | 0x1F300..=0x1F5FF | 0x1F680..=0x1F6FF | 0x2600..=0x26FF => 1.0,
28        // Hiragana, Katakana, CJK Unified, Hangul Syllables
29        0x3040..=0x30FF | 0x4E00..=0x9FFF | 0xAC00..=0xD7AF => 1.5,
30        // Arabic, Devanagari, Thai
31        0x0600..=0x06FF | 0x0900..=0x097F | 0x0E00..=0x0E7F => 2.0,
32        // Cyrillic (Russian, Ukrainian, Bulgarian, etc.)
33        0x0400..=0x04FF => 2.0,
34        // Greek and Coptic
35        0x0370..=0x03FF => 2.0,
36        // Hebrew
37        0x0590..=0x05FF => 2.0,
38        _ => DEFAULT_CPT,
39    }
40}
41
42/// Default chars-per-token for Latin/basic ASCII text.
43const DEFAULT_CPT: f32 = 4.0;
44
45/// Token weight contribution for whitespace (roughly 1 token per 4 spaces).
46const WS_WEIGHT: f32 = 0.25;
47
48/// Estimate the number of tokens in a string using Unicode-script heuristics.
49///
50/// This is a rough estimate (±20%) suitable for budget management.
51///
52/// # Example
53///
54/// ```
55/// use llm_kernel::tokens::estimate_tokens;
56/// // ASCII is roughly 1 token per 4 characters; CJK carries more weight.
57/// assert!(estimate_tokens("hello world") > 0);
58/// assert!(estimate_tokens("知識グラフ") > 0);
59/// assert_eq!(estimate_tokens(""), 0);
60/// ```
61pub fn estimate_tokens(text: &str) -> usize {
62    if text.is_empty() {
63        return 0;
64    }
65
66    let mut total_weight: f32 = 0.0;
67
68    for ch in text.chars() {
69        // Whitespace first: `\n`, `\t`, and `\r` are ASCII control chars too,
70        // and skipping them undercounts newline-dense text (code, markdown).
71        if ch.is_whitespace() {
72            total_weight += WS_WEIGHT;
73            continue;
74        }
75        if ch.is_ascii_control() {
76            continue;
77        }
78        total_weight += 1.0 / char_cpt(ch);
79    }
80
81    if total_weight == 0.0 {
82        return 0;
83    }
84
85    // Text with any visible content is at least one token — rounding alone
86    // reports 0 for short strings ("a" weighs 0.25), and a 0-token budget
87    // entry reads as "nothing to send".
88    (total_weight.round() as usize).max(1)
89}
90
91/// Estimate tokens for a single string, returning at least `min`.
92pub fn estimate_tokens_min(text: &str, min: usize) -> usize {
93    estimate_tokens(text).max(min)
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn empty_string() {
102        assert_eq!(estimate_tokens(""), 0);
103    }
104
105    #[test]
106    fn ascii_text() {
107        let tokens = estimate_tokens("Hello, world! This is a test.");
108        // ~30 chars / 4 cpt ≈ 7-8 tokens
109        assert!(tokens > 3 && tokens < 15, "got {tokens}");
110    }
111
112    #[test]
113    fn cjk_text() {
114        let tokens = estimate_tokens("こんにちは世界");
115        // 7 chars / 1.5 cpt ≈ 4-5 tokens
116        assert!(tokens > 2 && tokens < 10, "got {tokens}");
117    }
118
119    #[test]
120    fn mixed_scripts() {
121        let tokens = estimate_tokens("Hello こんにちは مرحبا");
122        assert!(tokens > 0);
123    }
124
125    #[test]
126    fn emoji() {
127        let tokens = estimate_tokens("🎉🚀👍");
128        assert!(tokens >= 2, "got {tokens}");
129    }
130
131    #[test]
132    fn min_clamp() {
133        assert_eq!(estimate_tokens_min("", 5), 5);
134    }
135
136    #[test]
137    fn long_text_proportional() {
138        let short = estimate_tokens("Hello world");
139        let long = estimate_tokens("Hello world Hello world Hello world");
140        assert!(long > short, "long={long} should be > short={short}");
141    }
142
143    #[test]
144    fn cyrillic_text() {
145        let tokens = estimate_tokens("Привет мир");
146        // 8 non-space Cyrillic chars / 2.0 cpt ≈ 4 tokens + whitespace
147        assert!(tokens > 2 && tokens < 10, "got {tokens}");
148    }
149
150    #[test]
151    fn greek_text() {
152        let tokens = estimate_tokens("Γεια σου κόσμε");
153        assert!(tokens > 0 && tokens < 10, "got {tokens}");
154    }
155
156    #[test]
157    fn hebrew_text() {
158        let tokens = estimate_tokens("שלום עולם");
159        assert!(tokens > 0 && tokens < 10, "got {tokens}");
160    }
161
162    #[test]
163    fn whitespace_contributes_tokens() {
164        let no_space = estimate_tokens("abcdef");
165        let with_space = estimate_tokens("a b c d e f");
166        // Whitespace should add some token weight, not zero
167        assert!(
168            with_space > no_space / 2,
169            "with_space={with_space} should not be negligible vs no_space={no_space}"
170        );
171    }
172}
llm_kernel/tokens/mod.rs

llm_kernel/tokens/
mod.rs
