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/// Format a token count with a compact decimal suffix for transcript labels.
50///
51/// Counts below one thousand stay exact. Larger counts are rounded to one
52/// decimal place using `k`, `M`, or `B` suffixes.
53pub fn format_token_count(tokens: usize) -> String {
54    let (divisor, suffix) = match tokens {
55        999_500_000.. => (1_000_000_000, "B"),
56        999_500.. => (1_000_000, "M"),
57        1_000.. => (1_000, "k"),
58        _ => return tokens.to_string(),
59    };
60
61    let whole = tokens / divisor;
62    let decimal_step = divisor / 10;
63    // Divide into tenths before rounding; multiplying a billion-scale
64    // remainder by ten would overflow on 32-bit targets.
65    let tenths = (tokens % divisor + decimal_step / 2) / decimal_step;
66    let (whole, tenths) = if tenths >= 10 {
67        (whole.saturating_add(1), 0)
68    } else {
69        (whole, tenths)
70    };
71    if tenths == 0 {
72        format!("{whole}{suffix}")
73    } else {
74        format!("{whole}.{tenths}{suffix}")
75    }
76}
77
78/// Truncate `text` to at most `max_tokens` tokens.
79///
80/// Decodes the truncated token sequence back to text so the result is always
81/// valid UTF-8 with no mid-token corruption. Falls back to byte-level
82/// truncation if BPE decode fails (should not happen in practice).
83pub fn truncate_to_tokens(text: &str, max_tokens: usize) -> String {
84    if max_tokens == 0 || text.is_empty() {
85        return String::new();
86    }
87    // Byte-level fallback used when BPE is unavailable or decode fails.
88    let byte_truncate = || {
89        let end = text.floor_char_boundary(max_tokens.saturating_mul(4));
90        let mut result = text[..end].to_string();
91        result.push_str("...");
92        result
93    };
94    let Some(bpe) = bpe() else {
95        return byte_truncate();
96    };
97    let tokens = bpe.encode_with_special_tokens(text);
98    if tokens.len() <= max_tokens {
99        return text.to_string();
100    }
101    bpe.decode_to_string(&tokens[..max_tokens]).unwrap_or_else(|_| byte_truncate())
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn empty_string_returns_zero() {
110        assert_eq!(estimate_tokens(""), 0);
111        assert_eq!(truncate_to_tokens("", 10), "");
112    }
113
114    #[test]
115    fn count_is_reasonable() {
116        let count = estimate_tokens("Hello, how are you today?");
117        assert!((4..=12).contains(&count), "count={count}");
118    }
119
120    #[test]
121    fn truncate_respects_limit() {
122        let text = "the quick brown fox jumps over the lazy dog";
123        let truncated = truncate_to_tokens(text, 5);
124        let count = estimate_tokens(&truncated);
125        assert!(count <= 5 + 1, "count={count} should be <= 6");
126    }
127
128    #[test]
129    fn truncate_zero_returns_empty() {
130        assert_eq!(truncate_to_tokens("hello", 0), "");
131    }
132
133    #[test]
134    fn code_and_prose_tokenize() {
135        let code = "fn main() { println!(\"hello\"); }";
136        let prose = "the main function prints hello to console";
137        assert!(estimate_tokens(code) > 0);
138        assert!(estimate_tokens(prose) > 0);
139    }
140
141    #[test]
142    fn json_tokenizes() {
143        let json = r#"{"name":"test","value":123,"nested":{"key":"value"}}"#;
144        let count = estimate_tokens(json);
145        assert!((10..=40).contains(&count), "json count={count}");
146    }
147
148    #[test]
149    fn token_count_format_is_compact_and_rounds_across_suffix_boundaries() {
150        assert_eq!(format_token_count(999), "999");
151        assert_eq!(format_token_count(1_000), "1k");
152        assert_eq!(format_token_count(1_050), "1.1k");
153        assert_eq!(format_token_count(6_100), "6.1k");
154        assert_eq!(format_token_count(999_500), "1M");
155        assert_eq!(format_token_count(15_600_000), "15.6M");
156        assert_eq!(format_token_count(999_500_000), "1B");
157        assert_eq!(format_token_count(1_550_000_000), "1.6B");
158        assert_eq!(format_token_count(1_549_999_999), "1.5B");
159    }
160}
vtcode_commons/tokens.rs

vtcode_commons/
tokens.rs
