Skip to main content

whycodes_core/
tokens.rs

1//! Token estimation — single source for `session` and `llm` (issue #52 C2).
2//!
3//! `session::estimate_tokens` and `llm::token_counter::chars_to_tokens_fallback`
4//! were two copies of `chars.div_ceil(4)` with a subtle mismatch: session
5//! returned `0` for empty input while llm returned `max(1)`. Unifying here
6//! makes the policy explicit and testable.
7
8/// ~4 Unicode scalars per token (matches `whycodes_llm` fallback family).
9///
10/// ASCII uses byte length (same as scalar count); non-ASCII walks `chars`.
11/// Returns `0` for empty input — callers that need at least 1 should use
12/// [`estimate_tokens_at_least_one`].
13pub fn estimate_tokens(text: &str) -> usize {
14    let n = if text.is_ascii() {
15        text.len()
16    } else {
17        text.chars().count()
18    };
19    n.div_ceil(4)
20}
21
22/// Same as [`estimate_tokens`] but at least `1` (for `llm::count_tokens`
23/// where an empty prompt still costs one token in provider accounting).
24pub fn estimate_tokens_at_least_one(text: &str) -> usize {
25    estimate_tokens(text).max(1)
26}
27
28#[cfg(test)]
29mod tests {
30    use super::*;
31
32    #[test]
33    fn empty_is_zero() {
34        assert_eq!(estimate_tokens(""), 0);
35        assert_eq!(estimate_tokens_at_least_one(""), 1);
36    }
37
38    #[test]
39    fn ascii_fast_path() {
40        assert_eq!(estimate_tokens("abcd"), 1);
41        assert_eq!(estimate_tokens("abcde"), 2);
42        assert_eq!(estimate_tokens("a"), 1); // 1.div_ceil(4)=1
43        // empty already tested
44    }
45
46    #[test]
47    fn unicode_counts_scalars() {
48        // 4 scalars -> 1 token, 5 -> 2
49        assert_eq!(estimate_tokens("🦀🦀🦀🦀"), 1);
50        assert_eq!(estimate_tokens("🦀🦀🦀🦀🦀"), 2);
51    }
52
53    #[test]
54    fn llm_wrapper_max_one() {
55        assert_eq!(estimate_tokens_at_least_one("a"), 1);
56        assert_eq!(estimate_tokens_at_least_one("abcd"), 1);
57        assert_eq!(estimate_tokens_at_least_one("abcde"), 2);
58    }
59}