vtcode_commons/
tokens.rs

1#![expect(
2    clippy::indexing_slicing,
3    clippy::string_slice,
4    reason = "Token and UTF-8 boundaries come from tokenizer output and the byte truncation helper."
5)]
6
7//! Token counting via tiktoken BPE tokenizer.
8//!
9//! All token estimation goes through [`tiktoken`]'s `cl100k_base` encoding
10//! (GPT-4, GPT-3.5-turbo). BPE tokenizers are similar enough across providers
11//! that this gives reasonable accuracy for Anthropic, Gemini, and others.
12//!
13//! Provider-reported exact token counts (from API responses) should always be
14//! preferred when available. This module is for pre-call budget estimation and
15//! offline token sizing where no provider response exists yet.
16
17use std::sync::OnceLock;
18use tiktoken::CoreBpe;
19
20/// Return the process-global `cl100k_base` BPE instance, if it could be loaded.
21///
22/// Loaded once on first call; all subsequent calls return the same reference.
23/// Returns `None` only if the builtin encoding fails to load, in which case
24/// callers fall back to a character-based heuristic rather than panicking.
25fn bpe() -> Option<&'static CoreBpe> {
26    static BPE: OnceLock<Option<&'static CoreBpe>> = OnceLock::new();
27    *BPE.get_or_init(|| tiktoken::get_encoding("cl100k_base"))
28}
29
30/// Approximate token count from character length (~4 chars per token).
31fn heuristic_token_count(text: &str) -> usize {
32    text.len().div_ceil(4)
33}
34
35/// Count the number of tokens in `text` using tiktoken BPE.
36///
37/// Returns 0 for empty strings. Falls back to a character-based heuristic if
38/// the BPE tokenizer is unavailable.
39pub fn estimate_tokens(text: &str) -> usize {
40    if text.is_empty() {
41        return 0;
42    }
43    match bpe() {
44        Some(bpe) => bpe.count(text),
45        None => heuristic_token_count(text),
46    }
47}
48
49/// Truncate `text` to at most `max_tokens` tokens.
50///
51/// Decodes the truncated token sequence back to text so the result is always
52/// valid UTF-8 with no mid-token corruption. Falls back to byte-level
53/// truncation if BPE decode fails (should not happen in practice).
54pub fn truncate_to_tokens(text: &str, max_tokens: usize) -> String {
55    if max_tokens == 0 || text.is_empty() {
56        return String::new();
57    }
58    // Byte-level fallback used when BPE is unavailable or decode fails.
59    let byte_truncate = || {
60        let end = (max_tokens * 4).min(text.len());
61        let mut end = end;
62        while end > 0 && !text.is_char_boundary(end) {
63            end -= 1;
64        }
65        let mut result = text[..end].to_string();
66        result.push_str("...");
67        result
68    };
69    let Some(bpe) = bpe() else {
70        return byte_truncate();
71    };
72    let tokens = bpe.encode_with_special_tokens(text);
73    if tokens.len() <= max_tokens {
74        return text.to_string();
75    }
76    bpe.decode_to_string(&tokens[..max_tokens]).unwrap_or_else(|_| byte_truncate())
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn empty_string_returns_zero() {
85        assert_eq!(estimate_tokens(""), 0);
86        assert_eq!(truncate_to_tokens("", 10), "");
87    }
88
89    #[test]
90    fn count_is_reasonable() {
91        let count = estimate_tokens("Hello, how are you today?");
92        assert!((4..=12).contains(&count), "count={count}");
93    }
94
95    #[test]
96    fn truncate_respects_limit() {
97        let text = "the quick brown fox jumps over the lazy dog";
98        let truncated = truncate_to_tokens(text, 5);
99        let count = estimate_tokens(&truncated);
100        assert!(count <= 5 + 1, "count={count} should be <= 6");
101    }
102
103    #[test]
104    fn truncate_zero_returns_empty() {
105        assert_eq!(truncate_to_tokens("hello", 0), "");
106    }
107
108    #[test]
109    fn code_and_prose_tokenize() {
110        let code = "fn main() { println!(\"hello\"); }";
111        let prose = "the main function prints hello to console";
112        assert!(estimate_tokens(code) > 0);
113        assert!(estimate_tokens(prose) > 0);
114    }
115
116    #[test]
117    fn json_tokenizes() {
118        let json = r#"{"name":"test","value":123,"nested":{"key":"value"}}"#;
119        let count = estimate_tokens(json);
120        assert!((10..=40).contains(&count), "json count={count}");
121    }
122}
vtcode_commons/tokens.rs

vtcode_commons/
tokens.rs
