lean_ctx/core/context_bundle/
limit.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Output budget: a size limit and the unit it is measured in.
3
4use std::fmt;
5
6/// What a bundle limit counts.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
8pub enum Unit {
9    /// Unicode scalar values — what chat input boxes count.
10    #[default]
11    Chars,
12    /// `o200k_base` BPE tokens, the tokenizer used for all lean-ctx accounting.
13    Tokens,
14}
15
16impl Unit {
17    pub fn parse(raw: &str) -> Result<Self, String> {
18        match raw.trim().to_ascii_lowercase().as_str() {
19            "chars" | "char" | "characters" => Ok(Self::Chars),
20            "tokens" | "token" => Ok(Self::Tokens),
21            other => Err(format!("unknown unit '{other}' (expected chars or tokens)")),
22        }
23    }
24
25    /// Size of `text` in this unit.
26    pub(crate) fn measure(self, text: &str) -> usize {
27        match self {
28            Self::Chars => text.chars().count(),
29            Self::Tokens => crate::core::tokens::count_tokens(text),
30        }
31    }
32}
33
34impl fmt::Display for Unit {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        f.write_str(match self {
37            Self::Chars => "chars",
38            Self::Tokens => "tokens",
39        })
40    }
41}
42
43/// Parse a limit such as `128000`, `128k`, `2M` or `1.5m`.
44///
45/// Suffixes are decimal (`k` = 1 000, `m` = 1 000 000) because chat products
46/// advertise their limits that way. Zero and values that overflow are refused.
47pub fn parse_limit(raw: &str) -> Result<usize, String> {
48    let trimmed = raw.trim().replace('_', "");
49    let (number, multiplier) = match trimmed.char_indices().last() {
50        Some((idx, 'k' | 'K')) => (&trimmed[..idx], 1_000.0),
51        Some((idx, 'm' | 'M')) => (&trimmed[..idx], 1_000_000.0),
52        _ => (trimmed.as_str(), 1.0),
53    };
54    let value: f64 = number
55        .parse()
56        .map_err(|_| format!("invalid limit '{raw}' (examples: 128000, 128k, 2M)"))?;
57    let scaled = value * multiplier;
58    // NaN and infinities fall outside the range too.
59    if !(1.0..=1e12).contains(&scaled) {
60        return Err(format!("limit '{raw}' is out of range"));
61    }
62    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
63    Ok(scaled.round() as usize)
64}
65
66#[cfg(test)]
67mod tests {
68    use super::{Unit, parse_limit};
69
70    #[test]
71    fn limits_accept_plain_and_suffixed_forms() {
72        assert_eq!(parse_limit("128000"), Ok(128_000));
73        assert_eq!(parse_limit("128k"), Ok(128_000));
74        assert_eq!(parse_limit("128K"), Ok(128_000));
75        assert_eq!(parse_limit("2M"), Ok(2_000_000));
76        assert_eq!(parse_limit("1.5m"), Ok(1_500_000));
77        assert_eq!(parse_limit("100_000"), Ok(100_000));
78    }
79
80    #[test]
81    fn limits_refuse_zero_garbage_and_overflow() {
82        assert!(parse_limit("0").is_err());
83        assert!(parse_limit("-5k").is_err());
84        assert!(parse_limit("lots").is_err());
85        assert!(parse_limit("k").is_err());
86        assert!(parse_limit("1e20").is_err());
87    }
88
89    #[test]
90    fn chars_count_scalar_values_not_bytes() {
91        assert_eq!(Unit::Chars.measure("héllo"), 5);
92        assert_eq!(Unit::Chars.measure("日本"), 2);
93    }
94
95    #[test]
96    fn unit_parse_is_strict() {
97        assert_eq!(Unit::parse("Tokens"), Ok(Unit::Tokens));
98        assert_eq!(Unit::parse("chars"), Ok(Unit::Chars));
99        assert!(Unit::parse("bytes").is_err());
100    }
101}
lean_ctx/core/context_bundle/limit.rs

lean_ctx/core/context_bundle/
limit.rs
