token-estimate 0.1.0

//! Estimate LLM token counts from text without loading a tokenizer.
//!
//! Real tokenizers (BPE, SentencePiece) need a vocabulary file of several
//! megabytes and a non-trivial dependency tree. That is the right tool when you
//! need an exact count. It is the wrong tool when you only need to answer
//! "will this roughly fit in the context window", which is the common case in
//! CLI tooling, log processing, and pre-flight budget checks.
//!
//! This crate gives a fast, dependency-free estimate and is explicit about how
//! wrong it can be.
//!
//! # Accuracy, stated honestly
//!
//! For ordinary English prose the character-based estimate is typically within
//! about 10-15% of a BPE tokenizer. It degrades in predictable ways:
//!
//! - **Code and markup** tokenize denser than prose (more punctuation, more
//!   short symbols), so the estimate tends to run low. Use
//!   [`Profile::Code`].
//! - **Non-Latin scripts** often use one token per character or worse, so a
//!   chars/4 heuristic can underestimate severely. [`estimate`] detects
//!   non-ASCII density and adjusts, but treat CJK figures as a rough floor.
//! - **Long runs of whitespace or repeated punctuation** compress well and the
//!   estimate runs high.
//!
//! If you need an exact count, use a real tokenizer. This is for budgeting.
//!
//! # Example
//!
//! ```
//! use token_estimate::{estimate, Profile};
//!
//! let prose = "The quick brown fox jumps over the lazy dog.";
//! let n = estimate(prose, Profile::Prose);
//! assert!(n > 0);
//!
//! // Code is denser, so the same byte count yields more tokens.
//! let code = "fn main() { let x: Vec<u8> = vec![1,2,3]; }";
//! assert!(estimate(code, Profile::Code) > estimate(code, Profile::Prose));
//! ```

/// Text kind, which changes the characters-per-token assumption.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Profile {
    /// Ordinary prose. ~4 characters per token.
    Prose,
    /// Source code or markup. Denser: ~3 characters per token.
    Code,
    /// Mixed content, e.g. prose with fenced code blocks.
    Mixed,
}

impl Profile {
    /// Characters per token assumed for this profile.
    pub fn chars_per_token(self) -> f64 {
        match self {
            Profile::Prose => 4.0,
            Profile::Code => 3.0,
            Profile::Mixed => 3.5,
        }
    }
}

/// Estimated token count for `text` under `profile`.
///
/// Returns 0 for empty or whitespace-only input. The result is a ceiling, so a
/// non-empty string always estimates at least 1 token.
pub fn estimate(text: &str, profile: Profile) -> usize {
    if text.trim().is_empty() {
        return 0;
    }
    // Collapse whitespace runs: tokenizers merge them, so counting raw bytes
    // overestimates heavily on indented or padded text.
    let mut effective = 0usize;
    let mut non_ascii = 0usize;
    let mut prev_ws = false;
    for ch in text.chars() {
        if ch.is_whitespace() {
            if !prev_ws {
                effective += 1;
            }
            prev_ws = true;
            continue;
        }
        prev_ws = false;
        effective += 1;
        if !ch.is_ascii() {
            non_ascii += 1;
        }
    }
    if effective == 0 {
        return 0;
    }
    let total_chars = text.chars().count();
    let non_ascii_ratio = if total_chars == 0 { 0.0 } else { non_ascii as f64 / total_chars as f64 };

    // Non-Latin scripts approach one token per character. Blend the profile
    // ratio toward 1.0 as non-ASCII density rises.
    let ratio = profile.chars_per_token();
    let adjusted = ratio * (1.0 - non_ascii_ratio) + 1.0 * non_ascii_ratio;
    let est = (effective as f64 / adjusted).ceil() as usize;
    est.max(1)
}

/// Estimate with the profile inferred from the text itself.
///
/// Heuristic: a high density of code-ish punctuation (`{}()<>;=` and friends)
/// selects [`Profile::Code`]; some but not much selects [`Profile::Mixed`].
pub fn estimate_auto(text: &str) -> usize {
    estimate(text, infer_profile(text))
}

