From 2ca3143b34b10a203b414faea6366e945df74f99 Mon Sep 17 00:00:00 2001 From: Mihidum Hettiyahandi <55163074+mihidumh@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:44:41 +1000 Subject: [PATCH] fix(proxy): floor budget reservation for input_audio requests at the serialised fallback Before #38459 an input_audio request raised inside token_counter and the budget reservation counted it from the serialised messages (base64 payload included). Making the blocks countable turned that into a size-derived estimate at a deliberately low bitrate, priced at the text rate, so a large or highly compressed audio payload could be admitted against a budget more cheaply than before. Floor audio-bearing requests at the serialised fallback, the same compensate-in-the-caller pattern the batch limiter uses; text-only requests are untouched. --- litellm/proxy/spend_tracking/input_tokens.py | 18 +++++++-- .../proxy/spend_tracking/test_input_tokens.py | 39 +++++++++++++++++++ 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 6c7083fb6db..a80bbb79d4a 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -19,6 +19,7 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.token_counter import messages_contain_input_audio_content_blocks from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.token_counter import ( @@ -128,15 +129,26 @@ def _count_input_tokens_for_models( def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: try: if "messages" in request_body: + messages: Final = request_body.get("messages") try: - return litellm.token_counter( + counted: Final = litellm.token_counter( model=model, - messages=request_body.get("messages") or (), + messages=messages or (), tools=request_body.get("tools"), tool_choice=request_body.get("tool_choice"), ) except ValueError: - return _count_text_tokens(model=model, text=request_body.get("messages")) + return _count_text_tokens(model=model, text=messages) + # An ``input_audio`` block counts as a size-derived estimate at a + # deliberately low assumed bitrate, and this reservation prices it + # at the text rate. Before such blocks were countable (#38459) an + # audio request RAISED in token_counter and reserved from the + # serialised-messages fallback above; floor it there again so a + # large or highly compressed audio payload cannot be admitted + # against a budget more cheaply than it was before. + if messages_contain_input_audio_content_blocks(messages): + return max(counted, _count_text_tokens(model=model, text=messages)) + return counted if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: diff --git a/tests/test_litellm/proxy/spend_tracking/test_input_tokens.py b/tests/test_litellm/proxy/spend_tracking/test_input_tokens.py index a1da6c79cd7..3439e0a16af 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_input_tokens.py +++ b/tests/test_litellm/proxy/spend_tracking/test_input_tokens.py @@ -2,6 +2,7 @@ from __future__ import annotations +import base64 from types import MappingProxyType from typing import Final @@ -9,6 +10,7 @@ import pytest from litellm.proxy.spend_tracking.input_tokens import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS, + _count_text_tokens, count_input_tokens, count_input_tokens_for_model, ) @@ -27,3 +29,40 @@ async def test_large_input_is_still_counted() -> None: assert counts[CL100K_MODEL] == count_input_tokens_for_model(request_body=request_body, model=CL100K_MODEL) assert isinstance(counts, MappingProxyType) + + +def test_input_audio_requests_reserve_at_least_the_serialised_fallback() -> None: + """ + Budget reservation counts an ``input_audio`` block as a size-derived + estimate at a deliberately low bitrate, priced at the text rate. Before + #38459 the same request raised inside ``token_counter`` and reserved + from the serialised-messages fallback, which tokenises the base64 + payload itself. Floor audio-bearing requests at that fallback so a + caller cannot be admitted against a budget more cheaply than before + the blocks became countable (compressed audio carries far more duration + per byte than the estimate assumes). + """ + model: Final = "gpt-4o-audio-preview" + audio_b64: Final = base64.b64encode(bytes(range(256)) * 400).decode() + audio_messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Transcribe this recording."}, + {"type": "input_audio", "input_audio": {"data": audio_b64, "format": "mp3"}}, + ], + } + ] + text_messages: Final = [{"role": "user", "content": "Transcribe this recording."}] + + audio_count: Final = count_input_tokens_for_model(request_body={"messages": audio_messages}, model=model) + fallback: Final = _count_text_tokens(model=model, text=audio_messages) + text_count: Final = count_input_tokens_for_model(request_body={"messages": text_messages}, model=model) + + assert audio_count is not None, "an audio-bearing request must still be countable" + assert fallback > 0, "the serialised fallback must see the base64 payload" + assert audio_count >= fallback, ( + f"audio request reserved {audio_count} tokens, below the pre-#38459 fallback of {fallback}" + ) + assert text_count is not None + assert text_count < fallback, "a text-only request must not be floored"