scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

feat(rust): count tiktoken cl100k_base admission tokens in Rust

Browse source 
The Rust admission token counter only had the Anthropic tokenizer, so every
other model (OpenAI gpt-4 family, Azure, Gemini, Bedrock non-Claude, Mistral)
tokenized with tiktoken on the Python inference worker.

Add an exact cl100k_base counter to litellm-token-counter: the vendored rank
file (base64 token / rank lines, the bytes Python's tiktoken uses) is parsed
into a byte-level BPE model and the cl100k split pattern is a handwritten
scanner over the shared Unicode classes, so no regex engine runs per request.
Both tokenizers share the message, tool and reply-priming accounting.

The PyO3 TokenCounter gains a from_cl100k_ranks constructor; Python reads the
rank file and passes it in, the way claude_json_str already works. The bridge
selects the counter through the same predicates litellm.token_counter uses
(huggingface_tokenizer_kind, openai_tokenizer_encoding), declines o200k_base,
downloaded HuggingFace and custom tokenizers to Python, and budget reservation
counts once per distinct tokenizer a request names.

The legacy gpt-3.5-turbo-0301 message accounting (4 per message, -1 per name)
stays in Python: the selector declines it through the predicate token_counter
itself uses.
This commit is contained in:
yassin 2026-09-11 20:11:29 +00:00
parent 3f81ba3d30
commit b1373d2456
21 changed files with 4500 additions and 195 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-rust/Cargo.lock generated
View file

@ -1656,11 +1656,13 @@ dependencies = [
name = "litellm-token-counter"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"criterion",
"indexmap",
"itoa",
"rand 0.8.7",
"rstest",
"rustc-hash",
"serde",
"serde_json",
"thiserror 2.0.19",

27
litellm-rust/crates/python-bridge/src/token_counter.rs
View file

@ -30,12 +30,12 @@ struct TokenCounter {
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())),
})
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
}
#[staticmethod]
fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
}
fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
@ -58,6 +58,19 @@ impl TokenCounter {
}
}
impl TokenCounter {
fn load(
py: Python<'_>,
load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
) -> PyResult<Self> {
let inner = release_gil(py, load).map_err(token_count_error_to_pyerr)?;
Ok(Self {
inner: Arc::new(inner),
encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
})
}
}
fn encode_parallelism() -> usize {
available_parallelism().map_or(TOKEN_COUNT_FALLBACK_PARALLELISM, NonZero::get)
}
@ -70,7 +83,7 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount
fn token_count_error_to_pyerr(error: Error) -> PyErr {
let message = error.to_string();
match error {
Error::Load(_) => PyValueError::new_err(message),
Error::Load(_) | Error::Ranks(_) | Error::UnicodeClasses => PyValueError::new_err(message),
Error::RequestParse(_)
| Error::MissingInput
| Error::FloatText

2
litellm-rust/crates/token-counter/Cargo.toml
View file

@ -6,8 +6,10 @@ license.workspace = true
repository.workspace = true
[dependencies]
base64.workspace = true
indexmap = { version = "2.14.0", features = ["serde"] }
itoa = "1.0"
rustc-hash = "2.1.3"
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true

29
litellm-rust/crates/token-counter/src/byte_level.rs
View file

@ -12,7 +12,7 @@ use tokenizers::pre_tokenizers::PreTokenizerWrapper;
use tokenizers::{Model, Tokenizer};
use unicode_normalization_alignments::{IsNormalized, UnicodeNormalization, is_nfkc_quick};
use super::unicode_classes::UnicodeClasses;
use super::unicode_classes::{Class, UnicodeClasses, class, run_len};
const CONTRACTIONS: [&str; 7] = ["'s", "'t", "'re", "'ve", "'m", "'ll", "'d"];
@ -110,27 +110,6 @@ fn mapped_len(piece: &str) -> usize {
.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,
@ -169,12 +148,6 @@ fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize
}
}
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 {

196
litellm-rust/crates/token-counter/src/cl100k.rs Normal file
View file

@ -0,0 +1,196 @@
//! Exact token counting for tiktoken's `cl100k_base`. A scanner reproduces the
//! piece boundaries of the encoding's split regex,
//! `'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}++|\p{N}{1,3}+| ?[^\s\p{L}\p{N}]++[\r\n]*+|\s++$|\s*[\r\n]|\s+(?!\S)|\s`,
//! and each piece is merged with the rank file. Special tokens are ordinary
//! text, as with `encode(text, disallowed_special=())`.
use std::iter;
use super::tiktoken::{MergeRanks, MergeScratch};
use super::unicode_classes::{Class, UnicodeClasses, class, run_len};
use crate::Error;
const MAX_DIGITS_PER_PIECE: usize = 3;
pub(super) struct Cl100kCounter {
ranks: MergeRanks,
unicode_classes: &'static UnicodeClasses,
}
impl Cl100kCounter {
pub(super) fn from_ranks(rank_file: &str) -> Result<Self, Error> {
Ok(Self {
ranks: MergeRanks::parse(rank_file)?,
unicode_classes: UnicodeClasses::get().ok_or(Error::UnicodeClasses)?,
})
}
pub(super) fn count(&self, text: &str) -> usize {
let mut scratch = MergeScratch::default();
pieces(text, self.unicode_classes)
.map(|piece| self.ranks.count_piece(piece.as_bytes(), &mut scratch))
.sum()
}
}
/// 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)))
}
/// The alternatives in regex order; the possessive quantifiers mean an
/// alternative that starts matching and runs out of input fails as a whole.
fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize {
if first == '\''
&& let Some(len) = contraction_len(&text[1..])
{
return 1 + len;
}
let first_class = class(first, unicode_classes);
match first_class {
Class::Letter => return run_len(text, Class::Letter, unicode_classes),
Class::Number => return digit_run_len(text, unicode_classes),
Class::Space | Class::Other => {}
}
let rest = &text[first.len_utf8()..];
let second_class = rest
.chars()
.next()
.map(|character| class(character, unicode_classes));
if !is_newline(first) && second_class == Some(Class::Letter) {
return first.len_utf8() + run_len(rest, Class::Letter, unicode_classes);
}
if first_class == Class::Other {
return symbol_run_len(text, unicode_classes);
}
if first == ' ' && second_class == Some(Class::Other) {
return 1 + symbol_run_len(rest, unicode_classes);
}
space_run_len(text, unicode_classes)
}
/// `(?i:[sdmt]|ll|ve|re)` after the apostrophe. Simple case folding also maps
/// U+017F (long s) onto `s`.
fn contraction_len(rest: &str) -> Option<usize> {
let mut characters = rest.chars();
let first = characters.next()?;
match first {
's' | 'S' | '\u{17F}' | 'd' | 'D' | 'm' | 'M' | 't' | 'T' => Some(first.len_utf8()),
'l' | 'L' => matches!(characters.next(), Some('l' | 'L')).then_some(2),
'v' | 'V' | 'r' | 'R' => matches!(characters.next(), Some('e' | 'E')).then_some(2),
_ => None,
}
}
fn is_newline(character: char) -> bool {
matches!(character, '\r' | '\n')
}
/// `\p{N}{1,3}+`
fn digit_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
text.chars()
.take(MAX_DIGITS_PER_PIECE)
.take_while(|character| class(*character, unicode_classes) == Class::Number)
.map(char::len_utf8)
.sum()
}
/// `[^\s\p{L}\p{N}]++[\r\n]*+`
fn symbol_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
let symbols = run_len(text, Class::Other, unicode_classes);
symbols
+ text[symbols..]
.bytes()
.take_while(|byte| matches!(byte, b'\r' | b'\n'))
.count()
}
/// `\s++$|\s*[\r\n]|\s+(?!\S)|\s`: whitespace to the end of the text is one
/// piece; otherwise the piece ends at the last newline of the run, or leaves
/// the run's last character for the next piece's optional leading space.
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;
}
if let Some(newline) = text[..run].rfind(['\r', '\n']) {
return newline + 1;
}
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 rstest::rstest;
use super::*;
fn split(text: &str) -> Vec<&str> {
pieces(
text,
UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"),
)
.collect()
}
#[rstest]
#[case("", &[])]
#[case("Hello world", &["Hello", " world"])]
#[case("don't I'LL you'Ve we'RE he'd I'm", &["don", "'t", " I", "'LL", " you", "'Ve", " we", "'RE", " he", "'d", " I", "'m"])]
#[case("IT'SOK it'Dbe 'Sx 'Tx", &["IT", "'S", "OK", " it", "'D", "be", " '", "Sx", " '", "Tx"])]
#[case("'Sx'Tx'Mx'LLx'VEx'REx'Dx", &["'S", "x", "'T", "x", "'M", "x", "'LL", "x", "'VE", "x", "'RE", "x", "'D", "x"])]
#[case("'ſ 'lx", &["'ſ", " '", "lx"])]
#[case("12345 6", &["123", "45", " ", "6"])]
#[case("!abc !!abc", &["!abc", " !!", "abc"])]
#[case(" !!!\r\n\r\nx", &[" !!!\r\n\r\n", "x"])]
#[case("a b \n\n c", &["a", " ", " b", " \n\n", " ", " c"])]
#[case("a\nb\r\nc\n\nd \n e", &["a", "\n", "b", "\r\n", "c", "\n\n", "d", " \n", " e"])]
#[case("x \t\n \t y\n", &["x", " \t\n", " \t", " y", "\n"])]
#[case("end ", &["end", " "])]
#[case("\u{a0}abc\u{a0}!", &["\u{a0}abc", "\u{a0}", "!"])]
#[case("<|endoftext|>", &["<|", "endoftext", "|>"])]
#[case("e\u{301}a", &["e", "\u{301}a"])]
#[case("日本語 ١٢٣٤", &["日本語", " ", "١٢٣", "٤"])]
fn scanner_splits_like_the_regex(#[case] text: &str, #[case] expected: &[&str]) {
assert_eq!(split(text), expected);
}
#[derive(serde::Deserialize)]
struct TextFixture {
text: String,
pieces: Vec<String>,
}
#[test]
fn scanner_splits_the_fixture_corpus_like_tiktoken_regex() {
let fixtures: Vec<TextFixture> = include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/cl100k/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(|fixture| split(&fixture.text) != fixture.pieces)
.map(|fixture| (&fixture.text, split(&fixture.text), &fixture.pieces))
.collect();
assert!(mismatches.is_empty(), "{mismatches:#?}");
}
}