/// Infer a [`Profile`] from punctuation density.
pub fn infer_profile(text: &str) -> Profile {
    let total = text.chars().filter(|c| !c.is_whitespace()).count();
    if total == 0 {
        return Profile::Prose;
    }
    let codeish = text
        .chars()
        .filter(|c| matches!(c, '{' | '}' | '(' | ')' | '<' | '>' | ';' | '=' | '[' | ']' | '|' | '_'))
        .count();
    let density = codeish as f64 / total as f64;
    if density > 0.08 {
        Profile::Code
    } else if density > 0.03 {
        Profile::Mixed
    } else {
        Profile::Prose
    }
}

/// Whether `text` is estimated to fit in `budget` tokens.
pub fn fits(text: &str, budget: usize, profile: Profile) -> bool {
    estimate(text, profile) <= budget
}

/// Remaining budget after `text`, saturating at zero.
pub fn remaining(text: &str, budget: usize, profile: Profile) -> usize {
    budget.saturating_sub(estimate(text, profile))
}

/// A conservative upper bound, inflating the estimate by `margin_pct`.
///
/// Useful when the cost of overflowing a context window is higher than the cost
/// of sending less. A margin of 20 is a reasonable default for prose.
pub fn estimate_with_margin(text: &str, profile: Profile, margin_pct: u32) -> usize {
    let base = estimate(text, profile) as f64;
    (base * (1.0 + margin_pct as f64 / 100.0)).ceil() as usize
}

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

    #[test]
    fn empty_is_zero() {
        assert_eq!(estimate("", Profile::Prose), 0);
        assert_eq!(estimate("   \n\t ", Profile::Prose), 0);
    }

    #[test]
    fn non_empty_is_at_least_one() {
        assert_eq!(estimate("a", Profile::Prose), 1);
        assert_eq!(estimate("hi", Profile::Code), 1);
    }

    #[test]
    fn code_estimates_higher_than_prose() {
        let s = "fn main() { let x: Vec<u8> = vec![1,2,3]; }";
        assert!(estimate(s, Profile::Code) > estimate(s, Profile::Prose));
    }

    #[test]
    fn whitespace_runs_are_collapsed() {
        let tight = "alpha beta gamma";
        let padded = "alpha        beta        gamma";
        assert_eq!(estimate(tight, Profile::Prose), estimate(padded, Profile::Prose));
    }

    #[test]
    fn non_ascii_raises_the_estimate() {
        let ascii = "aaaaaaaaaaaaaaaaaaaa";
        let cjk = "説明説明説明説明説明説明説明説明説明説明";
        assert!(estimate(cjk, Profile::Prose) > estimate(ascii, Profile::Prose));
    }

    #[test]
    fn profile_inference() {
        assert_eq!(infer_profile("A plain sentence about nothing at all."), Profile::Prose);
        assert_eq!(infer_profile("if (a == b) { c[d] = e; }"), Profile::Code);
    }

    #[test]
    fn auto_matches_inferred_profile() {
        let s = "let y = f(x);";
        assert_eq!(estimate_auto(s), estimate(s, infer_profile(s)));
    }

    #[test]
    fn budget_helpers() {
        let s = "a short line of prose";
        let n = estimate(s, Profile::Prose);
        assert!(fits(s, n, Profile::Prose));
        assert!(!fits(s, n - 1, Profile::Prose));
        assert_eq!(remaining(s, n + 5, Profile::Prose), 5);
        assert_eq!(remaining(s, 0, Profile::Prose), 0);
    }

    #[test]
    fn margin_inflates() {
        let s = "some prose to measure against a margin";
        let base = estimate(s, Profile::Prose);
        assert!(estimate_with_margin(s, Profile::Prose, 20) > base);
        assert_eq!(estimate_with_margin(s, Profile::Prose, 0), base);
    }

    #[test]
    fn prose_estimate_is_in_a_sane_range() {
        // ~4 chars/token: a 44-char sentence should land near 11 tokens.
        let s = "The quick brown fox jumps over the lazy dog.";
        let n = estimate(s, Profile::Prose);
        assert!((8..=14).contains(&n), "got {n}");
    }
}