tokenfold_core/
token_estimator.rs

1use crate::report::EstimatorInfo;
2
3pub trait TokenEstimator {
4    /// Estimator provenance for `CompressionReport.estimator`.
5    fn info(&self) -> EstimatorInfo;
6    fn count_bytes(&self, bytes: &[u8]) -> usize;
7}
8
9/// Fast pre-filter ONLY; under-counts dense JSON/code. Never used to claim exact savings.
10#[derive(Debug, Clone, Copy, Default)]
11pub struct ByteHeuristicEstimator;
12
13impl TokenEstimator for ByteHeuristicEstimator {
14    fn info(&self) -> EstimatorInfo {
15        EstimatorInfo {
16            backend: "heuristic".to_string(),
17            model: None,
18            is_exact: false,
19        }
20    }
21
22    fn count_bytes(&self, bytes: &[u8]) -> usize {
23        if bytes.is_empty() {
24            0
25        } else {
26            bytes.len().div_ceil(4)
27        }
28    }
29}
30
31#[cfg(feature = "tiktoken")]
32pub struct TiktokenEstimator {
33    model: &'static str,
34    bpe: tiktoken_rs::CoreBPE,
35}
36
37#[cfg(feature = "tiktoken")]
38impl TiktokenEstimator {
39    pub fn o200k_base() -> Result<Self, crate::errors::TokenFoldError> {
40        let bpe = tiktoken_rs::o200k_base()
41            .map_err(|e| crate::errors::TokenFoldError::EstimatorError(e.to_string()))?;
42        Ok(Self {
43            model: "o200k_base",
44            bpe,
45        })
46    }
47
48    pub fn cl100k_base() -> Result<Self, crate::errors::TokenFoldError> {
49        let bpe = tiktoken_rs::cl100k_base()
50            .map_err(|e| crate::errors::TokenFoldError::EstimatorError(e.to_string()))?;
51        Ok(Self {
52            model: "cl100k_base",
53            bpe,
54        })
55    }
56}
57
58#[cfg(feature = "tiktoken")]
59impl TokenEstimator for TiktokenEstimator {
60    fn info(&self) -> EstimatorInfo {
61        EstimatorInfo {
62            backend: "tiktoken".to_string(),
63            model: Some(self.model.to_string()),
64            is_exact: true,
65        }
66    }
67
68    fn count_bytes(&self, bytes: &[u8]) -> usize {
69        let text = String::from_utf8_lossy(bytes);
70        self.bpe.encode_with_special_tokens(&text).len()
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn heuristic_rounds_up_for_dense_input() {
80        let est = ByteHeuristicEstimator;
81        assert_eq!(est.count_bytes(b""), 0);
82        assert_eq!(est.count_bytes(b"a"), 1); // div_ceil(1, 4) == 1
83        assert_eq!(est.count_bytes(b"abcd"), 1);
84        assert_eq!(est.count_bytes(b"abcde"), 2);
85    }
86
87    #[test]
88    fn heuristic_info_is_labeled_not_exact() {
89        let info = ByteHeuristicEstimator.info();
90        assert_eq!(info.backend, "heuristic");
91        assert_eq!(info.model, None);
92        assert!(!info.is_exact);
93    }
94
95    #[cfg(feature = "tiktoken")]
96    #[test]
97    fn tiktoken_smoke_counts_known_string() {
98        let est = TiktokenEstimator::o200k_base().expect("o200k_base should load");
99        let info = est.info();
100        assert_eq!(info.backend, "tiktoken");
101        assert_eq!(info.model.as_deref(), Some("o200k_base"));
102        assert!(info.is_exact);
103
104        assert_eq!(est.count_bytes(b""), 0);
105        assert!(est.count_bytes(b"hello world") > 0);
106    }
107}
tokenfold_core/token_estimator.rs

tokenfold_core/
token_estimator.rs