59
litellm-rust/crates/token-counter/src/counter.rs
View file

@ -2,6 +2,7 @@ use serde::Serialize;
use crate::Error;
use crate::byte_level::ByteLevelCounter;
use crate::cl100k::Cl100kCounter;
use crate::python_json;
use crate::tools::format_function_definitions;
use crate::types::{
@ -23,12 +24,19 @@ pub struct InputTokenCount {
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.
enum Encoder {
HuggingFace {
tokenizer: Box<tokenizers::Tokenizer>,
byte_level: Option<ByteLevelCounter>,
},
Cl100k(Cl100kCounter),
}
/// 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 {
tokenizer: tokenizers::Tokenizer,
byte_level: Option<ByteLevelCounter>,
encoder: Encoder,
}
impl TokenCounter {
@ -39,23 +47,40 @@ impl TokenCounter {
.map_err(Error::Load)?;
let byte_level = ByteLevelCounter::detect(&tokenizer);
Ok(Self {
tokenizer,
byte_level,
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, Error> {
Ok(Self {
encoder: Encoder::Cl100k(Cl100kCounter::from_ranks(rank_file)?),
})
}
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);
match &self.encoder {
Encoder::Cl100k(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.tokenizer
.encode_fast(text, true)
.map(|encoding| encoding.len())
.map_err(Error::Encode)
}
/// Mirrors the host's key precedence: `messages`, then `prompt`, then

4
litellm-rust/crates/token-counter/src/error.rs
View file

@ -6,6 +6,10 @@ use thiserror::Error as 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("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")]

2
litellm-rust/crates/token-counter/src/lib.rs
View file

@ -5,9 +5,11 @@
#![forbid(unsafe_code)]
mod byte_level;
mod cl100k;
mod counter;
mod error;
mod python_json;
mod tiktoken;
mod tools;
mod types;
mod unicode_classes;

215
litellm-rust/crates/token-counter/src/tiktoken.rs Normal file
View file

@ -0,0 +1,215 @@
//! 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<Box<[u8]>, Rank>);
impl MergeRanks {
pub(super) fn parse(text: &str) -> Result<Self, Error> {
let ranks = text
.lines()
.filter(|line| !line.is_empty())
.map(parse_line)
.collect::<Result<FxHashMap<_, _>, _>>()?;
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<usize>,
prev: Vec<usize>,
rank: Vec<Rank>,
heap: BinaryHeap<Reverse<(Rank, usize)>>,
}
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<u8> = (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_stay_cheap() {
let ranks = ranks();
let mut scratch = MergeScratch::default();
let piece = vec![b' '; 1 << 20];
let started = std::time::Instant::now();
let count = ranks.count_piece(&piece, &mut scratch);
assert!(count > 0);
assert!(started.elapsed().as_secs() < 5, "{:?}", started.elapsed());
}
#[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());
}
}

36
litellm-rust/crates/token-counter/src/unicode_classes.rs
View file

@ -59,15 +59,45 @@ impl UnicodeClasses {
CLASSES.as_ref()
}
pub(super) fn is_letter(&self, character: char) -> bool {
fn is_letter(&self, character: char) -> bool {
self.letters.contains(character)
}
pub(super) fn is_number(&self, character: char) -> bool {
fn is_number(&self, character: char) -> bool {
self.numbers.contains(character)
}
pub(super) fn is_space(&self, character: char) -> bool {
fn is_space(&self, character: char) -> bool {
self.spaces.contains(character)
}
}
/// `\p{L}`, `\p{N}`, `\s` and everything else, the character classes the
/// split regexes are written in.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Class {
Letter,
Number,
Space,
Other,
}
pub(super) 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,
}
}
/// Byte length of the leading run of `run_class` characters.
pub(super) 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)
}

380
litellm-rust/crates/token-counter/tests/fixtures/cl100k/generate.py vendored Normal file
View file

@ -0,0 +1,380 @@
"""Pin tiktoken cl100k_base reference counts for the Rust parity tests.
Run from the repository root with the project environment:
uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/cl100k/generate.py
`texts.jsonl` holds `{"text", "tokens", "pieces"}` lines: `tokens` counted with
`tiktoken.get_encoding("cl100k_base").encode(text, disallowed_special=())`, the
same call `litellm.token_counter` makes, and `pieces` the installed encoding's
split pattern applied with the `regex` module tiktoken itself uses, so a scanner
that splits differently fails even where BPE would count the same. `requests.jsonl` holds
`{"body", "input_tokens"}` lines, `body` being the exact request bytes as a JSON
string, counted with the proxy's admission counter
(`_count_input_tokens(body, "gpt-4")`). Every message in the 50k-token body is
shorter than the Python chunk size so the chunked Python count equals the exact
whole-text tiktoken count the Rust counter produces.
"""
import itertools
import json
import random
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import regex
import tiktoken
from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens
HERE: Final = Path(__file__).resolve().parent
ENCODING: Final = tiktoken.get_encoding("cl100k_base")
SPLIT_PATTERN: Final = regex.compile(ENCODING._pat_str) # pyright: ignore[reportPrivateUsage] # tiktoken has no public accessor
# Mirrors ALPHABET in src/byte_level.rs, plus the pieces the cl100k pattern treats differently.
ALPHABET: Final = (
"a",
"Z",
"e",
"s",
"t",
"d",
"m",
"'",
"'s",
"'re",
"'ll",
"'S",
"0",
"9",
" ",
" ",
"\t",
"\n",
"\r\n",
"\x0b",
".",
",",
"!",
"-",
"(",
'"',
"\xa0",
"\x85",
"\u2028",
"\u3000",
"\u200b",
"\u200d",
"é",
"e\u0301",
"ß",
"漢",
"字",
"ع",
"३",
"½",
"Ⅳ",
"🙂",
"👍🏽",
"Ａ",
"fi",
"㍿",
"㋿",
"ꟲ",
"𐞁",
"a\u030a",
"\u1e0b\u0323",
"<",
">",
"EOT",
"<EOT>",
"<META_START>",
"'D",
"'M",
"'T",
"'VE",
"'Re",
"'ſ",
"ſ",
"12345678",
"٣٤٥٦",
"<|endoftext|>",
"<|fim_prefix|>",
"\r",
"\r\n\r\n",
" \n",
"!!",
"#$%",
"\u00ad",
"\u0301",
"\U0001f600\U0001f3fd",
"İ",
"Dž",
)
CORPUS: Final = (
"",
"Hello, how are you today?",
"I'm sure they're right, we'll see. WE'LL SEE, I'M SURE THEY'RE RIGHT, IT'S HERS AND IT'D BE 'D",
"don't Don'T DON'T won'T i've I'VE i'Ve you'RE 'S 'T 'M 'D 'LL 'VE 'RE 'ſ 'x",
"1234567890 123 12 1 0000000 ٣٤٥٦٧٨ ३४५६ 1,234,567.89 2026-09-11T18:00:00Z",
"$abc %def &ghi @jkl _mno #pqr ~stu ^vwx |yz \\a /b :c ;d ?e !f (g )h [i ]j {k }l <m >n =o +p *q",
"foo bar baz \t qux\t\tquux \n\nline\r\nline\r\n\r\n \n\t\r\n x ",
"trailing spaces ",
"trailing tabs\t\t",
"trailing newline\n",
"\n\n\n",
"\r\n\r\n\r\n",
" ",
"😀😃😄 👍🏽 🇺🇸 👨‍👩‍👧‍👦 ✈️ ❤️‍🔥 ٭ ※ ⌘ ⏎",
"漢字かな交じり文、東京都千代田区。日本語のテキストです。中文测试。한국어 텍스트",
"مرحبا بالعالم، هذا نص عربي مع أرقام ١٢٣٤٥٦٧ و علامات ترقيم!",
"Zürich, façade, naïve, Ærøskøbing, Ελληνικά, Русский текст, עברית, हिन्दी, ไทย",
"e\u0301 a\u030a \u1e0b\u0323 \u0301\u0301 combining\u0308 marks\u0301!",
"ΣΊΣΥΦΟΣ Džungla İstanbul file flow Ａｂｃ ㍿ ㋿ ꟲ 𐞁",
"<|endoftext|> <|fim_prefix|>code<|fim_middle|>more<|fim_suffix|> <|endofprompt|> <|im_start|>",
"<EOT> <META_START> <s> </s> [INST] [/INST] <<SYS>>",
"def f(x):\n return {'a': x ** 2, \"b\": [1, 2, 3]} # comment\n\nprint(f(10))\n",
'{"model":"gpt-4","messages":[{"role":"user","content":"hi\\n"}],"temperature":0.7}',
"https://example.com/path?query=1&other=two#fragment user@example.com 192.168.0.1",
"a" * 3000,
" " * 3000,
"." * 3000,
"ab" * 1500,
"\n" * 3000,
"0" * 3000,
"!" * 3000,
"😀" * 1000,
"漢" * 1000,
"\u00a0abc\u00a0! \u2028x \u3000y \u200bz \u200d\u200d q",
"x\u0085y \x0b\x0c z",
"\x00\x01\x02 \x7f \ufffd",
"tab\tseparated\tvalues\n1\t2\t3\n",
"MiXeD cAsE wOrDs AND ACRONYMS like NASA, HTTP/2, gRPC, iOS, macOS",
"snake_case_identifier camelCaseIdentifier PascalCaseIdentifier SCREAMING_SNAKE_CASE kebab-case",
"x'sy x'ty x'rey x'vey x'my x'lly x'dy x'S x'T x'RE x'VE x'M x'LL x'D x'sS x'llL",
"IT'SOK it'Dbe x'Sy x'Ty x'My x'Dy x'LLy x'VEy x'REy x'Ly x'Vy x'Ry 'Sx'Tx'Mx'LLx'VEx'REx'Dx",
"'s't're've'm'll'd 'S'T'RE'VE'M'LL'D ''s '''s",
"9'9 9's a'9 '9 ' 's' ' 's",
"١٢٣٤ ½⅓¼ ⅣⅤ 𝟘𝟙𝟚𝟛𝟜𝟝𝟞𝟟𝟠𝟡 ①②③",
)
WORDS: Final = (
"the",
"quick",
"brown",
"fox",
"jumps",
"over",
"lazy",
"dog",
"while",
"counting",
"tokens",
"for",
"budget",
"reservation",
"before",
"admission",
"on",
"the",
"gateway",
"and",
"every",
"request",
"body",
"is",
"scanned",
"exactly",
"once",
"with",
"a",
"hand",
"written",
"piece",
"scanner",
"that",
"mirrors",
"tiktoken's",
"regex",
"boundaries",
"It's",
"faster",
"because",
"there's",
"no",
"backtracking",
"engine",
"involved",
"so",
"we'll",
"keep",
"it",
"that",
"way",
"Zürich",
"café",
"naïve",
"東京",
"مرحبا",
"🙂",
"42",
"1999",
"3.14159",
"$1,234.56",
"100%",
"user@example.com",
"https://example.com/a/b?c=d",
"C++",
"F#",
"node.js",
"v1.2.3",
"(parens)",
"[brackets]",
"{braces}",
"<tags>",
'"quotes"',
"'single'",
"don't",
"WON'T",
"I'M",
"They'RE",
)
def random_text(rng: random.Random) -> str:
return "".join(rng.choice(ALPHABET) for _ in range(rng.randrange(0, 40)))
def paragraph(rng: random.Random, words: int) -> str:
return " ".join(rng.choice(WORDS) for _ in range(words))
def short_paragraphs(rng: random.Random) -> Iterator[str]:
while True:
content = paragraph(rng, rng.randrange(60, 140))
if len(content) < TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS:
yield content
def chat_body(rng: random.Random, target_tokens: int) -> dict[str, object]:
candidates: Final = tuple(itertools.islice(short_paragraphs(rng), 2000))
running: Final = tuple(itertools.accumulate(len(ENCODING.encode(content)) + 3 for content in candidates))
turns: Final = next(index for index, total in enumerate(running) if total >= target_tokens) + 1
contents: Final = candidates[: turns + (turns % 2)]
return {
"model": "gpt-4",
"messages": [
{"role": "system", "content": "You are a helpful assistant. Answer precisely and cite sources."},
*(
{"role": "user" if index % 2 == 0 else "assistant", "content": content}
for index, content in enumerate(contents)
),
{"role": "user", "content": "Summarise the conversation so far in three sentences."},
],
}
TOOLS: Final = [
{
"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"}},
]
SMALL_REQUESTS: Final = (
{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello, how are you today?"}]},
{
"model": "gpt-4",
"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",
],
},
{"role": "assistant", "content": [{"type": "text", "text": "Sure."}]},
],
},
{
"model": "gpt-4",
"messages": [{"role": "user", "content": "weather?"}],
"tools": TOOLS,
"tool_choice": {"type": "function", "function": {"name": "get_weather"}},
},
{
"model": "gpt-4",
"messages": [{"role": "system", "content": "sys"}, {"role": "user", "content": "weather?"}],
"tools": TOOLS,
"tool_choice": "none",
},
{"model": "gpt-4", "prompt": "Write a haiku about ships."},
{"model": "gpt-4", "prompt": ["first prompt", "second prompt"]},
{
"model": "gpt-4",
"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",
},
{"model": "gpt-4", "input": [[101, 2023, 5], [7]], "encoding_format": "float"},
{
"model": "gpt-4",
"query": "best harbour",
"documents": [
"doc one",
{"text": "doc two", "title": "T", "n": 3, "ok": True, "none": None, "tags": ["a", "b"]},
],
},
)
def main() -> None:
rng: Final = random.Random(2026)
texts: Final = tuple(CORPUS) + tuple(random_text(rng) for _ in range(3000))
with (HERE / "texts.jsonl").open("w", encoding="utf-8") as handle:
for text in texts:
tokens = len(ENCODING.encode(text, disallowed_special=()))
pieces = SPLIT_PATTERN.findall(text)
handle.write(json.dumps({"text": text, "tokens": tokens, "pieces": pieces}, ensure_ascii=False) + "\n")
bodies: Final = tuple(SMALL_REQUESTS) + (chat_body(rng, 50_000),)
with (HERE / "requests.jsonl").open("w", encoding="utf-8") as handle:
for body in bodies:
input_tokens = _count_input_tokens(dict(body), "gpt-4")
assert input_tokens is not None
handle.write(json.dumps({"body": json.dumps(body), "input_tokens": input_tokens}) + "\n")
if __name__ == "__main__":
main()

10
litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl vendored Normal file
View file

File diff suppressed because one or more lines are too long

3044
litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl vendored Normal file
View file

File diff suppressed because one or more lines are too long

111
litellm-rust/crates/token-counter/tests/token_counter.rs
View file

@ -1,4 +1,5 @@
use rstest::rstest;
use serde::Deserialize;
use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter};
@ -174,3 +175,113 @@ fn tool_choice_and_system_discount_change_the_count() {
fn loading_a_bad_tokenizer_is_a_load_error() {
assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_))));
}
fn cl100k_counter() -> TokenCounter {
let path = concat!(
env!("CARGO_MANIFEST_DIR"),
"/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
);
let ranks = std::fs::read_to_string(path).expect("cl100k rank file is in the repo");
TokenCounter::from_cl100k_ranks(&ranks).expect("cl100k ranks load")
}
fn cl100k_fixture(name: &str) -> String {
let path = format!(
"{}/tests/fixtures/cl100k/{name}",
env!("CARGO_MANIFEST_DIR")
);
std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/cl100k/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("cl100k_base")`; see
/// `tests/fixtures/cl100k/generate.py`.
#[test]
fn cl100k_text_counts_match_tiktoken() {
let counter = cl100k_counter();
let fixtures: Vec<TextFixture> = cl100k_fixture("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, "gpt-4")`), so this pins the shared message,
/// tool and reply-priming accounting on the cl100k path as well.
#[test]
fn cl100k_request_counts_match_python_admission_counter() {
let counter = cl100k_counter();
let fixtures: Vec<RequestFixture> = cl100k_fixture("requests.jsonl")
.lines()
.map(|line| serde_json::from_str(line).expect("fixture line is json"))
.collect();
let counts: Vec<usize> = 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("gpt-4"));
assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body);
count.input_tokens
})
.collect();
assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000));
}
#[test]
fn cl100k_shares_the_message_accounting_with_the_anthropic_path() {
let counter = cl100k_counter();
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) {
assert!(matches!(
TokenCounter::from_cl100k_ranks(rank_file),
Err(Error::Ranks(_))
));
}

