roma_core/
tokenizer.rs

1//! Token counting.
2//!
3//! Providers count tokens differently, so we abstract via [`Tokenizer`].
4//! Two default implementations:
5//!
6//! * [`TiktokenTokenizer`] wraps `tiktoken-rs`'s `cl100k_base` encoder
7//!   (close enough for OpenAI-family models and most estimation work).
8//! * [`ApproxTokenizer`] is a char-count / 4 heuristic used when no
9//!   native tokenizer exists (e.g. Anthropic). It systematically
10//!   under-counts; callers should apply a safety margin.
11
12use std::sync::Arc;
13
14use tiktoken_rs::CoreBPE;
15use tiktoken_rs::cl100k_base;
16
17/// A stateless token counter.
18pub trait Tokenizer: Send + Sync {
19    /// Count tokens in `text`.
20    fn count(&self, text: &str) -> u32;
21    /// Human-readable identifier for logging.
22    fn name(&self) -> &str;
23}
24
25/// Tiktoken-backed counter (cl100k_base).
26pub struct TiktokenTokenizer {
27    bpe: Arc<CoreBPE>,
28    name: &'static str,
29}
30
31impl TiktokenTokenizer {
32    /// Build a cl100k_base counter. Used by OpenAI GPT-4 family and a
33    /// reasonable default for most OpenAI-compatible APIs.
34    ///
35    /// Returns an `Internal` error if the tokenizer data fails to load
36    /// (tiktoken-rs keeps its BPE merges in-memory, so this only fails on
37    /// allocation failure or version mismatch).
38    pub fn cl100k() -> Result<Self, crate::ClassifiedError> {
39        let bpe = cl100k_base()
40            .map_err(|e| crate::ClassifiedError::Runtime(format!("tiktoken cl100k_base: {e}")))?;
41        Ok(Self {
42            bpe: Arc::new(bpe),
43            name: "tiktoken.cl100k",
44        })
45    }
46}
47
48impl Tokenizer for TiktokenTokenizer {
49    fn count(&self, text: &str) -> u32 {
50        self.bpe
51            .encode_with_special_tokens(text)
52            .len()
53            .try_into()
54            .unwrap_or(u32::MAX / 2)
55    }
56
57    fn name(&self) -> &str {
58        self.name
59    }
60}
61
62/// Character-based heuristic (`chars / 4`).
63///
64/// Under-counts in practice; use as a last resort (e.g. for Anthropic
65/// where no stable public tokenizer exists) and add a 1.3x safety factor
66/// when comparing against a hard context window.
67#[derive(Debug, Default, Clone, Copy)]
68pub struct ApproxTokenizer;
69
70impl Tokenizer for ApproxTokenizer {
71    fn count(&self, text: &str) -> u32 {
72        let chars = text.chars().count();
73        (chars as u32).div_ceil(4)
74    }
75
76    fn name(&self) -> &str {
77        "approx.chars_over_4"
78    }
79}
80
81#[cfg(test)]
82#[allow(clippy::expect_used, clippy::unwrap_used)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn approx_counts_nonzero_for_ascii() {
88        let t = ApproxTokenizer;
89        assert_eq!(t.count(""), 0);
90        assert_eq!(t.count("abcd"), 1);
91        assert_eq!(t.count("abcde"), 2);
92    }
93
94    #[test]
95    fn approx_counts_multibyte_by_char_not_byte() {
96        let t = ApproxTokenizer;
97        // "中" is one char, 3 bytes; should be counted as 1 char → 1 token.
98        assert_eq!(t.count("中"), 1);
99        // Four Chinese characters → 1 token (4 / 4 = 1).
100        assert_eq!(t.count("中国加油"), 1);
101    }
102
103    #[test]
104    fn tiktoken_counts_standard_phrase() {
105        let t = TiktokenTokenizer::cl100k().unwrap();
106        // This is one of the canonical prefix examples.
107        let n = t.count("Hello world");
108        assert!(n > 0 && n < 10, "expected a small token count, got {n}");
109    }
110
111    #[test]
112    fn tiktoken_empty_is_zero() {
113        let t = TiktokenTokenizer::cl100k().unwrap();
114        assert_eq!(t.count(""), 0);
115    }
116}
roma_core/tokenizer.rs

roma_core/
tokenizer.rs
