From 8cfcc70ad16a1223b3d9c1eb6defe820671024f6 Mon Sep 17 00:00:00 2001 From: harehare Date: Mon, 13 Jul 2026 21:25:31 +0900 Subject: [PATCH 1/2] =?UTF-8?q?=E2=9C=A8=20feat(mq-lang,mq-run):=20add=20t?= =?UTF-8?q?oken=5Fcount(text,=20model)=20builtin?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a native token_count(text, model) builtin for LLM context-window budgeting (e.g. select(token_count(.) < N)), pairing with the existing toon.mq module for minimizing token usage. Two-tier: a dependency-free chars-per-token heuristic (adjusted by detected script — CJK tokenizes much denser than Latin) ships by default and costs nothing in binary size, including for mq-wasm. Exact counting via tiktoken-rs is opt-in behind a new `tiktoken` Cargo feature (mq-lang and, forwarded, mq-run), since resolving `model` at runtime makes every vendored encoding reachable and unremovable by dead-code elimination. --- Cargo.lock | 27 ++++ crates/mq-check/src/builtin.rs | 3 + crates/mq-lang/Cargo.toml | 2 + crates/mq-lang/src/eval/builtin.rs | 32 +++++ crates/mq-lang/src/eval/builtin/tokenizer.rs | 140 +++++++++++++++++++ crates/mq-lang/tests/integration_tests.rs | 6 + crates/mq-run/Cargo.toml | 1 + 7 files changed, 211 insertions(+) create mode 100644 crates/mq-lang/src/eval/builtin/tokenizer.rs diff --git a/Cargo.lock b/Cargo.lock index 0d931f2f0..0ed4700d1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1309,6 +1309,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fancy-regex" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fantoccini" version = "0.22.1" @@ -2635,6 +2646,7 @@ dependencies = [ "strsim", "tempfile", "thiserror 2.0.18", + "tiktoken-rs", "toml 1.1.2+spec-1.1.0", "toon-format", "ureq", @@ -4630,6 +4642,21 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "tiktoken-rs" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "027853bbf8c7763b77c5c595f1c271c7d536ced7d6f83452911b944621e57fc2" +dependencies = [ + "anyhow", + "base64", + "bstr", + "fancy-regex", + "lazy_static", + "regex", + "rustc-hash", +] + [[package]] name = "time" version = "0.3.53" diff --git a/crates/mq-check/src/builtin.rs b/crates/mq-check/src/builtin.rs index cd9e254a4..e1b19b1ad 100644 --- a/crates/mq-check/src/builtin.rs +++ b/crates/mq-check/src/builtin.rs @@ -312,6 +312,7 @@ 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_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 +1547,8 @@ 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_number("token_count(42, \"gpt-4\")", false)] // Should fail: wrong type #[case::capture("capture(\"hello 42\", \"(?P\\\\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)] diff --git a/crates/mq-lang/Cargo.toml b/crates/mq-lang/Cargo.toml index cf643660d..150cc21d5 100644 --- a/crates/mq-lang/Cargo.toml +++ b/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} diff --git a/crates/mq-lang/src/eval/builtin.rs b/crates/mq-lang/src/eval/builtin.rs index 4363dfff0..45ed3afd4 100644 --- a/crates/mq-lang/src/eval/builtin.rs +++ b/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,27 @@ fn utf8bytelen_impl(_: &Ident, _: &RuntimeValue, args: Args, _: &SharedEnv) -> R } } +/// Counts (or, with the `tiktoken` Cargo feature, exactly counts) the LLM tokens `text` would +/// consume for `model`. See [`tokenizer`] for the heuristic/exact two-tier design. +#[mq_macros::mq_fn(name = "token_count", params = Fixed(2))] +fn token_count_impl(ident: &Ident, _: &RuntimeValue, mut args: Args, _: &SharedEnv) -> Result { + match args.as_mut_slice() { + [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, b] => Err(Error::InvalidTypes( + ident.to_string(), + vec![std::mem::take(a), std::mem::take(b)], + )), + _ => unreachable!("token_count should always receive exactly two arguments"), + } +} + #[mq_macros::mq_fn(name = "rindex", params = Fixed(2))] fn rindex_impl(ident: &Ident, _: &RuntimeValue, mut args: Args, _: &SharedEnv) -> Result { match args.as_mut_slice() { @@ -4172,6 +4194,7 @@ mq_macros::builtin_dispatch! { INDEX, LEN, UTF8BYTELEN, + TOKEN_COUNT, RINDEX, RANGE, DEL, @@ -5358,6 +5381,13 @@ pub static BUILTIN_FUNCTION_DOC: LazyLock params: &["value", "needle"], }, ); + map.insert( + SmolStr::new("token_count"), + BuiltinFunctionDoc { + description: "Estimates how many LLM tokens the given text would consume for `model` (e.g. \"gpt-5\"), 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.", + params: &["text", "model"], + }, + ); map.insert( SmolStr::new("join"), BuiltinFunctionDoc { @@ -6485,6 +6515,8 @@ 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("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())))] diff --git a/crates/mq-lang/src/eval/builtin/tokenizer.rs b/crates/mq-lang/src/eval/builtin/tokenizer.rs new file mode 100644 index 000000000..82a2daee6 --- /dev/null +++ b/crates/mq-lang/src/eval/builtin/tokenizer.rs @@ -0,0 +1,140 @@ +//! `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. + +use super::Error; + +/// Kept in its own module (instead of `#[cfg(not(feature = "tiktoken"))]`-ing it out) so it's +/// still compiled and tested under `--all-features`, even though [`token_count`] only calls it +/// when `tiktoken` is off. +#[cfg_attr(feature = "tiktoken", allow(dead_code))] +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 { + tiktoken_rs::bpe_for_model(model) + .map(|bpe| bpe.encode_ordinary(text).len()) + .map_err(|e| Error::Runtime(format!("token_count: {e}"))) +} + +pub(super) fn token_count(text: &str, model: &str) -> Result { + #[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); + } +} diff --git a/crates/mq-lang/tests/integration_tests.rs b/crates/mq-lang/tests/integration_tests.rs index a84ceda59..1c8a1fbcf 100644 --- a/crates/mq-lang/tests/integration_tests.rs +++ b/crates/mq-lang/tests/integration_tests.rs @@ -2731,6 +2731,10 @@ 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#"to_md_text("Hello, world!") | token_count("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::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 +3677,8 @@ fn test_eval(mut engine: Engine, #[case] program: &str, #[case] input: Vec Date: Mon, 13 Jul 2026 23:35:08 +0900 Subject: [PATCH 2/2] =?UTF-8?q?=E2=9C=A8=20feat(mq-lang,mq-check):=20make?= =?UTF-8?q?=20model=20optional=20in=20token=5Fcount?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Support token_count(text) alongside token_count(text, model), so callers who only want the heuristic estimate aren't forced to name a model just to satisfy arity. Since builtin dispatch resolves pipe injection purely by argument count, this collides with the previous text | token_count(model) idiom (1 explicit arg); that form is replaced by explicit two-arg calls, e.g. token_count(md, "gpt-4"). --- crates/mq-check/src/builtin.rs | 2 ++ crates/mq-lang/src/eval/builtin.rs | 23 +++++++++++++++----- crates/mq-lang/src/eval/builtin/tokenizer.rs | 16 ++++++++++---- crates/mq-lang/tests/integration_tests.rs | 7 +++++- 4 files changed, 38 insertions(+), 10 deletions(-) diff --git a/crates/mq-check/src/builtin.rs b/crates/mq-check/src/builtin.rs index e1b19b1ad..26766563f 100644 --- a/crates/mq-check/src/builtin.rs +++ b/crates/mq-check/src/builtin.rs @@ -312,6 +312,7 @@ 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 @@ -1548,6 +1549,7 @@ mod tests { #[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\\\\w+)\")", true)] #[case::is_regex_match("is_regex_match(\"hello123\", \"[0-9]+\")", true)] diff --git a/crates/mq-lang/src/eval/builtin.rs b/crates/mq-lang/src/eval/builtin.rs index 45ed3afd4..3210d872f 100644 --- a/crates/mq-lang/src/eval/builtin.rs +++ b/crates/mq-lang/src/eval/builtin.rs @@ -1528,10 +1528,20 @@ fn utf8bytelen_impl(_: &Ident, _: &RuntimeValue, args: Args, _: &SharedEnv) -> R } /// Counts (or, with the `tiktoken` Cargo feature, exactly counts) the LLM tokens `text` would -/// consume for `model`. See [`tokenizer`] for the heuristic/exact two-tier design. -#[mq_macros::mq_fn(name = "token_count", params = Fixed(2))] +/// 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 { 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())) } @@ -1540,11 +1550,12 @@ fn token_count_impl(ident: &Ident, _: &RuntimeValue, mut args: Args, _: &SharedE .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 exactly two arguments"), + _ => unreachable!("token_count should always receive one or two arguments"), } } @@ -5384,8 +5395,8 @@ pub static BUILTIN_FUNCTION_DOC: LazyLock map.insert( SmolStr::new("token_count"), BuiltinFunctionDoc { - description: "Estimates how many LLM tokens the given text would consume for `model` (e.g. \"gpt-5\"), 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.", - params: &["text", "model"], + 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( @@ -6517,6 +6528,8 @@ mod tests { #[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())))] diff --git a/crates/mq-lang/src/eval/builtin/tokenizer.rs b/crates/mq-lang/src/eval/builtin/tokenizer.rs index 82a2daee6..beb7f6367 100644 --- a/crates/mq-lang/src/eval/builtin/tokenizer.rs +++ b/crates/mq-lang/src/eval/builtin/tokenizer.rs @@ -2,13 +2,14 @@ //! 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; -/// Kept in its own module (instead of `#[cfg(not(feature = "tiktoken"))]`-ing it out) so it's -/// still compiled and tested under `--all-features`, even though [`token_count`] only calls it -/// when `tiktoken` is off. -#[cfg_attr(feature = "tiktoken", allow(dead_code))] 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; @@ -64,6 +65,13 @@ fn exact_token_count(text: &str, model: &str) -> Result { .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 { #[cfg(feature = "tiktoken")] { diff --git a/crates/mq-lang/tests/integration_tests.rs b/crates/mq-lang/tests/integration_tests.rs index 1c8a1fbcf..1c803479f 100644 --- a/crates/mq-lang/tests/integration_tests.rs +++ b/crates/mq-lang/tests/integration_tests.rs @@ -2733,8 +2733,12 @@ fn engine() -> DefaultEngine { #[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#"to_md_text("Hello, world!") | token_count("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()))] @@ -3679,6 +3683,7 @@ fn test_eval(mut engine: Engine, #[case] program: &str, #[case] input: Vec