mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
feat(rust_bridge): count budget-check input tokens in Rust on all LLM routes (#40381)
* feat(rust_bridge): count budget-check input tokens in Rust on all LLM routes Rust counts input tokens from the raw JSON body with the GIL released inside the existing budget reservation, covering every LLM route the auth dependency guards. It only fires for models on the Anthropic tokenizer when a budget is set, and Python counts whenever Rust is off, missing, or declines a body shape. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * perf(rust): count byte-level BPE tokens without the GPT-2 split regex (#40594) The oniguruma run of the ByteLevel pre-tokenizer regex is about 90% of encode_fast on a 100k token body (100 ms of the ~110 ms Rust admission count in the gateway pod). A hand-written scanner that yields the same pieces, then feeds the model directly, counts the same text in 10 ms. It only engages for tokenizers with the Anthropic shape (optional NFKC, ByteLevel without prefix space, no post-processor) and falls back to the full encoder when the text contains an added token. Parity with encode_fast is tested on random texts, the pieces are compared with the real pre-tokenizer, and the \p{L}/\p{N}/\s tables are checked against oniguruma for every code point. NFKC runs through unicode-normalization-alignments, the crate and Unicode tables NormalizedString::nfkc already uses, so the fast path normalizes exactly what the full encoder would. Using the newer unicode-normalization crate changed the count for 171 code points that gained compatibility decompositions after Unicode 9 (U+32FF, U+A7F1..). The fast normalizer is compared with the tokenizer's for every scalar value and on random texts. The scanner is built without mutable state: byte_char and mapped_len replace the const table builders and the reusable mapped buffer, and iter::successors replaces the stateful piece iterator. byte_chars_match_the_byte_level_alphabet checks the byte mapping against ByteLevel for every scalar value. Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(rust_bridge): bound concurrent token-count encodes and share the Anthropic tokenizer predicate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
0e35c8fee9
commit
46a185d3cd
30 changed files with 3196 additions and 46 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
422
litellm-rust/Cargo.lock
generated
422
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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://<host>/v1/realtime?model=<model>` (WebSocket)
|
||||
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
2
litellm-rust/crates/python-bridge/src/constants.rs
Normal file
2
litellm-rust/crates/python-bridge/src/constants.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
/// Concurrent token-count encodes allowed when the core count is unavailable.
|
||||
pub(crate) const TOKEN_COUNT_FALLBACK_PARALLELISM: usize = 1;
|
||||
|
|
@ -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<T, F>(
|
||||
pub(crate) fn run_sync<T, E, F>(
|
||||
py: Python<'_>,
|
||||
future: F,
|
||||
map_error: fn(Error) -> PyErr,
|
||||
map_error: fn(E) -> PyErr,
|
||||
) -> PyResult<Py<PyAny>>
|
||||
where
|
||||
T: Serialize + Send + 'static,
|
||||
F: Future<Output = Result<T, Error>> + Send + 'static,
|
||||
E: Send + 'static,
|
||||
F: Future<Output = Result<T, E>> + Send + 'static,
|
||||
{
|
||||
run_sync_on(
|
||||
py,
|
||||
|
|
@ -28,15 +28,16 @@ where
|
|||
)
|
||||
}
|
||||
|
||||
fn run_sync_on<T, F>(
|
||||
fn run_sync_on<T, E, F>(
|
||||
py: Python<'_>,
|
||||
runtime: &Runtime,
|
||||
future: F,
|
||||
map_error: fn(Error) -> PyErr,
|
||||
map_error: fn(E) -> PyErr,
|
||||
) -> PyResult<Py<PyAny>>
|
||||
where
|
||||
T: Serialize + Send + 'static,
|
||||
F: Future<Output = Result<T, Error>> + Send + 'static,
|
||||
E: Send + 'static,
|
||||
F: Future<Output = Result<T, E>> + 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<T, F>(
|
||||
pub(crate) fn run_async<T, E, F>(
|
||||
py: Python<'_>,
|
||||
future: F,
|
||||
map_error: fn(Error) -> PyErr,
|
||||
map_error: fn(E) -> PyErr,
|
||||
) -> PyResult<Bound<'_, PyAny>>
|
||||
where
|
||||
T: Serialize + Send + 'static,
|
||||
F: Future<Output = Result<T, Error>> + Send + 'static,
|
||||
E: Send + 'static,
|
||||
F: Future<Output = Result<T, E>> + 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<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
|
||||
fn map_core_result<T, E>(result: Result<T, E>, map_error: fn(E) -> PyErr) -> PyResult<T> {
|
||||
match result {
|
||||
Ok(value) => Ok(value),
|
||||
Err(error) => Err(
|
||||
|
|
@ -75,9 +77,9 @@ fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -
|
|||
}
|
||||
}
|
||||
|
||||
async fn catch_future_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
|
||||
async fn catch_future_panic<T, E, F>(future: F) -> PyResult<Result<T, E>>
|
||||
where
|
||||
F: Future<Output = Result<T, Error>>,
|
||||
F: Future<Output = Result<T, E>>,
|
||||
{
|
||||
AssertUnwindSafe(future)
|
||||
.catch_unwind()
|
||||
|
|
@ -85,9 +87,9 @@ where
|
|||
.map_err(panic_to_pyerr)
|
||||
}
|
||||
|
||||
async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
|
||||
async fn wait_for_sync_result<T, E, F>(future: F) -> PyResult<Result<T, E>>
|
||||
where
|
||||
F: Future<Output = Result<T, Error>>,
|
||||
F: Future<Output = Result<T, E>>,
|
||||
{
|
||||
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::<bool, _>(py, async { Ok(true) }, runtime_error)
|
||||
run_sync::<bool, Error, _>(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::<bool, _>(
|
||||
let error = run_sync::<bool, Error, _>(
|
||||
py,
|
||||
poll_fn(|_| -> Poll<Result<bool, Error>> { 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::<bool, _>(
|
||||
let error = run_sync::<bool, Error, _>(
|
||||
py,
|
||||
async { Err(Error::InvalidRequest("invalid".to_string())) },
|
||||
panicking_error_mapper,
|
||||
|
|
|
|||
|
|
@ -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::ResponsesWebSocketConnection>()?;
|
||||
super::token_counter::register(module)?;
|
||||
super::diagnostics::register(module)
|
||||
}
|
||||
}
|
||||
|
|
@ -106,6 +109,7 @@ mod tests {
|
|||
"chat_completions",
|
||||
"achat_completions",
|
||||
"ResponsesWebSocketConnection",
|
||||
"TokenCounter",
|
||||
"gil_stats",
|
||||
];
|
||||
|
||||
|
|
|
|||
87
litellm-rust/crates/python-bridge/src/token_counter.rs
Normal file
87
litellm-rust/crates/python-bridge/src/token_counter.rs
Normal file
|
|
@ -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<CoreTokenCounter>,
|
||||
encode_slots: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl TokenCounter {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
|
||||
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<Bound<'py, PyAny>> {
|
||||
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<InputTokenCount, Error> {
|
||||
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::<TokenCounter>()
|
||||
}
|
||||
28
litellm-rust/crates/token-counter/Cargo.toml
Normal file
28
litellm-rust/crates/token-counter/Cargo.toml
Normal file
|
|
@ -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
|
||||
128
litellm-rust/crates/token-counter/benches/allocations.rs
Normal file
128
litellm-rust/crates/token-counter/benches/allocations.rs
Normal file
|
|
@ -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,
|
||||
},
|
||||
);
|
||||
}
|
||||
100
litellm-rust/crates/token-counter/benches/token_counter.rs
Normal file
100
litellm-rust/crates/token-counter/benches/token_counter.rs
Normal file
|
|
@ -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::<Tokenizer>()
|
||||
.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);
|
||||
650
litellm-rust/crates/token-counter/src/byte_level.rs
Normal file
650
litellm-rust/crates/token-counter/src/byte_level.rs
Normal file
|
|
@ -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<String>,
|
||||
unicode_classes: &'static UnicodeClasses,
|
||||
}
|
||||
|
||||
impl ByteLevelCounter {
|
||||
pub(super) fn detect(tokenizer: &Tokenizer) -> Option<Self> {
|
||||
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<usize> {
|
||||
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<Item = &'a str> {
|
||||
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",
|
||||
"<EOT>",
|
||||
"<META_START>",
|
||||
];
|
||||
|
||||
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 <EOT> 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<char> = (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::<PreTokenizerWrapper>);
|
||||
}
|
||||
"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 Ⅳ", "<EOT> 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::<NormalizerWrapper>)
|
||||
.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漢字🙂",
|
||||
" <EOT> 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);
|
||||
}
|
||||
}
|
||||
194
litellm-rust/crates/token-counter/src/counter.rs
Normal file
194
litellm-rust/crates/token-counter/src/counter.rs
Normal file
|
|
@ -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<String>,
|
||||
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<ByteLevelCounter>,
|
||||
}
|
||||
|
||||
impl TokenCounter {
|
||||
/// Load a HuggingFace `tokenizer.json` document. The host reads the file.
|
||||
pub fn from_json(tokenizer_json: &str) -> Result<Self, Error> {
|
||||
let tokenizer = tokenizer_json
|
||||
.parse::<tokenizers::Tokenizer>()
|
||||
.map_err(Error::Load)?;
|
||||
let byte_level = ByteLevelCounter::detect(&tokenizer);
|
||||
Ok(Self {
|
||||
tokenizer,
|
||||
byte_level,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn count_text(&self, text: &str) -> Result<usize, Error> {
|
||||
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<InputTokenCount, Error> {
|
||||
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<usize, Error> {
|
||||
let message_tokens = messages
|
||||
.iter()
|
||||
.map(|message| self.count_message(message))
|
||||
.sum::<Result<usize, _>>()?;
|
||||
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<usize, Error> {
|
||||
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<usize, Error> {
|
||||
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::<Result<usize, _>>(),
|
||||
TextValue::Object(_) => self.count_text(&python_json::dumps(value)?),
|
||||
}
|
||||
}
|
||||
|
||||
fn count_message(&self, message: &Message) -> Result<usize, Error> {
|
||||
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::<Result<usize, _>>()?,
|
||||
None => 0,
|
||||
};
|
||||
Ok(TOKENS_PER_MESSAGE + role_tokens + name_tokens + content_tokens)
|
||||
}
|
||||
|
||||
fn count_content_item(&self, item: &ContentItem) -> Result<usize, Error> {
|
||||
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<usize, Error> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
31
litellm-rust/crates/token-counter/src/error.rs
Normal file
31
litellm-rust/crates/token-counter/src/error.rs
Normal file
|
|
@ -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),
|
||||
}
|
||||
17
litellm-rust/crates/token-counter/src/lib.rs
Normal file
17
litellm-rust/crates/token-counter/src/lib.rs
Normal file
|
|
@ -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;
|
||||
155
litellm-rust/crates/token-counter/src/python_json.rs
Normal file
155
litellm-rust/crates/token-counter/src/python_json.rs
Normal file
|
|
@ -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<String, Error> {
|
||||
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<usize, Error> {
|
||||
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<W>(&mut self, writer: &mut W, first: bool) -> io::Result<()>
|
||||
where
|
||||
W: ?Sized + Write,
|
||||
{
|
||||
if !first {
|
||||
writer.write_all(b", ")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn begin_object_key<W>(&mut self, writer: &mut W, first: bool) -> io::Result<()>
|
||||
where
|
||||
W: ?Sized + Write,
|
||||
{
|
||||
if !first {
|
||||
writer.write_all(b", ")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn begin_object_value<W>(&mut self, writer: &mut W) -> io::Result<()>
|
||||
where
|
||||
W: ?Sized + Write,
|
||||
{
|
||||
writer.write_all(b": ")
|
||||
}
|
||||
|
||||
fn write_string_fragment<W>(&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"
|
||||
);
|
||||
}
|
||||
}
|
||||
304
litellm-rust/crates/token-counter/src/tools.rs
Normal file
304
litellm-rust/crates/token-counter/src/tools.rs
Normal file
|
|
@ -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<String, Error> {
|
||||
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<String, Error> {
|
||||
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<ToolDefinition> {
|
||||
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)
|
||||
));
|
||||
}
|
||||
}
|
||||
238
litellm-rust/crates/token-counter/src/types.rs
Normal file
238
litellm-rust/crates/token-counter/src/types.rs
Normal file
|
|
@ -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<String>,
|
||||
#[serde(default, deserialize_with = "present_messages")]
|
||||
pub(crate) messages: Option<Vec<Message>>,
|
||||
pub(crate) tools: Option<Vec<ToolDefinition>>,
|
||||
pub(crate) tool_choice: Option<ToolChoice>,
|
||||
#[serde(default, deserialize_with = "present_text")]
|
||||
pub(crate) prompt: Option<TextValue>,
|
||||
#[serde(default, deserialize_with = "present_text")]
|
||||
pub(crate) input: Option<TextValue>,
|
||||
#[serde(default, deserialize_with = "present_text")]
|
||||
pub(crate) query: Option<TextValue>,
|
||||
#[serde(default, deserialize_with = "present_text")]
|
||||
pub(crate) documents: Option<TextValue>,
|
||||
}
|
||||
|
||||
impl CountableRequest {
|
||||
pub fn parse(body: &[u8]) -> Result<Self, Error> {
|
||||
serde_json::from_slice(body).map_err(Error::RequestParse)
|
||||
}
|
||||
}
|
||||
|
||||
fn present_messages<'de, D: Deserializer<'de>>(
|
||||
deserializer: D,
|
||||
) -> Result<Option<Vec<Message>>, D::Error> {
|
||||
Option::<Vec<Message>>::deserialize(deserializer)
|
||||
.map(|messages| Some(messages.unwrap_or_default()))
|
||||
}
|
||||
|
||||
fn present_text<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<TextValue>, 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<TextValue>),
|
||||
Object(IndexMap<String, TextValue>),
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for TextValue {
|
||||
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
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<E>(self) -> Result<Self::Value, E> {
|
||||
Ok(TextValue::Null)
|
||||
}
|
||||
|
||||
fn visit_none<E>(self) -> Result<Self::Value, E> {
|
||||
Ok(TextValue::Null)
|
||||
}
|
||||
|
||||
fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
|
||||
Ok(TextValue::Bool(value))
|
||||
}
|
||||
|
||||
fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
|
||||
Ok(TextValue::Number(value.into()))
|
||||
}
|
||||
|
||||
fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
|
||||
Ok(TextValue::Number(value.into()))
|
||||
}
|
||||
|
||||
fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
|
||||
where
|
||||
E: serde::de::Error,
|
||||
{
|
||||
Number::from_f64(value)
|
||||
.map(TextValue::Number)
|
||||
.ok_or_else(|| E::custom("non-finite JSON number"))
|
||||
}
|
||||
|
||||
fn visit_str<E>(self, value: &str) -> Result<Self::Value, E> {
|
||||
Ok(TextValue::Text(value.to_owned()))
|
||||
}
|
||||
|
||||
fn visit_string<E>(self, value: String) -> Result<Self::Value, E> {
|
||||
Ok(TextValue::Text(value))
|
||||
}
|
||||
|
||||
fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
|
||||
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<A>(self, mut map: A) -> Result<Self::Value, A::Error>
|
||||
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<String>,
|
||||
pub(crate) name: Option<String>,
|
||||
pub(crate) content: Option<MessageContent>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, PartialEq)]
|
||||
#[serde(untagged)]
|
||||
pub(crate) enum MessageContent {
|
||||
Text(String),
|
||||
Blocks(Vec<ContentItem>),
|
||||
}
|
||||
|
||||
#[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<String> },
|
||||
/// 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<FunctionDefinition>,
|
||||
pub(crate) name: Option<String>,
|
||||
pub(crate) description: Option<String>,
|
||||
pub(crate) input_schema: Option<Schema>,
|
||||
pub(crate) parameters: Option<Schema>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, PartialEq)]
|
||||
pub(crate) struct FunctionDefinition {
|
||||
pub(crate) name: Option<String>,
|
||||
pub(crate) description: Option<String>,
|
||||
pub(crate) parameters: Option<Schema>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
|
||||
pub(crate) struct Schema {
|
||||
#[serde(rename = "type")]
|
||||
pub(crate) schema_type: Option<SchemaType>,
|
||||
pub(crate) description: Option<String>,
|
||||
#[serde(rename = "enum")]
|
||||
pub(crate) enum_values: Option<Vec<EnumValue>>,
|
||||
pub(crate) items: Option<Box<Schema>>,
|
||||
pub(crate) properties: Option<IndexMap<String, Schema>>,
|
||||
pub(crate) required: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, PartialEq)]
|
||||
#[serde(untagged)]
|
||||
pub(crate) enum SchemaType {
|
||||
Name(String),
|
||||
Union(Vec<String>),
|
||||
}
|
||||
|
||||
#[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,
|
||||
}
|
||||
73
litellm-rust/crates/token-counter/src/unicode_classes.rs
Normal file
73
litellm-rust/crates/token-counter/src/unicode_classes.rs
Normal file
|
|
@ -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<Option<UnicodeClasses>> = 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<Self> {
|
||||
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::<Option<Box<[_]>>>()?;
|
||||
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)
|
||||
}
|
||||
}
|
||||
176
litellm-rust/crates/token-counter/tests/token_counter.rs
Normal file
176
litellm-rust/crates/token-counter/tests/token_counter.rs
Normal file
|
|
@ -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(_))));
|
||||
}
|
||||
|
|
@ -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),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
79
litellm/rust_bridge/token_counter.py
Normal file
79
litellm/rust_bridge/token_counter.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
242
tests/test_litellm/rust_bridge/test_token_counter.py
Normal file
242
tests/test_litellm/rust_bridge/test_token_counter.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue