From 0290573ffd452ac6197d53560e78c40561c1e264 Mon Sep 17 00:00:00 2001 From: Mihidum Hettiyahandi <55163074+mihidumh@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:09:48 +1100 Subject: [PATCH] fix(token_counter): support OpenAI 'input_audio' content blocks token_counter raised 'Invalid content item type: input_audio' on audio understanding payloads ({'type': 'input_audio', ...}). Callers that swallow the error fail open: Router._pre_call_checks returns the unfiltered deployment list (context-window guard skipped), the prompt-caching deployment check skips (no cache-affinity pinning), trim_messages returns over-budget conversations untrimmed, and /utils/token_counter 500s outright. parallel_request_limiter_v3 already documents the gap and strips/estimates audio blocks locally; every other call site still hit the raise. Estimate audio blocks the same way that limiter does -- decoded-base64 byte count at a conservative low bitrate, floored per block -- using the shared constants in litellm_core_utils/audio_utils/utils.py, which the limiter now imports instead of keeping its own copies. Batch rows and budget reservations that carry input_audio stay floored at the size-based fallback they used while the block raised, so an audio payload cannot slip under a TPM or budget limit. Fixes #38459 Co-Authored-By: Claude Opus 5.5 --- litellm/litellm_core_utils/token_counter.py | 68 ++++++++++++- litellm/proxy/hooks/batch_rate_limiter.py | 22 ++++- .../hooks/parallel_request_limiter_v3.py | 32 +++---- litellm/proxy/spend_tracking/input_tokens.py | 20 +++- .../litellm_core_utils/test_token_counter.py | 77 +++++++++++++++ .../proxy/hooks/test_batch_file_validation.py | 96 +++++++++++++++++++ tests/unit/proxy/hooks/test_tpm_concurrent.py | 6 +- .../proxy/spend_tracking/test_input_tokens.py | 38 ++++++++ 8 files changed, 331 insertions(+), 28 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 323af72940a..84c1b9d2f09 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -17,7 +17,7 @@ import anyio import anyio.lowlevel import httpx import tiktoken -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from tiktoken.model import MODEL_TO_ENCODING from typing_extensions import ParamSpec, TypeVar @@ -47,6 +47,7 @@ from litellm.constants import ( TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.audio_utils.utils import AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( ChatTemplateToolSection, @@ -1674,6 +1675,67 @@ def _count_anthropic_content( _TEXT_FIELD_BY_BLOCK_TYPE: Final = MappingProxyType( {"text": "text", "thinking": "thinking", "redacted_thinking": "thinking", "compaction": "content"} ) +_OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT_LIST: Final = TypeAdapter(list[object]) + + +def _object_mapping(value: object) -> Mapping[str, object] | None: + try: + return _OBJECT_MAPPING.validate_python(value) + except ValidationError: + return None + + +def _object_list(value: object) -> Sequence[object] | None: + try: + return _OBJECT_LIST.validate_python(value) + except ValidationError: + return None + + +def _count_input_audio_content_block(c: Mapping[str, object]) -> int: + """ + Estimate tokens for an OpenAI ``input_audio`` content block (audio + understanding), e.g. {"type": "input_audio", "input_audio": {"data": + "", "format": "wav"}}. The real token cost is the provider's + server-side audio tokenization and cannot be derived exactly client-side; + when the block carries a base64 payload, derive the estimate from the + decoded byte count at a conservative low bitrate, otherwise use the flat + per-block floor -- the same numbers ``parallel_request_limiter_v3`` + already uses for its audio reservations (issue #38459). + """ + input_audio: Final = _object_mapping(c.get("input_audio")) + b64_data: Final = input_audio.get("data") if input_audio is not None else None + if isinstance(b64_data, str) and b64_data: + decoded_bytes: Final = len(b64_data) * 3 // 4 + return max(decoded_bytes // AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE) + return DEFAULT_AUDIO_TOKEN_ESTIMATE + + +def messages_contain_input_audio_content_blocks(messages: object) -> bool: + """ + True when any message carries an OpenAI ``input_audio`` content block. + + Callers that use ``token_counter`` for reservations or rate limits must + check this first: an audio block's contribution is a size-derived + ESTIMATE at a deliberately low assumed bitrate (see + ``_count_input_audio_content_block``), not a measurement, so the count + for such a message must not be trusted as an upper bound. + """ + for message in _object_list(messages) or (): + message_mapping = _object_mapping(message) + content = _object_list(message_mapping.get("content")) if message_mapping is not None else None + for content_item in content or (): + item_mapping = _object_mapping(content_item) + if item_mapping is not None and item_mapping.get("type") == "input_audio": + return True + return False + + +def batch_entry_contains_input_audio_content_blocks(entry: Mapping[str, object]) -> bool: + """True when a batch input row's ``body.messages`` carries an ``input_audio`` block.""" + body: Final = _object_mapping(entry.get("body")) + return body is not None and messages_contain_input_audio_content_blocks(body.get("messages")) def _count_content_list( @@ -1746,13 +1808,15 @@ def _count_content_list( tool_name = str(c.get("tool_name") or "") if tool_name: num_tokens += count_function(tool_name) + elif c["type"] == "input_audio": + num_tokens += _count_input_audio_content_block(c) else: content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__ raise UnsupportedContentTypeError( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field " f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, " - f"compaction, tool_reference, search_result)." + f"compaction, tool_reference, search_result, input_audio)." ) return num_tokens except Exception as e: diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 747dda100bf..515739f45f4 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -39,6 +39,7 @@ from litellm.batches.batch_utils import ( from litellm.constants import BATCH_TPD_DESCRIPTOR_SUFFIX, BATCH_TPD_WINDOW_SECONDS from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.token_counter import batch_entry_contains_input_audio_content_blocks from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, @@ -125,6 +126,25 @@ def _batch_create_routing_model(data: Mapping[str, object]) -> str | None: IncrementAmounts: TypeAlias = dict[Literal["requests", "tokens"], int] +def _count_entry_tokens_with_audio_floor(entry: Mapping[str, object], raw_line: bytes) -> int: + """ + Count a batch row, flooring a row that carries ``input_audio`` at the + size-based estimate. + + An ``input_audio`` block's contribution is a size-derived estimate at a + deliberately low assumed bitrate (see + ``token_counter._count_input_audio_content_block``), far below the + raw-bytes fallback. Before such blocks were countable (#38459) an audio row + RAISED inside token_counter and fell back to that size-based estimate; + floor it there again so a row carrying a large base64 audio payload cannot + slide the batch under the TPM limit. + """ + counted: Final = count_entry_tokens(entry) + if batch_entry_contains_input_audio_content_blocks(entry): + return max(counted, estimate_batch_entry_tokens(raw_line)) + return counted + + class BatchFileUsage(LiteLLMBaseModel): """ Internal model for batch file usage tracking, used for batch rate limiting @@ -1014,7 +1034,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): output_tokens += entry_output_tokens try: - entry_total_tokens = count_entry_tokens(entry) + entry_total_tokens = _count_entry_tokens_with_audio_floor(entry, raw_line) except Exception: entry_total_tokens = estimate_batch_entry_tokens(raw_line) total_tokens += entry_total_tokens diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c83551a815c..2bdb5c1769f 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -55,7 +55,11 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, ) -from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens +from litellm.litellm_core_utils.audio_utils.utils import ( + AUDIO_BYTES_PER_TOKEN, + DEFAULT_AUDIO_TOKEN_ESTIMATE, + estimate_audio_tokens, +) from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, @@ -531,18 +535,6 @@ GOOGLE_GENAI_NATIVE_CALL_TYPES: Final = ( CallTypes.agenerate_content_stream.value, ) RESPONSES_API_MIN_OUTPUT_TOKENS: Final = 16 -# 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. When the block carries a base64 -# payload, the estimate is derived from the decoded byte count; when the -# block is a reference without a payload (or the payload is missing), this -# flat per-block floor is used instead. -DEFAULT_AUDIO_TOKEN_ESTIMATE: Final = 300 -# Conservative bytes-per-token assumption for size-based audio estimation: -# equivalent to 8 kHz mono PCM-16 (16 000 bytes/s) at 10 tokens/s. Choosing -# the lowest reasonable bitrate means we never under-reserve for higher- -# quality audio recorded at the same wall-clock duration. -_AUDIO_BYTES_PER_TOKEN: Final = 1600 # Descriptor "key" values for project-scoped ITPM/OTPM. Distinct from # "model_per_project" (the combined-TPM descriptor) so both can be enforced # on the same project+model simultaneously without colliding on cache keys. @@ -3925,7 +3917,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Token estimate for one ``input_audio`` content block. When the block carries a base64 ``data`` payload, the estimate comes - from the decoded byte count (``len(b64) * 3 // 4 // _AUDIO_BYTES_PER_TOKEN``), + from the decoded byte count (``len(b64) * 3 // 4 // AUDIO_BYTES_PER_TOKEN``), assuming the lowest reasonable audio bitrate so we never under-reserve for higher-quality recordings of the same duration. @@ -3938,7 +3930,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): b64_data: Final = input_audio.get("data") if isinstance(input_audio, dict) else None if b64_data and isinstance(b64_data, str): decoded_bytes: Final = len(b64_data) * 3 // 4 - return max(decoded_bytes // _AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE) + return max(decoded_bytes // AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE) return DEFAULT_AUDIO_TOKEN_ESTIMATE @classmethod @@ -3964,11 +3956,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _strip_audio_content_blocks(messages: object) -> object: """ 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. + ``token_counter``. It counts them with the same size-derived estimate + since #38459 (it used to raise); stripping keeps the audio contribution + added exactly once, by ``_estimate_audio_content_tokens``, so the rest + of the message (text/images/tools) is counted accurately without ever + double-counting the audio. """ if not isinstance(messages, list): return messages diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 9f37d2965bc..b2de1e24792 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -24,6 +24,7 @@ from litellm.constants import OPENAI_CHAT_PDF_FILE_TOKENS, PDF_DATA_URL_PREFIX from litellm.litellm_core_utils.deployment_pricing import deployment_pricing_models, str_keyed from litellm.litellm_core_utils.token_counter import ( deployment_openai_model, + messages_contain_input_audio_content_blocks, pdf_page_pricing, reads_pdf_pages, unresolved_tokenizer_model, @@ -221,18 +222,33 @@ def _served_by_openai(deployment: object) -> bool: return served_provider == "openai" +def _as_object(value: object) -> object: + return value + + def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: try: if "messages" in request_body: + messages: Final = _as_object(request_body.get("messages")) try: - return litellm.token_counter( + counted: Final = litellm.token_counter( model=model, messages=request_body.get("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/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 697ac76b4e5..d58e2e0eba1 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -3570,3 +3570,80 @@ def test_openai_model_check_holds_while_openai_models_are_registered_concurrentl registering.result() finally: litellm.open_ai_chat_completion_models.difference_update(registered) + + +def test_token_counter_with_input_audio_content_block(): + """ + Regression test for issue #38459: a message containing an OpenAI + `input_audio` content block (audio understanding) must NOT raise from + token_counter. Before the fix the raise poisoned the whole message and + every caller that swallows counter errors failed open (router + context-window pre-call check, prompt-caching deployment check), while + /utils/token_counter returned HTTP 500. + + The estimate mirrors parallel_request_limiter_v3's audio reservation: + decoded-base64 byte count at AUDIO_BYTES_PER_TOKEN, floored at + DEFAULT_AUDIO_TOKEN_ESTIMATE per block. + """ + from litellm.litellm_core_utils.audio_utils.utils import AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE + + small_b64 = "UklGRiQAAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YQAAAAA=" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What does the audio say?"}, + {"type": "input_audio", "input_audio": {"data": small_b64, "format": "wav"}}, + ], + } + ] + + tokens = token_counter_new(model="gpt-4o-audio-preview", messages=messages) + assert tokens >= DEFAULT_AUDIO_TOKEN_ESTIMATE, f"Expected at least the per-block floor, got {tokens}" + + # a large payload must scale the estimate (decoded bytes / AUDIO_BYTES_PER_TOKEN) + big_b64 = "A" * (AUDIO_BYTES_PER_TOKEN * 4000) # decoded ~3000 * AUDIO_BYTES_PER_TOKEN bytes + tokens_big = token_counter_new( + model="gpt-4o-audio-preview", + messages=[ + { + "role": "user", + "content": [{"type": "input_audio", "input_audio": {"data": big_b64, "format": "wav"}}], + } + ], + ) + assert tokens_big >= 2900, f"large audio payload must scale the estimate, got {tokens_big}" + assert tokens_big > tokens + + # a payload-less (reference-only) block must not raise and gets the floor + tokens_bare = token_counter_new( + model="gpt-4o-audio-preview", + messages=[{"role": "user", "content": [{"type": "input_audio", "input_audio": {"format": "wav"}}]}], + ) + assert tokens_bare >= DEFAULT_AUDIO_TOKEN_ESTIMATE + + +def test_trim_messages_with_input_audio_content_block(): + """Companion to issue #38459, same shape as the `file`-block case + (#28409): trim_messages swallows the token_counter error and silently + returns an over-budget conversation UNTRIMMED. With the fix, trimming + must actually happen.""" + small_b64 = "UklGRiQAAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YQAAAAA=" + messages = [ + {"role": "user", "content": "filler message " * 200}, + {"role": "user", "content": "filler message " * 200}, + { + "role": "user", + "content": [ + {"type": "text", "text": "What does the audio say?"}, + {"type": "input_audio", "input_audio": {"data": small_b64, "format": "wav"}}, + ], + }, + ] + + trimmed = litellm.utils.trim_messages(messages, model="gpt-4o-audio-preview", max_tokens=500) + + assert trimmed is not None + assert len(trimmed) < len(messages), ( + "trim_messages must actually trim an over-budget conversation containing an input_audio block" + ) diff --git a/tests/unit/proxy/hooks/test_batch_file_validation.py b/tests/unit/proxy/hooks/test_batch_file_validation.py index 8eeee2e6468..27bbe74b161 100644 --- a/tests/unit/proxy/hooks/test_batch_file_validation.py +++ b/tests/unit/proxy/hooks/test_batch_file_validation.py @@ -2286,3 +2286,99 @@ async def test_disable_flag_still_skips_batch_processing_with_enqueued_limits(): assert result is data afile_content_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_input_audio_row_reserves_size_based_floor(): + """A chat row carrying an OpenAI `input_audio` content block must reserve + at least the size-based estimate (serialized bytes / 4). + + token_counter's audio contribution is a size-derived estimate at a + deliberately low assumed bitrate (decoded bytes / AUDIO_BYTES_PER_TOKEN), + far below the raw-bytes fallback. Before audio blocks were countable + (#38459) such a row RAISED inside token_counter and fell back to the + size-based estimate; the floor restores exactly that conservatism so a + large base64 audio payload cannot slide the batch under the TPM limit. + """ + import json as _json + + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + prl = MagicMock() + prl.no_max_tokens_output_floor.return_value = 0 + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=prl, + ) + + blob = "A" * 400_000 + audio_row_bytes = _json.dumps( + { + "custom_id": "row-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o-audio-preview", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What does the audio say?"}, + {"type": "input_audio", "input_audio": {"data": blob, "format": "wav"}}, + ], + } + ], + }, + } + ).encode("utf-8") + fake_content = MagicMock() + fake_content.content = audio_row_bytes + + with patch( # test-quality-ok: mirrors the file's established harness — the download, not an HTTP boundary + "litellm.afile_content", + new=AsyncMock(return_value=fake_content), + ): + usage = await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=None, + ) + + size_based_floor = len(audio_row_bytes) // 4 + assert usage.request_count == 1 + assert usage.total_tokens >= size_based_floor, ( + f"audio-block row must reserve at least the size-based estimate " + f"({size_based_floor} tokens for {len(audio_row_bytes)} bytes), got " + f"{usage.total_tokens} — a large audio payload would evade TPM limits" + ) + + # Control: a plain-text row must NOT be floored at its serialized size — + # measured text rows keep the (smaller) real token count. + text_row_bytes = _json.dumps( + { + "custom_id": "row-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o-audio-preview", + "messages": [{"role": "user", "content": "What does the audio say?"}], + }, + } + ).encode("utf-8") + fake_text_content = MagicMock() + fake_text_content.content = text_row_bytes + + with patch( # test-quality-ok: mirrors the file's established harness — the download, not an HTTP boundary + "litellm.afile_content", + new=AsyncMock(return_value=fake_text_content), + ): + text_usage = await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=None, + ) + + assert text_usage.request_count == 1 + assert text_usage.total_tokens < len(text_row_bytes) // 4, ( + "plain-text rows must keep the measured token count, not the size-based floor" + ) diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index e43a72541c7..23ddb9ea1c0 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -29,11 +29,11 @@ from PIL import Image from litellm import token_counter from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.audio_utils.utils import AUDIO_BYTES_PER_TOKEN from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, estimate_image_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( - _AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE, PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, @@ -2284,7 +2284,7 @@ def test_audio_token_estimate_scales_with_payload_size(): to exhaust ITPM quota while reserving almost nothing. The estimate must now grow proportionally with the base64 payload size - (len(b64) * 3 // 4 // _AUDIO_BYTES_PER_TOKEN), floored at + (len(b64) * 3 // 4 // AUDIO_BYTES_PER_TOKEN), floored at DEFAULT_AUDIO_TOKEN_ESTIMATE so reference-only blocks and genuinely short clips still get a non-trivial reservation. @@ -2322,7 +2322,7 @@ def test_audio_token_estimate_scales_with_payload_size(): f"Large payload ({large_estimate}) must reserve more than small payload " f"({small_estimate}); flat-rate bug is back" ) - assert very_large_estimate == len(very_large_b64) * 3 // 4 // _AUDIO_BYTES_PER_TOKEN + assert very_large_estimate == len(very_large_b64) * 3 // 4 // AUDIO_BYTES_PER_TOKEN assert very_large_estimate > 6_000 assert no_data_estimate >= 300, ( f"Reference-only block (no data) must use the DEFAULT_AUDIO_TOKEN_ESTIMATE floor; got {no_data_estimate}" diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index 40aee7bb00d..bfd2140311d 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -11,6 +11,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, models_counted_for, @@ -280,3 +281,40 @@ def test_a_compaction_item_counts_only_its_summary(model: str) -> None: assert count_input_tokens_for_model(request_body=padded_body, model=model) == count_input_tokens_for_model( request_body=plain_body, model=model ) + + +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"