diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 62bde5806a4..af3a31ddbfa 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2257,6 +2257,7 @@ dependencies = [ "serde", "serde_json", "thiserror 2.0.19", + "tokenizers", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 109a2edaf4d..4250fa81d26 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -48,6 +48,7 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std", sha2 = "0.10" subtle = "2" thiserror = "2.0" +tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } 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"] } diff --git a/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py b/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py index 1bfbdf00218..2ca853cbd62 100644 --- a/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py +++ b/litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py @@ -2,8 +2,8 @@ Run from the repository root with the project environment, once per encoding: - uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py cl100k_base - uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py o200k_base + uv run --no-sync python litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py cl100k_base + uv run --no-sync python litellm-rust/crates/token-counter-fast/tests/fixtures/generate.py o200k_base `/texts.jsonl` holds `{"text", "tokens", "pieces"}` lines: `tokens` counted with `tiktoken.get_encoding(name).encode(text, disallowed_special=())`, diff --git a/litellm-rust/crates/token-counter/Cargo.toml b/litellm-rust/crates/token-counter/Cargo.toml index 59e4f9a6a1d..67e5bd60537 100644 --- a/litellm-rust/crates/token-counter/Cargo.toml +++ b/litellm-rust/crates/token-counter/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [features] -default = ["huggingface", "tiktoken"] +default = ["fast", "huggingface", "tiktoken"] fast = ["dep:litellm-token-counter-fast"] huggingface = ["dep:litellm-token-counter-huggingface"] tiktoken = ["dep:litellm-token-counter-tiktoken"] @@ -25,6 +25,7 @@ thiserror.workspace = true criterion.workspace = true rand.workspace = true rstest.workspace = true +tokenizers.workspace = true [[bench]] name = "token_counter" diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md index 3ad21a00697..a6b2b50aac0 100644 --- a/litellm-rust/crates/token-counter/README.md +++ b/litellm-rust/crates/token-counter/README.md @@ -8,7 +8,7 @@ The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through t 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. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend +All three backends are enabled by default. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. 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 diff --git a/litellm-rust/crates/token-counter/tests/token_counter.rs b/litellm-rust/crates/token-counter/tests/token_counter.rs index 378a16d5d72..ebbe2f00ffc 100644 --- a/litellm-rust/crates/token-counter/tests/token_counter.rs +++ b/litellm-rust/crates/token-counter/tests/token_counter.rs @@ -1,4 +1,4 @@ -#![cfg(feature = "huggingface")] +#![cfg(any(feature = "fast", feature = "huggingface"))] use rstest::rstest; @@ -6,13 +6,15 @@ use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCount /// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)` /// so this test also guards Python parity. -fn counter() -> TokenCounter { +type JsonLoader = fn(&str) -> Result; + +fn counter(load: JsonLoader) -> TokenCounter { let path = concat!( env!("CARGO_MANIFEST_DIR"), "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" ); let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo"); - TokenCounter::from_json(&json).expect("anthropic tokenizer loads") + load(&json).expect("anthropic tokenizer loads") } const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#; @@ -63,21 +65,15 @@ const EMBEDDINGS_TOKEN_IDS: &str = const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#; -/// Expected counts are pinned from -/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`. -#[rstest] -#[case::text_only(SIMPLE, 14)] -#[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)] -#[case::openai_tools_named_choice(TOOLS_OPENAI, 123)] -#[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)] -#[case::completions_prompt(COMPLETIONS_PROMPT, 7)] -#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)] -#[case::responses_input_items(RESPONSES_INPUT, 62)] -#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)] -#[case::rerank_query_and_documents(RERANK, 41)] -fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) { +fn assert_count_request_matches_python_token_counter( + load: JsonLoader, + body: &str, + expected: usize, +) { let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter().count_request(&request).expect("fixture counts"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); assert_eq!( count, InputTokenCount { @@ -87,63 +83,24 @@ fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expect ); } -#[rstest] -#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)] -#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] -#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] -#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] -fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { +fn assert_key_presence_follows_python(load: JsonLoader, body: &str, expected: usize) { let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); - let count = counter().count_request(&request).expect("fixture counts"); + let count = counter(load) + .count_request(&request) + .expect("fixture counts"); assert_eq!(count.input_tokens, expected); } -#[rstest] -#[case::not_json(b"not json" as &[u8])] -#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)] -#[case::message_with_tool_calls( - br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"# -)] -#[case::dict_content( - br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"# -)] -#[case::float_enum( - br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"# -)] -#[case::anthropic_tool_choice_without_function( - br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"# -)] -fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { - assert!(matches!( - CountableRequest::parse(body), - Err(Error::RequestParse(_)) - )); -} - -#[rstest] -#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] -#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] -#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] -#[case::image_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# -)] -#[case::tool_result_block( - br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# -)] -#[case::array_without_items( - br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# -)] -fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { +fn assert_shapes_outside_the_mirror_are_declined_at_count(load: JsonLoader, body: &[u8]) { let request = CountableRequest::parse(body).expect("shape parses"); assert!(matches!( - counter().count_request(&request), + counter(load).count_request(&request), Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems) )); } -#[test] -fn tool_choice_and_system_discount_change_the_count() { - let counter = counter(); +fn assert_tool_choice_and_system_discount_change_the_count(load: JsonLoader) { + let counter = counter(load); let count = |body: &str| { counter .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses")) @@ -172,9 +129,104 @@ fn tool_choice_and_system_discount_change_the_count() { ); } -#[test] -fn loading_a_bad_tokenizer_is_a_load_error() { - assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_)))); +fn assert_loading_a_bad_tokenizer_is_a_load_error(load: JsonLoader) { + assert!(matches!(load("{}"), Err(Error::Load(_)))); +} + +macro_rules! json_backend_tests { + ($loader:path) => { + #[rstest] + #[case::text_only(super::SIMPLE, 14)] + #[case::content_blocks_name_and_system(super::BLOCKS_AND_SYSTEM, 45)] + #[case::openai_tools_named_choice(super::TOOLS_OPENAI, 123)] + #[case::anthropic_tools_system_discount_choice_none(super::TOOLS_ANTHROPIC_SYSTEM, 53)] + #[case::completions_prompt(super::COMPLETIONS_PROMPT, 7)] + #[case::completions_prompt_list(super::COMPLETIONS_PROMPT_LIST, 4)] + #[case::responses_input_items(super::RESPONSES_INPUT, 62)] + #[case::embeddings_token_ids(super::EMBEDDINGS_TOKEN_IDS, 5)] + #[case::rerank_query_and_documents(super::RERANK, 41)] + fn count_request_matches_python_token_counter( + #[case] body: &str, + #[case] expected: usize, + ) { + super::assert_count_request_matches_python_token_counter($loader, body, expected); + } + + #[rstest] + #[case::null_messages_win_over_prompt( + r#"{"model":"m","messages":null,"prompt":"ignored"}"#, + 3 + )] + #[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] + #[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] + #[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] + fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { + super::assert_key_presence_follows_python($loader, body, expected); + } + + #[rstest] + #[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])] + #[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)] + #[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)] + #[case::image_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"# + )] + #[case::tool_result_block( + br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"# + )] + #[case::array_without_items( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"# + )] + fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) { + super::assert_shapes_outside_the_mirror_are_declined_at_count($loader, body); + } + + #[test] + fn tool_choice_and_system_discount_change_the_count() { + super::assert_tool_choice_and_system_discount_change_the_count($loader); + } + + #[test] + fn loading_a_bad_tokenizer_is_a_load_error() { + super::assert_loading_a_bad_tokenizer_is_a_load_error($loader); + } + }; +} + +#[rstest] +#[case::not_json(b"not json" as &[u8])] +#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)] +#[case::message_with_tool_calls( + br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"# +)] +#[case::dict_content( + br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"# +)] +#[case::float_enum( + br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"# +)] +#[case::anthropic_tool_choice_without_function( + br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"# +)] +fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { + assert!(matches!( + CountableRequest::parse(body), + Err(Error::RequestParse(_)) + )); +} + +#[cfg(feature = "fast")] +mod fast_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json_fast); +} + +#[cfg(feature = "huggingface")] +mod huggingface_json { + use super::*; + + json_backend_tests!(TokenCounter::from_json); } #[cfg(feature = "fast")] @@ -222,7 +274,8 @@ mod fast { env!("CARGO_MANIFEST_DIR"), encoding.fixtures ); - std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py") + std::fs::read_to_string(&path) + .expect("fixture generated by token-counter-fast/tests/fixtures/generate.py") } #[derive(Deserialize)] @@ -238,7 +291,7 @@ mod fast { } /// Reference counts come from `tiktoken.get_encoding(name)`; see - /// `tests/fixtures/generate.py`. + /// `token-counter-fast/tests/fixtures/generate.py`. #[rstest] #[case::cl100k(CL100K)] #[case::o200k(O200K)]