diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2720cf01f2e..62bde5806a4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -61,6 +61,12 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + [[package]] name = "arc-swap" version = "1.9.2" @@ -598,6 +604,17 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "bstr" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f" +dependencies = [ + "memchr", + "regex-automata", + "serde_core", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -1181,6 +1198,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fancy-regex" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fancy-regex" version = "0.19.2" @@ -1926,6 +1954,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.186" @@ -2083,7 +2117,7 @@ dependencies = [ name = "litellm-core-utils" version = "0.1.0" dependencies = [ - "fancy-regex", + "fancy-regex 0.19.2", "litellm-types", "rstest", "serde", @@ -2212,10 +2246,24 @@ dependencies = [ name = "litellm-token-counter" version = "0.1.0" dependencies = [ - "base64 0.22.1", "criterion", "indexmap 2.14.0", "itoa", + "litellm-token-counter-fast", + "litellm-token-counter-huggingface", + "litellm-token-counter-tiktoken", + "rand 0.8.7", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", +] + +[[package]] +name = "litellm-token-counter-fast" +version = "0.1.0" +dependencies = [ + "base64 0.22.1", "rand 0.8.7", "rstest", "rustc-hash", @@ -2226,6 +2274,22 @@ dependencies = [ "unicode-normalization-alignments", ] +[[package]] +name = "litellm-token-counter-huggingface" +version = "0.1.0" +dependencies = [ + "thiserror 2.0.19", + "tokenizers", +] + +[[package]] +name = "litellm-token-counter-tiktoken" +version = "0.1.0" +dependencies = [ + "thiserror 2.0.19", + "tiktoken-rs", +] + [[package]] name = "litellm-types" version = "0.1.0" @@ -3794,6 +3858,21 @@ dependencies = [ "syn 3.0.0", ] +[[package]] +name = "tiktoken-rs" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "027853bbf8c7763b77c5c595f1c271c7d536ced7d6f83452911b944621e57fc2" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bstr", + "fancy-regex 0.17.0", + "lazy_static", + "regex", + "rustc-hash", +] + [[package]] name = "time" version = "0.3.53" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a6185632871..109a2edaf4d 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -24,6 +24,9 @@ litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-token-counter = { path = "crates/token-counter" } +litellm-token-counter-fast = { path = "crates/token-counter-fast" } +litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" } +litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" } litellm-host-python = { path = "crates/host-python" } bytes = "1" @@ -45,6 +48,7 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std", sha2 = "0.10" subtle = "2" thiserror = "2.0" +tiktoken-rs = "0.12.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 3a4a579efa3..2c868ebc769 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,10 +10,13 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["abi3"] +default = ["abi3", "token-counter-huggingface", "token-counter-tiktoken"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] +token-counter-fast = ["litellm-token-counter/fast"] +token-counter-huggingface = ["litellm-token-counter/huggingface"] +token-counter-tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true @@ -26,7 +29,7 @@ litellm-http.workspace = true litellm-llms.workspace = true litellm-types.workspace = true litellm-host-python.workspace = true -litellm-token-counter.workspace = true +litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index 7dc86b78ad6..a7c6bce4d0b 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/token_counter.rs @@ -1,6 +1,19 @@ -use std::{num::NonZero, sync::Arc, thread::available_parallelism}; +use std::sync::Arc; -use litellm_host_python::{release_gil, run_async}; +#[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" +))] +use std::{num::NonZero, thread::available_parallelism}; + +#[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" +))] +use litellm_host_python::release_gil; +use litellm_host_python::run_async; use litellm_token_counter::{ CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, }; @@ -28,17 +41,47 @@ pub(crate) struct TokenCounter { impl TokenCounter { #[new] fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { - Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) + #[cfg(feature = "token-counter-huggingface")] + { + Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) + } + #[cfg(not(feature = "token-counter-huggingface"))] + { + let _ = (py, tokenizer_json); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the token-counter-huggingface feature", + )) + } } #[staticmethod] - fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { - Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file)) + fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult { + #[cfg(feature = "token-counter-fast")] + { + Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) + } + #[cfg(not(feature = "token-counter-fast"))] + { + let _ = (py, tokenizer_json); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the token-counter-fast feature", + )) + } } #[staticmethod] - fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { - Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file)) + fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { + #[cfg(feature = "token-counter-tiktoken")] + { + Self::load(py, || CoreTokenCounter::from_tiktoken(encoding)) + } + #[cfg(not(feature = "token-counter-tiktoken"))] + { + let _ = (py, encoding); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the token-counter-tiktoken feature", + )) + } } fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult> { @@ -62,6 +105,11 @@ impl TokenCounter { } impl TokenCounter { + #[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" + ))] fn load( py: Python<'_>, load: impl FnOnce() -> Result + Send, @@ -74,6 +122,11 @@ impl TokenCounter { } } +#[cfg(any( + feature = "token-counter-fast", + feature = "token-counter-huggingface", + feature = "token-counter-tiktoken" +))] fn encode_parallelism() -> usize { available_parallelism().map_or(1, NonZero::get) } @@ -86,7 +139,10 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result PyErr { let message = error.to_string(); match error { - Error::Load(_) | Error::Ranks(_) | Error::UnicodeClasses => PyValueError::new_err(message), + Error::Load(_) + | Error::Ranks(_) + | Error::UnicodeClasses + | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message), Error::RequestParse(_) | Error::MissingInput | Error::FloatText diff --git a/litellm-rust/crates/token-counter-fast/Cargo.toml b/litellm-rust/crates/token-counter-fast/Cargo.toml new file mode 100644 index 00000000000..4950277f5b0 --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "litellm-token-counter-fast" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +base64.workspace = true +rustc-hash = "2.1.3" +thiserror.workspace = true +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } +unicode-normalization-alignments = "0.1.12" + +[dev-dependencies] +rand.workspace = true +rstest.workspace = true +serde.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/token-counter/src/byte_level.rs b/litellm-rust/crates/token-counter-fast/src/byte_level.rs similarity index 97% rename from litellm-rust/crates/token-counter/src/byte_level.rs rename to litellm-rust/crates/token-counter-fast/src/byte_level.rs index ec6134a252e..6fc7d9146b9 100644 --- a/litellm-rust/crates/token-counter/src/byte_level.rs +++ b/litellm-rust/crates/token-counter-fast/src/byte_level.rs @@ -473,13 +473,13 @@ mod tests { _ => unreachable!(), } assert!(ByteLevelCounter::detect(&anthropic_tokenizer).is_none()); - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); for text in ["", "Hello WORLD! ＡＢ fi Ⅳ", " stop"] { assert_eq!( - counter.count_text(text).expect("count"), + counter.count_tokens(text).expect("count"), reference_count(&anthropic_tokenizer, text) ); } @@ -545,7 +545,7 @@ mod tests { .rstrip(rstrip)]) .expect("add token"); let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported"); - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); @@ -557,7 +557,7 @@ mod tests { ] { assert_eq!(fast.count(&anthropic_tokenizer, text), None); assert_eq!( - counter.count_text(text).expect("count"), + counter.count_tokens(text).expect("count"), reference_count(&anthropic_tokenizer, text) ); } @@ -571,17 +571,17 @@ mod tests { 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")) + crate::FastTokenizer::from_json(&tokenizer.to_string(false).expect("serialize")) .expect("load"); assert!(matches!( - counter.count_text("hello"), + counter.count_tokens("hello"), Err(crate::Error::Encode(_)) )); } #[rstest] fn shared_counter_matches_encoder_across_threads(anthropic_tokenizer: Tokenizer) { - let counter = crate::TokenCounter::from_json( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); @@ -598,7 +598,7 @@ mod tests { scope.spawn(move || { for _ in 0..100 { for (text, count) in inputs.iter().zip(expected) { - assert_eq!(counter.count_text(text).expect("count"), count); + assert_eq!(counter.count_tokens(text).expect("count"), count); } } }); @@ -614,10 +614,10 @@ mod tests { 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( + let counter = crate::FastTokenizer::from_json( &anthropic_tokenizer.to_string(false).expect("serialize"), ) .expect("load"); - assert_eq!(counter.count_text("ABCD EFGH").expect("count"), 1); + assert_eq!(counter.count_tokens("ABCD EFGH").expect("count"), 1); } } diff --git a/litellm-rust/crates/token-counter/src/cl100k.rs b/litellm-rust/crates/token-counter-fast/src/cl100k.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/cl100k.rs rename to litellm-rust/crates/token-counter-fast/src/cl100k.rs diff --git a/litellm-rust/crates/token-counter-fast/src/error.rs b/litellm-rust/crates/token-counter-fast/src/error.rs new file mode 100644 index 00000000000..e63ccf3ad39 --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/src/error.rs @@ -0,0 +1,13 @@ +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("failed to load tokenizer: tiktoken rank file: {0}")] + Ranks(String), + #[error("failed to load tokenizer: Unicode character classes are unavailable")] + UnicodeClasses, + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), +} diff --git a/litellm-rust/crates/token-counter-fast/src/lib.rs b/litellm-rust/crates/token-counter-fast/src/lib.rs new file mode 100644 index 00000000000..ce91af642ea --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/src/lib.rs @@ -0,0 +1,70 @@ +#![forbid(unsafe_code)] + +mod byte_level; +mod cl100k; +mod error; +mod o200k; +mod scanner; +mod tiktoken; +mod unicode_classes; + +use byte_level::ByteLevelCounter; +use scanner::{SplitPattern, TiktokenCounter}; + +pub use error::Error; + +enum Encoder { + HuggingFace { + tokenizer: Box, + byte_level: Option, + }, + Tiktoken(TiktokenCounter), +} + +pub struct FastTokenizer(Encoder); + +impl FastTokenizer { + pub fn from_json(json: &str) -> Result { + let tokenizer = json.parse::().map_err(Error::Load)?; + let byte_level = ByteLevelCounter::detect(&tokenizer); + Ok(Self(Encoder::HuggingFace { + tokenizer: Box::new(tokenizer), + byte_level, + })) + } + + pub fn from_cl100k_ranks(ranks: &str) -> Result { + Self::from_ranks(SplitPattern::Cl100k, ranks) + } + + pub fn from_o200k_ranks(ranks: &str) -> Result { + Self::from_ranks(SplitPattern::O200k, ranks) + } + + fn from_ranks(split: SplitPattern, ranks: &str) -> Result { + TiktokenCounter::from_ranks(split, ranks) + .map(Encoder::Tiktoken) + .map(Self) + } + + pub fn count_tokens(&self, text: &str) -> Result { + match &self.0 { + Encoder::Tiktoken(counter) => Ok(counter.count(text)), + Encoder::HuggingFace { + tokenizer, + byte_level, + } => { + if let Some(count) = byte_level + .as_ref() + .and_then(|counter| counter.count(tokenizer, text)) + { + return Ok(count); + } + tokenizer + .encode_fast(text, true) + .map(|encoding| encoding.len()) + .map_err(Error::Encode) + } + } + } +} diff --git a/litellm-rust/crates/token-counter/src/o200k.rs b/litellm-rust/crates/token-counter-fast/src/o200k.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/o200k.rs rename to litellm-rust/crates/token-counter-fast/src/o200k.rs diff --git a/litellm-rust/crates/token-counter/src/scanner.rs b/litellm-rust/crates/token-counter-fast/src/scanner.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/scanner.rs rename to litellm-rust/crates/token-counter-fast/src/scanner.rs diff --git a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs new file mode 100644 index 00000000000..7a9e71ed587 --- /dev/null +++ b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs @@ -0,0 +1,222 @@ +//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and +//! the merge loop that turns one regex piece into tokens. The merge order is +//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is +//! identical, but pairs are tracked in a heap so a long piece costs +//! `O(n log n)` instead of tiktoken's `O(n^2)`. + +use std::cmp::Reverse; +use std::collections::BinaryHeap; + +use base64::Engine; +use base64::engine::general_purpose::STANDARD; +use rustc_hash::FxHashMap; + +use crate::Error; + +type Rank = u32; + +const NO_RANK: Rank = Rank::MAX; +const END: usize = usize::MAX; + +pub(super) struct MergeRanks(FxHashMap, Rank>); + +impl MergeRanks { + pub(super) fn parse(text: &str) -> Result { + let ranks = text + .lines() + .filter(|line| !line.is_empty()) + .map(parse_line) + .collect::, _>>()?; + if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) { + return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token"))); + } + Ok(Self(ranks)) + } + + fn rank(&self, bytes: &[u8]) -> Rank { + self.0.get(bytes).copied().unwrap_or(NO_RANK) + } + + /// Token count of one regex piece, as `encode_ordinary` would produce. + pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize { + if piece.len() < 2 || self.0.contains_key(piece) { + return 1; + } + scratch.reset(piece.len()); + for start in 0..piece.len() - 1 { + scratch.set_rank(start, self.rank(&piece[start..start + 2])); + } + let mut parts = piece.len(); + while let Some(Reverse((rank, start))) = scratch.heap.pop() { + if scratch.next[start] == END || scratch.rank[start] != rank { + continue; + } + let merged = scratch.next[start]; + let after = scratch.next[merged]; + scratch.next[merged] = END; + scratch.next[start] = after; + parts -= 1; + if after < piece.len() { + scratch.prev[after] = start; + scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)])); + } else { + scratch.rank[start] = NO_RANK; + } + let before = scratch.prev[start]; + if before != END { + scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)])); + } + } + parts + } +} + +fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> { + let (token, rank) = line + .split_once(' ') + .ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?; + let bytes = STANDARD + .decode(token) + .map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?; + let rank = rank + .parse() + .map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?; + Ok((bytes.into_boxed_slice(), rank)) +} + +/// Buffers reused across the pieces of one text. Parts are addressed by the +/// byte offset they start at, which also gives the leftmost-pair tie break. +#[derive(Default)] +pub(super) struct MergeScratch { + next: Vec, + prev: Vec, + rank: Vec, + heap: BinaryHeap>, +} + +impl MergeScratch { + fn reset(&mut self, len: usize) { + self.next.clear(); + self.next.extend(1..=len); + self.prev.clear(); + self.prev.push(END); + self.prev.extend(0..len - 1); + self.rank.clear(); + self.rank.resize(len, NO_RANK); + self.heap.clear(); + } + + fn end(&self, start: usize) -> usize { + self.next[start] + } + + fn set_rank(&mut self, start: usize, rank: Rank) { + self.rank[start] = rank; + if rank != NO_RANK { + self.heap.push(Reverse((rank, start))); + } + } +} + +#[cfg(test)] +mod tests { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + use super::*; + + fn ranks() -> MergeRanks { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4" + ); + MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo")) + .expect("rank file parses") + } + + /// tiktoken's `_byte_pair_merge`, transcribed, as the reference. + fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize { + if piece.len() < 2 || ranks.0.contains_key(piece) { + return 1; + } + let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1) + .map(|index| (index, ranks.rank(&piece[index..index + 2]))) + .chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)]) + .collect(); + let get_rank = |parts: &[(usize, Rank)], index: usize| { + if index + 3 < parts.len() { + ranks.rank(&piece[parts[index].0..parts[index + 3].0]) + } else { + NO_RANK + } + }; + loop { + let Some(index) = parts[..parts.len() - 1] + .iter() + .enumerate() + .filter(|(_, (_, rank))| *rank != NO_RANK) + .min_by_key(|(index, (_, rank))| (*rank, *index)) + .map(|(index, _)| index) + else { + return parts.len() - 1; + }; + if index > 0 { + parts[index - 1].1 = get_rank(&parts, index - 1); + } + parts[index].1 = get_rank(&parts, index); + parts.remove(index + 1); + } + } + + #[test] + fn every_byte_is_a_token() { + let ranks = ranks(); + assert_eq!(ranks.0.len(), 100_256); + assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK)); + } + + #[test] + fn heap_merge_matches_tiktokens_merge_loop() { + let ranks = ranks(); + let mut scratch = MergeScratch::default(); + let mut rng = StdRng::seed_from_u64(99); + let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123"; + for _ in 0..20_000 { + let piece: Vec = (0..rng.gen_range(1..24)) + .map(|_| alphabet[rng.gen_range(0..alphabet.len())]) + .collect(); + assert_eq!( + ranks.count_piece(&piece, &mut scratch), + reference_count(&ranks, &piece), + "piece {:?}", + String::from_utf8_lossy(&piece) + ); + } + } + + #[test] + fn long_repeated_runs_cost_close_to_linear() { + let ranks = ranks(); + let mut scratch = MergeScratch::default(); + let mut time = |len: usize| { + let piece = vec![b' '; len]; + let started = std::time::Instant::now(); + assert!(ranks.count_piece(&piece, &mut scratch) > 0); + started.elapsed() + }; + let small = (0..3).map(|_| time(1 << 14)).min().unwrap(); + let large = time(1 << 18); + assert!( + large < small * 64, + "{small:?} for 2^14 bytes, {large:?} for 2^18" + ); + } + + #[test] + fn malformed_rank_files_are_rejected() { + assert!(MergeRanks::parse("IQ==").is_err()); + assert!(MergeRanks::parse("IQ== x").is_err()); + assert!(MergeRanks::parse("!!! 1").is_err()); + assert!(MergeRanks::parse("IQ== 1").is_err()); + } +} diff --git a/litellm-rust/crates/token-counter/src/unicode_classes.rs b/litellm-rust/crates/token-counter-fast/src/unicode_classes.rs similarity index 100% rename from litellm-rust/crates/token-counter/src/unicode_classes.rs rename to litellm-rust/crates/token-counter-fast/src/unicode_classes.rs diff --git a/litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/requests.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/requests.jsonl diff --git a/litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/texts.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/cl100k/texts.jsonl diff --git a/litellm-rust/crates/token-counter/tests/fixtures/generate.py b/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/generate.py rename to litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py diff --git a/litellm-rust/crates/token-counter/tests/fixtures/o200k/requests.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/requests.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/o200k/requests.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/requests.jsonl diff --git a/litellm-rust/crates/token-counter/tests/fixtures/o200k/texts.jsonl b/litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/texts.jsonl similarity index 100% rename from litellm-rust/crates/token-counter/tests/fixtures/o200k/texts.jsonl rename to litellm-rust/crates/token-counter-fast/tests/fixtures/o200k/texts.jsonl diff --git a/litellm-rust/crates/token-counter-huggingface/Cargo.toml b/litellm-rust/crates/token-counter-huggingface/Cargo.toml new file mode 100644 index 00000000000..ad15d9f19d0 --- /dev/null +++ b/litellm-rust/crates/token-counter-huggingface/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "litellm-token-counter-huggingface" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +thiserror.workspace = true +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } diff --git a/litellm-rust/crates/token-counter-huggingface/src/lib.rs b/litellm-rust/crates/token-counter-huggingface/src/lib.rs new file mode 100644 index 00000000000..184956b5f39 --- /dev/null +++ b/litellm-rust/crates/token-counter-huggingface/src/lib.rs @@ -0,0 +1,29 @@ +#![forbid(unsafe_code)] + +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +pub enum Error { + #[error("failed to load tokenizer: {0}")] + Load(#[source] tokenizers::Error), + #[error("tokenization failed: {0}")] + Encode(#[source] tokenizers::Error), +} + +pub struct HuggingFaceTokenizer(Box); + +impl HuggingFaceTokenizer { + pub fn from_json(json: &str) -> Result { + json.parse::() + .map(Box::new) + .map(Self) + .map_err(Error::Load) + } + + pub fn count_tokens(&self, text: &str) -> Result { + self.0 + .encode_fast(text, true) + .map(|encoding| encoding.len()) + .map_err(Error::Encode) + } +} diff --git a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml new file mode 100644 index 00000000000..494a9233e69 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "litellm-token-counter-tiktoken" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +thiserror.workspace = true +tiktoken-rs.workspace = true diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs new file mode 100644 index 00000000000..b35c4876d00 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -0,0 +1,54 @@ +#![forbid(unsafe_code)] + +use thiserror::Error as ThisError; + +#[derive(Debug, ThisError)] +#[error("unsupported tokenizer: {0}")] +pub struct UnsupportedTokenizer(pub String); + +pub struct TiktokenTokenizer(&'static tiktoken_rs::CoreBPE); + +impl TiktokenTokenizer { + pub fn from_name(name: &str) -> Result { + let tokenizer = match name { + "cl100k_base" => tiktoken_rs::cl100k_base_singleton(), + "o200k_base" => tiktoken_rs::o200k_base_singleton(), + "o200k_harmony" => tiktoken_rs::o200k_harmony_singleton(), + "p50k_base" => tiktoken_rs::p50k_base_singleton(), + "p50k_edit" => tiktoken_rs::p50k_edit_singleton(), + "r50k_base" | "gpt2" => tiktoken_rs::r50k_base_singleton(), + _ => return Err(UnsupportedTokenizer(name.to_owned())), + }; + Ok(Self(tokenizer)) + } + + pub fn count_tokens(&self, text: &str) -> usize { + self.0.count_ordinary(text) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn special_tokens_are_counted_as_ordinary_text() { + let counter = TiktokenTokenizer::from_name("cl100k_base").unwrap(); + assert!(counter.count_tokens("<|endoftext|>") > 1); + } + + #[test] + fn all_python_tiktoken_encodings_are_available() { + for name in [ + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", + "gpt2", + ] { + assert!(TiktokenTokenizer::from_name(name).is_ok(), "{name}"); + } + } +} diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml index d0369631682..59e4f9a6a1d 100644 --- a/litellm-rust/crates/token-counter/Cargo.toml +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -5,16 +5,21 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +default = ["huggingface", "tiktoken"] +fast = ["dep:litellm-token-counter-fast"] +huggingface = ["dep:litellm-token-counter-huggingface"] +tiktoken = ["dep:litellm-token-counter-tiktoken"] + [dependencies] -base64.workspace = true indexmap = { version = "2.14.0", features = ["serde"] } itoa = "1.0" -rustc-hash = "2.1.3" +litellm-token-counter-fast = { workspace = true, optional = true } +litellm-token-counter-huggingface = { workspace = true, optional = true } +litellm-token-counter-tiktoken = { workspace = true, optional = true } 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 @@ -24,7 +29,9 @@ rstest.workspace = true [[bench]] name = "token_counter" harness = false +required-features = ["fast"] [[bench]] name = "allocations" harness = false +required-features = ["fast"] diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md new file mode 100644 index 00000000000..c8ed0737a7e --- /dev/null +++ b/litellm-rust/crates/token-counter/README.md @@ -0,0 +1,23 @@ +# Token counting + +`Tokenizer` is the text-counting interface. `TokenCounter` applies LiteLLM request, message, and tool accounting using any implementation of that interface + +The `fast` feature provides `fast::FastTokenizer` from `litellm-token-counter-fast`. `TokenCounter::from_json_fast`, `TokenCounter::from_cl100k_ranks`, and `TokenCounter::from_o200k_ranks` use this implementation + +The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through the upstream `tokenizers` library. `TokenCounter::from_json` uses this implementation + +The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2` + +The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend + +Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits + +Run the feature matrix with: + +```sh +cargo test -p litellm-token-counter +cargo test -p litellm-token-counter --no-default-features +cargo test -p litellm-token-counter --no-default-features --features fast +cargo test -p litellm-token-counter --no-default-features --features huggingface +cargo test -p litellm-token-counter --no-default-features --features tiktoken +``` diff --git a/litellm-rust/crates/token-counter/benches/allocations.rs b/litellm-rust/crates/token-counter/benches/allocations.rs index 343f815749d..7aebceb6e9c 100644 --- a/litellm-rust/crates/token-counter/benches/allocations.rs +++ b/litellm-rust/crates/token-counter/benches/allocations.rs @@ -87,7 +87,7 @@ fn main() { }, ); - let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("tokenizer loads"); + let counter = TokenCounter::from_json_fast(TOKENIZER_JSON).expect("tokenizer loads"); let object = CountableRequest::parse(OBJECT_BODY).expect("object request parses"); counter .count_request(&object) diff --git a/litellm-rust/crates/token-counter/benches/token_counter.rs b/litellm-rust/crates/token-counter/benches/token_counter.rs index a7c3177b88a..5e30cf6f77e 100644 --- a/litellm-rust/crates/token-counter/benches/token_counter.rs +++ b/litellm-rust/crates/token-counter/benches/token_counter.rs @@ -42,7 +42,7 @@ fn inputs(tokenizer: &Tokenizer) -> Vec<(&'static str, String)> { } fn token_counter(c: &mut Criterion) { - let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("token counter should load"); + let counter = TokenCounter::from_json_fast(TOKENIZER_JSON).expect("token counter should load"); let tokenizer = TOKENIZER_JSON .parse::() .expect("reference tokenizer should load"); diff --git a/litellm-rust/crates/token-counter/src/counter.rs b/litellm-rust/crates/token-counter/src/counter.rs index 7eedc449dd1..ce08e225be4 100644 --- a/litellm-rust/crates/token-counter/src/counter.rs +++ b/litellm-rust/crates/token-counter/src/counter.rs @@ -1,9 +1,7 @@ use serde::Serialize; use crate::Error; -use crate::byte_level::ByteLevelCounter; use crate::python_json; -use crate::scanner::{SplitPattern, TiktokenCounter}; use crate::tools::format_function_definitions; use crate::types::{ ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice, @@ -24,73 +22,22 @@ pub struct InputTokenCount { pub input_tokens: usize, } -enum Encoder { - HuggingFace { - tokenizer: Box, - byte_level: Option, - }, - Tiktoken(TiktokenCounter), -} - /// A loaded 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 { - encoder: Encoder, + encoder: Box, } impl TokenCounter { - /// Load a HuggingFace `tokenizer.json` document. The host reads the file. - pub fn from_json(tokenizer_json: &str) -> Result { - let tokenizer = tokenizer_json - .parse::() - .map_err(Error::Load)?; - let byte_level = ByteLevelCounter::detect(&tokenizer); - Ok(Self { - encoder: Encoder::HuggingFace { - tokenizer: Box::new(tokenizer), - byte_level, - }, - }) - } - - /// Load tiktoken's `cl100k_base` rank file (`base64(token) rank` lines). - /// The host reads the file. - pub fn from_cl100k_ranks(rank_file: &str) -> Result { - Self::from_tiktoken_ranks(SplitPattern::Cl100k, rank_file) - } - - /// Load tiktoken's `o200k_base` rank file (`base64(token) rank` lines). - /// The host reads the file. - pub fn from_o200k_ranks(rank_file: &str) -> Result { - Self::from_tiktoken_ranks(SplitPattern::O200k, rank_file) - } - - fn from_tiktoken_ranks(split: SplitPattern, rank_file: &str) -> Result { - Ok(Self { - encoder: Encoder::Tiktoken(TiktokenCounter::from_ranks(split, rank_file)?), - }) + pub fn new(tokenizer: impl crate::Tokenizer + 'static) -> Self { + Self { + encoder: Box::new(tokenizer), + } } pub fn count_text(&self, text: &str) -> Result { - match &self.encoder { - Encoder::Tiktoken(counter) => Ok(counter.count(text)), - Encoder::HuggingFace { - tokenizer, - byte_level, - } => { - if let Some(count) = byte_level - .as_ref() - .and_then(|counter| counter.count(tokenizer, text)) - { - return Ok(count); - } - tokenizer - .encode_fast(text, true) - .map(|encoding| encoding.len()) - .map_err(Error::Encode) - } - } + self.encoder.count_tokens(text) } /// Mirrors the host's key precedence: `messages`, then `prompt`, then diff --git a/litellm-rust/crates/token-counter/src/error.rs b/litellm-rust/crates/token-counter/src/error.rs index 6b8668fe182..b05ce007e46 100644 --- a/litellm-rust/crates/token-counter/src/error.rs +++ b/litellm-rust/crates/token-counter/src/error.rs @@ -4,8 +4,10 @@ use thiserror::Error as ThisError; #[derive(Debug, ThisError)] pub enum Error { + #[error("unsupported tokenizer: {0}")] + UnsupportedTokenizer(String), #[error("failed to load tokenizer: {0}")] - Load(#[source] tokenizers::Error), + Load(#[source] Box), #[error("failed to load tokenizer: tiktoken rank file: {0}")] Ranks(String), #[error("failed to load tokenizer: Unicode character classes are unavailable")] @@ -29,7 +31,7 @@ pub enum 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), + Encode(#[source] Box), #[error("token counting task failed: {0}")] Task(String), } diff --git a/litellm-rust/crates/token-counter/src/fast.rs b/litellm-rust/crates/token-counter/src/fast.rs new file mode 100644 index 00000000000..de3f86abd68 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/fast.rs @@ -0,0 +1,41 @@ +use litellm_token_counter_fast::Error as BackendError; +pub use litellm_token_counter_fast::FastTokenizer; + +use crate::{Error, TokenCounter, Tokenizer}; + +impl TokenCounter { + pub fn from_json_fast(tokenizer_json: &str) -> Result { + FastTokenizer::from_json(tokenizer_json) + .map(Self::new) + .map_err(Error::from) + } + + pub fn from_cl100k_ranks(rank_file: &str) -> Result { + FastTokenizer::from_cl100k_ranks(rank_file) + .map(Self::new) + .map_err(Error::from) + } + + pub fn from_o200k_ranks(rank_file: &str) -> Result { + FastTokenizer::from_o200k_ranks(rank_file) + .map(Self::new) + .map_err(Error::from) + } +} + +impl Tokenizer for FastTokenizer { + fn count_tokens(&self, text: &str) -> Result { + FastTokenizer::count_tokens(self, text).map_err(Error::from) + } +} + +impl From for Error { + fn from(error: BackendError) -> Self { + match error { + BackendError::Load(source) => Self::Load(source), + BackendError::Ranks(message) => Self::Ranks(message), + BackendError::UnicodeClasses => Self::UnicodeClasses, + BackendError::Encode(source) => Self::Encode(source), + } + } +} diff --git a/litellm-rust/crates/token-counter/src/huggingface.rs b/litellm-rust/crates/token-counter/src/huggingface.rs new file mode 100644 index 00000000000..fb7683b373e --- /dev/null +++ b/litellm-rust/crates/token-counter/src/huggingface.rs @@ -0,0 +1,27 @@ +use litellm_token_counter_huggingface::Error as BackendError; +pub use litellm_token_counter_huggingface::HuggingFaceTokenizer; + +use crate::{Error, TokenCounter, Tokenizer}; + +impl TokenCounter { + pub fn from_json(tokenizer_json: &str) -> Result { + HuggingFaceTokenizer::from_json(tokenizer_json) + .map(Self::new) + .map_err(Error::from) + } +} + +impl Tokenizer for HuggingFaceTokenizer { + fn count_tokens(&self, text: &str) -> Result { + HuggingFaceTokenizer::count_tokens(self, text).map_err(Error::from) + } +} + +impl From for Error { + fn from(error: BackendError) -> Self { + match error { + BackendError::Load(source) => Self::Load(source), + BackendError::Encode(source) => Self::Encode(source), + } + } +} diff --git a/litellm-rust/crates/token-counter/src/lib.rs b/litellm-rust/crates/token-counter/src/lib.rs index fa0014e2bad..446c91049de 100644 --- a/litellm-rust/crates/token-counter/src/lib.rs +++ b/litellm-rust/crates/token-counter/src/lib.rs @@ -4,18 +4,21 @@ #![forbid(unsafe_code)] -mod byte_level; -mod cl100k; mod counter; mod error; -mod o200k; mod python_json; -mod scanner; -mod tiktoken; +mod tokenizer; mod tools; mod types; -mod unicode_classes; + +#[cfg(feature = "fast")] +pub mod fast; +#[cfg(feature = "huggingface")] +pub mod huggingface; +#[cfg(feature = "tiktoken")] +pub mod tiktoken; pub use counter::{InputTokenCount, TokenCounter}; pub use error::Error; +pub use tokenizer::Tokenizer; pub use types::CountableRequest; diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index 7a9e71ed587..07c1c9f5b73 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -1,222 +1,24 @@ -//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and -//! the merge loop that turns one regex piece into tokens. The merge order is -//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is -//! identical, but pairs are tracked in a heap so a long piece costs -//! `O(n log n)` instead of tiktoken's `O(n^2)`. +pub use litellm_token_counter_tiktoken::TiktokenTokenizer; +use litellm_token_counter_tiktoken::UnsupportedTokenizer; -use std::cmp::Reverse; -use std::collections::BinaryHeap; +use crate::{Error, TokenCounter, Tokenizer}; -use base64::Engine; -use base64::engine::general_purpose::STANDARD; -use rustc_hash::FxHashMap; - -use crate::Error; - -type Rank = u32; - -const NO_RANK: Rank = Rank::MAX; -const END: usize = usize::MAX; - -pub(super) struct MergeRanks(FxHashMap, Rank>); - -impl MergeRanks { - pub(super) fn parse(text: &str) -> Result { - let ranks = text - .lines() - .filter(|line| !line.is_empty()) - .map(parse_line) - .collect::, _>>()?; - if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) { - return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token"))); - } - Ok(Self(ranks)) - } - - fn rank(&self, bytes: &[u8]) -> Rank { - self.0.get(bytes).copied().unwrap_or(NO_RANK) - } - - /// Token count of one regex piece, as `encode_ordinary` would produce. - pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize { - if piece.len() < 2 || self.0.contains_key(piece) { - return 1; - } - scratch.reset(piece.len()); - for start in 0..piece.len() - 1 { - scratch.set_rank(start, self.rank(&piece[start..start + 2])); - } - let mut parts = piece.len(); - while let Some(Reverse((rank, start))) = scratch.heap.pop() { - if scratch.next[start] == END || scratch.rank[start] != rank { - continue; - } - let merged = scratch.next[start]; - let after = scratch.next[merged]; - scratch.next[merged] = END; - scratch.next[start] = after; - parts -= 1; - if after < piece.len() { - scratch.prev[after] = start; - scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)])); - } else { - scratch.rank[start] = NO_RANK; - } - let before = scratch.prev[start]; - if before != END { - scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)])); - } - } - parts +impl TokenCounter { + pub fn from_tiktoken(encoding: &str) -> Result { + TiktokenTokenizer::from_name(encoding) + .map(Self::new) + .map_err(Error::from) } } -fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> { - let (token, rank) = line - .split_once(' ') - .ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?; - let bytes = STANDARD - .decode(token) - .map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?; - let rank = rank - .parse() - .map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?; - Ok((bytes.into_boxed_slice(), rank)) -} - -/// Buffers reused across the pieces of one text. Parts are addressed by the -/// byte offset they start at, which also gives the leftmost-pair tie break. -#[derive(Default)] -pub(super) struct MergeScratch { - next: Vec, - prev: Vec, - rank: Vec, - heap: BinaryHeap>, -} - -impl MergeScratch { - fn reset(&mut self, len: usize) { - self.next.clear(); - self.next.extend(1..=len); - self.prev.clear(); - self.prev.push(END); - self.prev.extend(0..len - 1); - self.rank.clear(); - self.rank.resize(len, NO_RANK); - self.heap.clear(); - } - - fn end(&self, start: usize) -> usize { - self.next[start] - } - - fn set_rank(&mut self, start: usize, rank: Rank) { - self.rank[start] = rank; - if rank != NO_RANK { - self.heap.push(Reverse((rank, start))); - } +impl Tokenizer for TiktokenTokenizer { + fn count_tokens(&self, text: &str) -> Result { + Ok(TiktokenTokenizer::count_tokens(self, text)) } } -#[cfg(test)] -mod tests { - use rand::rngs::StdRng; - use rand::{Rng, SeedableRng}; - - use super::*; - - fn ranks() -> MergeRanks { - let path = concat!( - env!("CARGO_MANIFEST_DIR"), - "/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4" - ); - MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo")) - .expect("rank file parses") - } - - /// tiktoken's `_byte_pair_merge`, transcribed, as the reference. - fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize { - if piece.len() < 2 || ranks.0.contains_key(piece) { - return 1; - } - let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1) - .map(|index| (index, ranks.rank(&piece[index..index + 2]))) - .chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)]) - .collect(); - let get_rank = |parts: &[(usize, Rank)], index: usize| { - if index + 3 < parts.len() { - ranks.rank(&piece[parts[index].0..parts[index + 3].0]) - } else { - NO_RANK - } - }; - loop { - let Some(index) = parts[..parts.len() - 1] - .iter() - .enumerate() - .filter(|(_, (_, rank))| *rank != NO_RANK) - .min_by_key(|(index, (_, rank))| (*rank, *index)) - .map(|(index, _)| index) - else { - return parts.len() - 1; - }; - if index > 0 { - parts[index - 1].1 = get_rank(&parts, index - 1); - } - parts[index].1 = get_rank(&parts, index); - parts.remove(index + 1); - } - } - - #[test] - fn every_byte_is_a_token() { - let ranks = ranks(); - assert_eq!(ranks.0.len(), 100_256); - assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK)); - } - - #[test] - fn heap_merge_matches_tiktokens_merge_loop() { - let ranks = ranks(); - let mut scratch = MergeScratch::default(); - let mut rng = StdRng::seed_from_u64(99); - let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123"; - for _ in 0..20_000 { - let piece: Vec = (0..rng.gen_range(1..24)) - .map(|_| alphabet[rng.gen_range(0..alphabet.len())]) - .collect(); - assert_eq!( - ranks.count_piece(&piece, &mut scratch), - reference_count(&ranks, &piece), - "piece {:?}", - String::from_utf8_lossy(&piece) - ); - } - } - - #[test] - fn long_repeated_runs_cost_close_to_linear() { - let ranks = ranks(); - let mut scratch = MergeScratch::default(); - let mut time = |len: usize| { - let piece = vec![b' '; len]; - let started = std::time::Instant::now(); - assert!(ranks.count_piece(&piece, &mut scratch) > 0); - started.elapsed() - }; - let small = (0..3).map(|_| time(1 << 14)).min().unwrap(); - let large = time(1 << 18); - assert!( - large < small * 64, - "{small:?} for 2^14 bytes, {large:?} for 2^18" - ); - } - - #[test] - fn malformed_rank_files_are_rejected() { - assert!(MergeRanks::parse("IQ==").is_err()); - assert!(MergeRanks::parse("IQ== x").is_err()); - assert!(MergeRanks::parse("!!! 1").is_err()); - assert!(MergeRanks::parse("IQ== 1").is_err()); +impl From for Error { + fn from(error: UnsupportedTokenizer) -> Self { + Self::UnsupportedTokenizer(error.0) } } diff --git a/litellm-rust/crates/token-counter/src/tokenizer.rs b/litellm-rust/crates/token-counter/src/tokenizer.rs new file mode 100644 index 00000000000..88c29c672a7 --- /dev/null +++ b/litellm-rust/crates/token-counter/src/tokenizer.rs @@ -0,0 +1,31 @@ +use crate::Error; + +pub trait Tokenizer: Send + Sync { + fn count_tokens(&self, text: &str) -> Result; +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{CountableRequest, TokenCounter}; + + struct Characters; + + impl Tokenizer for Characters { + fn count_tokens(&self, text: &str) -> Result { + Ok(text.chars().count()) + } + } + + #[test] + fn request_accounting_works_with_an_injected_backend() { + let counter = TokenCounter::new(Characters); + let request = + CountableRequest::parse(br#"{"messages":[{"role":"user","content":"hello"}]}"#) + .unwrap(); + assert_eq!( + counter.count_request(&request).unwrap().input_tokens, + 3 + 4 + 5 + 3 + ); + } +} diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index 12c59768952..378a16d5d72 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -1,5 +1,6 @@ +#![cfg(feature = "huggingface")] + use rstest::rstest; -use serde::Deserialize; use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter}; @@ -176,146 +177,155 @@ fn loading_a_bad_tokenizer_is_a_load_error() { assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_)))); } -/// A tiktoken encoding: its fixture directory, the vendored rank file Python -/// loads, the constructor, and the model `generate.py` counted the requests for. -#[derive(Clone, Copy)] -struct TiktokenEncoding { - fixtures: &'static str, - rank_file: &'static str, - load: fn(&str) -> Result, - model: &'static str, -} +#[cfg(feature = "fast")] +mod fast { + use super::*; + use serde::Deserialize; -const CL100K: TiktokenEncoding = TiktokenEncoding { - fixtures: "cl100k", - rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4", - load: TokenCounter::from_cl100k_ranks, - model: "gpt-4", -}; + /// A tiktoken encoding: its fixture directory, the vendored rank file Python + /// loads, the constructor, and the model `generate.py` counted the requests for. + #[derive(Clone, Copy)] + struct TiktokenEncoding { + fixtures: &'static str, + rank_file: &'static str, + load: fn(&str) -> Result, + model: &'static str, + } -const O200K: TiktokenEncoding = TiktokenEncoding { - fixtures: "o200k", - rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790", - load: TokenCounter::from_o200k_ranks, - model: "gpt-4o", -}; - -fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter { - let path = format!( - "{}/../../../litellm/litellm_core_utils/tokenizers/{}", - env!("CARGO_MANIFEST_DIR"), - encoding.rank_file - ); - let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo"); - (encoding.load)(&ranks).expect("ranks load") -} - -fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String { - let path = format!( - "{}/tests/fixtures/{}/{name}", - env!("CARGO_MANIFEST_DIR"), - encoding.fixtures - ); - std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py") -} - -#[derive(Deserialize)] -struct TextFixture { - text: String, - tokens: usize, -} - -#[derive(Deserialize)] -struct RequestFixture { - body: String, - input_tokens: usize, -} - -/// Reference counts come from `tiktoken.get_encoding(name)`; see -/// `tests/fixtures/generate.py`. -#[rstest] -#[case::cl100k(CL100K)] -#[case::o200k(O200K)] -fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) { - let counter = tiktoken_counter(encoding); - let fixtures: Vec = tiktoken_fixture(encoding, "texts.jsonl") - .lines() - .map(|line| serde_json::from_str(line).expect("fixture line is json")) - .collect(); - assert!(fixtures.len() > 3000); - let mismatches: Vec<_> = fixtures - .iter() - .filter_map(|fixture| { - let count = counter.count_text(&fixture.text).expect("text counts"); - (count != fixture.tokens).then(|| (fixture.text.clone(), fixture.tokens, count)) - }) - .collect(); - assert!( - mismatches.is_empty(), - "(text, tiktoken, rust): {mismatches:?}" - ); -} - -/// Reference counts come from the proxy's admission counter -/// (`_count_input_tokens(body, model)`), so this pins the shared message, -/// tool and reply-priming accounting on the tiktoken paths as well. -#[rstest] -#[case::cl100k(CL100K)] -#[case::o200k(O200K)] -fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) { - let counter = tiktoken_counter(encoding); - let fixtures: Vec = tiktoken_fixture(encoding, "requests.jsonl") - .lines() - .map(|line| serde_json::from_str(line).expect("fixture line is json")) - .collect(); - let counts: Vec = fixtures - .iter() - .map(|fixture| { - let request = CountableRequest::parse(fixture.body.as_bytes()).expect("fixture parses"); - let count = counter.count_request(&request).expect("fixture counts"); - assert_eq!(count.model.as_deref(), Some(encoding.model)); - assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body); - count.input_tokens - }) - .collect(); - assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000)); -} - -#[rstest] -#[case::cl100k(CL100K)] -#[case::o200k(O200K)] -fn tiktoken_shares_the_message_accounting_with_the_anthropic_path( - #[case] encoding: TiktokenEncoding, -) { - let counter = tiktoken_counter(encoding); - let count = |body: &str| { - counter - .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) - .expect("counts") - .input_tokens + const CL100K: TiktokenEncoding = TiktokenEncoding { + fixtures: "cl100k", + rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4", + load: TokenCounter::from_cl100k_ranks, + model: "gpt-4", }; - let text = |text: &str| counter.count_text(text).expect("counts"); - let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); - assert_eq!(base, 3 + text("user") + text("hi") + 3); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","name":"al","content":"hi"}]}"#), - base + text("al") + 1 - ); - assert_eq!( - count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#), - base + 1 - ); -} -#[rstest] -#[case::empty("")] -#[case::not_base64("!!!! 0")] -#[case::missing_rank("YQ==")] -#[case::rank_not_a_number("YQ== x")] -#[case::single_byte_tokens_missing("YWI= 0")] -fn loading_a_bad_rank_file_is_a_load_error( - #[case] rank_file: &str, - #[values(CL100K, O200K)] encoding: TiktokenEncoding, -) { - assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); + const O200K: TiktokenEncoding = TiktokenEncoding { + fixtures: "o200k", + rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790", + load: TokenCounter::from_o200k_ranks, + model: "gpt-4o", + }; + + fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter { + let path = format!( + "{}/../../../litellm/litellm_core_utils/tokenizers/{}", + env!("CARGO_MANIFEST_DIR"), + encoding.rank_file + ); + let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo"); + (encoding.load)(&ranks).expect("ranks load") + } + + fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String { + let path = format!( + "{}/../token-counter-fast/tests/fixtures/{}/{name}", + env!("CARGO_MANIFEST_DIR"), + encoding.fixtures + ); + std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py") + } + + #[derive(Deserialize)] + struct TextFixture { + text: String, + tokens: usize, + } + + #[derive(Deserialize)] + struct RequestFixture { + body: String, + input_tokens: usize, + } + + /// Reference counts come from `tiktoken.get_encoding(name)`; see + /// `tests/fixtures/generate.py`. + #[rstest] + #[case::cl100k(CL100K)] + #[case::o200k(O200K)] + fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) { + let counter = tiktoken_counter(encoding); + let fixtures: Vec = tiktoken_fixture(encoding, "texts.jsonl") + .lines() + .map(|line| serde_json::from_str(line).expect("fixture line is json")) + .collect(); + assert!(fixtures.len() > 3000); + let mismatches: Vec<_> = fixtures + .iter() + .filter_map(|fixture| { + let count = counter.count_text(&fixture.text).expect("text counts"); + (count != fixture.tokens).then(|| (fixture.text.clone(), fixture.tokens, count)) + }) + .collect(); + assert!( + mismatches.is_empty(), + "(text, tiktoken, rust): {mismatches:?}" + ); + } + + /// Reference counts come from the proxy's admission counter + /// (`_count_input_tokens(body, model)`), so this pins the shared message, + /// tool and reply-priming accounting on the tiktoken paths as well. + #[rstest] + #[case::cl100k(CL100K)] + #[case::o200k(O200K)] + fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) { + let counter = tiktoken_counter(encoding); + let fixtures: Vec = tiktoken_fixture(encoding, "requests.jsonl") + .lines() + .map(|line| serde_json::from_str(line).expect("fixture line is json")) + .collect(); + let counts: Vec = fixtures + .iter() + .map(|fixture| { + let request = + CountableRequest::parse(fixture.body.as_bytes()).expect("fixture parses"); + let count = counter.count_request(&request).expect("fixture counts"); + assert_eq!(count.model.as_deref(), Some(encoding.model)); + assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body); + count.input_tokens + }) + .collect(); + assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000)); + } + + #[rstest] + #[case::cl100k(CL100K)] + #[case::o200k(O200K)] + fn tiktoken_shares_the_message_accounting_with_the_anthropic_path( + #[case] encoding: TiktokenEncoding, + ) { + let counter = tiktoken_counter(encoding); + let count = |body: &str| { + counter + .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) + .expect("counts") + .input_tokens + }; + let text = |text: &str| counter.count_text(text).expect("counts"); + let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#); + assert_eq!(base, 3 + text("user") + text("hi") + 3); + assert_eq!( + count(r#"{"model":"m","messages":[{"role":"user","name":"al","content":"hi"}]}"#), + base + text("al") + 1 + ); + assert_eq!( + count( + r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"# + ), + base + 1 + ); + } + + #[rstest] + #[case::empty("")] + #[case::not_base64("!!!! 0")] + #[case::missing_rank("YQ==")] + #[case::rank_not_a_number("YQ== x")] + #[case::single_byte_tokens_missing("YWI= 0")] + fn loading_a_bad_rank_file_is_a_load_error( + #[case] rank_file: &str, + #[values(CL100K, O200K)] encoding: TiktokenEncoding, + ) { + assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_)))); + } } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index c0a06364261..7d20fc5a02f 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -97,9 +97,9 @@ class ResponsesWebSocketConnection: class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... @staticmethod - def from_cl100k_ranks(rank_file: str) -> TokenCounter: ... + def from_json_fast(tokenizer_json: str) -> TokenCounter: ... @staticmethod - def from_o200k_ranks(rank_file: str) -> TokenCounter: ... + def from_tiktoken(encoding: str) -> TokenCounter: ... def acount_request(self, body: bytes) -> Future[dict[str, object]]: ... def gil_stats() -> dict[str, int]: ... diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index d36234f56c1..e816f1ab388 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -11,14 +11,21 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting 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, huggingface_tokenizer_kind -RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"] +RustTokenizer = Literal[ + "anthropic", + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", +] class RustTokenCounter(Protocol): @@ -30,10 +37,7 @@ class RustTokenCounterFactory(Protocol): def __call__(self, tokenizer_json: str) -> RustTokenCounter: raise NotImplementedError - def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter: - raise NotImplementedError - - def from_o200k_ranks(self, rank_file: str) -> RustTokenCounter: + def from_tiktoken(self, encoding: RustTokenizer) -> RustTokenCounter: raise NotImplementedError @@ -63,9 +67,8 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: """The Rust counter for the tokenizer `litellm.token_counter` selects for `model`, `None` when Python must count. Mirrors `_select_tokenizer_helper`: the Anthropic tokenizer has a Rust port, the other HuggingFace - downloads do not, and of the tiktoken encodings `cl100k_base` and `o200k_base` do (p50k/r50k do not). Rust - prices every message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in - Python.""" + downloads do not. All tiktoken encodings used by Python are backed by tiktoken-rs. Rust prices every + message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in Python.""" if litellm.disable_token_counter is True: return None kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model) @@ -73,24 +76,19 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: return "anthropic" if kind is not None or uses_legacy_message_accounting(model): return None - match openai_tokenizer_encoding(model).name: - case "cl100k_base": - return "cl100k_base" - case "o200k_base": - return "o200k_base" - case _: - return None + encoding: Final = openai_tokenizer_encoding(model).name + if encoding in {"cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base"}: + return cast(RustTokenizer, encoding) + return None -@lru_cache(maxsize=4) +@lru_cache(maxsize=8) def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter: match tokenizer: case "anthropic": return factory(claude_json_str) - case "cl100k_base": - return factory.from_cl100k_ranks(cl100k_base_rank_file()) - case "o200k_base": - return factory.from_o200k_ranks(o200k_base_rank_file()) + case _: + return factory.from_tiktoken(tokenizer) async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None: diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 0e0025f5194..c785a7b657a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -256,7 +256,7 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`.""" + """Stands in for the native `TokenCounter` class.""" def __init__(self) -> None: self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = [] @@ -264,11 +264,8 @@ class _RecordingFactory: def __call__(self, tokenizer_json: str) -> _RecordingCounter: return _RecordingCounter(self, "anthropic") - def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: - return _RecordingCounter(self, "cl100k_base") - - def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: - return _RecordingCounter(self, "o200k_base") + def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _RecordingCounter: + return _RecordingCounter(self, encoding) class _DecliningCounter: @@ -280,10 +277,7 @@ class _DecliningFactory: def __call__(self, tokenizer_json: str) -> _DecliningCounter: return _DecliningCounter() - def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter: - return _DecliningCounter() - - def from_o200k_ranks(self, rank_file: str) -> _DecliningCounter: + def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _DecliningCounter: return _DecliningCounter() diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py index 71aa79cc4bb..b95e6a62e9c 100644 --- a/tests/test_litellm/rust_bridge/test_token_counter.py +++ b/tests/test_litellm/rust_bridge/test_token_counter.py @@ -8,7 +8,6 @@ cases need the extension and are skipped when it is not built. from __future__ import annotations import json -from types import MappingProxyType from typing import Final import pytest @@ -27,7 +26,6 @@ MODEL: Final = "claude-sonnet-4-5-20250929" CL100K_MODEL: Final = "gpt-4" O200K_MODEL: Final = "gpt-4o" TOKENIZERS: Final[tuple[bridge.RustTokenizer, ...]] = ("anthropic", "cl100k_base", "o200k_base") -RANK_FILE_LINES: Final = MappingProxyType({"cl100k_base": 100_256, "o200k_base": 199_998}) BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode() @@ -55,24 +53,20 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files.""" + """Stands in for the native `TokenCounter` class.""" def __init__(self) -> None: self.counters: list[_RecordingCounter] = [] - self.rank_files: list[str] = [] + self.encodings: list[str] = [] def __call__(self, tokenizer_json: str) -> _RecordingCounter: counter = _RecordingCounter(tokenizer_json) self.counters.append(counter) return counter - def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: - self.rank_files.append(rank_file) - return self("cl100k_base") - - def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: - self.rank_files.append(rank_file) - return self("o200k_base") + def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RecordingCounter: + self.encodings.append(encoding) + return self(encoding) class _RaisingCounter: @@ -92,10 +86,7 @@ class _RaisingFactory: def __call__(self, tokenizer_json: str) -> _RaisingCounter: return _RaisingCounter(self.error) - def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter: - return _RaisingCounter(self.error) - - def from_o200k_ranks(self, rank_file: str) -> _RaisingCounter: + def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RaisingCounter: return _RaisingCounter(self.error) @@ -140,7 +131,7 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No @pytest.mark.asyncio @pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base")) -async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokenizer: bridge.RustTokenizer) -> None: +async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridge.RustTokenizer) -> None: factory: Final = _RecordingFactory() litellm.rust(True) bridge.TOKEN_COUNTER.override(factory) @@ -149,9 +140,7 @@ async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokeni second: Final = await bridge.count_input_tokens(BODY, tokenizer) assert first == second == bridge.InputTokenCount(model=MODEL, input_tokens=42) - assert len(factory.rank_files) == 1 - assert factory.rank_files[0].startswith("IQ== 0\n") - assert factory.rank_files[0].count("\n") == RANK_FILE_LINES[tokenizer] + assert factory.encodings == [tokenizer] assert factory.counters[0].tokenizer_json == tokenizer assert factory.counters[0].bodies == [BODY, BODY] @@ -235,13 +224,13 @@ def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: ("model", "python_encoding"), (("text-davinci-003", "p50k_base"), ("gpt-oss-120b", "o200k_harmony")), ) -def test_rust_tokenizer_declines_tiktoken_encodings_rust_does_not_have( - monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: str +def test_rust_tokenizer_uses_every_tiktoken_encoding_supported_by_rust( + monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: bridge.RustTokenizer ) -> None: monkeypatch.setattr(litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {model}) assert openai_tokenizer_encoding(model).name == python_encoding - assert bridge.rust_tokenizer(model) is None + assert bridge.rust_tokenizer(model) == python_encoding def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None: