From 13036b72f8254451c335fd7a86bb46c13e4c4a41 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 9 Sep 2026 22:36:11 +0000 Subject: [PATCH] feat(rust_bridge): count prompt, input, query and documents bodies in Rust The Rust token counter now mirrors _count_input_tokens key precedence (messages, prompt, input, query/documents) so every LLM route that goes through budget reservation gets the GIL-free count, not only /v1/messages and /v1/chat/completions. Objects are serialised like json.dumps before tokenizing; floats and unknown shapes still decline to Python. The body model is optional so route-selected models can be matched by the caller Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/core/src/token_counter/mod.rs | 71 ++++++++++++++--- .../core/src/token_counter/python_json.rs | 79 +++++++++++++++++++ .../crates/core/src/token_counter/tests.rs | 53 ++++++++++++- .../crates/core/src/token_counter/types.rs | 49 ++++++++++-- litellm/rust_bridge/token_counter.py | 4 +- .../spend_tracking/test_budget_reservation.py | 19 ++++- .../rust_bridge/test_token_counter.py | 58 +++++++++----- 7 files changed, 290 insertions(+), 43 deletions(-) create mode 100644 litellm-rust/crates/core/src/token_counter/python_json.rs diff --git a/litellm-rust/crates/core/src/token_counter/mod.rs b/litellm-rust/crates/core/src/token_counter/mod.rs index 637691407a4..765b272979d 100644 --- a/litellm-rust/crates/core/src/token_counter/mod.rs +++ b/litellm-rust/crates/core/src/token_counter/mod.rs @@ -2,6 +2,7 @@ //! for the shapes it can count exactly. Everything else is declined so the host //! keeps its own counter as the reference. +mod python_json; mod tools; pub mod types; @@ -13,7 +14,9 @@ use crate::constants::{ TOOL_CHOICE_NONE_TOKENS, TOOL_DEFINITIONS_TOKENS, TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT, }; use tools::format_function_definitions; -use types::{ContentBlock, ContentItem, CountableRequest, Message, MessageContent, ToolChoice}; +use types::{ + ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice, +}; #[derive(Debug, ThisError, PartialEq, Eq)] pub enum TokenCountError { @@ -29,7 +32,7 @@ pub enum TokenCountError { #[derive(Clone, Debug, PartialEq, Eq, Serialize)] pub struct InputTokenCount { - pub model: String, + pub model: Option, pub input_tokens: usize, } @@ -56,14 +59,37 @@ impl TokenCounter { .map_err(|error| TokenCountError::Encode(error.to_string())) } + /// Mirrors the host's key precedence: `messages`, then `prompt`, then + /// `input`, then `query` plus `documents`. pub fn count_request( &self, request: &CountableRequest, ) -> Result { - let messages = request - .messages - .as_deref() - .ok_or_else(|| TokenCountError::Unsupported("request has no messages".to_string()))?; + let input_tokens = if let Some(messages) = &request.messages { + self.count_messages(request, messages)? + } else if let Some(prompt) = &request.prompt { + self.count_text_value(prompt)? + } else if let Some(input) = &request.input { + self.count_text_value(input)? + } else if request.query.is_some() || request.documents.is_some() { + self.count_optional_text_value(request.query.as_ref())? + + self.count_optional_text_value(request.documents.as_ref())? + } else { + return Err(TokenCountError::Unsupported( + "request has no countable input".to_string(), + )); + }; + Ok(InputTokenCount { + model: request.model.clone(), + input_tokens, + }) + } + + fn count_messages( + &self, + request: &CountableRequest, + messages: &[Message], + ) -> Result { let message_tokens = messages .iter() .map(|message| self.count_message(message)) @@ -76,10 +102,35 @@ impl TokenCounter { request.tool_choice.as_ref(), includes_system_message, )?; - Ok(InputTokenCount { - model: request.model.clone(), - input_tokens: message_tokens + extra_tokens, - }) + Ok(message_tokens + extra_tokens) + } + + fn count_optional_text_value( + &self, + value: Option<&TextValue>, + ) -> Result { + value.map_or(Ok(0), |value| self.count_text_value(value)) + } + + /// `str()` for scalars, `json.dumps()` for objects, lists flattened, nulls + /// skipped. Floats are declined because Python's `repr` and Rust's float + /// formatting disagree on exponents. + fn count_text_value(&self, value: &TextValue) -> Result { + match value { + TextValue::Null => Ok(0), + TextValue::Bool(true) => self.count_text("True"), + TextValue::Bool(false) => self.count_text("False"), + TextValue::Integer(number) => self.count_text(&number.to_string()), + TextValue::Float(_) => Err(TokenCountError::Unsupported( + "float text values are counted by the python path".to_string(), + )), + TextValue::Text(text) => self.count_text(text), + TextValue::List(items) => items + .iter() + .map(|item| self.count_text_value(item)) + .sum::>(), + TextValue::Object(_) => self.count_text(&python_json::dumps(value)?), + } } fn count_message(&self, message: &Message) -> Result { diff --git a/litellm-rust/crates/core/src/token_counter/python_json.rs b/litellm-rust/crates/core/src/token_counter/python_json.rs new file mode 100644 index 00000000000..3b93fa38c3b --- /dev/null +++ b/litellm-rust/crates/core/src/token_counter/python_json.rs @@ -0,0 +1,79 @@ +//! `json.dumps(value)` with Python's default arguments: `", "` and `": "` +//! separators, `ensure_ascii=True`, and keys in insertion order. + +use std::fmt::Write; + +use super::TokenCountError; +use super::types::TextValue; + +pub(super) fn dumps(value: &TextValue) -> Result { + let mut out = String::new(); + write_value(&mut out, value)?; + Ok(out) +} + +fn write_value(out: &mut String, value: &TextValue) -> Result<(), TokenCountError> { + match value { + TextValue::Null => out.push_str("null"), + TextValue::Bool(true) => out.push_str("true"), + TextValue::Bool(false) => out.push_str("false"), + TextValue::Integer(number) => write_number(out, number), + TextValue::Float(_) => { + return Err(TokenCountError::Unsupported( + "float repr is formatted by the python path".to_string(), + )); + } + TextValue::Text(text) => write_string(out, text), + TextValue::List(items) => { + out.push('['); + for (index, item) in items.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + write_value(out, item)?; + } + out.push(']'); + } + TextValue::Object(entries) => { + out.push('{'); + for (index, (key, item)) in entries.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + write_string(out, key); + out.push_str(": "); + write_value(out, item)?; + } + out.push('}'); + } + } + Ok(()) +} + +fn write_number(out: &mut String, number: &i64) { + // Writing an integer into a String cannot fail. + let _ = write!(out, "{number}"); +} + +fn write_string(out: &mut String, text: &str) { + out.push('"'); + for character in text.chars() { + match character { + '"' => out.push_str("\\\""), + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + '\u{08}' => out.push_str("\\b"), + '\u{0c}' => out.push_str("\\f"), + ' '..='~' => out.push(character), + _ => { + let mut units = [0u16; 2]; + for unit in character.encode_utf16(&mut units) { + let _ = write!(out, "\\u{unit:04x}"); + } + } + } + } + out.push('"'); +} diff --git a/litellm-rust/crates/core/src/token_counter/tests.rs b/litellm-rust/crates/core/src/token_counter/tests.rs index c171bb32bdc..67b2fc9fda7 100644 --- a/litellm-rust/crates/core/src/token_counter/tests.rs +++ b/litellm-rust/crates/core/src/token_counter/tests.rs @@ -45,23 +45,69 @@ const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5", "type":"object","properties":{"location":{"type":["string","null"]}},"required":["location"]}}], "tool_choice":"none"}"#; +const COMPLETIONS_PROMPT: &str = + r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#; + +const COMPLETIONS_PROMPT_LIST: &str = + r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#; + +const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[ + {"role":"user","content":[{"type":"input_text","text":"Summarise caf\u00e9 menus, na\u00efve \u2014 ok? \"quoted\"\n"}]}, + {"role":"assistant","content":"Sure."}],"instructions":"be terse"}"#; + +const EMBEDDINGS_TOKEN_IDS: &str = + r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#; + +const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour", + "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#; + +/// Expected counts are pinned from +/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`. #[rstest] #[case::text_only(SIMPLE, 14)] #[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)] #[case::openai_tools_named_choice(TOOLS_OPENAI, 123)] #[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)] +#[case::completions_prompt(COMPLETIONS_PROMPT, 7)] +#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)] +#[case::responses_input_items(RESPONSES_INPUT, 62)] +#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)] +#[case::rerank_query_and_documents(RERANK, 41)] fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) { let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); let count = counter().count_request(&request).expect("fixture counts"); assert_eq!( count, InputTokenCount { - model: "claude-sonnet-4-5".to_string(), + model: Some("claude-sonnet-4-5".to_string()), input_tokens: expected, } ); } +#[rstest] +#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)] +#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)] +#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)] +#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)] +fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) { + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let count = counter().count_request(&request).expect("fixture counts"); + assert_eq!(count.input_tokens, expected); +} + +#[test] +fn objects_dump_like_python_json_dumps() { + let body = r#"{"model":"m","input":{"text":"caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~","n":-3,"ok":true,"no":false,"none":null,"list":[1,"a",{"z":[]}],"empty":{}}}"#; + let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses"); + let dumped = python_json::dumps(request.input.as_ref().expect("input is present")) + .expect("fixture dumps"); + assert_eq!( + dumped, + r#"{"text": "caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~", "n": -3, "ok": true, "no": false, "none": null, "list": [1, "a", {"z": []}], "empty": {}}"# + ); +} + #[test] fn tool_definitions_render_like_python() { let request = CountableRequest::parse(TOOLS_OPENAI.as_bytes()).expect("fixture parses"); @@ -85,7 +131,6 @@ fn union_types_and_anthropic_schema_render_like_python() { #[rstest] #[case::not_json(b"not json" as &[u8])] -#[case::missing_model(br#"{"messages":[]}"#)] #[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":"{}"}}]}]}"# @@ -107,7 +152,9 @@ fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) { } #[rstest] -#[case::no_messages(br#"{"model":"m"}"# as &[u8])] +#[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=="}}]}]}"# )] diff --git a/litellm-rust/crates/core/src/token_counter/types.rs b/litellm-rust/crates/core/src/token_counter/types.rs index e3b58e9327d..022a6c68012 100644 --- a/litellm-rust/crates/core/src/token_counter/types.rs +++ b/litellm-rust/crates/core/src/token_counter/types.rs @@ -1,17 +1,30 @@ use indexmap::IndexMap; -use serde::Deserialize; +use serde::{Deserialize, Deserializer}; use super::TokenCountError; -/// The parts of a request body `litellm.token_counter` reads when a host counts -/// input tokens for budget checks. Anything outside this shape is declined so -/// the host can fall back to its own counter instead of silently miscounting. +/// The parts of a request body the host's budget counter reads. Chat and +/// Anthropic Messages bodies carry `messages`; completions carry `prompt`; +/// Responses and embeddings carry `input`; rerank carries `query` and +/// `documents`. The host checks key presence, not nullness, so an explicit +/// `null` is kept distinct from an absent key. Anything outside this shape is +/// declined so the host can fall back to its own counter instead of silently +/// miscounting. #[derive(Clone, Debug, Deserialize, PartialEq)] pub struct CountableRequest { - pub model: String, + pub model: Option, + #[serde(default, deserialize_with = "present_messages")] pub messages: Option>, pub tools: Option>, pub tool_choice: Option, + #[serde(default, deserialize_with = "present_text")] + pub prompt: Option, + #[serde(default, deserialize_with = "present_text")] + pub input: Option, + #[serde(default, deserialize_with = "present_text")] + pub query: Option, + #[serde(default, deserialize_with = "present_text")] + pub documents: Option, } impl CountableRequest { @@ -21,6 +34,32 @@ impl CountableRequest { } } +fn present_messages<'de, D: Deserializer<'de>>( + deserializer: D, +) -> Result>, D::Error> { + Option::>::deserialize(deserializer) + .map(|messages| Some(messages.unwrap_or_default())) +} + +fn present_text<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + TextValue::deserialize(deserializer).map(Some) +} + +/// Free-form JSON the host counts as text: strings and integers via `str()`, +/// objects via `json.dumps()`, lists flattened. Objects keep document order so +/// the dumped text matches Python byte for byte. +#[derive(Clone, Debug, Deserialize, PartialEq)] +#[serde(untagged)] +pub enum TextValue { + Null, + Bool(bool), + Integer(i64), + Float(f64), + Text(String), + List(Vec), + Object(IndexMap), +} + /// Python counts every string-valued key of a message, so any key beyond these /// makes the shape unsupported rather than silently uncounted. #[derive(Clone, Debug, Deserialize, PartialEq)] diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index 618ec7610ed..5a82ad716dd 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -28,7 +28,7 @@ class RustTokenCounterFactory(Protocol): @dataclass(frozen=True, slots=True) class InputTokenCount: - model: str + model: str | None input_tokens: int @@ -49,7 +49,7 @@ TOKEN_COUNTER: Final = NativeBinding("TokenCounter", validate=_as_factory) def uses_anthropic_tokenizer(model: str) -> bool: - if litellm.disable_hf_tokenizer_download is True: + if litellm.disable_token_counter is True or litellm.disable_hf_tokenizer_download is True: return False return model in litellm.anthropic_models and "claude-3" not in model 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 9918588a1b6..f272361062b 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -199,13 +199,26 @@ def rust_counter(monkeypatch: pytest.MonkeyPatch): @pytest.mark.asyncio -async def test_rust_count_replaces_python_tokenizing_for_anthropic_models(rust_counter: None) -> None: +@pytest.mark.parametrize( + ("route", "request_body"), + ( + ("/v1/messages", RUST_COUNTED_BODY), + ("/v1/chat/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "messages": ANTHROPIC_MESSAGES}), + ("/v1/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "prompt": "hi"}), + ("/v1/responses", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": "hi"}), + ("/v1/embeddings", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": ["hi"]}), + ("/v1/rerank", {"model": ANTHROPIC_TOKENIZER_MODEL, "query": "hi", "documents": ["a"]}), + ), +) +async def test_rust_count_replaces_python_tokenizing_on_every_llm_route( + rust_counter: None, route: str, request_body: dict +) -> None: litellm.rust(True) rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter) - raw_body: Final = json.dumps(RUST_COUNTED_BODY).encode() + raw_body: Final = json.dumps(request_body).encode() counts: Final = await count_request_input_tokens( - request_body=RUST_COUNTED_BODY, route="/v1/messages", llm_router=None, raw_body=raw_body + request_body=request_body, route=route, llm_router=None, raw_body=raw_body ) assert dict(counts) == {ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS} diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py index 4d29251dca0..2df91204390 100644 --- a/tests/test_litellm/rust_bridge/test_token_counter.py +++ b/tests/test_litellm/rust_bridge/test_token_counter.py @@ -13,6 +13,7 @@ from typing import Final import pytest import litellm +from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens from litellm.rust_bridge import bindings, configuration from litellm.rust_bridge import token_counter as bridge @@ -140,8 +141,9 @@ def test_uses_anthropic_tokenizer_mirrors_python_tokenizer_selection(model: str, assert bridge.uses_anthropic_tokenizer(model) is expected -def test_uses_anthropic_tokenizer_respects_hf_download_opt_out(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True) +@pytest.mark.parametrize("flag", ("disable_hf_tokenizer_download", "disable_token_counter")) +def test_uses_anthropic_tokenizer_respects_python_opt_outs(monkeypatch: pytest.MonkeyPatch, flag: str) -> None: + monkeypatch.setattr(litellm, flag, True) assert bridge.uses_anthropic_tokenizer(MODEL) is False @@ -182,12 +184,27 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = ( "model": MODEL, "messages": [{"role": "user", "content": "x " * 20_000}], }, + {"model": MODEL, "prompt": "Write a haiku about ships.", "max_tokens": 20}, + {"model": MODEL, "prompt": ["first prompt", "second prompt"]}, + { + "model": MODEL, + "instructions": "be terse", + "input": [ + {"role": "user", "content": [{"type": "input_text", "text": "Summarise caf\u00e9 menus \u2014 \"ok\"?\n"}]}, + {"role": "assistant", "content": "Sure."}, + ], + }, + {"model": MODEL, "input": "a single embedding string"}, + {"model": MODEL, "input": [[101, 2023, 5], [7]], "encoding_format": "float"}, + {"model": MODEL, "query": "best harbour", "documents": ["doc one", {"text": "doc two", "title": "T", "n": 3}]}, + {"model": MODEL, "messages": None, "prompt": "messages key wins even when null"}, + {"prompt": "model comes from the route"}, ) @pytest.mark.asyncio @pytest.mark.parametrize("request_body", PARITY_REQUESTS) -async def test_native_count_matches_python_token_counter( +async def test_native_count_matches_python_budget_counter( monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object] ) -> None: native: Final = pytest.importorskip("litellm.rust_bridge._native") @@ -195,30 +212,31 @@ async def test_native_count_matches_python_token_counter( litellm.rust(True) rust_count: Final = await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode()) - python_count: Final = litellm.token_counter( - model=MODEL, - messages=request_body["messages"], - tools=request_body.get("tools"), - tool_choice=request_body.get("tool_choice"), - ) + python_count: Final = _count_input_tokens(request_body=request_body, model=MODEL) assert rust_count is not None - assert rust_count.model == MODEL + assert rust_count.model == request_body.get("model") assert rust_count.input_tokens == python_count +DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = ( + { + "model": MODEL, + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}], + }, + {"model": MODEL, "prompt": 1.5}, + {"model": MODEL, "documents": [{"score": 0.5}]}, + {"model": MODEL, "file": "audio.mp3"}, +) + + @pytest.mark.asyncio -async def test_native_declines_image_content(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize("request_body", DECLINED_REQUESTS) +async def test_native_declines_shapes_python_prices_differently( + monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object] +) -> None: native: Final = pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) litellm.rust(True) - body: Final = json.dumps( - { - "model": MODEL, - "messages": [ - {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]} - ], - } - ).encode() - assert await bridge.count_anthropic_input_tokens(body) is None + assert await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode()) is None