9
litellm/litellm_core_utils/default_encoding.py
View file

@ -1,4 +1,5 @@
import os
from pathlib import Path
from typing import Final
import litellm
@ -14,6 +15,14 @@ except (ImportError, AttributeError):
filename = pkg_resources.resource_filename(__name__, "litellm_core_utils/tokenizers")
CL100K_BASE_RANK_FILE: Final = "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
def cl100k_base_rank_file() -> str:
"""The vendored tiktoken `cl100k_base` rank file (`base64(token) rank` lines)."""
return Path(filename, CL100K_BASE_RANK_FILE).read_text(encoding="ascii")
# Always default TIKTOKEN_CACHE_DIR to the bundled tokenizers directory
# unless the user explicitly overrides it via CUSTOM_TIKTOKEN_CACHE_DIR.
# This keeps tiktoken fully offline-capable by default (see #1071).

33
litellm/litellm_core_utils/token_counter.py
View file

@ -381,7 +381,7 @@ class _MessageCountParams:
from litellm.utils import print_verbose
actual_model: Final = _fix_model_name(model)
if actual_model == "gpt-3.5-turbo-0301":
if uses_legacy_message_accounting(model):
self.tokens_per_message = 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n
self.tokens_per_name = -1 # if there's a name, the role is omitted
elif actual_model in litellm.open_ai_chat_completion_models or actual_model in litellm.azure_llms:
@ -615,7 +615,7 @@ def _get_exact_count_function(
) -> TokenCounterFunction:
"""
Get the function to count tokens based on the model and custom tokenizer."""
from litellm.utils import _select_tokenizer, print_verbose
from litellm.utils import _select_tokenizer
if model is not None or custom_tokenizer is not None:
tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model)
@ -627,15 +627,7 @@ def _get_exact_count_function(
return count_tokens
elif tokenizer_json["type"] == "openai_tokenizer":
model_to_use: Final = _fix_model_name(model)
try:
if "gpt-4o" in model_to_use:
encoding = tiktoken.get_encoding("o200k_base")
else:
encoding = tiktoken.encoding_for_model(model_to_use)
except KeyError:
print_verbose("Warning: model not found. Using cl100k_base encoding.")
encoding = tiktoken.get_encoding("cl100k_base")
encoding: Final = openai_tokenizer_encoding(model)
def encode_length(text: str) -> int:
return len(encoding.encode(text, disallowed_special=()))
@ -651,6 +643,25 @@ def _get_exact_count_function(
return _get_tiktoken_count_function(encode_length)
def openai_tokenizer_encoding(model: str) -> tiktoken.Encoding:
"""The tiktoken encoding `token_counter` uses for a model on the `openai_tokenizer` path."""
from litellm.utils import print_verbose
model_to_use: Final = _fix_model_name(model)
if "gpt-4o" in model_to_use:
return tiktoken.get_encoding("o200k_base")
try:
return tiktoken.encoding_for_model(model_to_use)
except KeyError:
print_verbose("Warning: model not found. Using cl100k_base encoding.")
return tiktoken.get_encoding("cl100k_base")
def uses_legacy_message_accounting(model: str) -> bool:
"""Whether `token_counter` prices messages with the `gpt-3.5-turbo-0301` constants (4 per message, -1 per name)."""
return _fix_model_name(model) == "gpt-3.5-turbo-0301"
def _fix_model_name(model: str) -> str:
"""We normalize some model names to others"""
if model in litellm.azure_llms:

36
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -34,7 +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.rust_bridge.token_counter import RustTokenizer, count_input_tokens, rust_tokenizer
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
from litellm.types.router import DeploymentTypedDict
@ -1365,8 +1365,9 @@ async def count_request_input_tokens(
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.
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.
Models whose tokenizer the Rust bridge ports (Anthropic, tiktoken cl100k_base)
are counted from the raw body by the bridge when it is enabled, once per
distinct tokenizer, 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.
@ -1374,23 +1375,31 @@ async def count_request_input_tokens(
models: Final = _get_request_models(request_body=request_body, route=route, llm_router=llm_router)
if not models:
return MappingProxyType({})
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
tokenizers: Final[Mapping[str, RustTokenizer | None]] = MappingProxyType(
{model: rust_tokenizer(model) for model in models}
)
distinct_tokenizers: Final[tuple[RustTokenizer, ...]] = tuple(
dict.fromkeys(tokenizer for tokenizer in tokenizers.values() if tokenizer is not None)
)
rust_counts_by_tokenizer: Final[Mapping[RustTokenizer, int]] = MappingProxyType(
{
tokenizer: count.input_tokens
for tokenizer in distinct_tokenizers
if raw_body is not None and (count := await count_input_tokens(raw_body, tokenizer)) is not 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)
model: rust_counts_by_tokenizer[tokenizer]
for model, tokenizer in tokenizers.items()
if tokenizer is not None and tokenizer in rust_counts_by_tokenizer
}
)
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)
MappingProxyType({})
if not python_models
else _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,
@ -1398,6 +1407,7 @@ async def count_request_input_tokens(
models=python_models,
)
)
verbose_proxy_logger.debug("input token counts: rust=%s python=%s", dict(rust_counts), dict(python_counts))
return MappingProxyType({**rust_counts, **python_counts})

51
litellm/rust_bridge/token_counter.py
View file

@ -5,17 +5,20 @@ 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 typing import Final, Literal, Protocol, cast # noqa: TID251 # native extension exposes untyped callables
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.default_encoding import cl100k_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
from litellm.utils import uses_anthropic_tokenizer as _python_uses_anthropic_tokenizer
from litellm.utils import claude_json_str, huggingface_tokenizer_kind
RustTokenizer = Literal["anthropic", "cl100k_base"]
class RustTokenCounter(Protocol):
@ -27,6 +30,9 @@ class RustTokenCounterFactory(Protocol):
def __call__(self, tokenizer_json: str) -> RustTokenCounter:
raise NotImplementedError
def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter:
raise NotImplementedError
@dataclass(frozen=True, slots=True)
class InputTokenCount:
@ -50,18 +56,32 @@ def _as_factory(value: object) -> RustTokenCounterFactory | 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)
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 only `cl100k_base` does. 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)
if kind == "anthropic":
return "anthropic"
if kind is not None or uses_legacy_message_accounting(model):
return None
return "cl100k_base" if openai_tokenizer_encoding(model).name == "cl100k_base" else None
@lru_cache(maxsize=4)
def _anthropic_counter(factory: RustTokenCounterFactory) -> RustTokenCounter:
return factory(claude_json_str)
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())
async def count_anthropic_input_tokens(body: bytes) -> InputTokenCount | None:
async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None:
if not rust_enabled():
return None
factory: Final = TOKEN_COUNTER.load()
@ -69,11 +89,14 @@ async def count_anthropic_input_tokens(body: bytes) -> InputTokenCount | None:
return None
try:
attempt: Final = await aattempt(
native_call=lambda: _anthropic_counter(factory).acount_request(body),
native_call=lambda: _counter(factory, tokenizer).acount_request(body),
adapt=_INPUT_TOKEN_COUNT.validate_python,
context=BridgeErrorContext(route="token_counter", provider="anthropic", model=""),
context=BridgeErrorContext(route="token_counter", provider=tokenizer, model=""),
)
except (RuntimeError, ValueError) as error:
verbose_logger.debug("Rust token counter failed, counting in Python: %s", error)
verbose_logger.debug("Rust token counter (%s) failed, counting in Python: %s", tokenizer, error)
return None
return attempt.value if isinstance(attempt, RustHandled) else None
if not isinstance(attempt, RustHandled):
return None
verbose_logger.debug("Rust token counter (%s) counted %d input tokens", tokenizer, attempt.value.input_tokens)
return attempt.value

48
litellm/utils.py
View file

@ -2227,25 +2227,39 @@ 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:
HuggingFaceTokenizerKind = Literal["cohere", "anthropic", "llama2", "llama3"]
def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None:
"""Which HuggingFace tokenizer `token_counter` selects for a model; `None` means tiktoken."""
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 uses_anthropic_tokenizer(model):
claude_tokenizer: Final = Tokenizer.from_str(claude_json_str)
return {"type": "huggingface_tokenizer", "tokenizer": claude_tokenizer}
# llama2
elif "llama-2" in model.lower() or "replicate" in model.lower():
tokenizer = Tokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
# llama3
elif "llama-3" in model.lower():
tokenizer = Tokenizer.from_pretrained("Xenova/llama-3-tokenizer")
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
else:
return "cohere"
if uses_anthropic_tokenizer(model):
return "anthropic"
if "llama-2" in model.lower() or "replicate" in model.lower():
return "llama2"
if "llama-3" in model.lower():
return "llama3"
return None
def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None:
kind: Final = huggingface_tokenizer_kind(model)
if kind is None:
return None
return {"type": "huggingface_tokenizer", "tokenizer": _load_huggingface_tokenizer(kind)}
def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind) -> Tokenizer:
match kind:
case "cohere":
return Tokenizer.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer")
case "anthropic":
return Tokenizer.from_str(claude_json_str)
case "llama2":
return Tokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
case "llama3":
return Tokenizer.from_pretrained("Xenova/llama-3-tokenizer")
def encode(model="", text="", custom_tokenizer: dict | None = None):

