tokenfold-core 0.3.2

Token-aware compression for LLM payloads: shrink JSON tool-call bodies, command output, and diffs with exact tiktoken accounting and a typed safety report.
Documentation 
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

use crate::report::EstimatorInfo;

pub trait TokenEstimator {
    /// Estimator provenance for `CompressionReport.estimator`.
    fn info(&self) -> EstimatorInfo;
    fn count_bytes(&self, bytes: &[u8]) -> usize;
}

/// Fast pre-filter ONLY; under-counts dense JSON/code. Never used to claim exact savings.
#[derive(Debug, Clone, Copy, Default)]
pub struct ByteHeuristicEstimator;

impl TokenEstimator for ByteHeuristicEstimator {
    fn info(&self) -> EstimatorInfo {
        EstimatorInfo {
            backend: "heuristic".to_string(),
            model: None,
            is_exact: false,
        }
    }

    fn count_bytes(&self, bytes: &[u8]) -> usize {
        if bytes.is_empty() {
            0
        } else {
            bytes.len().div_ceil(4)
        }
    }
}

#[cfg(feature = "tiktoken")]
pub struct TiktokenEstimator {
    model: &'static str,
    bpe: tiktoken_rs::CoreBPE,
}

#[cfg(feature = "tiktoken")]
impl TiktokenEstimator {
    pub fn o200k_base() -> Result<Self, crate::errors::TokenFoldError> {
        let bpe = tiktoken_rs::o200k_base()
            .map_err(|e| crate::errors::TokenFoldError::EstimatorError(e.to_string()))?;
        Ok(Self {
            model: "o200k_base",
            bpe,
        })
    }

    pub fn cl100k_base() -> Result<Self, crate::errors::TokenFoldError> {
        let bpe = tiktoken_rs::cl100k_base()
            .map_err(|e| crate::errors::TokenFoldError::EstimatorError(e.to_string()))?;
        Ok(Self {
            model: "cl100k_base",
            bpe,
        })
    }
}

#[cfg(feature = "tiktoken")]
impl TokenEstimator for TiktokenEstimator {
    fn info(&self) -> EstimatorInfo {
        EstimatorInfo {
            backend: "tiktoken".to_string(),
            model: Some(self.model.to_string()),
            is_exact: true,
        }
    }

    fn count_bytes(&self, bytes: &[u8]) -> usize {
        let text = String::from_utf8_lossy(bytes);
        self.bpe.encode_with_special_tokens(&text).len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn heuristic_rounds_up_for_dense_input() {
        let est = ByteHeuristicEstimator;
        assert_eq!(est.count_bytes(b""), 0);
        assert_eq!(est.count_bytes(b"a"), 1); // div_ceil(1, 4) == 1
        assert_eq!(est.count_bytes(b"abcd"), 1);
        assert_eq!(est.count_bytes(b"abcde"), 2);
    }

    #[test]
    fn heuristic_info_is_labeled_not_exact() {
        let info = ByteHeuristicEstimator.info();
        assert_eq!(info.backend, "heuristic");
        assert_eq!(info.model, None);
        assert!(!info.is_exact);
    }

    #[cfg(feature = "tiktoken")]
    #[test]
    fn tiktoken_smoke_counts_known_string() {
        let est = TiktokenEstimator::o200k_base().expect("o200k_base should load");
        let info = est.info();
        assert_eq!(info.backend, "tiktoken");
        assert_eq!(info.model.as_deref(), Some("o200k_base"));
        assert!(info.is_exact);

        assert_eq!(est.count_bytes(b""), 0);
        assert!(est.count_bytes(b"hello world") > 0);
    }
}