cerberust/scanners/
token_limit.rs

1//! The token-limit scanner: reject input whose approximate token count exceeds a
2//! configured maximum.
3//!
4//! The token count here is a whitespace/heuristic estimate, not a real
5//! tokenizer: a request is split on whitespace and each whitespace-delimited
6//! chunk contributes a count proportional to its length, approximating how a
7//! BPE tokenizer fragments long runs. A model-accurate tokenizer (`tiktoken` /
8//! the provider's BPE) is a later refinement; it plugs in behind
9//! [`approx_token_count`] without changing the scanner's contract.
10//!
11//! [`Disposition::Block`], [`Direction::Input`]: the point is to reject an
12//! over-budget prompt before it reaches the model, so this never runs on output
13//! and never rewrites text.
14
15use crate::scanner::{
16    Direction, Disposition, HoldBack, ScanCtx, ScanResult, Scanner, ScannerId, Verdict,
17};
18
19/// Average characters per token a BPE tokenizer fragments a long run into. Four
20/// is the conventional English-text rule of thumb; a long whitespace-free run
21/// (a URL, a base64 blob) is counted as `ceil(len / 4)` tokens rather than one.
22const CHARS_PER_TOKEN: usize = 4;
23
24/// Approximate the token count of `text`: each whitespace-delimited chunk costs
25/// `ceil(chars / CHARS_PER_TOKEN)` tokens, summed. A pure-whitespace string is
26/// zero tokens.
27#[must_use]
28pub fn approx_token_count(text: &str) -> usize {
29    text.split_whitespace()
30        .map(|chunk| chunk.chars().count().div_ceil(CHARS_PER_TOKEN))
31        .sum()
32}
33
34/// Rejects input exceeding `max_tokens` approximate tokens.
35#[derive(Debug, Clone)]
36pub struct TokenLimitScanner {
37    max_tokens: usize,
38}
39
40impl TokenLimitScanner {
41    /// This scanner's id.
42    pub const ID: ScannerId = ScannerId("native:token-limit");
43
44    /// A scanner that rejects input over `max_tokens` approximate tokens.
45    #[must_use]
46    pub fn new(max_tokens: usize) -> Self {
47        Self { max_tokens }
48    }
49
50    /// The configured maximum.
51    #[must_use]
52    pub fn max_tokens(&self) -> usize {
53        self.max_tokens
54    }
55}
56
57impl Scanner for TokenLimitScanner {
58    fn id(&self) -> ScannerId {
59        Self::ID
60    }
61
62    fn direction(&self) -> Direction {
63        Direction::Input
64    }
65
66    fn disposition(&self) -> Disposition {
67        Disposition::Block
68    }
69
70    fn scan<'a>(&self, text: &'a str, _ctx: &mut ScanCtx) -> ScanResult<'a> {
71        let over = approx_token_count(text) > self.max_tokens;
72        let risk = if over { 1.0 } else { 0.0 };
73        Ok(Verdict::detected(text, risk, 1.0))
74    }
75
76    fn hold_back(&self) -> HoldBack {
77        // Input-only and block-only: it never runs on a stream, so the
78        // hold-back contract is moot. WholeStream is the conservative declaration
79        // for a scanner that judges complete text.
80        HoldBack::WholeStream
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    #![allow(
87        clippy::unwrap_used,
88        clippy::expect_used,
89        reason = "tests assert on known-good values"
90    )]
91    use super::*;
92
93    #[test]
94    fn under_limit_passes() {
95        let scanner = TokenLimitScanner::new(10);
96        let mut ctx = ScanCtx::new();
97        let verdict = scanner.scan("a short prompt", &mut ctx).unwrap();
98        assert!(verdict.valid);
99    }
100
101    #[test]
102    fn over_limit_blocks() {
103        let scanner = TokenLimitScanner::new(3);
104        let mut ctx = ScanCtx::new();
105        let verdict = scanner
106            .scan("one two three four five six", &mut ctx)
107            .unwrap();
108        assert!(!verdict.valid);
109        assert!((verdict.risk - 1.0).abs() < f32::EPSILON);
110    }
111
112    #[test]
113    fn long_whitespace_free_run_counts_multiple_tokens() {
114        // A 40-char run is ~10 tokens, not one.
115        assert_eq!(approx_token_count(&"x".repeat(40)), 10);
116    }
117
118    #[test]
119    fn whitespace_only_is_zero_tokens() {
120        assert_eq!(approx_token_count("   \t\n  "), 0);
121    }
122
123    #[test]
124    fn exactly_at_limit_passes() {
125        let scanner = TokenLimitScanner::new(4);
126        let mut ctx = ScanCtx::new();
127        // Four single-char-ish words → four tokens, not over.
128        assert!(scanner.scan("aa bb cc dd", &mut ctx).unwrap().valid);
129    }
130}
cerberust/scanners/token_limit.rs

cerberust/scanners/
token_limit.rs
