diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index b8d1f2db4d7..17856218e60 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,18 +1,19 @@ # AGENTS.md -litellm-rust has five crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. +litellm-rust has six crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. ## Crates | Crate | Role | |-------|------| | litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. | +| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. | | litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. | | litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. | | litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`. The interop foundation depends on no LiteLLM domain crate. +Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers, `litellm-token-counter`, and `litellm-python-interop`. The token counter and interop foundations depend on no LiteLLM domain crate. ## Where a route lives diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7addc4e3e45..0e8e6e09a21 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,20 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.5" @@ -418,7 +432,7 @@ checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" dependencies = [ "async-trait", "axum-core", - "base64", + "base64 0.22.1", "bytes", "futures-util", "http 1.4.2", @@ -468,6 +482,12 @@ dependencies = [ "tracing", ] +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + [[package]] name = "base64" version = "0.22.1" @@ -542,6 +562,15 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + [[package]] name = "cc" version = "1.3.0" @@ -644,6 +673,21 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + [[package]] name = "const-oid" version = "0.10.2" @@ -696,7 +740,7 @@ dependencies = [ "ciborium", "clap", "criterion-plot", - "itertools", + "itertools 0.13.0", "num-traits", "oorandom", "page_size", @@ -716,7 +760,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" dependencies = [ "cast", - "itertools", + "itertools 0.13.0", ] [[package]] @@ -778,6 +822,56 @@ dependencies = [ "cmov", ] +[[package]] +name = "daachorse" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5614204febbc33cc07a2806aa6440b904ac012b68eecc37f4493ea4a76455a3d" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.119", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + [[package]] name = "data-encoding" version = "2.11.0" @@ -790,6 +884,37 @@ version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn 2.0.119", +] + [[package]] name = "digest" version = "0.10.7" @@ -841,6 +966,12 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + [[package]] name = "fastrand" version = "2.5.0" @@ -964,6 +1095,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + [[package]] name = "getrandom" version = "0.4.3" @@ -973,7 +1116,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", + "r-efi 6.0.0", "rand_core 0.10.1", "wasm-bindgen", ] @@ -1220,7 +1363,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -1319,6 +1462,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "idna" version = "1.1.0" @@ -1348,6 +1497,8 @@ checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", "hashbrown", + "serde", + "serde_core", ] [[package]] @@ -1365,6 +1516,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1409,7 +1569,7 @@ name = "litellm-ai-gateway" version = "0.1.0" dependencies = [ "axum", - "base64", + "base64 0.22.1", "futures-channel", "futures-util", "litellm-config", @@ -1447,7 +1607,7 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", - "base64", + "base64 0.22.1", "rand 0.8.7", "reqwest", "rstest", @@ -1471,6 +1631,7 @@ dependencies = [ "litellm-ai-gateway", "litellm-core", "litellm-python-interop", + "litellm-token-counter", "pyo3", "pyo3-async-runtimes", "serde", @@ -1491,6 +1652,22 @@ dependencies = [ "serde_json", ] +[[package]] +name = "litellm-token-counter" +version = "0.1.0" +dependencies = [ + "criterion", + "indexmap", + "itoa", + "rand 0.8.7", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokenizers", + "unicode-normalization-alignments", +] + [[package]] name = "litemap" version = "0.8.2" @@ -1509,6 +1686,22 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "macro_rules_attribute" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3ae8f6d608c795738406608304d30a2dfbdc8e58e44f7ba43236da5208ded3c" +dependencies = [ + "macro_rules_attribute-proc_macro", + "pastey", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c" + [[package]] name = "matchit" version = "0.7.3" @@ -1537,6 +1730,12 @@ dependencies = [ "unicase", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "mio" version = "1.2.2" @@ -1548,6 +1747,38 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -1578,6 +1809,28 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "oorandom" version = "11.1.5" @@ -1606,6 +1859,18 @@ dependencies = [ "winapi", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1852,6 +2117,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r-efi" version = "6.0.0" @@ -1865,10 +2136,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", - "rand_chacha", + "rand_chacha 0.3.1", "rand_core 0.6.4", ] +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.2" @@ -1890,6 +2171,16 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + [[package]] name = "rand_core" version = "0.6.4" @@ -1899,6 +2190,15 @@ dependencies = [ "getrandom 0.2.17", ] +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + [[package]] name = "rand_core" version = "0.10.1" @@ -1924,6 +2224,17 @@ dependencies = [ "rayon-core", ] +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools 0.14.0", + "rayon", +] + [[package]] name = "rayon-core" version = "1.13.0" @@ -1981,7 +2292,7 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-core", @@ -2363,12 +2674,36 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom", + "serde", + "unicode-segmentation", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "subtle" version = "2.6.1" @@ -2537,6 +2872,39 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tokenizers" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7afbf6e88718afcc138bad01d6ccc3051dbbc3b2ce9793d8b8a3aeb610969cfc" +dependencies = [ + "ahash", + "compact_str", + "daachorse", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "itertools 0.14.0", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.9.5", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.19", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokio" version = "1.53.0" @@ -2775,6 +3143,27 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "untrusted" version = "0.9.0" @@ -2858,6 +3247,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "wasm-bindgen" version = "0.2.126" @@ -3083,6 +3481,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + [[package]] name = "writeable" version = "0.6.3" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 00fa6faaaef..f3e54e5b2aa 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "crates/core", + "crates/token-counter", "crates/config", "crates/ai-gateway", "crates/python-interop", @@ -18,6 +19,7 @@ repository = "https://github.com/BerriAI/litellm" tracing = "0.1" tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } litellm-core = { path = "crates/core" } +litellm-token-counter = { path = "crates/token-counter" } litellm-config = { path = "crates/config" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } litellm-python-interop = { path = "crates/python-interop" } @@ -40,6 +42,7 @@ tokio-tungstenite = { version = "0.24", default-features = false, features = ["c futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" url = "2.5.8" +criterion = "0.8.2" [profile.release] opt-level = 3 diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 9fef59a277d..cbcd8119546 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -6,17 +6,18 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. ## Crates -`litellm-rust` has five crates. A crate is a layer or shared foundation, not a route: +`litellm-rust` has six crates. A crate is a layer or shared foundation, not a route: | Crate | Role | |-------|------| | litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. | +| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. | | litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. | | litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | | litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. | -Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop. +Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers, token counter, and Python interop. - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs index fc0ab2b62a3..e7739fe7312 100644 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs @@ -1,6 +1,7 @@ -//! Enforcement: the litellm-rust workspace has exactly five crates. +//! Enforcement: the litellm-rust workspace has exactly six crates. //! -//! `core` (the Rust SDK), `config` (the config-loading boundary), +//! `core` (the Rust SDK), `token-counter` (standalone input token counting), +//! `config` (the config-loading boundary), //! `ai-gateway` (the HTTP/WebSocket host), //! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the //! PyO3 cdylib). Adding or removing a crate must be a @@ -20,6 +21,7 @@ use std::path::{Path, PathBuf}; /// workspace legitimately gains or loses a crate. const EXPECTED_MEMBERS: &[&str] = &[ "crates/core", + "crates/token-counter", "crates/config", "crates/ai-gateway", "crates/python-interop", @@ -29,6 +31,7 @@ const EXPECTED_MEMBERS: &[&str] = &[ /// The crate subdirectory names that must exist under `crates/`. const EXPECTED_CRATE_DIRS: &[&str] = &[ "core", + "token-counter", "config", "ai-gateway", "python-interop", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index bda09a7d840..337a1e8e5ac 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -24,16 +24,17 @@ trace-parity = [ futures-util.workspace = true tracing = { workspace = true, optional = true } litellm-core = { workspace = true, features = ["bedrock-auth"] } +litellm-token-counter.workspace = true litellm-ai-gateway = { workspace = true, default-features = false } litellm-python-interop.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true serde.workspace = true serde_json.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["sync"] } [dev-dependencies] -criterion = "0.8.2" +criterion.workspace = true tokio-tungstenite.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/constants.rs b/litellm-rust/crates/python-bridge/src/constants.rs new file mode 100644 index 00000000000..d5cf5749820 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/constants.rs @@ -0,0 +1,2 @@ +/// Concurrent token-count encodes allowed when the core count is unavailable. +pub(crate) const TOKEN_COUNT_FALLBACK_PARALLELISM: usize = 1; diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index f3648158cf6..b57197b9ddf 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -3,7 +3,6 @@ use std::panic::AssertUnwindSafe; use std::time::Duration; use futures_util::FutureExt; -use litellm_core::error::Error; use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil}; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; @@ -11,14 +10,15 @@ use serde::Serialize; use tokio::runtime::{Handle, Runtime}; use tokio::time::{self, MissedTickBehavior}; -pub(crate) fn run_sync( +pub(crate) fn run_sync( py: Python<'_>, future: F, - map_error: fn(Error) -> PyErr, + map_error: fn(E) -> PyErr, ) -> PyResult> where T: Serialize + Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { run_sync_on( py, @@ -28,15 +28,16 @@ where ) } -fn run_sync_on( +fn run_sync_on( py: Python<'_>, runtime: &Runtime, future: F, - map_error: fn(Error) -> PyErr, + map_error: fn(E) -> PyErr, ) -> PyResult> where T: Serialize + Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { if Handle::try_current().is_ok() { return Err(PyRuntimeError::new_err( @@ -49,14 +50,15 @@ where Pythonized(result).into_pyobject(py).map(Bound::unbind) } -pub(crate) fn run_async( +pub(crate) fn run_async( py: Python<'_>, future: F, - map_error: fn(Error) -> PyErr, + map_error: fn(E) -> PyErr, ) -> PyResult> where T: Serialize + Send + 'static, - F: Future> + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, { pyo3_async_runtimes::tokio::future_into_py(py, async move { let result = catch_future_panic(future).await?; @@ -65,7 +67,7 @@ where }) } -fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { +fn map_core_result(result: Result, map_error: fn(E) -> PyErr) -> PyResult { match result { Ok(value) => Ok(value), Err(error) => Err( @@ -75,9 +77,9 @@ fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) - } } -async fn catch_future_panic(future: F) -> PyResult> +async fn catch_future_panic(future: F) -> PyResult> where - F: Future>, + F: Future>, { AssertUnwindSafe(future) .catch_unwind() @@ -85,9 +87,9 @@ where .map_err(panic_to_pyerr) } -async fn wait_for_sync_result(future: F) -> PyResult> +async fn wait_for_sync_result(future: F) -> PyResult> where - F: Future>, + F: Future>, { let future = catch_future_panic(future); tokio::pin!(future); @@ -114,6 +116,7 @@ mod tests { use std::thread; use std::time::Instant; + use litellm_core::error::Error; use pyo3::panic::PanicException; use pyo3::types::{PyDict, PyModule}; use serde::Serializer; @@ -237,7 +240,7 @@ mod tests { let error = runtime.block_on(async { Python::attach(|py| { - run_sync::(py, async { Ok(true) }, runtime_error) + run_sync::(py, async { Ok(true) }, runtime_error) .expect_err("sync route should reject a nested Tokio runtime") }) }); @@ -273,7 +276,7 @@ mod tests { fn sync_runner_maps_a_panicked_future() { Python::initialize(); Python::attach(|py| { - let error = run_sync::( + let error = run_sync::( py, poll_fn(|_| -> Poll> { panic!("route future panicked") }), runtime_error, @@ -289,7 +292,7 @@ mod tests { fn sync_runner_maps_a_panicked_error_mapper() { Python::initialize(); Python::attach(|py| { - let error = run_sync::( + let error = run_sync::( py, async { Err(Error::InvalidRequest("invalid".to_string())) }, panicking_error_mapper, diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 384f0be5a1b..cf0450a1b30 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,3 +1,4 @@ +mod constants; mod diagnostics; mod errors; mod execution; @@ -5,6 +6,7 @@ mod execution; mod function_trace; mod marshal; mod routes; +mod token_counter; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use pyo3::prelude::*; @@ -71,6 +73,7 @@ mod _native { super::errors::register(module)?; super::routes::register(module)?; module.add_class::()?; + super::token_counter::register(module)?; super::diagnostics::register(module) } } @@ -106,6 +109,7 @@ mod tests { "chat_completions", "achat_completions", "ResponsesWebSocketConnection", + "TokenCounter", "gil_stats", ]; diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs new file mode 100644 index 00000000000..ee82c170d55 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -0,0 +1,87 @@ +use std::num::NonZero; +use std::sync::Arc; +use std::thread::available_parallelism; + +use litellm_python_interop::release_gil; +use litellm_token_counter::{ + CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, +}; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyAny; +use tokio::sync::Semaphore; + +use crate::constants::TOKEN_COUNT_FALLBACK_PARALLELISM; +use crate::errors::RustBridgeDeclined; +use crate::execution::run_async; + +/// Counts the input tokens of a raw request body off the Python event loop with +/// the GIL released. Python owns which requests get here and what to do with +/// the count. At most one encode per core runs at a time; the rest wait in the +/// async task, where a cancelled Python awaiter drops them before any blocking +/// work is scheduled. +#[pyclass(frozen)] +struct TokenCounter { + inner: Arc, + encode_slots: Arc, +} + +#[pymethods] +impl TokenCounter { + #[new] + fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { + let inner = release_gil(py, || CoreTokenCounter::from_json(tokenizer_json)) + .map_err(token_count_error_to_pyerr)?; + Ok(Self { + inner: Arc::new(inner), + encode_slots: Arc::new(Semaphore::new(encode_parallelism())), + }) + } + + fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult> { + let counter = Arc::clone(&self.inner); + let encode_slots = Arc::clone(&self.encode_slots); + let body = body.to_vec(); + run_async( + py, + async move { + let _slot = encode_slots + .acquire_owned() + .await + .map_err(|error| Error::Task(error.to_string()))?; + tokio::task::spawn_blocking(move || count_body(&counter, &body)) + .await + .map_err(|error| Error::Task(error.to_string()))? + }, + token_count_error_to_pyerr, + ) + } +} + +fn encode_parallelism() -> usize { + available_parallelism().map_or(TOKEN_COUNT_FALLBACK_PARALLELISM, NonZero::get) +} + +fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result { + let request = CountableRequest::parse(body)?; + counter.count_request(&request) +} + +fn token_count_error_to_pyerr(error: Error) -> PyErr { + let message = error.to_string(); + match error { + Error::Load(_) => PyValueError::new_err(message), + Error::RequestParse(_) + | Error::MissingInput + | Error::FloatText + | Error::ContentBlock + | Error::ArrayItems + | Error::JsonSerialization(_) + | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message), + Error::Encode(_) | Error::Task(_) => PyRuntimeError::new_err(message), + } +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::() +} diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml new file mode 100644 index 00000000000..61c9cf6e991 --- /dev/null +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "litellm-token-counter" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +indexmap = { version = "2.14.0", features = ["serde"] } +itoa = "1.0" +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +unicode-normalization-alignments = "0.1.12" + +[dev-dependencies] +criterion.workspace = true +rand.workspace = true +rstest.workspace = true + +[[bench]] +name = "token_counter" +harness = false + +[[bench]] +name = "allocations" +harness = false diff --git a/litellm-rust/crates/token-counter/benches/allocations.rs b/litellm-rust/crates/token-counter/benches/allocations.rs new file mode 100644 index 00000000000..343f815749d --- /dev/null +++ b/litellm-rust/crates/token-counter/benches/allocations.rs @@ -0,0 +1,128 @@ +use std::alloc::{GlobalAlloc, Layout, System}; +use std::hint::black_box; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use litellm_token_counter::{CountableRequest, TokenCounter}; + +struct CountingAllocator; + +static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0); +static BYTES: AtomicUsize = AtomicUsize::new(0); + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + BYTES.fetch_add(layout.size(), Ordering::Relaxed); + // SAFETY: This allocator delegates the unchanged layout to `System`. + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) { + // SAFETY: The pointer and layout came from the delegated `System` allocation. + unsafe { System.dealloc(pointer, layout) } + } + + unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + BYTES.fetch_add(size, Ordering::Relaxed); + // SAFETY: The pointer and layout came from `System`; the new size is unchanged. + unsafe { System.realloc(pointer, layout, size) } + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +#[derive(Clone, Copy)] +struct AllocationCount { + allocations: usize, + bytes: usize, +} + +impl AllocationCount { + fn assert_max(self, label: &str, maximum: Self) { + eprintln!( + "{label}: {} allocations, {} bytes", + self.allocations, self.bytes + ); + assert!( + self.allocations <= maximum.allocations, + "{label} allocation count exceeded {}", + maximum.allocations + ); + assert!( + self.bytes <= maximum.bytes, + "{label} allocated bytes exceeded {}", + maximum.bytes + ); + } +} + +const TOKENIZER_JSON: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" +)); +const OBJECT_BODY: &[u8] = br#"{"model":"claude-sonnet-4-5","input":{"text":"caf\u00e9","n":3,"ok":true,"list":[1,"a",{"z":[]}]}}"#; +const INTEGER_BODY: &[u8] = br#"{"input":[-9223372036854775808,0,18446744073709551615]}"#; + +fn measure(operation: impl FnOnce()) -> AllocationCount { + ALLOCATIONS.store(0, Ordering::Relaxed); + BYTES.store(0, Ordering::Relaxed); + operation(); + AllocationCount { + allocations: ALLOCATIONS.load(Ordering::Relaxed), + bytes: BYTES.load(Ordering::Relaxed), + } +} + +fn main() { + measure(|| { + black_box(CountableRequest::parse(OBJECT_BODY).expect("object request parses")); + }) + .assert_max( + "parse object request", + AllocationCount { + allocations: 16, + bytes: 1_900, + }, + ); + + let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("tokenizer loads"); + let object = CountableRequest::parse(OBJECT_BODY).expect("object request parses"); + counter + .count_request(&object) + .expect("object warmup succeeds"); + measure(|| { + black_box( + counter + .count_request(black_box(&object)) + .expect("object counts"), + ); + }) + .assert_max( + "count object request", + AllocationCount { + allocations: 74, + bytes: 2_200, + }, + ); + + let integers = CountableRequest::parse(INTEGER_BODY).expect("integer request parses"); + counter + .count_request(&integers) + .expect("integer warmup succeeds"); + measure(|| { + black_box( + counter + .count_request(black_box(&integers)) + .expect("integers count"), + ); + }) + .assert_max( + "count integer list", + AllocationCount { + allocations: 26, + bytes: 1_050, + }, + ); +} diff --git a/litellm-rust/crates/token-counter/benches/token_counter.rs b/litellm-rust/crates/token-counter/benches/token_counter.rs new file mode 100644 index 00000000000..a7c3177b88a --- /dev/null +++ b/litellm-rust/crates/token-counter/benches/token_counter.rs @@ -0,0 +1,100 @@ +use std::hint::black_box; +use std::time::Duration; + +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use litellm_token_counter::TokenCounter; +use tokenizers::Tokenizer; + +const TOKENIZER_JSON: &str = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" +)); +const FULL_CONTEXT_TOKENS: usize = 1_000_000; +const LARGE_PROMPT_UNIT: &str = "The quick brown fox jumps over the lazy dog. 0123456789\n"; + +fn full_context_input(tokenizer: &Tokenizer) -> String { + let unit_tokens = tokenizer + .encode_fast(LARGE_PROMPT_UNIT, false) + .expect("reference tokenizer should encode") + .len(); + let input = LARGE_PROMPT_UNIT.repeat(FULL_CONTEXT_TOKENS.div_ceil(unit_tokens)); + let actual_tokens = tokenizer + .encode_fast(input.as_str(), true) + .expect("reference tokenizer should encode") + .len(); + assert!((FULL_CONTEXT_TOKENS..FULL_CONTEXT_TOKENS + unit_tokens).contains(&actual_tokens)); + input +} + +fn inputs(tokenizer: &Tokenizer) -> Vec<(&'static str, String)> { + vec![ + ( + "ascii_chat", + "User: Summarize the benefits of statistical benchmarking.\nAssistant:".repeat(8), + ), + ( + "unicode_nfkc", + "Ａ quick café résumé: مرحبا 世界 🙂 fi Ⅳ.\n".repeat(16), + ), + ("large_prompt", LARGE_PROMPT_UNIT.repeat(256)), + ("full_context_1m_tokens", full_context_input(tokenizer)), + ] +} + +fn token_counter(c: &mut Criterion) { + let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("token counter should load"); + let tokenizer = TOKENIZER_JSON + .parse::() + .expect("reference tokenizer should load"); + let mut group = c.benchmark_group("anthropic_token_counter"); + + for (name, input) in inputs(&tokenizer) { + let expected = tokenizer + .encode_fast(input.as_str(), true) + .expect("reference tokenizer should encode") + .len(); + let actual = counter + .count_text(input.as_str()) + .expect("benchmark path should count"); + assert_eq!( + actual, expected, + "benchmark paths should produce the same count" + ); + group.throughput(Throughput::Bytes(input.len() as u64)); + group.bench_with_input( + BenchmarkId::new("byte_level_fast_path", name), + &input, + |b, input| { + b.iter(|| { + counter + .count_text(black_box(input.as_str())) + .expect("fast path should count") + }) + }, + ); + group.bench_with_input( + BenchmarkId::new("full_encoder", name), + &input, + |b, input| { + b.iter(|| { + tokenizer + .encode_fast(black_box(input.as_str()), true) + .expect("reference tokenizer should encode") + .len() + }) + }, + ); + } + + group.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(20) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(4)); + targets = token_counter +} +criterion_main!(benches); diff --git a/litellm-rust/crates/token-counter/src/byte_level.rs b/litellm-rust/crates/token-counter/src/byte_level.rs new file mode 100644 index 00000000000..2b435eb39bb --- /dev/null +++ b/litellm-rust/crates/token-counter/src/byte_level.rs @@ -0,0 +1,650 @@ +//! Exact token counting for a supported tokenizer configuration: optional +//! NFKC normalization, `ByteLevel` pre-tokenization with the GPT-2 split regex, +//! and no post-processing. A scanner reproduces the regex's piece boundaries +//! and hands each piece to the tokenizer's model. Unsupported configurations +//! and added-token inputs fall back to the full encoder. + +use std::borrow::Cow; +use std::iter; + +use tokenizers::normalizers::NormalizerWrapper; +use tokenizers::pre_tokenizers::PreTokenizerWrapper; +use tokenizers::{Model, Tokenizer}; +use unicode_normalization_alignments::{IsNormalized, UnicodeNormalization, is_nfkc_quick}; + +use super::unicode_classes::UnicodeClasses; + +const CONTRACTIONS: [&str; 7] = ["'s", "'t", "'re", "'ve", "'m", "'ll", "'d"]; + +pub(super) struct ByteLevelCounter { + nfkc: bool, + normalized_added_tokens: Vec, + unicode_classes: &'static UnicodeClasses, +} + +impl ByteLevelCounter { + pub(super) fn detect(tokenizer: &Tokenizer) -> Option { + let nfkc = match tokenizer.get_normalizer() { + None => false, + Some(NormalizerWrapper::NFKC(_)) => true, + Some(_) => return None, + }; + let Some(PreTokenizerWrapper::ByteLevel(byte_level)) = tokenizer.get_pre_tokenizer() else { + return None; + }; + let plain = !byte_level.add_prefix_space + && byte_level.use_regex + && tokenizer.get_post_processor().is_none() + && tokenizer.get_truncation().is_none() + && tokenizer.get_padding().is_none(); + if !plain { + return None; + } + let vocabulary = tokenizer.get_added_vocabulary(); + let normalized_added_tokens = vocabulary + .get_vocab() + .iter() + .filter_map(|(original, id)| { + vocabulary + .simple_id_to_token(*id) + .filter(|normalized| normalized != original) + }) + .collect(); + Some(Self { + nfkc, + normalized_added_tokens, + unicode_classes: UnicodeClasses::get()?, + }) + } + + /// `None` when the text contains an added token or the model rejects a + /// piece; the caller then runs the full encoder. + pub(super) fn count(&self, tokenizer: &Tokenizer, text: &str) -> Option { + let normalized = self.normalize(text); + let added_tokens = tokenizer.get_added_vocabulary().get_vocab(); + if added_tokens + .keys() + .chain(self.normalized_added_tokens.iter()) + .any(|token| text.contains(token.as_str()) || normalized.contains(token.as_str())) + { + return None; + } + let model = tokenizer.get_model(); + let mapped: String = normalized.bytes().map(byte_char).collect(); + pieces(&normalized, self.unicode_classes) + .try_fold((0, 0), |(start, total), piece| { + let end = start + mapped_len(piece); + let tokens = model.tokenize(&mapped[start..end]).ok()?; + Some((end, total + tokens.len())) + }) + .map(|(_, total)| total) + } + + /// Same crate and Unicode tables as `NormalizedString::nfkc`, so the + /// result is what the full encoder would have tokenized. + fn normalize<'a>(&self, text: &'a str) -> Cow<'a, str> { + if !self.nfkc || text.is_ascii() || is_nfkc_quick(text.chars()) == IsNormalized::Yes { + return Cow::Borrowed(text); + } + Cow::Owned(text.nfkc().map(|(character, _)| character).collect()) + } +} + +/// GPT-2 `bytes_to_unicode`: printable Latin-1 bytes map to themselves, the +/// rest to U+0100 onwards in byte order. +fn byte_char(byte: u8) -> char { + let code = match byte { + 0x21..=0x7E | 0xA1..=0xAC | 0xAE..=0xFF => u32::from(byte), + 0x00..=0x20 => 0x100 + u32::from(byte), + 0x7F..=0xA0 => 0x121 + u32::from(byte - 0x7F), + 0xAD => 0x143, + }; + char::from_u32(code).unwrap_or(char::REPLACEMENT_CHARACTER) +} + +fn mapped_len(piece: &str) -> usize { + piece.len() + + piece + .bytes() + .filter(|byte| !byte.is_ascii_graphic()) + .count() +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Class { + Letter, + Number, + Space, + Other, +} + +fn class(character: char, unicode_classes: &UnicodeClasses) -> Class { + match character { + 'A'..='Z' | 'a'..='z' => Class::Letter, + '0'..='9' => Class::Number, + '\t'..='\r' | ' ' => Class::Space, + _ if character.is_ascii() => Class::Other, + _ if unicode_classes.is_letter(character) => Class::Letter, + _ if unicode_classes.is_number(character) => Class::Number, + _ if unicode_classes.is_space(character) => Class::Space, + _ => Class::Other, + } +} + +/// The regex matches every character, so the pieces tile the text. +fn pieces<'a>( + text: &'a str, + unicode_classes: &'static UnicodeClasses, +) -> impl Iterator { + iter::successors(split_piece(text, unicode_classes), move |(_, rest)| { + split_piece(rest, unicode_classes) + }) + .map(|(piece, _)| piece) +} + +fn split_piece<'a>(text: &'a str, unicode_classes: &UnicodeClasses) -> Option<(&'a str, &'a str)> { + let first = text.chars().next()?; + Some(text.split_at(piece_len(text, first, unicode_classes))) +} + +fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize { + if let Some(contraction) = CONTRACTIONS.iter().find(|word| text.starts_with(**word)) { + return contraction.len(); + } + let first_class = class(first, unicode_classes); + if first_class != Class::Space { + return run_len(text, first_class, unicode_classes); + } + if first != ' ' { + return space_run_len(text, unicode_classes); + } + let after_space = &text[1..]; + match after_space + .chars() + .next() + .map(|character| class(character, unicode_classes)) + { + None | Some(Class::Space) => space_run_len(text, unicode_classes), + Some(run_class) => 1 + run_len(after_space, run_class, unicode_classes), + } +} + +fn run_len(text: &str, run_class: Class, unicode_classes: &UnicodeClasses) -> usize { + text.char_indices() + .find(|(_, character)| class(*character, unicode_classes) != run_class) + .map_or(text.len(), |(index, _)| index) +} + +/// `\s+(?!\S)|\s+`: whitespace followed by a non-space leaves its last +/// character to start the next piece (` ?` on the following alternatives). +fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize { + let run = run_len(text, Class::Space, unicode_classes); + if run == text.len() { + return run; + } + let last = text[..run].chars().next_back().map_or(0, char::len_utf8); + match run - last { + 0 => run, + shorter => shorter, + } +} + +#[cfg(test)] +mod tests { + use rand::rngs::StdRng; + use rand::seq::SliceRandom; + use rand::{Rng, SeedableRng}; + use rstest::{fixture, rstest}; + use tokenizers::normalizers::NFKC; + use tokenizers::pre_tokenizers::byte_level::ByteLevel; + use tokenizers::utils::SysRegex; + use tokenizers::{ + NormalizedString, Normalizer, OffsetReferential, OffsetType, PreTokenizedString, + PreTokenizer, + }; + + use super::*; + + #[fixture] + fn anthropic_tokenizer() -> Tokenizer { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" + ); + std::fs::read_to_string(path) + .expect("anthropic tokenizer json is in the repo") + .parse() + .expect("anthropic tokenizer loads") + } + + fn reference_count(tokenizer: &Tokenizer, text: &str) -> usize { + tokenizer.encode_fast(text, true).expect("encode").len() + } + + fn byte_level_counter(nfkc: bool) -> ByteLevelCounter { + ByteLevelCounter { + nfkc, + normalized_added_tokens: Vec::new(), + unicode_classes: UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"), + } + } + + const ALPHABET: &[&str] = &[ + "a", + "Z", + "e", + "s", + "t", + "d", + "m", + "'", + "'s", + "'re", + "'ll", + "'S", + "0", + "9", + " ", + " ", + "\t", + "\n", + "\r\n", + "\u{b}", + ".", + ",", + "!", + "-", + "(", + "\"", + "\u{a0}", + "\u{85}", + "\u{2028}", + "\u{3000}", + "\u{200b}", + "\u{200d}", + "é", + "e\u{301}", + "ß", + "漢", + "字", + "ع", + "३", + "½", + "Ⅳ", + "🙂", + "👍🏽", + "Ａ", + "fi", + "㍿", + "㋿", + "ꟲ", + "𐞁", + "a\u{30a}", + "\u{1e0b}\u{323}", + "<", + ">", + "EOT", + "", + "", + ]; + + fn random_text(rng: &mut StdRng) -> String { + let pieces = rng.gen_range(0..40); + (0..pieces) + .map(|_| *ALPHABET.choose(rng).expect("alphabet is not empty")) + .collect() + } + + #[rstest] + #[case::plain_text("Hello, how are you today?", true)] + #[case::added_token("stop here", false)] + #[case::normalized_added_token("stop ＜ＥＯＴ＞ here", false)] + fn anthropic_tokenizer_takes_the_fast_path( + anthropic_tokenizer: Tokenizer, + #[case] text: &str, + #[case] supported: bool, + ) { + let fast = + ByteLevelCounter::detect(&anthropic_tokenizer).expect("anthropic shape is supported"); + assert!(fast.nfkc); + let count = fast.count(&anthropic_tokenizer, text); + if supported { + assert_eq!(count, Some(reference_count(&anthropic_tokenizer, text))); + } else { + assert_eq!(count, None); + } + } + + #[rstest] + fn counts_match_the_full_encoder(anthropic_tokenizer: Tokenizer) { + let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); + let mut rng = StdRng::seed_from_u64(2026); + for _ in 0..4000 { + let text = random_text(&mut rng).replace('<', "("); + let expected = reference_count(&anthropic_tokenizer, &text); + assert_eq!( + fast.count(&anthropic_tokenizer, &text), + Some(expected), + "text {text:?}" + ); + } + } + + #[rstest] + fn nfkc_matches_the_tokenizer_normalizer_for_every_scalar_value() { + let fast = byte_level_counter(true); + let mut text = String::new(); + for character in (0..=0x10FFFFu32).filter_map(char::from_u32) { + text.clear(); + text.push(character); + let mut expected = NormalizedString::from(text.as_str()); + NFKC.normalize(&mut expected).expect("nfkc"); + assert_eq!( + fast.normalize(&text), + expected.get(), + "U+{:04X}", + u32::from(character) + ); + } + } + + #[rstest] + fn nfkc_matches_the_tokenizer_normalizer_on_random_texts() { + let fast = byte_level_counter(true); + let mut rng = StdRng::seed_from_u64(11); + for _ in 0..4000 { + let text = random_text(&mut rng); + let mut expected = NormalizedString::from(text.as_str()); + NFKC.normalize(&mut expected).expect("nfkc"); + assert_eq!(fast.normalize(&text), expected.get(), "text {text:?}"); + } + } + + #[rstest] + fn pieces_match_the_byte_level_pre_tokenizer() { + let byte_level = ByteLevel::new(false, true, true); + let mut rng = StdRng::seed_from_u64(7); + for _ in 0..4000 { + let text = byte_level_counter(true) + .normalize(&random_text(&mut rng)) + .into_owned(); + let mut pre_tokenized = PreTokenizedString::from(text.as_str()); + byte_level + .pre_tokenize(&mut pre_tokenized) + .expect("pre-tokenize"); + let expected: Vec<(String, (usize, usize))> = pre_tokenized + .get_splits(OffsetReferential::Original, OffsetType::Byte) + .into_iter() + .map(|(mapped, offsets, _)| (mapped.to_string(), offsets)) + .collect(); + let actual: Vec<(String, (usize, usize))> = pieces( + &text, + UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"), + ) + .map(|piece| { + let start = piece.as_ptr() as usize - text.as_ptr() as usize; + let mapped: String = piece.bytes().map(byte_char).collect(); + (mapped, (start, start + piece.len())) + }) + .collect(); + assert_eq!(actual, expected, "text {text:?}"); + } + } + + #[rstest] + fn byte_chars_match_the_byte_level_alphabet() { + let byte_level = ByteLevel::new(false, false, false); + let characters: Vec = (0..=0x10FFFFu32).filter_map(char::from_u32).collect(); + for chunk in characters.chunks(1024) { + let text: String = chunk.iter().collect(); + let mut pre_tokenized = PreTokenizedString::from(text.as_str()); + byte_level + .pre_tokenize(&mut pre_tokenized) + .expect("pre-tokenize"); + let expected: String = pre_tokenized + .get_splits(OffsetReferential::Original, OffsetType::Byte) + .into_iter() + .map(|(mapped, _, _)| mapped) + .collect(); + let actual: String = text.bytes().map(byte_char).collect(); + assert_eq!(actual.len(), mapped_len(&text)); + assert_eq!( + actual, + expected, + "chunk starting at U+{:04X}", + u32::from(chunk[0]) + ); + } + } + + #[rstest] + fn classes_match_oniguruma() { + let unicode_classes = UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"); + let letter = SysRegex::new(r"\p{L}").expect("regex"); + let number = SysRegex::new(r"\p{N}").expect("regex"); + let space = SysRegex::new(r"\s").expect("regex"); + let whole = + |regex: &SysRegex, text: &str| regex.find_iter(text).next() == Some((0, text.len())); + let mut text = String::new(); + for character in (0..=0x10FFFFu32).filter_map(char::from_u32) { + text.clear(); + text.push(character); + let expected = if whole(&letter, &text) { + Class::Letter + } else if whole(&number, &text) { + Class::Number + } else if whole(&space, &text) { + Class::Space + } else { + Class::Other + }; + assert_eq!( + class(character, unicode_classes), + expected, + "U+{:04X}", + u32::from(character) + ); + } + } + + #[rstest] + #[case("prefix")] + #[case("regex")] + #[case("normalizer")] + #[case("no_pre_tokenizer")] + #[case("other_pre_tokenizer")] + #[case("post_processor")] + #[case("truncation")] + #[case("padding")] + fn other_tokenizer_shapes_are_declined( + mut anthropic_tokenizer: Tokenizer, + #[case] shape: &str, + ) { + use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams}; + match shape { + "prefix" => { + anthropic_tokenizer.with_pre_tokenizer(Some(ByteLevel::new(true, true, true))); + } + "regex" => { + anthropic_tokenizer.with_pre_tokenizer(Some(ByteLevel::new(false, true, false))); + } + "normalizer" => { + anthropic_tokenizer + .with_normalizer(Some(tokenizers::normalizers::Lowercase)) + .expect("normalizer"); + } + "no_pre_tokenizer" => { + anthropic_tokenizer.with_pre_tokenizer(None::); + } + "other_pre_tokenizer" => { + anthropic_tokenizer + .with_pre_tokenizer(Some(tokenizers::pre_tokenizers::whitespace::Whitespace)); + } + "post_processor" => { + anthropic_tokenizer.with_post_processor(Some(ByteLevel::default())); + } + "truncation" => { + anthropic_tokenizer + .with_truncation(Some(TruncationParams { + max_length: 2, + ..Default::default() + })) + .expect("truncation"); + } + "padding" => { + anthropic_tokenizer.with_padding(Some(PaddingParams { + strategy: PaddingStrategy::Fixed(32), + ..Default::default() + })); + } + _ => unreachable!(), + } + assert!(ByteLevelCounter::detect(&anthropic_tokenizer).is_none()); + let counter = crate::TokenCounter::from_json( + &anthropic_tokenizer.to_string(false).expect("serialize"), + ) + .expect("load"); + for text in ["", "Hello WORLD! ＡＢ fi Ⅳ", " stop"] { + assert_eq!( + counter.count_text(text).expect("count"), + reference_count(&anthropic_tokenizer, text) + ); + } + } + + #[rstest] + #[case(false)] + #[case(true)] + fn arbitrary_unicode_and_long_inputs_use_fast_path( + mut anthropic_tokenizer: Tokenizer, + #[case] nfkc: bool, + ) { + if !nfkc { + anthropic_tokenizer + .with_normalizer(None::) + .expect("normalizer"); + } + let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); + let mut rng = StdRng::seed_from_u64(314159); + for _ in 0..1000 { + let text: String = (0..64) + .filter_map(|_| char::from_u32(rng.gen_range(0..=0x10ffff))) + .collect(); + assert_eq!( + fast.count(&anthropic_tokenizer, &text), + Some(reference_count(&anthropic_tokenizer, &text)), + "text {text:?}" + ); + } + for text in [ + "", + "'s't're've'm'll'd'S'RE", + " a \t\r\n b\u{85}\u{a0}c ", + "\0é漢🙂", + "a\u{30a}\u{301}", + "ＡfiⅣ", + ] { + let text = text.repeat(2048); + assert_eq!( + fast.count(&anthropic_tokenizer, &text), + Some(reference_count(&anthropic_tokenizer, &text)) + ); + } + } + + #[rstest] + #[case(false, false, false, false)] + #[case(true, false, false, false)] + #[case(false, true, false, false)] + #[case(false, false, true, false)] + #[case(false, false, false, true)] + fn added_token_options_fall_back( + mut anthropic_tokenizer: Tokenizer, + #[case] special: bool, + #[case] single_word: bool, + #[case] lstrip: bool, + #[case] rstrip: bool, + ) { + anthropic_tokenizer + .add_tokens([tokenizers::AddedToken::from("custom token", special) + .single_word(single_word) + .lstrip(lstrip) + .rstrip(rstrip)]) + .expect("add token"); + let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); + let counter = crate::TokenCounter::from_json( + &anthropic_tokenizer.to_string(false).expect("serialize"), + ) + .expect("load"); + for text in [ + "custom token", + "a custom token b", + "acustom tokenb", + " custom token ", + ] { + assert_eq!(fast.count(&anthropic_tokenizer, text), None); + assert_eq!( + counter.count_text(text).expect("count"), + reference_count(&anthropic_tokenizer, text) + ); + } + } + + #[test] + fn model_errors_reach_public_caller() { + let mut tokenizer = Tokenizer::new(tokenizers::models::wordpiece::WordPiece::default()); + tokenizer.with_pre_tokenizer(Some(ByteLevel::new(false, true, true))); + let fast = ByteLevelCounter::detect(&tokenizer).expect("supported"); + assert_eq!(fast.count(&tokenizer, "hello"), None); + assert!(tokenizer.encode_fast("hello", true).is_err()); + let counter = + crate::TokenCounter::from_json(&tokenizer.to_string(false).expect("serialize")) + .expect("load"); + assert!(matches!( + counter.count_text("hello"), + Err(crate::Error::Encode(_)) + )); + } + + #[rstest] + fn shared_counter_matches_encoder_across_threads(anthropic_tokenizer: Tokenizer) { + let counter = crate::TokenCounter::from_json( + &anthropic_tokenizer.to_string(false).expect("serialize"), + ) + .expect("load"); + let inputs = [ + "hello world", + "ＡＢ fi\n漢字🙂", + " stop", + "\t 're \r\n", + ]; + let expected = inputs.map(|text| reference_count(&anthropic_tokenizer, text)); + std::thread::scope(|scope| { + for _ in 0..8 { + let counter = &counter; + scope.spawn(move || { + for _ in 0..100 { + for (text, count) in inputs.iter().zip(expected) { + assert_eq!(counter.count_text(text).expect("count"), count); + } + } + }); + } + }); + } + + #[rstest] + fn normalized_added_token_spelling_declines_fast_path(mut anthropic_tokenizer: Tokenizer) { + anthropic_tokenizer + .add_tokens([tokenizers::AddedToken::from("ＡＢＣＤ　ＥＦＧＨ", false)]) + .expect("add token"); + let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); + assert_eq!(reference_count(&anthropic_tokenizer, "ABCD EFGH"), 1); + assert_eq!(fast.count(&anthropic_tokenizer, "ABCD EFGH"), None); + let counter = crate::TokenCounter::from_json( + &anthropic_tokenizer.to_string(false).expect("serialize"), + ) + .expect("load"); + assert_eq!(counter.count_text("ABCD EFGH").expect("count"), 1); + } +} diff --git a/litellm-rust/crates/token-counter/src/counter.rs b/litellm-rust/crates/token-counter/src/counter.rs new file mode 100644 index 00000000000..a3943515a64 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/counter.rs @@ -0,0 +1,194 @@ +use serde::Serialize; + +use crate::Error; +use crate::byte_level::ByteLevelCounter; +use crate::python_json; +use crate::tools::format_function_definitions; +use crate::types::{ + ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice, + ToolDefinition, +}; + +const TOKENS_PER_MESSAGE: usize = 3; +const TOKENS_PER_NAME: usize = 1; +const REPLY_PRIMING_TOKENS: usize = 3; +const TOOL_DEFINITIONS_TOKENS: usize = 9; +const TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT: usize = 4; +const TOOL_CHOICE_NONE_TOKENS: usize = 1; +const NAMED_TOOL_CHOICE_TOKENS: usize = 7; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct InputTokenCount { + pub model: Option, + pub input_tokens: usize, +} + +/// A loaded HuggingFace tokenizer plus the message accounting Python applies on +/// top of it. Encoding is CPU-bound and synchronous; hosts run it off their +/// event loop. +pub struct TokenCounter { + tokenizer: tokenizers::Tokenizer, + byte_level: Option, +} + +impl TokenCounter { + /// Load a HuggingFace `tokenizer.json` document. The host reads the file. + pub fn from_json(tokenizer_json: &str) -> Result { + let tokenizer = tokenizer_json + .parse::() + .map_err(Error::Load)?; + let byte_level = ByteLevelCounter::detect(&tokenizer); + Ok(Self { + tokenizer, + byte_level, + }) + } + + pub fn count_text(&self, text: &str) -> Result { + if let Some(count) = self + .byte_level + .as_ref() + .and_then(|counter| counter.count(&self.tokenizer, text)) + { + return Ok(count); + } + self.tokenizer + .encode_fast(text, true) + .map(|encoding| encoding.len()) + .map_err(Error::Encode) + } + + /// Mirrors the host's key precedence: `messages`, then `prompt`, then + /// `input`, then `query` plus `documents`. + pub fn count_request(&self, request: &CountableRequest) -> Result { + let input_tokens = if let Some(messages) = &request.messages { + self.count_messages(request, messages)? + } else if let Some(prompt) = &request.prompt { + self.count_text_value(prompt)? + } else if let Some(input) = &request.input { + self.count_text_value(input)? + } else if request.query.is_some() || request.documents.is_some() { + self.count_optional_text_value(request.query.as_ref())? + + self.count_optional_text_value(request.documents.as_ref())? + } else { + return Err(Error::MissingInput); + }; + Ok(InputTokenCount { + model: request.model.clone(), + input_tokens, + }) + } + + fn count_messages( + &self, + request: &CountableRequest, + messages: &[Message], + ) -> Result { + let message_tokens = messages + .iter() + .map(|message| self.count_message(message)) + .sum::>()?; + let includes_system_message = messages + .iter() + .any(|message| message.role.as_deref() == Some("system")); + let extra_tokens = self.count_extra( + request.tools.as_deref().unwrap_or_default(), + request.tool_choice.as_ref(), + includes_system_message, + )?; + Ok(message_tokens + extra_tokens) + } + + fn count_optional_text_value(&self, value: Option<&TextValue>) -> Result { + value.map_or(Ok(0), |value| self.count_text_value(value)) + } + + /// `str()` for scalars, `json.dumps()` for objects, lists flattened, nulls + /// skipped. Floats are declined because Python's `repr` and Rust's float + /// formatting disagree on exponents. + fn count_text_value(&self, value: &TextValue) -> Result { + match value { + TextValue::Null => Ok(0), + TextValue::Bool(true) => self.count_text("True"), + TextValue::Bool(false) => self.count_text("False"), + TextValue::Number(number) => match (number.as_i64(), number.as_u64()) { + (Some(number), _) => self.count_text(itoa::Buffer::new().format(number)), + (_, Some(number)) => self.count_text(itoa::Buffer::new().format(number)), + _ => Err(Error::FloatText), + }, + TextValue::Text(text) => self.count_text(text), + TextValue::List(items) => items + .iter() + .map(|item| self.count_text_value(item)) + .sum::>(), + TextValue::Object(_) => self.count_text(&python_json::dumps(value)?), + } + } + + fn count_message(&self, message: &Message) -> Result { + let role_tokens = match &message.role { + Some(role) => self.count_text(role)?, + None => 0, + }; + let name_tokens = match &message.name { + Some(name) => self.count_text(name)? + TOKENS_PER_NAME, + None => 0, + }; + let content_tokens = match &message.content { + Some(MessageContent::Text(text)) => self.count_text(text)?, + Some(MessageContent::Blocks(items)) => items + .iter() + .map(|item| self.count_content_item(item)) + .sum::>()?, + None => 0, + }; + Ok(TOKENS_PER_MESSAGE + role_tokens + name_tokens + content_tokens) + } + + fn count_content_item(&self, item: &ContentItem) -> Result { + match item { + ContentItem::Text(text) => self.count_text(text), + ContentItem::Block(ContentBlock::Text { text }) => self.count_text(text), + ContentItem::Block(ContentBlock::Thinking { thinking }) => { + if thinking.is_empty() { + return Ok(0); + } + self.count_text(thinking) + } + ContentItem::Block(ContentBlock::ToolReference { tool_name }) => { + match tool_name.as_deref().filter(|name| !name.is_empty()) { + Some(name) => self.count_text(name), + None => Ok(0), + } + } + ContentItem::Block(ContentBlock::Unsupported) => Err(Error::ContentBlock), + } + } + + fn count_extra( + &self, + tools: &[ToolDefinition], + tool_choice: Option<&ToolChoice>, + includes_system_message: bool, + ) -> Result { + let tool_tokens = if tools.is_empty() { + 0 + } else { + let definitions = self.count_text(&format_function_definitions(tools)?)?; + let discount = if includes_system_message { + TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT + } else { + 0 + }; + definitions + TOOL_DEFINITIONS_TOKENS - discount + }; + let choice_tokens = match tool_choice { + Some(ToolChoice::Mode(mode)) if mode == "none" => TOOL_CHOICE_NONE_TOKENS, + Some(ToolChoice::Mode(_)) | None => 0, + Some(ToolChoice::Named(named)) => { + NAMED_TOOL_CHOICE_TOKENS + self.count_text(&named.function.name)? + } + }; + Ok(REPLY_PRIMING_TOKENS + tool_tokens + choice_tokens) + } +} diff --git a/litellm-rust/crates/token-counter/src/error.rs b/litellm-rust/crates/token-counter/src/error.rs new file mode 100644 index 00000000000..dc590093f64 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/error.rs @@ -0,0 +1,31 @@ +use std::string::FromUtf8Error; + +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("unsupported by the rust token counter: request body could not be parsed: {0}")] + RequestParse(#[source] serde_json::Error), + #[error("unsupported by the rust token counter: request has no countable input")] + MissingInput, + #[error( + "unsupported by the rust token counter: float text values are counted by the python path" + )] + FloatText, + #[error( + "unsupported by the rust token counter: content block type is counted by the python path" + )] + ContentBlock, + #[error("unsupported by the rust token counter: array parameter without items")] + ArrayItems, + #[error("unsupported by the rust token counter: text value could not be serialized: {0}")] + JsonSerialization(#[source] serde_json::Error), + #[error("unsupported by the rust token counter: serialized text value is not UTF-8: {0}")] + JsonUtf8(#[source] FromUtf8Error), + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), + #[error("token counting task failed: {0}")] + Task(String), +} diff --git a/litellm-rust/crates/token-counter/src/lib.rs b/litellm-rust/crates/token-counter/src/lib.rs new file mode 100644 index 00000000000..eb602b3cade --- /dev/null +++ b/litellm-rust/crates/token-counter/src/lib.rs @@ -0,0 +1,17 @@ +//! Input token counting for a request body, mirroring `litellm.token_counter` +//! for the shapes it can count exactly. Everything else is declined so the host +//! keeps its own counter as the reference. + +#![forbid(unsafe_code)] + +mod byte_level; +mod counter; +mod error; +mod python_json; +mod tools; +mod types; +mod unicode_classes; + +pub use counter::{InputTokenCount, TokenCounter}; +pub use error::Error; +pub use types::CountableRequest; diff --git a/litellm-rust/crates/token-counter/src/python_json.rs b/litellm-rust/crates/token-counter/src/python_json.rs new file mode 100644 index 00000000000..00a55ad300c --- /dev/null +++ b/litellm-rust/crates/token-counter/src/python_json.rs @@ -0,0 +1,155 @@ +//! `json.dumps(value)` with Python's default arguments: `", "` and `": "` +//! separators, `ensure_ascii=True`, and keys in insertion order. + +use std::io::{self, Write}; + +use serde::Serialize; +use serde_json::ser::{Formatter, Serializer}; + +use super::Error; +use super::types::TextValue; + +pub(super) fn dumps(value: &TextValue) -> Result { + let mut output = Vec::with_capacity(serialized_len(value)?); + value + .serialize(&mut Serializer::with_formatter( + &mut output, + PythonFormatter, + )) + .map_err(Error::JsonSerialization)?; + debug_assert_eq!(output.len(), output.capacity()); + String::from_utf8(output).map_err(Error::JsonUtf8) +} + +fn serialized_len(value: &TextValue) -> Result { + match value { + TextValue::Null => Ok(4), + TextValue::Bool(true) => Ok(4), + TextValue::Bool(false) => Ok(5), + TextValue::Number(number) => match (number.as_i64(), number.as_u64()) { + (Some(number), _) => Ok(unsigned_len(number.unsigned_abs()) + usize::from(number < 0)), + (_, Some(number)) => Ok(unsigned_len(number)), + _ => Err(Error::FloatText), + }, + TextValue::Text(text) => Ok(quoted_len(text)), + TextValue::List(items) => items + .iter() + .try_fold(2 + items.len().saturating_sub(1) * 2, |len, item| { + Ok(len + serialized_len(item)?) + }), + TextValue::Object(entries) => entries.iter().try_fold( + 2 + entries.len().saturating_sub(1) * 2, + |len, (key, value)| Ok(len + quoted_len(key) + 2 + serialized_len(value)?), + ), + } +} + +fn unsigned_len(number: u64) -> usize { + if number == 0 { + 1 + } else { + number.ilog10() as usize + 1 + } +} + +fn quoted_len(value: &str) -> usize { + value.chars().fold(2, |len, character| { + len + match character { + '"' | '\\' | '\u{0008}' | '\u{000c}' | '\n' | '\r' | '\t' => 2, + '\u{0000}'..='\u{001f}' => 6, + ' '..='~' => 1, + _ => character.len_utf16() * 6, + } + }) +} + +struct PythonFormatter; + +impl Formatter for PythonFormatter { + fn begin_array_value(&mut self, writer: &mut W, first: bool) -> io::Result<()> + where + W: ?Sized + Write, + { + if !first { + writer.write_all(b", ")?; + } + Ok(()) + } + + fn begin_object_key(&mut self, writer: &mut W, first: bool) -> io::Result<()> + where + W: ?Sized + Write, + { + if !first { + writer.write_all(b", ")?; + } + Ok(()) + } + + fn begin_object_value(&mut self, writer: &mut W) -> io::Result<()> + where + W: ?Sized + Write, + { + writer.write_all(b": ") + } + + fn write_string_fragment(&mut self, writer: &mut W, fragment: &str) -> io::Result<()> + where + W: ?Sized + Write, + { + for character in fragment.chars() { + if (' '..='~').contains(&character) { + write!(writer, "{character}")?; + continue; + } + let mut units = [0u16; 2]; + for unit in character.encode_utf16(&mut units) { + write!(writer, "\\u{unit:04x}")?; + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::dumps; + + #[rstest] + #[case::null("null", "null")] + #[case::boolean("true", "true")] + #[case::signed_integer("-3", "-3")] + #[case::unsigned_integer("18446744073709551615", "18446744073709551615")] + #[case::empty_array("[]", "[]")] + #[case::empty_object("{}", "{}")] + #[case::nested( + r#"{"first":1,"second":{"ok":true,"none":null},"third":[false,2]}"#, + r#"{"first": 1, "second": {"ok": true, "none": null}, "third": [false, 2]}"# + )] + #[case::string_escaping( + r#""caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~ /""#, + r#""caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~ /""# + )] + #[case::short_control_escapes(r#""\b\f\r""#, r#""\b\f\r""#)] + fn matches_python_json_dumps(#[case] input: &str, #[case] expected: &str) { + let value = serde_json::from_str(input).expect("fixture parses"); + + assert_eq!(dumps(&value).expect("fixture dumps"), expected); + } + + #[rstest] + #[case::top_level("1.5")] + #[case::array("[1,2.5]")] + #[case::object(r#"{"nested":{"value":-0.25}}"#)] + fn rejects_floats(#[case] input: &str) { + let value = serde_json::from_str(input).expect("fixture parses"); + let error = dumps(&value).expect_err("floats are declined"); + + assert_eq!( + error.to_string(), + "unsupported by the rust token counter: float text values are counted by the python path" + ); + } +} diff --git a/litellm-rust/crates/token-counter/src/tools.rs b/litellm-rust/crates/token-counter/src/tools.rs new file mode 100644 index 00000000000..22150a4e4b1 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/tools.rs @@ -0,0 +1,304 @@ +//! Renders tool definitions the way `litellm.token_counter` does before +//! tokenizing them (the TypeScript-like namespace OpenAI appears to use). + +use std::fmt::Write; + +use super::Error; +use super::types::{EnumValue, Schema, SchemaType, ToolDefinition}; + +pub(super) fn format_function_definitions(tools: &[ToolDefinition]) -> Result { + ToolFormatter::new(tools.len()).format(tools) +} + +struct ToolFormatter { + output: String, +} + +impl ToolFormatter { + fn new(tool_count: usize) -> Self { + Self { + output: String::with_capacity(tool_count.saturating_mul(128).saturating_add(48)), + } + } + + fn format(mut self, tools: &[ToolDefinition]) -> Result { + self.output.push_str("namespace functions {\n\n"); + + for tool in tools { + self.write_function(tool)?; + } + + self.output.push_str("} // namespace functions"); + Ok(self.output) + } + + fn write_function(&mut self, tool: &ToolDefinition) -> Result<(), Error> { + let (name, description, parameters) = resolve_function(tool); + let Some(name) = name.filter(|name| !name.is_empty()) else { + return Ok(()); + }; + + if let Some(description) = description.filter(|description| !description.is_empty()) { + self.output.push_str("// "); + self.output.push_str(description); + self.output.push('\n'); + } + + match parameters.filter(|parameters| { + parameters + .properties + .as_ref() + .is_some_and(|properties| !properties.is_empty()) + }) { + Some(parameters) => { + self.output.push_str("type "); + self.output.push_str(name); + self.output.push_str(" = (_: {\n"); + self.write_object_parameters(parameters, 0)?; + self.output.push_str("\n}) => any;\n\n"); + } + _ => { + self.output.push_str("type "); + self.output.push_str(name); + self.output.push_str(" = () => any;\n\n"); + } + } + + Ok(()) + } + + fn write_object_parameters(&mut self, parameters: &Schema, indent: usize) -> Result<(), Error> { + let Some(properties) = parameters + .properties + .as_ref() + .filter(|properties| !properties.is_empty()) + else { + return Ok(()); + }; + let required = parameters.required.as_deref().unwrap_or_default(); + for (index, (key, props)) in properties.iter().enumerate() { + if index > 0 { + self.output.push('\n'); + } + if let Some(description) = props + .description + .as_deref() + .filter(|description| !description.is_empty()) + { + self.write_indent(indent); + self.output.push_str("// "); + self.output.push_str(description); + self.output.push('\n'); + } + + self.write_indent(indent); + self.output.push_str(key); + if !required.iter().any(|required| required == key) { + self.output.push('?'); + } + self.output.push_str(": "); + self.write_type(props, indent)?; + self.output.push(','); + } + + Ok(()) + } + + fn write_type(&mut self, props: &Schema, indent: usize) -> Result<(), Error> { + let Some(SchemaType::Name(schema_type)) = &props.schema_type else { + self.output.push_str("any"); + return Ok(()); + }; + + match schema_type.as_str() { + "string" | "integer" | "number" => match &props.enum_values { + Some(values) => self.write_enum(values), + None if schema_type == "string" => self.output.push_str("string"), + None => self.output.push_str("number"), + }, + "array" => { + let items = props.items.as_deref().ok_or(Error::ArrayItems)?; + self.write_type(items, indent)?; + self.output.push_str("[]"); + } + "object" => { + self.output.push_str("{\n"); + self.write_object_parameters(props, indent + 2)?; + self.output.push_str("\n}"); + } + "boolean" => self.output.push_str("boolean"), + "null" => self.output.push_str("null"), + _ => self.output.push_str("any"), + } + + Ok(()) + } + + fn write_enum(&mut self, values: &[EnumValue]) { + for (index, value) in values.iter().enumerate() { + if index > 0 { + self.output.push_str(" | "); + } + self.output.push('"'); + match value { + EnumValue::Text(text) => self.output.push_str(text), + EnumValue::Integer(number) => { + write!(self.output, "{number}").expect("writing to a String cannot fail"); + } + } + self.output.push('"'); + } + } + + fn write_indent(&mut self, indent: usize) { + for _ in 0..indent { + self.output.push(' '); + } + } +} + +fn resolve_function(tool: &ToolDefinition) -> (Option<&str>, Option<&str>, Option<&Schema>) { + match &tool.function { + Some(function) => ( + function.name.as_deref(), + function.description.as_deref(), + function.parameters.as_ref(), + ), + None => ( + tool.name.as_deref(), + tool.description.as_deref(), + tool.input_schema.as_ref().or(tool.parameters.as_ref()), + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn parse_tools(json: &str) -> Vec { + serde_json::from_str(json).expect("tool fixture parses") + } + + #[test] + fn empty_tool_list_renders_an_empty_namespace() { + assert_eq!( + format_function_definitions(&[]).expect("empty tool list renders"), + "namespace functions {\n\n} // namespace functions" + ); + } + + #[test] + fn unnamed_tools_are_skipped_and_function_shape_takes_precedence() { + let tools = parse_tools( + r#"[ + {}, + {"name":""}, + {"name":"ignored","function":{"description":"missing a function name"}}, + {"name":"ping","description":""} + ]"#, + ); + + assert_eq!( + format_function_definitions(&tools).expect("tools render"), + "namespace functions {\n\ntype ping = () => any;\n\n} // namespace functions" + ); + } + + #[test] + fn empty_or_missing_properties_render_a_no_argument_function() { + let tools = parse_tools( + r#"[ + {"name":"missing","input_schema":{"type":"object"}}, + {"name":"empty","input_schema":{"type":"object","properties":{}}} + ]"#, + ); + + assert_eq!( + format_function_definitions(&tools).expect("tools render"), + concat!( + "namespace functions {\n\n", + "type missing = () => any;\n\n", + "type empty = () => any;\n\n", + "} // namespace functions" + ) + ); + } + + #[test] + fn anthropic_parameters_render_all_supported_types() { + let tools = parse_tools( + r#"[{ + "name":"inspect", + "description":"Inspect a value", + "input_schema":{ + "type":"object", + "properties":{ + "text":{"type":"string"}, + "count":{"type":"integer","description":"Number of attempts"}, + "ratio":{"type":"number"}, + "enabled":{"type":"boolean"}, + "nothing":{"type":"null"}, + "unknown":{"type":"custom"}, + "union":{"type":["string","null"]}, + "labels":{"type":"array","items":{"type":"string"}}, + "config":{"type":"object","properties":{"retries":{"type":"integer"}},"required":["retries"]}, + "mode":{"type":"string","enum":["fast",2]} + }, + "required":["text"] + } + }]"#, + ); + + assert_eq!( + format_function_definitions(&tools).expect("tool renders"), + concat!( + "namespace functions {\n\n", + "// Inspect a value\n", + "type inspect = (_: {\n", + "text: string,\n", + "// Number of attempts\n", + "count?: number,\n", + "ratio?: number,\n", + "enabled?: boolean,\n", + "nothing?: null,\n", + "unknown?: any,\n", + "union?: any,\n", + "labels?: string[],\n", + "config?: {\n", + " retries: number,\n", + "},\n", + "mode?: \"fast\" | \"2\",\n", + "}) => any;\n\n", + "} // namespace functions" + ) + ); + } + + #[test] + fn input_schema_takes_precedence_over_legacy_parameters() { + let tools = parse_tools( + r#"[{ + "name":"choose", + "input_schema":{"type":"object","properties":{"current":{"type":"string"}}}, + "parameters":{"type":"object","properties":{"legacy":{"type":"string"}}} + }]"#, + ); + + let rendered = format_function_definitions(&tools).expect("tool renders"); + assert!(rendered.contains("current?: string,")); + assert!(!rendered.contains("legacy")); + } + + #[test] + fn array_without_items_returns_an_error() { + let tools = parse_tools( + r#"[{"name":"broken","parameters":{"type":"object","properties":{"values":{"type":"array"}}}}]"#, + ); + + assert!(matches!( + format_function_definitions(&tools), + Err(Error::ArrayItems) + )); + } +} diff --git a/litellm-rust/crates/token-counter/src/types.rs b/litellm-rust/crates/token-counter/src/types.rs new file mode 100644 index 00000000000..d25554beaac --- /dev/null +++ b/litellm-rust/crates/token-counter/src/types.rs @@ -0,0 +1,238 @@ +use std::fmt; + +use indexmap::IndexMap; +use serde::de::{MapAccess, SeqAccess, Visitor}; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::Number; + +use super::Error; + +/// The parts of a request body the host's budget counter reads. Chat and +/// Anthropic Messages bodies carry `messages`; completions carry `prompt`; +/// Responses and embeddings carry `input`; rerank carries `query` and +/// `documents`. The host checks key presence, not nullness, so an explicit +/// `null` is kept distinct from an absent key. Anything outside this shape is +/// declined so the host can fall back to its own counter instead of silently +/// miscounting. +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub struct CountableRequest { + pub(crate) model: Option, + #[serde(default, deserialize_with = "present_messages")] + pub(crate) messages: Option>, + pub(crate) tools: Option>, + pub(crate) tool_choice: Option, + #[serde(default, deserialize_with = "present_text")] + pub(crate) prompt: Option, + #[serde(default, deserialize_with = "present_text")] + pub(crate) input: Option, + #[serde(default, deserialize_with = "present_text")] + pub(crate) query: Option, + #[serde(default, deserialize_with = "present_text")] + pub(crate) documents: Option, +} + +impl CountableRequest { + pub fn parse(body: &[u8]) -> Result { + serde_json::from_slice(body).map_err(Error::RequestParse) + } +} + +fn present_messages<'de, D: Deserializer<'de>>( + deserializer: D, +) -> Result>, D::Error> { + Option::>::deserialize(deserializer) + .map(|messages| Some(messages.unwrap_or_default())) +} + +fn present_text<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + TextValue::deserialize(deserializer).map(Some) +} + +/// Free-form JSON the host counts as text: strings and integers via `str()`, +/// objects via `json.dumps()`, lists flattened. Objects keep document order so +/// the dumped text matches Python byte for byte. +#[derive(Clone, Debug, PartialEq, Serialize)] +#[serde(untagged)] +pub(crate) enum TextValue { + Null, + Bool(bool), + Number(Number), + Text(String), + List(Vec), + Object(IndexMap), +} + +impl<'de> Deserialize<'de> for TextValue { + fn deserialize>(deserializer: D) -> Result { + deserializer.deserialize_any(TextValueVisitor) + } +} + +struct TextValueVisitor; + +impl<'de> Visitor<'de> for TextValueVisitor { + type Value = TextValue; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a JSON value") + } + + fn visit_unit(self) -> Result { + Ok(TextValue::Null) + } + + fn visit_none(self) -> Result { + Ok(TextValue::Null) + } + + fn visit_bool(self, value: bool) -> Result { + Ok(TextValue::Bool(value)) + } + + fn visit_i64(self, value: i64) -> Result { + Ok(TextValue::Number(value.into())) + } + + fn visit_u64(self, value: u64) -> Result { + Ok(TextValue::Number(value.into())) + } + + fn visit_f64(self, value: f64) -> Result + where + E: serde::de::Error, + { + Number::from_f64(value) + .map(TextValue::Number) + .ok_or_else(|| E::custom("non-finite JSON number")) + } + + fn visit_str(self, value: &str) -> Result { + Ok(TextValue::Text(value.to_owned())) + } + + fn visit_string(self, value: String) -> Result { + Ok(TextValue::Text(value)) + } + + fn visit_seq(self, mut sequence: A) -> Result + where + A: SeqAccess<'de>, + { + let mut items = Vec::with_capacity(sequence.size_hint().unwrap_or(0)); + while let Some(item) = sequence.next_element()? { + items.push(item); + } + Ok(TextValue::List(items)) + } + + fn visit_map(self, mut map: A) -> Result + where + A: MapAccess<'de>, + { + let mut entries = IndexMap::with_capacity(map.size_hint().unwrap_or(0)); + while let Some((key, value)) = map.next_entry()? { + entries.insert(key, value); + } + Ok(TextValue::Object(entries)) + } +} + +/// Python counts every string-valued key of a message, so any key beyond these +/// makes the shape unsupported rather than silently uncounted. +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(deny_unknown_fields)] +pub(crate) struct Message { + pub(crate) role: Option, + pub(crate) name: Option, + pub(crate) content: Option, +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum MessageContent { + Text(String), + Blocks(Vec), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum ContentItem { + Text(String), + Block(ContentBlock), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(tag = "type")] +pub(crate) enum ContentBlock { + #[serde(rename = "text")] + Text { text: String }, + #[serde(rename = "thinking")] + Thinking { thinking: String }, + #[serde(rename = "tool_reference")] + ToolReference { tool_name: Option }, + /// Images, documents, files and tool use/result blocks price through + /// Python-only helpers, so they stay on the Python counter. + #[serde(other)] + Unsupported, +} + +/// Either the OpenAI `{"type": "function", "function": {...}}` shape or the +/// Anthropic `{"name", "description", "input_schema"}` shape. +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub(crate) struct ToolDefinition { + pub(crate) function: Option, + pub(crate) name: Option, + pub(crate) description: Option, + pub(crate) input_schema: Option, + pub(crate) parameters: Option, +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub(crate) struct FunctionDefinition { + pub(crate) name: Option, + pub(crate) description: Option, + pub(crate) parameters: Option, +} + +#[derive(Clone, Debug, Default, Deserialize, PartialEq)] +pub(crate) struct Schema { + #[serde(rename = "type")] + pub(crate) schema_type: Option, + pub(crate) description: Option, + #[serde(rename = "enum")] + pub(crate) enum_values: Option>, + pub(crate) items: Option>, + pub(crate) properties: Option>, + pub(crate) required: Option>, +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum SchemaType { + Name(String), + Union(Vec), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum EnumValue { + Text(String), + Integer(i64), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub(crate) enum ToolChoice { + Mode(String), + Named(NamedToolChoice), +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub(crate) struct NamedToolChoice { + pub(crate) function: NamedFunction, +} + +#[derive(Clone, Debug, Deserialize, PartialEq)] +pub(crate) struct NamedFunction { + pub(crate) name: String, +} diff --git a/litellm-rust/crates/token-counter/src/unicode_classes.rs b/litellm-rust/crates/token-counter/src/unicode_classes.rs new file mode 100644 index 00000000000..405cee34949 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/unicode_classes.rs @@ -0,0 +1,73 @@ +use std::cmp::Ordering; +use std::sync::LazyLock; + +use tokenizers::utils::SysRegex; + +struct Ranges(Box<[(u32, u32)]>); + +pub(super) struct UnicodeClasses { + letters: Ranges, + numbers: Ranges, + spaces: Ranges, +} + +static CLASSES: LazyLock> = LazyLock::new(|| { + let scalars: String = (0..=u32::from(char::MAX)) + .filter_map(char::from_u32) + .collect(); + Some(UnicodeClasses { + letters: Ranges::load(r"\p{L}+", &scalars)?, + numbers: Ranges::load(r"\p{N}+", &scalars)?, + spaces: Ranges::load(r"\s+", &scalars)?, + }) +}); + +impl Ranges { + fn load(pattern: &str, scalars: &str) -> Option { + let regex = SysRegex::new(pattern).ok()?; + let ranges = regex + .find_iter(scalars) + .map(|(start, end)| { + let matched = scalars.get(start..end)?; + Some(( + u32::from(matched.chars().next()?), + u32::from(matched.chars().next_back()?), + )) + }) + .collect::>>()?; + Some(Self(ranges)) + } + + fn contains(&self, character: char) -> bool { + let code = u32::from(character); + self.0 + .binary_search_by(|(low, high)| { + if *high < code { + Ordering::Less + } else if *low > code { + Ordering::Greater + } else { + Ordering::Equal + } + }) + .is_ok() + } +} + +impl UnicodeClasses { + pub(super) fn get() -> Option<&'static Self> { + CLASSES.as_ref() + } + + pub(super) fn is_letter(&self, character: char) -> bool { + self.letters.contains(character) + } + + pub(super) fn is_number(&self, character: char) -> bool { + self.numbers.contains(character) + } + + pub(super) fn is_space(&self, character: char) -> bool { + self.spaces.contains(character) + } +} diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs new file mode 100644 index 00000000000..0e71a5c03cf --- /dev/null +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -0,0 +1,176 @@ +use rstest::rstest; + +use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter}; + +/// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)` +/// so this test also guards Python parity. +fn counter() -> TokenCounter { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" + ); + let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo"); + TokenCounter::from_json(&json).expect("anthropic tokenizer loads") +} + +const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#; + +const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[ + {"role":"system","content":"You are a terse assistant."}, + {"role":"user","name":"alice","content":[ + {"type":"text","text":"Summarise this paragraph about ships and harbours."}, + "plain string item", + {"type":"thinking","thinking":"pondering"}, + {"type":"tool_reference","tool_name":"get_weather"}]}, + {"role":"assistant","content":[{"type":"text","text":"Sure.","cache_control":{"type":"ephemeral"}}]}]}"#; + +const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}], + "tools":[ + {"type":"function","function":{"name":"get_weather","description":"Get weather","parameters":{ + "type":"object", + "properties":{ + "location":{"type":"string","description":"City name"}, + "unit":{"type":"string","enum":["celsius","fahrenheit"]}, + "days":{"type":"integer"}, + "tags":{"type":"array","items":{"type":"string"}}, + "opts":{"type":"object","properties":{"verbose":{"type":"boolean"},"level":{"type":"integer","enum":[1,2]}},"required":["verbose"]}, + "anything":{}}, + "required":["location"]}}}, + {"type":"function","function":{"name":"noop"}}], + "tool_choice":{"type":"function","function":{"name":"get_weather"}}}"#; + +const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5", + "messages":[{"role":"system","content":"sys"},{"role":"user","content":"weather?"}], + "tools":[{"name":"get_weather","description":"Get weather","input_schema":{ + "type":"object","properties":{"location":{"type":["string","null"]}},"required":["location"]}}], + "tool_choice":"none"}"#; + +const COMPLETIONS_PROMPT: &str = + r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#; + +const COMPLETIONS_PROMPT_LIST: &str = + r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#; + +const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[ + {"role":"user","content":[{"type":"input_text","text":"Summarise caf\u00e9 menus, na\u00efve \u2014 ok? \"quoted\"\n"}]}, + {"role":"assistant","content":"Sure."}],"instructions":"be terse"}"#; + +const EMBEDDINGS_TOKEN_IDS: &str = + r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#; + +const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", + "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#; + +/// Expected counts are pinned from +/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`. +#[rstest] +#[case::text_only(SIMPLE, 14)] +#[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)] +#[case::openai_tools_named_choice(TOOLS_OPENAI, 123)] +#[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)] +#[case::completions_prompt(COMPLETIONS_PROMPT, 7)] +#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)] +#[case::responses_input_items(RESPONSES_INPUT, 62)] +#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)] +#[case::rerank_query_and_documents(RERANK, 41)] +fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) { + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let count = counter().count_request(&request).expect("fixture counts"); + assert_eq!( + count, + InputTokenCount { + model: Some("claude-sonnet-4-5".to_string()), + input_tokens: expected, + } + ); +} + +#[rstest] +#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)] +#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] +#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] +#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] +fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let count = counter().count_request(&request).expect("fixture counts"); + assert_eq!(count.input_tokens, expected); +} + +#[rstest] +#[case::not_json(b"not json" as &[u8])] +#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)] +#[case::message_with_tool_calls( + br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"# +)] +#[case::dict_content( + br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"# +)] +#[case::float_enum( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"# +)] +#[case::anthropic_tool_choice_without_function( + br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"# +)] +fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { + assert!(matches!( + CountableRequest::parse(body), + Err(Error::RequestParse(_)) + )); +} + +#[rstest] +#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] +#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] +#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] +#[case::image_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# +)] +#[case::tool_result_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# +)] +#[case::array_without_items( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# +)] +fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { + let request = CountableRequest::parse(body).expect("shape parses"); + assert!(matches!( + counter().count_request(&request), + Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) + )); +} + +#[test] +fn tool_choice_and_system_discount_change_the_count() { + let counter = counter(); + let count = |body: &str| { + counter + .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) + .expect("counts") + .input_tokens + }; + let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#), + base + 1 + ); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"#), + base + ); + let with_tools = count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#, + ); + let with_tools_and_system = count( + r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#, + ); + assert_eq!(with_tools - with_tools_and_system, 4); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#), + base + ); +} + +#[test] +fn loading_a_bad_tokenizer_is_a_load_error() { + assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_)))); +} diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 64c0da1c28f..20ab9904f46 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -91,6 +91,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_query_params, _safe_set_request_parsed_body, populate_request_with_path_params, + read_raw_json_body, ) from litellm.proxy.common_utils.model_listing_utils import claude_code_requested_group from litellm.proxy.common_utils.realtime_utils import _realtime_request_body @@ -2689,6 +2690,7 @@ async def _run_centralized_common_checks( await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, + request=request, request_data=request_data, route=route, llm_router=llm_router, @@ -2724,6 +2726,7 @@ async def _reserve_budget_after_common_checks( general_settings: dict, end_user_id: str | None = None, end_user_object: LiteLLM_EndUserTable | None = None, + request: Request | None = None, ) -> None: user_api_key_auth_obj.budget_reservation = None if skip_budget_checks: @@ -2749,6 +2752,7 @@ async def _reserve_budget_after_common_checks( end_user_object=end_user_object, apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, + raw_body=await read_raw_json_body(request=request), ) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 635cbf3ba9b..845589aee7a 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -213,6 +213,18 @@ async def _read_request_body(request: Request | None) -> dict: return {} +async def read_raw_json_body(request: Request | None) -> bytes | None: + if request is None or _safe_get_request_parsed_body(request=request) is None: + return None + content_type: Final = _safe_get_request_headers(request=request).get("content-type", "") + if _is_form_content_type(content_type): + return None + try: + return await request.body() + except RuntimeError: + return None + + def _safe_get_request_parsed_body(request: Request | None) -> dict | None: if request is None: return None diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 8d4d72d9974..ef21551ca93 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -34,6 +34,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router +from litellm.rust_bridge.token_counter import count_anthropic_input_tokens, uses_anthropic_tokenizer from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.types.router import DeploymentTypedDict @@ -210,6 +211,7 @@ async def reserve_budget_for_request( end_user_object: object = None, apply_user_budget_to_team_keys: bool = False, fail_closed_budget_enforcement: bool = False, + raw_body: bytes | None = None, ) -> dict | None: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None @@ -237,6 +239,7 @@ async def reserve_budget_for_request( request_body=request_body, route=route, llm_router=llm_router, + raw_body=raw_body, ) current_spend_by_counter_key: Final[dict[str, float]] = {} @@ -1356,24 +1359,46 @@ async def count_request_input_tokens( request_body: dict, route: str, llm_router: Router | None, + raw_body: bytes | None = None, ) -> Mapping[str, int]: """Input-token count per candidate model, counted once per request. Tokenizing is the reservation path's dominant CPU cost and is O(prompt), so counting a large prompt inline stalls every other request on the worker. - Large prompts are counted in a worker thread, and the counts are reused by - both the max-cost and the input-cost estimate. + Models on the Anthropic tokenizer are counted from the raw body by the Rust + bridge when it is enabled, which parses and tokenizes with the GIL released. + Everything it declines is counted in Python, large prompts in a worker + thread. The counts are reused by both the max-cost and the input-cost + estimate. """ models: Final = _get_request_models(request_body=request_body, route=route, llm_router=llm_router) if not models: return MappingProxyType({}) - if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: - return _count_input_tokens_for_models(request_body=request_body, models=models) - return await asyncio.to_thread( - _count_input_tokens_for_models, - request_body=request_body, - models=models, + rust_count: Final = ( + await count_anthropic_input_tokens(raw_body) + if raw_body is not None and any(uses_anthropic_tokenizer(model) for model in models) + else None ) + rust_counts: Final = MappingProxyType( + { + model: rust_count.input_tokens + for model in models + if rust_count is not None and uses_anthropic_tokenizer(model) + } + ) + python_models: Final = tuple(model for model in models if model not in rust_counts) + if not python_models: + return rust_counts + python_counts: Final = ( + _count_input_tokens_for_models(request_body=request_body, models=python_models) + if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS + else await asyncio.to_thread( + _count_input_tokens_for_models, + request_body=request_body, + models=python_models, + ) + ) + return MappingProxyType({**rust_counts, **python_counts}) def _count_input_tokens_for_models( diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py new file mode 100644 index 00000000000..ee755b317d8 --- /dev/null +++ b/litellm/rust_bridge/token_counter.py @@ -0,0 +1,79 @@ +"""Thin Python wrapper for the native Rust input token counter.""" + +from __future__ import annotations + +from collections.abc import Awaitable +from dataclasses import dataclass +from functools import lru_cache +from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables + +from pydantic import TypeAdapter + +import litellm +from litellm._logging import verbose_logger +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.configuration import rust_enabled +from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt +from litellm.utils import claude_json_str +from litellm.utils import uses_anthropic_tokenizer as _python_uses_anthropic_tokenizer + + +class RustTokenCounter(Protocol): + def acount_request(self, body: bytes) -> Awaitable[object]: + raise NotImplementedError + + +class RustTokenCounterFactory(Protocol): + def __call__(self, tokenizer_json: str) -> RustTokenCounter: + raise NotImplementedError + + +@dataclass(frozen=True, slots=True) +class InputTokenCount: + model: str | None + input_tokens: int + + +_INPUT_TOKEN_COUNT: Final = TypeAdapter(InputTokenCount) + + +def _as_factory(value: object) -> RustTokenCounterFactory | None: + return ( + cast( # cast-ok: native extension protocol is runtime-defined + RustTokenCounterFactory, value + ) + if callable(value) + else None + ) + + +TOKEN_COUNTER: Final = NativeBinding("TokenCounter", validate=_as_factory) + + +def uses_anthropic_tokenizer(model: str) -> bool: + if litellm.disable_token_counter is True or litellm.disable_hf_tokenizer_download is True: + return False + return _python_uses_anthropic_tokenizer(model) + + +@lru_cache(maxsize=4) +def _anthropic_counter(factory: RustTokenCounterFactory) -> RustTokenCounter: + return factory(claude_json_str) + + +async def count_anthropic_input_tokens(body: bytes) -> InputTokenCount | None: + if not rust_enabled(): + return None + factory: Final = TOKEN_COUNTER.load() + if factory is None: + return None + try: + attempt: Final = await aattempt( + native_call=lambda: _anthropic_counter(factory).acount_request(body), + adapt=_INPUT_TOKEN_COUNT.validate_python, + context=BridgeErrorContext(route="token_counter", provider="anthropic", model=""), + ) + except (RuntimeError, ValueError) as error: + verbose_logger.debug("Rust token counter failed, counting in Python: %s", error) + return None + return attempt.value if isinstance(attempt, RustHandled) else None diff --git a/litellm/utils.py b/litellm/utils.py index 917af2b89d4..504de08fa91 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2197,13 +2197,17 @@ def _return_openai_tokenizer(model: str) -> SelectTokenizerResponse: return {"type": "openai_tokenizer", "tokenizer": _get_default_encoding()} +def uses_anthropic_tokenizer(model: str) -> bool: + return model in litellm.anthropic_models and "claude-3" not in model + + def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None: if model in litellm.cohere_models and "command-r" in model: # cohere cohere_tokenizer: Final = Tokenizer.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer") return {"type": "huggingface_tokenizer", "tokenizer": cohere_tokenizer} # anthropic - elif model in litellm.anthropic_models and "claude-3" not in model: + elif uses_anthropic_tokenizer(model): claude_tokenizer: Final = Tokenizer.from_str(claude_json_str) return {"type": "huggingface_tokenizer", "tokenizer": claude_tokenizer} # llama2 diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index 011571a37e0..bc4e756eb65 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -26,9 +26,58 @@ from litellm.proxy.common_utils.http_parsing_utils import ( get_tags_from_request_body, numeric_form_fields, populate_request_with_path_params, + read_raw_json_body, ) +def _starlette_request(body: bytes, content_type: str) -> Request: + scope = { + "type": "http", + "method": "POST", + "path": "/v1/messages", + "headers": [(b"content-type", content_type.encode())], + "query_string": b"", + } + chunks = iter((body,)) + + async def receive(): + return {"type": "http.request", "body": next(chunks, b""), "more_body": False} + + return Request(scope, receive) + + +@pytest.mark.asyncio +async def test_read_raw_json_body_returns_the_bytes_the_parsed_body_came_from(): + body = b'{"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}]}' + request = _starlette_request(body, "application/json") + + assert await _read_request_body(request) == orjson.loads(body) + assert await read_raw_json_body(request) == body + + +@pytest.mark.asyncio +async def test_read_raw_json_body_is_none_until_the_body_has_been_parsed(): + request = _starlette_request(b'{"model": "claude-sonnet-4-5"}', "application/json") + + assert await read_raw_json_body(request) is None + assert await read_raw_json_body(None) is None + + +@pytest.mark.asyncio +async def test_read_raw_json_body_is_none_for_form_bodies(): + request = _starlette_request(b"model=claude-sonnet-4-5", "application/x-www-form-urlencoded") + + assert await _read_request_body(request) == {"model": "claude-sonnet-4-5"} + assert await read_raw_json_body(request) is None + + +@pytest.mark.asyncio +async def test_read_raw_json_body_is_none_for_a_request_that_only_mocks_the_parsed_body_path(): + mock_request = MagicMock() + + assert await read_raw_json_body(mock_request) is None + + @pytest.mark.asyncio async def test_request_body_caching(): """ diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index e7dc5af9bfd..ee062b14b96 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -1,16 +1,23 @@ +import json import math from typing import Final import pytest import litellm -import litellm.proxy.proxy_server as proxy_server from litellm.caching import DualCache +from litellm.proxy import proxy_server from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.spend_tracking.budget_reservation import estimate_request_max_cost, reserve_budget_for_request +from litellm.proxy.spend_tracking.budget_reservation import ( + count_request_input_tokens, + estimate_request_max_cost, + reserve_budget_for_request, +) from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge import token_counter as rust_token_counter from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo TOKEN_COUNTING_ROUTES: Final = ( @@ -197,3 +204,128 @@ def test_deployment_pricing_update_invalidates_cached_estimate() -> None: after: Final = estimate_request_max_cost(request_body=TIERED_BODY, route="/chat/completions", llm_router=router) assert after is not None assert math.isclose(after, before * 1000) + + +ANTHROPIC_TOKENIZER_MODEL: Final = "claude-sonnet-4-5-20250929" +RUST_COUNTED_BODY: Final = {"model": ANTHROPIC_TOKENIZER_MODEL, "max_tokens": 16, "messages": ANTHROPIC_MESSAGES} +RUST_INPUT_TOKENS: Final = 4_321 + + +class _FakeDeclined(Exception): + pass + + +class _FakeUpstream(Exception): + pass + + +class _FakeNative: + RustBridgeDeclined = _FakeDeclined + RustUpstreamError = _FakeUpstream + + +class _RecordingCounter: + bodies: Final[list[bytes]] = [] + + def __init__(self, tokenizer_json: str) -> None: + pass + + async def acount_request(self, body: bytes) -> object: + self.bodies.append(body) + return {"model": ANTHROPIC_TOKENIZER_MODEL, "input_tokens": RUST_INPUT_TOKENS} + + +class _DecliningCounter: + def __init__(self, tokenizer_json: str) -> None: + pass + + async def acount_request(self, body: bytes) -> object: + raise _FakeDeclined("unsupported content block") + + +@pytest.fixture +def rust_counter(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) + rust_token_counter._anthropic_counter.cache_clear() + configuration.reset_rust_configuration() + _RecordingCounter.bodies.clear() + yield + rust_token_counter.TOKEN_COUNTER.reset() + rust_token_counter._anthropic_counter.cache_clear() + configuration.reset_rust_configuration() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "request_body"), + ( + ("/v1/messages", RUST_COUNTED_BODY), + ("/v1/chat/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "messages": ANTHROPIC_MESSAGES}), + ("/v1/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "prompt": "hi"}), + ("/v1/responses", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": "hi"}), + ("/v1/embeddings", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": ["hi"]}), + ("/v1/rerank", {"model": ANTHROPIC_TOKENIZER_MODEL, "query": "hi", "documents": ["a"]}), + ), +) +async def test_rust_count_replaces_python_tokenizing_on_every_llm_route( + rust_counter: None, route: str, request_body: dict +) -> None: + litellm.rust(True) + rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter) + raw_body: Final = json.dumps(request_body).encode() + + counts: Final = await count_request_input_tokens( + request_body=request_body, route=route, llm_router=None, raw_body=raw_body + ) + + assert dict(counts) == {ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS} + assert _RecordingCounter.bodies == [raw_body] + + +@pytest.mark.asyncio +async def test_rust_decline_falls_back_to_python_count(rust_counter: None) -> None: + litellm.rust(True) + rust_token_counter.TOKEN_COUNTER.override(_DecliningCounter) + python_counts: Final = await count_request_input_tokens( + request_body=RUST_COUNTED_BODY, route="/v1/messages", llm_router=None + ) + + counts: Final = await count_request_input_tokens( + request_body=RUST_COUNTED_BODY, + route="/v1/messages", + llm_router=None, + raw_body=json.dumps(RUST_COUNTED_BODY).encode(), + ) + + assert dict(counts) == dict(python_counts) + assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS + + +@pytest.mark.asyncio +async def test_disabled_rust_never_sees_the_raw_body(rust_counter: None) -> None: + litellm.rust(False) + rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter) + + counts: Final = await count_request_input_tokens( + request_body=RUST_COUNTED_BODY, + route="/v1/messages", + llm_router=None, + raw_body=json.dumps(RUST_COUNTED_BODY).encode(), + ) + + assert _RecordingCounter.bodies == [] + assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS + + +@pytest.mark.asyncio +async def test_non_anthropic_tokenizer_models_stay_in_python(rust_counter: None) -> None: + litellm.rust(True) + rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter) + body: Final = {"model": "gpt-4o", "messages": ANTHROPIC_MESSAGES} + + counts: Final = await count_request_input_tokens( + request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=json.dumps(body).encode() + ) + + assert _RecordingCounter.bodies == [] + assert counts["gpt-4o"] != RUST_INPUT_TOKENS diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py new file mode 100644 index 00000000000..2df91204390 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_token_counter.py @@ -0,0 +1,242 @@ +"""Tests for the Rust input token counter bridge. + +The native factory is dependency-injected through ``TOKEN_COUNTER.override`` +so the fallback cases run without the compiled extension present. The parity +cases need the extension and are skipped when it is not built. +""" + +from __future__ import annotations + +import json +from typing import Final + +import pytest + +import litellm +from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens +from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge import token_counter as bridge + +MODEL: Final = "claude-sonnet-4-5-20250929" +BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode() + + +class _FakeDeclined(Exception): + pass + + +class _FakeUpstream(Exception): + pass + + +class _FakeNative: + RustBridgeDeclined = _FakeDeclined + RustUpstreamError = _FakeUpstream + + +class _RecordingCounter: + def __init__(self, tokenizer_json: str) -> None: + self.tokenizer_json = tokenizer_json + self.bodies: list[bytes] = [] + + async def acount_request(self, body: bytes) -> object: + self.bodies.append(body) + return {"model": MODEL, "input_tokens": 42} + + +class _DecliningCounter: + def __init__(self, tokenizer_json: str) -> None: + pass + + async def acount_request(self, body: bytes) -> object: + raise _FakeDeclined("request has no messages") + + +class _FailingCounter: + def __init__(self, tokenizer_json: str) -> None: + pass + + async def acount_request(self, body: bytes) -> object: + raise RuntimeError("encode failed") + + +@pytest.fixture(autouse=True) +def _reset_bridge(monkeypatch: pytest.MonkeyPatch): + bridge.TOKEN_COUNTER.reset() + bridge._anthropic_counter.cache_clear() + configuration.reset_rust_configuration() + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) + yield + bridge.TOKEN_COUNTER.reset() + bridge._anthropic_counter.cache_clear() + configuration.reset_rust_configuration() + + +@pytest.mark.asyncio +async def test_disabled_bridge_never_constructs_a_counter() -> None: + constructed: list[str] = [] + + def factory(tokenizer_json: str) -> _RecordingCounter: + constructed.append(tokenizer_json) + return _RecordingCounter(tokenizer_json) + + litellm.rust(False) + bridge.TOKEN_COUNTER.override(factory) + + assert await bridge.count_anthropic_input_tokens(BODY) is None + assert constructed == [] + + +@pytest.mark.asyncio +async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> None: + counters: list[_RecordingCounter] = [] + + def factory(tokenizer_json: str) -> _RecordingCounter: + counter = _RecordingCounter(tokenizer_json) + counters.append(counter) + return counter + + litellm.rust(True) + bridge.TOKEN_COUNTER.override(factory) + + first: Final = await bridge.count_anthropic_input_tokens(BODY) + second: Final = await bridge.count_anthropic_input_tokens(BODY) + + assert first == bridge.InputTokenCount(model=MODEL, input_tokens=42) + assert second == first + assert len(counters) == 1 + assert counters[0].bodies == [BODY, BODY] + assert json.loads(counters[0].tokenizer_json)["model"]["type"] == "BPE" + + +@pytest.mark.asyncio +async def test_missing_native_module_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: + litellm.rust(True) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) + + assert await bridge.count_anthropic_input_tokens(BODY) is None + + +@pytest.mark.asyncio +async def test_declined_request_falls_back() -> None: + litellm.rust(True) + bridge.TOKEN_COUNTER.override(_DecliningCounter) + + assert await bridge.count_anthropic_input_tokens(BODY) is None + + +@pytest.mark.asyncio +async def test_runtime_failure_falls_back() -> None: + litellm.rust(True) + bridge.TOKEN_COUNTER.override(_FailingCounter) + + assert await bridge.count_anthropic_input_tokens(BODY) is None + + +@pytest.mark.parametrize( + ("model", "expected"), + ((MODEL, True), ("claude-3-5-sonnet-20241022", False), ("gpt-4o", False), ("my-router-alias", False)), +) +def test_uses_anthropic_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: bool) -> None: + assert bridge.uses_anthropic_tokenizer(model) is expected + + +@pytest.mark.parametrize("flag", ("disable_hf_tokenizer_download", "disable_token_counter")) +def test_uses_anthropic_tokenizer_respects_python_opt_outs(monkeypatch: pytest.MonkeyPatch, flag: str) -> None: + monkeypatch.setattr(litellm, flag, True) + + assert bridge.uses_anthropic_tokenizer(MODEL) is False + + +PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = ( + {"model": MODEL, "messages": [{"role": "user", "content": "Hello, how are you today?"}]}, + { + "model": MODEL, + "messages": [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "name": "bob", "content": [{"type": "text", "text": "Summarize this."}]}, + {"role": "assistant", "content": "Sure."}, + ], + }, + { + "model": MODEL, + "messages": [{"role": "user", "content": "weather in sf?"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City"}, + "unit": {"type": "string", "enum": ["c", "f"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": {"type": "function", "function": {"name": "get_weather"}}, + }, + { + "model": MODEL, + "messages": [{"role": "user", "content": "x " * 20_000}], + }, + {"model": MODEL, "prompt": "Write a haiku about ships.", "max_tokens": 20}, + {"model": MODEL, "prompt": ["first prompt", "second prompt"]}, + { + "model": MODEL, + "instructions": "be terse", + "input": [ + {"role": "user", "content": [{"type": "input_text", "text": "Summarise caf\u00e9 menus \u2014 \"ok\"?\n"}]}, + {"role": "assistant", "content": "Sure."}, + ], + }, + {"model": MODEL, "input": "a single embedding string"}, + {"model": MODEL, "input": [[101, 2023, 5], [7]], "encoding_format": "float"}, + {"model": MODEL, "query": "best harbour", "documents": ["doc one", {"text": "doc two", "title": "T", "n": 3}]}, + {"model": MODEL, "messages": None, "prompt": "messages key wins even when null"}, + {"prompt": "model comes from the route"}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_body", PARITY_REQUESTS) +async def test_native_count_matches_python_budget_counter( + monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object] +) -> None: + native: Final = pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + litellm.rust(True) + + rust_count: Final = await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode()) + python_count: Final = _count_input_tokens(request_body=request_body, model=MODEL) + + assert rust_count is not None + assert rust_count.model == request_body.get("model") + assert rust_count.input_tokens == python_count + + +DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = ( + { + "model": MODEL, + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}], + }, + {"model": MODEL, "prompt": 1.5}, + {"model": MODEL, "documents": [{"score": 0.5}]}, + {"model": MODEL, "file": "audio.mp3"}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_body", DECLINED_REQUESTS) +async def test_native_declines_shapes_python_prices_differently( + monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object] +) -> None: + native: Final = pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) + litellm.rust(True) + + assert await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode()) is None