telos_agent/model/
tokens.rs

1//! Token counting using OpenAI-compatible `cl100k_base` BPE tokenizer.
2//!
3//! All providers supported by tiny-agent-core (DeepSeek, Kimi, etc.) use
4//! tokenizers that are highly compatible with `cl100k_base`. This module
5//! provides a shared, lazily-loaded counter that replaces the old 4-char-per-
6//! token heuristic with an accurate (±5%) token count.
7
8use std::sync::OnceLock;
9use tiktoken_rs::CoreBPE;
10
11/// Global `cl100k_base` BPE instance — loaded once and reused.
12static CL100K: OnceLock<Option<CoreBPE>> = OnceLock::new();
13
14fn get_bpe() -> Option<&'static CoreBPE> {
15    CL100K
16        .get_or_init(|| {
17            let bpe = tiktoken_rs::cl100k_base().ok();
18            if bpe.is_none() {
19                tracing::warn!(
20                    "tiktoken-rs failed to load cl100k_base vocabulary; \
21                     falling back to char/4 heuristic"
22                );
23            }
24            bpe
25        })
26        .as_ref()
27}
28
29/// Count tokens in `text` using the `cl100k_base` tokenizer.
30///
31/// If the tokenizer fails to load (e.g. corrupted vocabulary file), falls back
32/// to the legacy heuristic: `ceil(text.len() / 4)`.
33pub fn count_tokens(text: &str) -> usize {
34    if text.is_empty() {
35        return 0;
36    }
37    match get_bpe() {
38        Some(bpe) => bpe.encode_ordinary(text).len(),
39        None => {
40            // Legacy fallback — keep the agent running even if the tokenizer
41            // vocabulary is somehow unavailable.
42            (text.len() as f64 / 4.0_f64).ceil() as usize
43        }
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50
51    #[test]
52    fn empty_string_is_zero() {
53        assert_eq!(count_tokens(""), 0);
54    }
55
56    #[test]
57    fn english_text_is_counted() {
58        // "hello world" should be ~2 tokens with cl100k_base
59        let count = count_tokens("hello world");
60        assert!(count > 0, "token count should be > 0");
61        assert!(count <= 10, "short text shouldn't be many tokens");
62    }
63
64    #[test]
65    fn chinese_text_is_counted() {
66        // Chinese text: each character is typically 1-2 tokens in cl100k_base,
67        // NOT 0.25 as the old chars/4 heuristic would give.
68        let count = count_tokens("你好世界");
69        assert!(count > 0);
70        // Old heuristic: 4 chars / 4 = 1 token — wildly underestimates.
71        // Actual should be 2-8 tokens.
72        assert!(count >= 2, "Chinese chars should be > 1 token (old heuristic was wrong)");
73    }
74
75    #[test]
76    fn code_text_is_counted() {
77        let count = count_tokens("fn main() { println!(\"hi\"); }");
78        assert!(count > 0);
79        assert!(count <= 50);
80    }
81
82    #[test]
83    fn larger_text_is_reasonable() {
84        let text = "This is a longer piece of text that should be counted. ".repeat(100);
85        let count = count_tokens(&text);
86        let chars = text.len();
87        // Token count should be materially less than chars
88        assert!(count < chars);
89        // And should be roughly in the expected range
90        assert!(count > chars / 5, "token count should be within reasonable range");
91    }
92
93    #[test]
94    fn count_is_stable() {
95        let text = "The quick brown fox jumps over the lazy dog.";
96        let a = count_tokens(text);
97        let b = count_tokens(text);
98        assert_eq!(a, b, "same text should always give same count");
99    }
100}
telos_agent/model/tokens.rs

telos_agent/model/
tokens.rs