136
tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
View file

@ -1,5 +1,8 @@
from __future__ import annotations
import json
import math
from types import MappingProxyType
from typing import Final
import pytest
@ -207,8 +210,10 @@ def test_deployment_pricing_update_invalidates_cached_estimate() -> None:
ANTHROPIC_TOKENIZER_MODEL: Final = "claude-sonnet-4-5-20250929"
CL100K_MODEL: Final = "gpt-4"
RUST_COUNTED_BODY: Final = {"model": ANTHROPIC_TOKENIZER_MODEL, "max_tokens": 16, "messages": ANTHROPIC_MESSAGES}
RUST_INPUT_TOKENS: Final = 4_321
RUST_INPUT_TOKENS_BY_TOKENIZER: Final = MappingProxyType({"anthropic": RUST_INPUT_TOKENS, "cl100k_base": 1_234})
class _FakeDeclined(Exception):
@ -225,33 +230,51 @@ class _FakeNative:
class _RecordingCounter:
bodies: Final[list[bytes]] = []
"""Stands in for one native counter; records `(tokenizer, body)` on the shared factory."""
def __init__(self, tokenizer_json: str) -> None:
pass
def __init__(self, factory: _RecordingFactory, tokenizer: rust_token_counter.RustTokenizer) -> None:
self.factory = factory
self.tokenizer = tokenizer
async def acount_request(self, body: bytes) -> object:
self.bodies.append(body)
return {"model": ANTHROPIC_TOKENIZER_MODEL, "input_tokens": RUST_INPUT_TOKENS}
self.factory.calls.append((self.tokenizer, body))
return {"model": "", "input_tokens": RUST_INPUT_TOKENS_BY_TOKENIZER[self.tokenizer]}
class _RecordingFactory:
"""Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_cl100k_ranks`."""
def __init__(self) -> None:
self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = []
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")
class _DecliningCounter:
def __init__(self, tokenizer_json: str) -> None:
pass
async def acount_request(self, body: bytes) -> object:
raise _FakeDeclined("unsupported content block")
class _DecliningFactory:
def __call__(self, tokenizer_json: str) -> _DecliningCounter:
return _DecliningCounter()
def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter:
return _DecliningCounter()
@pytest.fixture
def rust_counter(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative())
rust_token_counter._anthropic_counter.cache_clear()
rust_token_counter._counter.cache_clear()
configuration.reset_rust_configuration()
_RecordingCounter.bodies.clear()
yield
rust_token_counter.TOKEN_COUNTER.reset()
rust_token_counter._anthropic_counter.cache_clear()
rust_token_counter._counter.cache_clear()
configuration.reset_rust_configuration()
@ -270,8 +293,9 @@ def rust_counter(monkeypatch: pytest.MonkeyPatch):
async def test_rust_count_replaces_python_tokenizing_on_every_llm_route(
rust_counter: None, route: str, request_body: dict
) -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
rust_token_counter.TOKEN_COUNTER.override(factory)
raw_body: Final = json.dumps(request_body).encode()
counts: Final = await count_request_input_tokens(
@ -279,53 +303,105 @@ async def test_rust_count_replaces_python_tokenizing_on_every_llm_route(
)
assert dict(counts) == {ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS}
assert _RecordingCounter.bodies == [raw_body]
assert factory.calls == [("anthropic", raw_body)]
@pytest.mark.asyncio
async def test_rust_decline_falls_back_to_python_count(rust_counter: None) -> None:
@pytest.mark.parametrize("model", (CL100K_MODEL, "azure/gpt-35-turbo", "gemini/gemini-2.5-pro", "my-router-alias"))
async def test_tiktoken_cl100k_models_are_counted_by_rust(rust_counter: None, model: str) -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(_DecliningCounter)
rust_token_counter.TOKEN_COUNTER.override(factory)
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
raw_body: Final = json.dumps(body).encode()
counts: Final = await count_request_input_tokens(
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
)
assert dict(counts) == {model: RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"]}
assert factory.calls == [("cl100k_base", raw_body)]
@pytest.mark.asyncio
async def test_multi_model_request_counts_once_per_tokenizer_and_python_for_the_rest(rust_counter: None) -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(factory)
models: Final = (CL100K_MODEL, ANTHROPIC_TOKENIZER_MODEL, "gemini/gemini-2.5-pro", "gpt-4o")
body: Final = {"model": list(models), "messages": ANTHROPIC_MESSAGES}
raw_body: Final = json.dumps(body).encode()
python_counts: Final = await count_request_input_tokens(
request_body=RUST_COUNTED_BODY, route="/v1/messages", llm_router=None
request_body=body, route="/v1/chat/completions", llm_router=None
)
counts: Final = await count_request_input_tokens(
request_body=RUST_COUNTED_BODY,
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
)
assert factory.calls == [("cl100k_base", raw_body), ("anthropic", raw_body)]
assert dict(counts) == {
CL100K_MODEL: RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"],
"gemini/gemini-2.5-pro": RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"],
ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS,
"gpt-4o": python_counts["gpt-4o"],
}
assert counts["gpt-4o"] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
@pytest.mark.asyncio
@pytest.mark.parametrize("model", (ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL))
async def test_rust_decline_falls_back_to_python_count(rust_counter: None, model: str) -> None:
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(_DecliningFactory())
body: Final = {**RUST_COUNTED_BODY, "model": model}
python_counts: Final = await count_request_input_tokens(request_body=body, route="/v1/messages", llm_router=None)
counts: Final = await count_request_input_tokens(
request_body=body,
route="/v1/messages",
llm_router=None,
raw_body=json.dumps(RUST_COUNTED_BODY).encode(),
raw_body=json.dumps(body).encode(),
)
assert dict(counts) == dict(python_counts)
assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS
assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
@pytest.mark.asyncio
async def test_disabled_rust_never_sees_the_raw_body(rust_counter: None) -> None:
factory: Final = _RecordingFactory()
litellm.rust(False)
rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
rust_token_counter.TOKEN_COUNTER.override(factory)
body: Final = {"model": [ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL], "messages": ANTHROPIC_MESSAGES}
counts: Final = await count_request_input_tokens(
request_body=RUST_COUNTED_BODY,
route="/v1/messages",
request_body=body,
route="/v1/chat/completions",
llm_router=None,
raw_body=json.dumps(RUST_COUNTED_BODY).encode(),
raw_body=json.dumps(body).encode(),
)
assert _RecordingCounter.bodies == []
assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS
assert factory.calls == []
assert set(counts) == {ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL}
assert not set(counts.values()) & set(RUST_INPUT_TOKENS_BY_TOKENIZER.values())
@pytest.mark.asyncio
async def test_non_anthropic_tokenizer_models_stay_in_python(rust_counter: None) -> None:
@pytest.mark.parametrize("model", ("gpt-4o", "o3", "replicate/meta/llama-2-70b-chat"))
async def test_models_without_a_rust_tokenizer_stay_in_python(rust_counter: None, model: str) -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
body: Final = {"model": "gpt-4o", "messages": ANTHROPIC_MESSAGES}
rust_token_counter.TOKEN_COUNTER.override(factory)
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
python_counts: Final = await count_request_input_tokens(
request_body=body, route="/v1/chat/completions", llm_router=None
)
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
assert factory.calls == []
assert dict(counts) == dict(python_counts)
assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()

265
tests/test_litellm/rust_bridge/test_token_counter.py
View file

@ -11,13 +11,18 @@ import json
from typing import Final
import pytest
import tiktoken
from tokenizers import Tokenizer
import litellm
from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
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
from litellm.utils import claude_json_str
MODEL: Final = "claude-sonnet-4-5-20250929"
CL100K_MODEL: Final = "gpt-4"
BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode()
@ -44,69 +49,111 @@ class _RecordingCounter:
return {"model": MODEL, "input_tokens": 42}
class _DecliningCounter:
def __init__(self, tokenizer_json: str) -> None:
pass
class _RecordingFactory:
"""Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_cl100k_ranks` for ranks."""
def __init__(self) -> None:
self.counters: list[_RecordingCounter] = []
self.rank_files: 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")
class _RaisingCounter:
def __init__(self, error: Exception) -> None:
self.error = error
async def acount_request(self, body: bytes) -> object:
raise _FakeDeclined("request has no messages")
raise self.error
class _FailingCounter:
def __init__(self, tokenizer_json: str) -> None:
pass
class _RaisingFactory:
"""Every counter it builds, for either tokenizer, raises `error` on count."""
async def acount_request(self, body: bytes) -> object:
raise RuntimeError("encode failed")
def __init__(self, error: Exception) -> None:
self.error = error
def __call__(self, tokenizer_json: str) -> _RaisingCounter:
return _RaisingCounter(self.error)
def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter:
return _RaisingCounter(self.error)
@pytest.fixture(autouse=True)
def _reset_bridge(monkeypatch: pytest.MonkeyPatch):
bridge.TOKEN_COUNTER.reset()
bridge._anthropic_counter.cache_clear()
bridge._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()
bridge._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)
@pytest.mark.parametrize("tokenizer", ("anthropic", "cl100k_base"))
async def test_disabled_bridge_never_constructs_a_counter(tokenizer: bridge.RustTokenizer) -> None:
factory: Final = _RecordingFactory()
litellm.rust(False)
bridge.TOKEN_COUNTER.override(factory)
assert await bridge.count_anthropic_input_tokens(BODY) is None
assert constructed == []
assert await bridge.count_input_tokens(BODY, tokenizer) is None
assert factory.counters == []
@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
factory: Final = _RecordingFactory()
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)
first: Final = await bridge.count_input_tokens(BODY, "anthropic")
second: Final = await bridge.count_input_tokens(BODY, "anthropic")
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"
assert len(factory.counters) == 1
assert factory.counters[0].bodies == [BODY, BODY]
assert json.loads(factory.counters[0].tokenizer_json)["model"]["type"] == "BPE"
@pytest.mark.asyncio
async def test_cl100k_counter_is_built_from_the_vendored_rank_file_once() -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
bridge.TOKEN_COUNTER.override(factory)
first: Final = await bridge.count_input_tokens(BODY, "cl100k_base")
second: Final = await bridge.count_input_tokens(BODY, "cl100k_base")
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") == 100_256
assert factory.counters[0].bodies == [BODY, BODY]
@pytest.mark.asyncio
async def test_each_tokenizer_gets_its_own_cached_counter() -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
bridge.TOKEN_COUNTER.override(factory)
await bridge.count_input_tokens(BODY, "anthropic")
await bridge.count_input_tokens(BODY, "cl100k_base")
await bridge.count_input_tokens(BODY, "anthropic")
assert [counter.tokenizer_json == "cl100k_base" for counter in factory.counters] == [False, True]
assert [len(counter.bodies) for counter in factory.counters] == [2, 1]
@pytest.mark.asyncio
@ -114,38 +161,105 @@ async def test_missing_native_module_falls_back(monkeypatch: pytest.MonkeyPatch)
litellm.rust(True)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
assert await bridge.count_anthropic_input_tokens(BODY) is None
assert await bridge.count_input_tokens(BODY, "anthropic") is None
assert await bridge.count_input_tokens(BODY, "cl100k_base") is None
@pytest.mark.asyncio
async def test_declined_request_falls_back() -> None:
@pytest.mark.parametrize("tokenizer", ("anthropic", "cl100k_base"))
async def test_declined_request_falls_back(tokenizer: bridge.RustTokenizer) -> None:
litellm.rust(True)
bridge.TOKEN_COUNTER.override(_DecliningCounter)
bridge.TOKEN_COUNTER.override(_RaisingFactory(_FakeDeclined("request has no messages")))
assert await bridge.count_anthropic_input_tokens(BODY) is None
assert await bridge.count_input_tokens(BODY, tokenizer) is None
@pytest.mark.asyncio
async def test_runtime_failure_falls_back() -> None:
@pytest.mark.parametrize("tokenizer", ("anthropic", "cl100k_base"))
async def test_runtime_failure_falls_back(tokenizer: bridge.RustTokenizer) -> None:
litellm.rust(True)
bridge.TOKEN_COUNTER.override(_FailingCounter)
bridge.TOKEN_COUNTER.override(_RaisingFactory(RuntimeError("encode failed")))
assert await bridge.count_anthropic_input_tokens(BODY) is None
assert await bridge.count_input_tokens(BODY, tokenizer) is None
@pytest.mark.parametrize(
("model", "expected"),
((MODEL, True), ("claude-3-5-sonnet-20241022", False), ("gpt-4o", False), ("my-router-alias", False)),
(
(MODEL, "anthropic"),
("claude-3-5-sonnet-20241022", "cl100k_base"),
("gpt-4", "cl100k_base"),
("gpt-4-turbo", "cl100k_base"),
("gpt-3.5-turbo", "cl100k_base"),
("azure/gpt-35-turbo", "cl100k_base"),
("gemini/gemini-2.5-pro", "cl100k_base"),
("mistral/mistral-large-latest", "cl100k_base"),
("my-router-alias", "cl100k_base"),
("azure/gpt-4o", "cl100k_base"),
("command-r-plus", "cl100k_base"),
("gpt-4o", None),
("gpt-4o-mini", None),
("gpt-4.1", None),
("gpt-5", None),
("o3", None),
("replicate/meta/llama-2-70b-chat", None),
("meta-llama/Llama-3-8b", None),
),
)
def test_uses_anthropic_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: bool) -> None:
assert bridge.uses_anthropic_tokenizer(model) is expected
def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: bridge.RustTokenizer | None) -> None:
assert bridge.rust_tokenizer(model) == 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)
def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "cohere_models", litellm.cohere_models | {"command-r-plus"})
assert bridge.uses_anthropic_tokenizer(MODEL) is False
assert bridge.rust_tokenizer("command-r-plus") is None
@pytest.mark.parametrize("legacy_model", ("gpt-3.5-turbo-0301", "gpt-35-turbo-0301"))
def test_rust_tokenizer_declines_legacy_message_accounting_python_prices_differently(
monkeypatch: pytest.MonkeyPatch, legacy_model: str
) -> None:
monkeypatch.setattr(
litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {"gpt-3.5-turbo-0301"}
)
monkeypatch.setattr(litellm, "azure_llms", {**litellm.azure_llms, "gpt-35-turbo-0301": "azure"})
messages: Final = [{"role": "user", "name": "bob", "content": "hello there"}]
assert litellm.token_counter(model=legacy_model, messages=messages) != litellm.token_counter(
model=CL100K_MODEL, messages=messages
)
assert bridge.rust_tokenizer(legacy_model) is None
assert bridge.rust_tokenizer(CL100K_MODEL) == "cl100k_base"
@pytest.mark.parametrize("model", (MODEL, "gpt-4", "gpt-4o"))
def test_rust_tokenizer_names_the_encoding_python_actually_counts_with(model: str) -> None:
text: Final = "Hello, world! \u00e9\u00e8 12345 \u3053\u3093\u306b\u3061\u306f <|endoftext|>\r\n" * 10
python_count: Final = litellm.token_counter(model=model, text=text)
cl100k_count: Final = len(tiktoken.get_encoding("cl100k_base").encode(text, disallowed_special=()))
match bridge.rust_tokenizer(model):
case "cl100k_base":
assert python_count == cl100k_count
case "anthropic":
assert python_count == len(Tokenizer.from_str(claude_json_str).encode(text).ids) != cl100k_count
case None:
assert python_count != cl100k_count
def test_disabled_hf_download_routes_anthropic_models_to_cl100k_like_python(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
assert bridge.rust_tokenizer(MODEL) == "cl100k_base"
assert bridge.rust_tokenizer("meta-llama/Llama-3-8b") == "cl100k_base"
assert bridge.rust_tokenizer("gpt-4o") is None
def test_disabled_token_counter_declines_every_model(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_token_counter", True)
assert bridge.rust_tokenizer(MODEL) is None
assert bridge.rust_tokenizer(CL100K_MODEL) is None
PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
@ -182,7 +296,16 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
},
{
"model": MODEL,
"messages": [{"role": "user", "content": "x " * 20_000}],
"messages": [{"role": "user", "content": "x " * 500}],
},
{
"model": MODEL,
"messages": [
{
"role": "user",
"content": "I'VE got 1234567 things; it's \"fine\"...\r\n\r\n caf\u00e9 \u0645\u0631\u062d\u0628\u0627 \U0001f600 <|endoftext|>",
}
],
},
{"model": MODEL, "prompt": "Write a haiku about ships.", "max_tokens": 20},
{"model": MODEL, "prompt": ["first prompt", "second prompt"]},
@ -190,7 +313,7 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
"model": MODEL,
"instructions": "be terse",
"input": [
{"role": "user", "content": [{"type": "input_text", "text": "Summarise caf\u00e9 menus \u2014 \"ok\"?\n"}]},
{"role": "user", "content": [{"type": "input_text", "text": 'Summarise caf\u00e9 menus \u2014 "ok"?\n'}]},
{"role": "assistant", "content": "Sure."},
],
},
@ -202,27 +325,58 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
)
PARITY_MODELS: Final[tuple[tuple[str, bridge.RustTokenizer], ...]] = (
(MODEL, "anthropic"),
(CL100K_MODEL, "cl100k_base"),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(("model", "tokenizer"), PARITY_MODELS)
@pytest.mark.parametrize("request_body", PARITY_REQUESTS)
async def test_native_count_matches_python_budget_counter(
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object]
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object], model: str, tokenizer: bridge.RustTokenizer
) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
litellm.rust(True)
body: Final = json.dumps(request_body).replace(MODEL, model)
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)
rust_count: Final = await bridge.count_input_tokens(body.encode(), tokenizer)
python_count: Final = _count_input_tokens(request_body=json.loads(body), model=model)
assert rust_count is not None
assert rust_count.model == request_body.get("model")
assert rust_count.model == json.loads(body).get("model")
assert rust_count.input_tokens == python_count
@pytest.mark.asyncio
async def test_cl100k_counts_long_text_exactly_where_python_chunks(monkeypatch: pytest.MonkeyPatch) -> None:
"""Python encodes tiktoken text in fixed-size chunks (drift of up to one token per chunk boundary); Rust does not."""
native: Final = pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
litellm.rust(True)
text: Final = "x " * 20_000
body: Final = {"model": CL100K_MODEL, "messages": [{"role": "user", "content": text}]}
encoding: Final = tiktoken.get_encoding("cl100k_base")
exact: Final = 3 + len(encoding.encode("user")) + len(encoding.encode(text)) + 3
chunks: Final = -(-len(text) // TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS)
rust_count: Final = await bridge.count_input_tokens(json.dumps(body).encode(), "cl100k_base")
python_count: Final = _count_input_tokens(request_body=body, model=CL100K_MODEL)
assert rust_count is not None
assert rust_count.input_tokens == exact
assert python_count is not None
assert exact < python_count <= exact + chunks
DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
{
"model": MODEL,
"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}],
"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}]},
@ -231,12 +385,13 @@ DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
@pytest.mark.asyncio
@pytest.mark.parametrize("tokenizer", ("anthropic", "cl100k_base"))
@pytest.mark.parametrize("request_body", DECLINED_REQUESTS)
async def test_native_declines_shapes_python_prices_differently(
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object]
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object], tokenizer: bridge.RustTokenizer
) -> 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
assert await bridge.count_input_tokens(json.dumps(request_body).encode(), tokenizer) is None