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.
51pub fn estimate_tokens(text: &str) -> usize {
52    if text.is_empty() {
53        return 0;
54    }
55
56    let mut total_weight: f32 = 0.0;
57
58    for ch in text.chars() {
59        if ch.is_ascii_control() {
60            continue;
61        }
62        if ch.is_whitespace() {
63            total_weight += WS_WEIGHT;
64            continue;
65        }
66        total_weight += 1.0 / char_cpt(ch);
67    }
68
69    if total_weight == 0.0 {
70        return 0;
71    }
72
73    total_weight.round() as usize
74}
75
76/// Estimate tokens for a single string, returning at least `min`.
77pub fn estimate_tokens_min(text: &str, min: usize) -> usize {
78    estimate_tokens(text).max(min)
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    #[test]
86    fn empty_string() {
87        assert_eq!(estimate_tokens(""), 0);
88    }
89
90    #[test]
91    fn ascii_text() {
92        let tokens = estimate_tokens("Hello, world! This is a test.");
93        // ~30 chars / 4 cpt ≈ 7-8 tokens
94        assert!(tokens > 3 && tokens < 15, "got {tokens}");
95    }
96
97    #[test]
98    fn cjk_text() {
99        let tokens = estimate_tokens("こんにちは世界");
100        // 7 chars / 1.5 cpt ≈ 4-5 tokens
101        assert!(tokens > 2 && tokens < 10, "got {tokens}");
102    }
103
104    #[test]
105    fn mixed_scripts() {
106        let tokens = estimate_tokens("Hello こんにちは مرحبا");
107        assert!(tokens > 0);
108    }
109
110    #[test]
111    fn emoji() {
112        let tokens = estimate_tokens("🎉🚀👍");
113        assert!(tokens >= 2, "got {tokens}");
114    }
115
116    #[test]
117    fn min_clamp() {
118        assert_eq!(estimate_tokens_min("", 5), 5);
119    }
120
121    #[test]
122    fn long_text_proportional() {
123        let short = estimate_tokens("Hello world");
124        let long = estimate_tokens("Hello world Hello world Hello world");
125        assert!(long > short, "long={long} should be > short={short}");
126    }
127
128    #[test]
129    fn cyrillic_text() {
130        let tokens = estimate_tokens("Привет мир");
131        // 8 non-space Cyrillic chars / 2.0 cpt ≈ 4 tokens + whitespace
132        assert!(tokens > 2 && tokens < 10, "got {tokens}");
133    }
134
135    #[test]
136    fn greek_text() {
137        let tokens = estimate_tokens("Γεια σου κόσμε");
138        assert!(tokens > 0 && tokens < 10, "got {tokens}");
139    }
140
141    #[test]
142    fn hebrew_text() {
143        let tokens = estimate_tokens("שלום עולם");
144        assert!(tokens > 0 && tokens < 10, "got {tokens}");
145    }
146
147    #[test]
148    fn whitespace_contributes_tokens() {
149        let no_space = estimate_tokens("abcdef");
150        let with_space = estimate_tokens("a b c d e f");
151        // Whitespace should add some token weight, not zero
152        assert!(
153            with_space > no_space / 2,
154            "with_space={with_space} should not be negligible vs no_space={no_space}"
155        );
156    }
157}
llm_kernel/tokens/mod.rs

llm_kernel/tokens/
mod.rs
