1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
//! BPE token counting (bd-cv653.7.1).
//!
//! Real O200k (OpenAI) + Cl100k (Anthropic-approx) BPE counting via
//! tiktoken-rs replaces the chars/4 heuristic on the estimation path.
//! Hybrid accounting is preserved: measured API usage still wins when
//! present; BPE replaces ONLY the heuristic path; chars/4 stays as the
//! final fallback when the `bpe-tokens` feature is off (minimal builds).
//!
//! Table selection: anthropic → Cl100k-class, everything else → O200k
//! (documented approximation for non-OpenAI providers — Cl100k and O200k
//! diverge mostly on code-token frequencies, so O200k is the safer default
//! for OpenAI-compatible hosts).
/// Token table families.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TokenTable {
O200k,
Cl100k,
}
impl TokenTable {
pub const fn as_str(self) -> &'static str {
match self {
Self::O200k => "o200k",
Self::Cl100k => "cl100k",
}
}
}
/// Pick the table for a provider id (anthropic → Cl100k-class; everything
/// else → O200k, the documented approximation).
#[must_use]
pub const fn table_for_provider(provider: &str) -> TokenTable {
if provider.eq_ignore_ascii_case("anthropic") {
TokenTable::Cl100k
} else {
TokenTable::O200k
}
}
/// Counting surface (tests inject a deterministic stub).
pub trait TokenCounter: Send + Sync {
fn count(&self, text: &str, table: TokenTable) -> u64;
}
/// Real BPE counting (feature `bpe-tokens`).
#[cfg(feature = "bpe-tokens")]
pub struct BpeCounter;
#[cfg(feature = "bpe-tokens")]
impl TokenCounter for BpeCounter {
fn count(&self, text: &str, table: TokenTable) -> u64 {
let bpe = match table {
TokenTable::O200k => tiktoken_rs::o200k_base_singleton(),
TokenTable::Cl100k => tiktoken_rs::cl100k_base_singleton(),
};
bpe.encode_with_special_tokens(text).len() as u64
}
}
/// chars/4 fallback (feature-off builds and the final fallback).
pub struct HeuristicCounter;
impl TokenCounter for HeuristicCounter {
fn count(&self, text: &str, _table: TokenTable) -> u64 {
(text.len() / 4) as u64
}
}
/// The active counter for this build.
#[must_use]
pub fn active_counter() -> &'static dyn TokenCounter {
#[cfg(feature = "bpe-tokens")]
{
static COUNTER: BpeCounter = BpeCounter;
&COUNTER
}
#[cfg(not(feature = "bpe-tokens"))]
{
static COUNTER: HeuristicCounter = HeuristicCounter;
&COUNTER
}
}
/// Count text with the active counter for a provider family.
#[must_use]
pub fn count_tokens(text: &str, provider: &str) -> u64 {
active_counter().count(text, table_for_provider(provider))
}
/// Per-table counts for `pi token` output.
#[must_use]
pub fn count_all_tables(text: &str) -> Vec<(TokenTable, u64)> {
[TokenTable::O200k, TokenTable::Cl100k]
.into_iter()
.map(|table| (table, active_counter().count(text, table)))
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn table_selection() {
assert_eq!(table_for_provider("anthropic"), TokenTable::Cl100k);
assert_eq!(table_for_provider("Anthropic"), TokenTable::Cl100k);
assert_eq!(table_for_provider("openai"), TokenTable::O200k);
assert_eq!(table_for_provider("ollama"), TokenTable::O200k);
}
#[cfg(feature = "bpe-tokens")]
#[test]
fn bpe_counts_reference_vectors() {
// Reference vectors (tiktoken oracle):
// "hello world" is 2 tokens in both families.
let bpe = BpeCounter;
assert_eq!(bpe.count("hello world", TokenTable::O200k), 2);
assert_eq!(bpe.count("hello world", TokenTable::Cl100k), 2);
// Code-heavy text counts materially above the chars/4 heuristic
// on symbol-dense input (the motivation for the swap).
let code = "fn main() { println!(\"{}\", 1 + 1); }";
let bpe_count = bpe.count(code, TokenTable::O200k);
let heuristic = HeuristicCounter.count(code, TokenTable::O200k);
assert!(bpe_count > 0);
assert!(bpe_count != heuristic, "BPE should diverge from chars/4");
}
#[test]
fn heuristic_counts_quarters() {
let counter = HeuristicCounter;
assert_eq!(counter.count("abcd", TokenTable::O200k), 1);
assert_eq!(counter.count("abcdefgh", TokenTable::Cl100k), 2);
}
#[cfg(feature = "bpe-tokens")]
#[test]
fn counting_1mb_stays_linear_rather_than_merely_fast() {
let text = "lorem ipsum dolor sit amet ".repeat(40_000); // ~1.08 MB
let bpe = BpeCounter;
let start = std::time::Instant::now();
let count = bpe.count(&text, TokenTable::O200k);
let elapsed = start.elapsed();
eprintln!("1MB BPE count: {count} tokens in {elapsed:?}");
assert!(count > 100_000);
// A catastrophe guard, not a performance measurement. It exists to
// catch tokenisation going super-linear, which on 1 MB would take
// minutes rather than seconds; nothing in tests/perf covers
// tokenisation, so this is the only guard there is.
//
// The bound was 2s and failed at 2.97s in a full `--lib` run on a host
// at load average 50. That is a debug build sharing 14 cores with
// ~10,000 other tests, so a two-second wall clock is a coin flip
// rather than a property of the code. 30s keeps the regression it is
// actually for and stops it reporting host load as a defect.
assert!(
elapsed < std::time::Duration::from_secs(30),
"1MB count took {elapsed:?}; tokenisation is super-linear, not merely slow"
);
}
}