✨ feat(mq-lang,mq-run): add token_count(text, model) builtin

crates/mq-check/src/builtin.rs

@@ -312,6 +312,8 @@ fn register_string(ctx: &mut InferenceContext) {
     register_binary(ctx, "ends_with", Type::String, Type::String, Type::Bool);
     register_binary(ctx, "index", Type::String, Type::String, Type::Number);
     register_binary(ctx, "rindex", Type::String, Type::String, Type::Number);
+    register_unary(ctx, "token_count", Type::String, Type::Number);
+    register_binary(ctx, "token_count", Type::String, Type::String, Type::Number);
 
     // String manipulation
     register_ternary(ctx, "replace", Type::String, Type::String, Type::String, Type::String);
@@ -1546,6 +1548,9 @@ mod tests {
     #[case::ascii_downcase_number("ascii_downcase(42)", false)] // Should fail: wrong type
     #[case::rtrim("rtrim(\"  hello  \")", true)]
     #[case::rindex("rindex(\"hello world hello\", \"hello\")", true)]
+    #[case::token_count("token_count(\"hello world\", \"gpt-4\")", true)]
+    #[case::token_count_no_model("token_count(\"hello world\")", true)]
+    #[case::token_count_number("token_count(42, \"gpt-4\")", false)] // Should fail: wrong type
     #[case::capture("capture(\"hello 42\", \"(?P<word>\\\\w+)\")", true)]
     #[case::is_regex_match("is_regex_match(\"hello123\", \"[0-9]+\")", true)]
     #[case::is_not_regex_match("is_not_regex_match(\"hello123\", \"[0-9]+\")", true)]

crates/mq-lang/Cargo.toml

@@ -45,6 +45,7 @@ url = {workspace = true}
 toon-format = { version = "0.5", default-features = false }
 quick-xml = "0.41.0"
 similar = "3.0.0"
+tiktoken-rs = { version = "0.12", optional = true }
 ureq = { workspace = true, optional = true }
 uuid = {workspace = true, features = ["v4", "v7"]}
 
@@ -64,6 +65,7 @@ sync = []
 http-import = []
 http = ["dep:ureq"]
 http-import-ureq = ["http-import", "http"]
+tiktoken = ["dep:tiktoken-rs"]
 
 [dev-dependencies]
 divan = {workspace = true}

crates/mq-lang/src/eval/builtin.rs

@@ -8,6 +8,7 @@ pub(super) mod path;
 mod random;
 mod range;
 mod regex;
+pub(super) mod tokenizer;
 
 use crate::arena::Arena;
 use crate::ast::{constants, node as ast};
@@ -1526,6 +1527,38 @@ fn utf8bytelen_impl(_: &Ident, _: &RuntimeValue, args: Args, _: &SharedEnv) -> R
     }
 }
 
