From cfde33ac5149bd8ccf2ee102f4cc4f9afad6ef90 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Wed, 22 Jul 2026 23:59:29 +0530 Subject: [PATCH] fix(proxy): Non-text input bypasses upfront ITPM enforcement --- .../hooks/parallel_request_limiter_v3.py | 78 +++++++++++++++++-- .../proxy/hooks/test_tpm_concurrent.py | 54 +++++++++++++ 2 files changed, 125 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 9fcf085967a..fb9ddf26260 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -290,6 +290,14 @@ DEFAULT_CHARS_PER_TOKEN = 4 # (baseline floor) and to the smallest configured TPM limit (capped floor for # small per-tenant TPM caps). _TPM_FLOOR_FRACTION = 4 +# litellm.token_counter has no per-type handling for "input_audio" content +# blocks (unlike images, which use use_default_image_token_count) -- it +# silently contributes 0 tokens for them. This is a conservative flat +# estimate added per audio block on top of token_counter's result for the +# project ITPM reservation, so an audio-heavy burst can't reserve close to +# nothing. Not a billing-accurate count, just enough to avoid grossly +# under-reserving. +DEFAULT_AUDIO_TOKEN_ESTIMATE = int(os.getenv("LITELLM_DEFAULT_AUDIO_TOKEN_ESTIMATE", 300)) # Stash for the reserved-token count on the request data dict so success/ # failure callbacks can reconcile against the upfront reservation. TPM_RESERVED_TOKENS_KEY = "_litellm_tpm_reserved_tokens" @@ -2558,6 +2566,53 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): llm_provider=llm_provider, ) + @staticmethod + def _count_audio_content_blocks(messages: Any) -> int: + """ + Number of ``input_audio`` content blocks across ``messages``. + ``litellm.token_counter`` has no per-type handling for these (unlike + images) -- it raises ``ValueError`` on them -- so + ``_estimate_precise_input_tokens`` strips them before counting and + adds a flat estimate per block back on top; see + ``DEFAULT_AUDIO_TOKEN_ESTIMATE``. + """ + if not isinstance(messages, list): + return 0 + count = 0 + for message in messages: + content = message.get("content") if isinstance(message, dict) else None + if not isinstance(content, list): + continue + count += sum(1 for block in content if isinstance(block, dict) and block.get("type") == "input_audio") + return count + + @staticmethod + def _strip_audio_content_blocks(messages: Any) -> Any: + """ + Drop ``input_audio`` content blocks before passing ``messages`` to + ``token_counter``, which raises ``ValueError`` on them (no per-type + handling, unlike images). The audio contribution is added back + separately via ``DEFAULT_AUDIO_TOKEN_ESTIMATE`` so the rest of the + message (text/images/tools) still gets counted accurately instead of + the whole call falling back to the cheap char-count estimate. + """ + if not isinstance(messages, list): + return messages + sanitized = [] + for message in messages: + if not isinstance(message, dict): + sanitized.append(message) + continue + content = message.get("content") + if not isinstance(content, list): + sanitized.append(message) + continue + filtered_content = [ + block for block in content if not (isinstance(block, dict) and block.get("type") == "input_audio") + ] + sanitized.append({**message, "content": filtered_content}) + return sanitized + def _estimate_precise_input_tokens(self, data: dict, model: str | None) -> int: """ Model-aware input token estimate for the project ITPM reservation, @@ -2565,23 +2620,31 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): deployment-level itpm/otpm check uses in ``io_token_rate_limit_check.py``. Unlike the cheap char-count estimate the combined-TPM path uses, this accounts for image/tool - content so a burst of multimodal or tool-heavy requests can't each - reserve only the one-token floor and blow past ITPM before post-call - reconciliation catches up. + content (plus a conservative flat add-on per audio block, which + ``token_counter`` can't count at all) so a burst of multimodal, + tool-heavy, or audio-heavy requests can't each reserve only the + one-token floor and blow past ITPM before post-call reconciliation + catches up. Falls back to the cheap char-count estimate if ``token_counter`` can't resolve a tokenizer for this model (e.g. an unrecognized - custom model name) or otherwise raises. + custom model name) or otherwise raises -- the audio add-on still + applies on top of the fallback. """ from litellm import token_counter + messages = data.get("messages") + audio_blocks = self._count_audio_content_blocks(messages) + audio_token_estimate = audio_blocks * DEFAULT_AUDIO_TOKEN_ESTIMATE + countable_messages = self._strip_audio_content_blocks(messages) if audio_blocks else messages + try: - return max( + estimate = max( 0, int( token_counter( model=model or "", - messages=data.get("messages"), + messages=countable_messages, text=data.get("prompt") or data.get("input"), tools=data.get("tools"), tool_choice=data.get("tool_choice"), @@ -2589,9 +2652,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ), ) + return estimate + audio_token_estimate except Exception: # noqa: BLE001 - any tokenizer/model-resolution failure degrades to the cheap estimate, never a 500 estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data) - return estimated_input_tokens + return estimated_input_tokens + audio_token_estimate async def _reserve_project_io_tokens_or_raise( self, diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 143c369c360..a7907f39b46 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -1727,5 +1727,59 @@ async def test_responses_api_usage_reconciles_using_input_output_tokens_fields(r ) +@pytest.mark.asyncio +async def test_itpm_reservation_accounts_for_audio_content_not_just_text(rate_limiter): + """ + Regression for the audio half of a Medium-severity review finding: + litellm.token_counter has no per-type handling for `input_audio` + content blocks (unlike images, which it does count via + use_default_image_token_count), so it silently contributes 0 tokens for + them. Without DEFAULT_AUDIO_TOKEN_ESTIMATE, a burst of audio-heavy + requests with minimal text would each reserve only the one-token floor + and blow past the project ITPM limit before post-call reconciliation. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-audio-itpm") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-audio", + project_metadata={ + # Tighter than DEFAULT_AUDIO_TOKEN_ESTIMATE (300), but far bigger + # than the handful of tokens the bare text "hi" would cost. + "model_itpm_limit": {"bedrock_mantle/claude-opus": 50}, + }, + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + { + "type": "input_audio", + "input_audio": {"data": "base64-audio-bytes", "format": "wav"}, + }, + ], + } + ], + } + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + assert getattr(exc_info.value, "status_code", None) == 429, ( + "Expected the audio content to push the ITPM reservation over the " + "50-token limit; if this doesn't raise, audio content isn't being " + "counted again." + ) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"])