av_core/
tokens.rs

1//! Deterministic approximate tokenizer.
2//!
3//! Used for budgets, token-velocity tracking, and compression ratios. This is
4//! an *approximation* (documented, deliberate): exact counts differ per
5//! provider/model and arrive with responses; we record those separately when
6//! present. Properties guaranteed (and property-tested):
7//!
8//! - deterministic;
9//! - monotone: appending text never lowers the count;
10//! - Unicode-safe (multi-byte chars never split or panic);
11//! - zero for the empty string.
12//!
13//! Heuristic: ASCII words contribute `ceil(len/4)` tokens (the ~4 chars/token
14//! BPE rule of thumb), each punctuation/symbol char is one token, and each CJK
15//! or other non-ASCII alphabetic char is one token.
16
17/// Approximate token count for `text`.
18pub fn approx_tokens(text: &str) -> u64 {
19    let mut tokens: u64 = 0;
20    let mut ascii_run: u64 = 0;
21    for ch in text.chars() {
22        if ch.is_ascii_alphanumeric() {
23            ascii_run += 1;
24        } else {
25            tokens += ascii_run.div_ceil(4);
26            ascii_run = 0;
27            if ch.is_whitespace() {
28                continue;
29            }
30            // Punctuation, symbols, CJK, emoji: one token each.
31            tokens += 1;
32        }
33    }
34    tokens + ascii_run.div_ceil(4)
35}
36
37/// Approximate token count for a serialized JSON value.
38///
39/// Reuses a thread-local byte buffer for JSON serialization, so on repeated
40/// calls no per-call allocation happens (the buffer grows to the peak size
41/// once and then stays). Called on every request through `prepare_chat`.
42pub fn approx_tokens_json(value: &serde_json::Value) -> u64 {
43    use std::cell::RefCell;
44    thread_local! {
45        static BUF: RefCell<Vec<u8>> = const { RefCell::new(Vec::new()) };
46    }
47    BUF.with(|cell| {
48        let mut buf = cell.borrow_mut();
49        buf.clear();
50        if serde_json::to_writer(&mut *buf, value).is_err() {
51            return 0;
52        }
53        match std::str::from_utf8(&buf) {
54            Ok(s) => approx_tokens(s),
55            Err(_) => 0,
56        }
57    })
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63    use proptest::prelude::*;
64
65    #[test]
66    fn empty_is_zero() {
67        assert_eq!(approx_tokens(""), 0);
68    }
69
70    #[test]
71    fn simple_words() {
72        // "hello world" => hello(2) + world(2)
73        assert_eq!(approx_tokens("hello world"), 4);
74    }
75
76    #[test]
77    fn punctuation_counts() {
78        assert_eq!(approx_tokens("a,b"), 3); // a(1) , (1) b(1)
79    }
80
81    #[test]
82    fn cjk_one_per_char() {
83        assert_eq!(approx_tokens("日本語"), 3);
84    }
85
86    #[test]
87    fn emoji_do_not_panic() {
88        assert!(approx_tokens("🎉🎉🎉") >= 3);
89    }
90
91    #[test]
92    fn approx_tokens_json_reflects_payload_size() {
93        // Catches `approx_tokens_json -> 0 / 1` stubs: a non-empty payload
94        // must produce more tokens than an empty one.
95        let empty = approx_tokens_json(&serde_json::json!({}));
96        let small = approx_tokens_json(&serde_json::json!({"k": "v"}));
97        let bigger = approx_tokens_json(&serde_json::json!({
98            "messages": [
99                {"role": "user", "content": "hello world"},
100                {"role": "assistant", "content": "hi back"},
101            ]
102        }));
103        assert!(small > empty, "small={small} not > empty={empty}");
104        assert!(bigger > small, "bigger={bigger} not > small={small}");
105        assert!(bigger >= 4);
106    }
107
108    proptest! {
109        #[test]
110        fn monotone_under_append(a in ".{0,200}", b in ".{0,200}") {
111            let joined = format!("{a}{b}");
112            prop_assert!(approx_tokens(&joined) >= approx_tokens(&a));
113        }
114
115        #[test]
116        fn never_panics(s in "\\PC{0,500}") {
117            let _ = approx_tokens(&s);
118        }
119
120        #[test]
121        fn nonempty_ascii_word_is_positive(s in "[a-zA-Z0-9]{1,100}") {
122            prop_assert!(approx_tokens(&s) > 0);
123        }
124    }
125}
av_core/tokens.rs

av_core/
tokens.rs
