roma-core 0.1.0

//! Token counting.
//!
//! Providers count tokens differently, so we abstract via [`Tokenizer`].
//! Two default implementations:
//!
//! * [`TiktokenTokenizer`] wraps `tiktoken-rs`'s `cl100k_base` encoder
//!   (close enough for OpenAI-family models and most estimation work).
//! * [`ApproxTokenizer`] is a char-count / 4 heuristic used when no
//!   native tokenizer exists (e.g. Anthropic). It systematically
//!   under-counts; callers should apply a safety margin.

use std::sync::Arc;

use tiktoken_rs::CoreBPE;
use tiktoken_rs::cl100k_base;

/// A stateless token counter.
pub trait Tokenizer: Send + Sync {
    /// Count tokens in `text`.
    fn count(&self, text: &str) -> u32;
    /// Human-readable identifier for logging.
    fn name(&self) -> &str;
}

/// Tiktoken-backed counter (cl100k_base).
pub struct TiktokenTokenizer {
    bpe: Arc<CoreBPE>,
    name: &'static str,
}

impl TiktokenTokenizer {
    /// Build a cl100k_base counter. Used by OpenAI GPT-4 family and a
    /// reasonable default for most OpenAI-compatible APIs.
    ///
    /// Returns an `Internal` error if the tokenizer data fails to load
    /// (tiktoken-rs keeps its BPE merges in-memory, so this only fails on
    /// allocation failure or version mismatch).
    pub fn cl100k() -> Result<Self, crate::ClassifiedError> {
        let bpe = cl100k_base()
            .map_err(|e| crate::ClassifiedError::Runtime(format!("tiktoken cl100k_base: {e}")))?;
        Ok(Self {
            bpe: Arc::new(bpe),
            name: "tiktoken.cl100k",
        })
    }
}

impl Tokenizer for TiktokenTokenizer {
    fn count(&self, text: &str) -> u32 {
        self.bpe
            .encode_with_special_tokens(text)
            .len()
            .try_into()
            .unwrap_or(u32::MAX / 2)
    }

    fn name(&self) -> &str {
        self.name
    }
}

/// Character-based heuristic (`chars / 4`).
///
/// Under-counts in practice; use as a last resort (e.g. for Anthropic
/// where no stable public tokenizer exists) and add a 1.3x safety factor
/// when comparing against a hard context window.
#[derive(Debug, Default, Clone, Copy)]
pub struct ApproxTokenizer;

impl Tokenizer for ApproxTokenizer {
    fn count(&self, text: &str) -> u32 {
        let chars = text.chars().count();
        (chars as u32).div_ceil(4)
    }

    fn name(&self) -> &str {
        "approx.chars_over_4"
    }
}

#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)]
mod tests {
    use super::*;

    #[test]
    fn approx_counts_nonzero_for_ascii() {
        let t = ApproxTokenizer;
        assert_eq!(t.count(""), 0);
        assert_eq!(t.count("abcd"), 1);
        assert_eq!(t.count("abcde"), 2);
    }

    #[test]
    fn approx_counts_multibyte_by_char_not_byte() {
        let t = ApproxTokenizer;
        // "中" is one char, 3 bytes; should be counted as 1 char → 1 token.
        assert_eq!(t.count("中"), 1);
        // Four Chinese characters → 1 token (4 / 4 = 1).
        assert_eq!(t.count("中国加油"), 1);
    }

    #[test]
    fn tiktoken_counts_standard_phrase() {
        let t = TiktokenTokenizer::cl100k().unwrap();
        // This is one of the canonical prefix examples.
        let n = t.count("Hello world");
        assert!(n > 0 && n < 10, "expected a small token count, got {n}");
    }

    #[test]
    fn tiktoken_empty_is_zero() {
        let t = TiktokenTokenizer::cl100k().unwrap();
        assert_eq!(t.count(""), 0);
    }
}