+/// Counts (or, with the `tiktoken` Cargo feature, exactly counts) the LLM tokens `text` would
+/// consume. `model` is optional: without it, `text` is always run through the heuristic
+/// estimate, regardless of the `tiktoken` feature. See [`tokenizer`] for the heuristic/exact
+/// two-tier design.
+#[mq_macros::mq_fn(name = "token_count", params = Range(1, 2))]
+fn token_count_impl(ident: &Ident, _: &RuntimeValue, mut args: Args, _: &SharedEnv) -> Result<RuntimeValue, Error> {
+    match args.as_mut_slice() {
+        [RuntimeValue::String(text)] => Ok(RuntimeValue::Number(tokenizer::token_count_estimate(text).into())),
+        [node @ RuntimeValue::Markdown(_, _)] => Ok(RuntimeValue::Number(
+            node.markdown_node()
+                .map(|md| tokenizer::token_count_estimate(md.value().as_str()))
+                .unwrap_or(0)
+                .into(),
+        )),
+        [RuntimeValue::None] => Ok(RuntimeValue::Number(0.into())),
+        [RuntimeValue::String(text), RuntimeValue::String(model)] => {
+            tokenizer::token_count(text, model).map(|n| RuntimeValue::Number(n.into()))
+        }
+        [node @ RuntimeValue::Markdown(_, _), RuntimeValue::String(model)] => node
+            .markdown_node()
+            .map(|md| tokenizer::token_count(md.value().as_str(), model).map(|n| RuntimeValue::Number(n.into())))
+            .unwrap_or(Ok(RuntimeValue::Number(0.into()))),
+        [RuntimeValue::None, RuntimeValue::String(_)] => Ok(RuntimeValue::Number(0.into())),
+        [a] => Err(Error::InvalidTypes(ident.to_string(), vec![std::mem::take(a)])),
+        [a, b] => Err(Error::InvalidTypes(
+            ident.to_string(),
+            vec![std::mem::take(a), std::mem::take(b)],
+        )),
+        _ => unreachable!("token_count should always receive one or two arguments"),
+    }
+}
+
 #[mq_macros::mq_fn(name = "rindex", params = Fixed(2))]
 fn rindex_impl(ident: &Ident, _: &RuntimeValue, mut args: Args, _: &SharedEnv) -> Result<RuntimeValue, Error> {
     match args.as_mut_slice() {
@@ -4172,6 +4205,7 @@ mq_macros::builtin_dispatch! {
     INDEX,
     LEN,
     UTF8BYTELEN,
+    TOKEN_COUNT,
     RINDEX,
     RANGE,
     DEL,
@@ -5358,6 +5392,13 @@ pub static BUILTIN_FUNCTION_DOC: LazyLock<FxHashMap<SmolStr, BuiltinFunctionDoc>
             params: &["value", "needle"],
         },
     );
+    map.insert(
+        SmolStr::new("token_count"),
+        BuiltinFunctionDoc {
+            description: "Estimates how many LLM tokens the given text would consume, for context-window budgeting. Uses a lightweight chars-per-token heuristic by default; built with the `tiktoken` Cargo feature, counts exactly via tiktoken-rs instead when `model` (e.g. \"gpt-5\") is given. `model` is optional; without it, the heuristic estimate is always used.",
+            params: &["text", "model?"],
+        },
+    );
     map.insert(
         SmolStr::new("join"),
         BuiltinFunctionDoc {
@@ -6485,6 +6526,10 @@ mod tests {
     #[rstest]
     #[case("type", vec![RuntimeValue::String("test".into())], Ok(RuntimeValue::String("string".into())))]
     #[case("len", vec![RuntimeValue::String("test".into())], Ok(RuntimeValue::Number(4.into())))]
+    #[case("token_count", vec![RuntimeValue::String("Hello, world!".into()), RuntimeValue::String("gpt-4".into())], Ok(RuntimeValue::Number(4.into())))]
+    #[case("token_count", vec![RuntimeValue::String("".into()), RuntimeValue::String("gpt-4".into())], Ok(RuntimeValue::Number(0.into())))]
+    #[case("token_count", vec![RuntimeValue::String("Hello, world!".into())], Ok(RuntimeValue::Number(4.into())))]
+    #[case("token_count", vec![RuntimeValue::String("".into())], Ok(RuntimeValue::Number(0.into())))]
     #[case("abs", vec![RuntimeValue::Number((-10).into())], Ok(RuntimeValue::Number(10.into())))]
     #[case("ceil", vec![RuntimeValue::Number(3.2.into())], Ok(RuntimeValue::Number(4.0.into())))]
     #[case("floor", vec![RuntimeValue::Number(3.8.into())], Ok(RuntimeValue::Number(3.0.into())))]

crates/mq-lang/src/eval/builtin/tokenizer.rs

@@ -0,0 +1,148 @@
+//! `token_count` builtin. Two-tier: a dependency-free chars-per-token heuristic by default, or
+//! exact counts via `tiktoken-rs` behind the opt-in `tiktoken` feature. The exact tier pulls in
+//! several MB of vendored vocabulary data once `model` is resolved at runtime (every encoding
+//! becomes reachable, so none can be dead-code-eliminated), hence the feature gate.
+//!
+//! The `model` argument is optional: `token_count(text)` always uses the heuristic estimate
+//! (via [`token_count_estimate`]), so callers who don't care about an exact, model-specific
+//! count aren't forced to name one just to satisfy the arity, whether or not `tiktoken` is
+//! compiled in.
+
+use super::Error;
+
+mod heuristic {
+    // Rough chars-per-token by script; CJK tokenizes much denser than Latin under typical BPE.
+    const CJK_CHARS_PER_TOKEN: f64 = 1.6;
+    const LATIN_CHARS_PER_TOKEN: f64 = 4.0;
+    const OTHER_CHARS_PER_TOKEN: f64 = 2.5;
+
+    pub(super) fn is_cjk(c: char) -> bool {
+        matches!(c as u32,
+            0x3040..=0x30FF   // Hiragana, Katakana
+            | 0x3130..=0x318F // Hangul Compatibility Jamo
+            | 0x3400..=0x4DBF // CJK Unified Ideographs Extension A
+            | 0x4E00..=0x9FFF // CJK Unified Ideographs
+            | 0xAC00..=0xD7A3 // Hangul Syllables
+            | 0xF900..=0xFAFF // CJK Compatibility Ideographs
+            | 0xFF00..=0xFFEF // Halfwidth and Fullwidth Forms
+        )
+    }
+
+    pub(super) fn is_latin(c: char) -> bool {
+        (c as u32) < 0x0250 // Basic Latin, Latin-1 Supplement, Latin Extended-A/B
+    }
+
+    pub(super) fn token_count(text: &str) -> usize {
+        if text.is_empty() {
+            return 0;
+        }
+
+        let estimate: f64 = text
+            .chars()
+            .map(|c| {
+                if is_cjk(c) {
+                    1.0 / CJK_CHARS_PER_TOKEN
+                } else if is_latin(c) {
+                    1.0 / LATIN_CHARS_PER_TOKEN
+                } else {
+                    1.0 / OTHER_CHARS_PER_TOKEN
+                }
+            })
+            .sum();
+
+        estimate.ceil() as usize
+    }
+}
+
+/// `model` must be one of tiktoken-rs's recognized names (e.g. `"gpt-5"`, `"gpt-4"`,
+/// `"text-embedding-3-small"`); unrecognized names error rather than falling back silently.
+/// Uses `encode_ordinary` so special-token syntax in `text` (e.g. `<|endoftext|>`) is counted
+/// as plain text, since `text` is arbitrary document content, not a trusted prompt.
+#[cfg(feature = "tiktoken")]
+fn exact_token_count(text: &str, model: &str) -> Result<usize, Error> {
+    tiktoken_rs::bpe_for_model(model)
+        .map(|bpe| bpe.encode_ordinary(text).len())
+        .map_err(|e| Error::Runtime(format!("token_count: {e}")))
+}
+
+/// No-model overload: always the dependency-free heuristic estimate, regardless of the
+/// `tiktoken` feature. Lets callers who don't need an exact, model-specific count avoid
+/// naming one at all.
+pub(super) fn token_count_estimate(text: &str) -> usize {
+    heuristic::token_count(text)
+}
+
+pub(super) fn token_count(text: &str, model: &str) -> Result<usize, Error> {
+    #[cfg(feature = "tiktoken")]
+    {
+        exact_token_count(text, model)
+    }
+    #[cfg(not(feature = "tiktoken"))]
+    {
+        let _ = model;
+        Ok(heuristic::token_count(text))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::heuristic::{is_cjk, is_latin, token_count as heuristic_token_count};
+    #[cfg(feature = "tiktoken")]
+    use super::*;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case::empty("", 0)]
+    #[case::single_ascii_char("a", 1)]
+    #[case::short_english("Hello, world!", 4)]
+    fn test_heuristic_token_count(#[case] text: &str, #[case] expected: usize) {
+        assert_eq!(heuristic_token_count(text), expected);
+    }
+
+    #[test]
+    fn test_heuristic_token_count_cjk_denser_than_latin() {
+        let english = "hello world";
+        let japanese = "こんにちは世界";
+        assert!(english.chars().count() >= japanese.chars().count());
+        assert!(heuristic_token_count(japanese) > heuristic_token_count(english) / 2);
+    }
+
+    #[test]
+    fn test_is_cjk() {
+        assert!(is_cjk('あ'));
+        assert!(is_cjk('ア'));
+        assert!(is_cjk('漢'));
+        assert!(is_cjk('한'));
+        assert!(!is_cjk('a'));
+        assert!(!is_cjk('1'));
+    }
+
+    #[test]
+    fn test_is_latin() {
+        assert!(is_latin('a'));
+        assert!(is_latin('Z'));
+        assert!(is_latin('é'));
+        assert!(!is_latin('漢'));
+        assert!(!is_latin('Ж')); // Cyrillic
+    }
+
+    #[cfg(feature = "tiktoken")]
+    #[rstest]
+    #[case::empty("", "gpt-4", 0)]
+    #[case::simple("Hello, world!", "gpt-4", 4)]
+    fn test_exact_token_count(#[case] text: &str, #[case] model: &str, #[case] expected: usize) {
+        assert_eq!(exact_token_count(text, model).unwrap(), expected);
+    }
+
+    #[cfg(feature = "tiktoken")]
+    #[test]
+    fn test_exact_token_count_unknown_model() {
+        assert!(exact_token_count("hello", "not-a-real-model").is_err());
+    }
+
+    #[cfg(feature = "tiktoken")]
+    #[test]
+    fn test_exact_token_count_ignores_special_tokens_in_text() {
+        assert!(exact_token_count("<|endoftext|>", "gpt-4").unwrap() > 0);
+    }
+}

crates/mq-lang/tests/integration_tests.rs

@@ -2731,6 +2731,14 @@ fn engine() -> DefaultEngine {
 #[case::sort_array("sort([3, 1, 2])", vec![RuntimeValue::None], Ok(vec![RuntimeValue::Array(vec![RuntimeValue::Number(1.into()), RuntimeValue::Number(2.into()), RuntimeValue::Number(3.into())])].into()))]
 #[case::utf8bytelen_simple(r##"utf8bytelen("あ")"##, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(3.into())].into()))]
 #[case::rindex_simple(r##"rindex("banana", "a")"##, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(5.into())].into()))]
+#[case::token_count_empty(r#"token_count("", "gpt-4")"#, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(0.into())].into()))]
+#[case::token_count_simple(r#"token_count("Hello, world!", "gpt-4")"#, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(4.into())].into()))]
+#[case::token_count_markdown(r#"token_count(to_md_text("Hello, world!"), "gpt-4")"#, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(4.into())].into()))]
+#[case::token_count_none(r#"token_count(None, "gpt-4")"#, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(0.into())].into()))]
+#[case::token_count_no_model_empty(r#"token_count("")"#, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(0.into())].into()))]
+#[case::token_count_no_model_simple(r#"token_count("Hello, world!")"#, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(4.into())].into()))]
+#[case::token_count_no_model_markdown(r#"to_md_text("Hello, world!") | token_count()"#, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(4.into())].into()))]
+#[case::token_count_no_model_none(r#"token_count(None)"#, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Number(0.into())].into()))]
 #[case::explode_simple(r##"explode("abc")"##, vec![RuntimeValue::None], Ok(vec![RuntimeValue::Array(vec![RuntimeValue::Number(97.into()), RuntimeValue::Number(98.into()), RuntimeValue::Number(99.into())])].into()))]
 #[case::implode_simple(r##"implode([97, 98, 99])"##, vec![RuntimeValue::None], Ok(vec![RuntimeValue::String("abc".to_string())].into()))]
 #[case::intern_simple(r##"intern("foo")"##, vec![RuntimeValue::None], Ok(vec![RuntimeValue::String("foo".to_string())].into()))]
@@ -3673,6 +3681,9 @@ fn test_eval(mut engine: Engine, #[case] program: &str, #[case] input: Vec<Runti
 #[case::html_unescape_non_string("html_unescape(42)", vec![RuntimeValue::None],)]
 // strip_tags: non-string/non-markdown/non-none → type error
 #[case::strip_tags_non_string("strip_tags(42)", vec![RuntimeValue::None],)]
+// token_count: non-string/non-markdown text → type error
+#[case::token_count_non_string(r#"token_count(42, "gpt-4")"#, vec![RuntimeValue::None],)]
+#[case::token_count_no_model_non_string(r#"token_count(42)"#, vec![RuntimeValue::None],)]
 // range: multi-char string with step → error
 #[case::range_multichar_with_step(r#"range("aa", "zz", 2)"#, vec![RuntimeValue::None],)]
 // range: invalid type → error

crates/mq-run/Cargo.toml

@@ -17,6 +17,7 @@ debugger = ["mq-lang/debugger", "dep:rustyline", "dep:strum", "dep:regex-lite",
 default = ["std", "use_mimalloc", "http-import"]
 http-import = ["mq-lang/http-import-ureq"]
 std = []
+tiktoken = ["mq-lang/tiktoken"]
 use_mimalloc = ["mimalloc"]
 
 [dependencies]