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

llm_kernel/tokens/
mod.rs
