diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 858b078d626..9ce50cac005 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -695,6 +695,48 @@ def _count_anthropic_content( return tokens +def _count_file_content_block(count_function: TokenCounterFunction, c: Mapping[str, object]) -> int: + """ + Count an OpenAI file content block (document understanding), e.g. + {"type": "file", "file": {"file_id" | "file_data" | "filename", ...}}. + The real token cost is the provider's server-side extraction of the + document and cannot be derived client-side; counting the base64 file_data + blob would wildly overcount. Count only the lightweight textual fields so + token_counter/trim_messages accept valid payloads instead of raising + (issue #28409). + """ + num_tokens = 0 # rebind-ok: accumulates over the file fields below + file_obj: Final = c.get("file") + if isinstance(file_obj, dict): + for file_field in ("file_id", "filename", "format"): + file_field_value = file_obj.get(file_field) + if isinstance(file_field_value, str) and file_field_value: + num_tokens += count_function(file_field_value) + return num_tokens + + +def messages_contain_file_content_blocks(messages: object) -> bool: + """ + True when any message carries an OpenAI ``file`` content block. + + Callers that use ``token_counter`` for reservations or rate limits must + check this first: a file block's document payload is opaque to the counter + (only the ``file_id``/``filename``/``format`` fields are counted -- see + ``_count_file_content_block``), so the count for such a message is a floor, + not a measurement, and it must not be trusted as an upper bound. + """ + if not isinstance(messages, list): + return False + for message in messages: + content = message.get("content") if isinstance(message, dict) else None + if not isinstance(content, list): + continue + for content_item in content: + if isinstance(content_item, dict) and content_item.get("type") == "file": + return True + return False + + def _count_content_list( count_function: TokenCounterFunction, content_list: OpenAIMessageContent, @@ -738,11 +780,13 @@ def _count_content_list( tool_name = str(c.get("tool_name") or "") if tool_name: num_tokens += count_function(tool_name) + elif c["type"] == "file": + num_tokens += _count_file_content_block(count_function, c) else: content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__ raise ValueError( f"Invalid content item type: {content_type}. " - f"Expected str or dict with 'type' field (text, image_url, tool_use, tool_result, thinking, tool_reference)." + f"Expected str or dict with 'type' field (text, image_url, tool_use, tool_result, thinking, tool_reference, file)." ) 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 5b814ad28fd..8672580e194 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -35,6 +35,7 @@ from litellm.batches.batch_utils import ( ) from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.token_counter import messages_contain_file_content_blocks from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, @@ -853,6 +854,15 @@ class _PROXY_BatchRateLimiter(CustomLogger): try: entry_total_tokens = _count_entry_tokens(entry) + # A `file` content block's document payload is opaque to + # token_counter -- only its filename/id fields are counted + # (see token_counter._count_file_content_block) -- so a row + # carrying a large base64 `file_data` would count as a few + # tokens and slide the whole batch under the TPM limit. + # Floor such rows at the size-based estimate, which does + # cover the payload bytes. + if messages_contain_file_content_blocks((entry.get("body") or {}).get("messages")): + entry_total_tokens = max(entry_total_tokens, _estimate_batch_entry_tokens(raw_line)) except Exception: entry_total_tokens = _estimate_batch_entry_tokens(raw_line) total_tokens += entry_total_tokens diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 7cad3f0a022..a755b1c1a76 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -15,6 +15,7 @@ from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate +from litellm.litellm_core_utils.token_counter import messages_contain_file_content_blocks from litellm.proxy._types import ( Litellm_EntityType, LiteLLM_TeamMembership, @@ -1279,9 +1280,20 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int: def _count_input_tokens(request_body: dict, model: str) -> int | None: try: if "messages" in request_body: + messages: Final = request_body.get("messages") or [] + # A `file` content block's real token cost is the provider's + # server-side extraction of the referenced document and cannot be + # derived client-side — token_counter only counts its filename/id + # fields, which would under-reserve arbitrarily large uploads. + # Return None so the caller takes the conservative max_input_tokens + # reservation — an explicit return, because falling through would + # let a stray `prompt` / `input` / `query` on the same body count + # as the whole request. + if messages_contain_file_content_blocks(messages): + return None return litellm.token_counter( model=model, - messages=request_body.get("messages") or [], + messages=messages, tools=request_body.get("tools"), tool_choice=request_body.get("tool_choice"), ) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index a2590dbca2d..e116c39edcc 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1160,3 +1160,90 @@ def test_count_content_list_rejects_unknown_type(): message = str(exc_info.value) assert "Invalid content item type: totally_unknown_block" in message assert "tool_reference" in message + assert "file" in message + + +def test_token_counter_with_file_content_block(): + """ + Regression test for issue #28409: a message containing an OpenAI `file` + content block (document understanding) must NOT raise from token_counter + or trim_messages. + + The real token cost of a file is the provider's server-side extraction and + cannot be derived client-side; only the lightweight textual fields + (file_id / filename / format) are counted — the opaque base64 `file_data` + blob must not inflate the count. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this document."}, + { + "type": "file", + "file": { + "file_id": "file-abc123", + "filename": "report.pdf", + "format": "application/pdf", + }, + }, + ], + } + ] + + tokens = token_counter_new(model="gpt-4o", messages=messages) + assert tokens > 0, f"Expected positive token count, got {tokens}" + + # base64 file_data is opaque — it must not blow up the count + small_blob = dict(messages[0]["content"][1]["file"], file_data="data:application/pdf;base64,AAAA") + big_blob = dict(messages[0]["content"][1]["file"], file_data="data:application/pdf;base64," + "A" * 100_000) + tokens_small = token_counter_new( + model="gpt-4o", + messages=[{"role": "user", "content": [{"type": "file", "file": small_blob}]}], + ) + tokens_big = token_counter_new( + model="gpt-4o", + messages=[{"role": "user", "content": [{"type": "file", "file": big_blob}]}], + ) + assert tokens_small == tokens_big, "opaque file_data blob must not be counted" + + # a bare file block (e.g. only file_data) must also not raise + tokens_bare = token_counter_new( + model="gpt-4o", + messages=[ + { + "role": "user", + "content": [{"type": "file", "file": {"file_data": "data:application/pdf;base64,AAAA"}}], + } + ], + ) + assert tokens_bare >= 0 + + +def test_trim_messages_with_file_content_block(): + """The original repro from issue #28409: trim_messages on a document + understanding payload raised ValueError from token_counter (1.83.14) — + current versions swallow that error and silently return the messages + UNTRIMMED instead. With the fix, trimming must actually happen.""" + messages = [ + {"role": "user", "content": "filler message " * 200}, + {"role": "user", "content": "filler message " * 200}, + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this file?"}, + { + "type": "file", + "file": {"file_id": "file-abc123", "filename": "report.pdf"}, + }, + ], + }, + ] + + trimmed = litellm.utils.trim_messages(messages, model="gpt-4o", max_tokens=120) + + assert trimmed is not None + assert len(trimmed) < len(messages), ( + "trim_messages must actually trim an over-budget conversation containing " + f"a file content block — got {len(trimmed)} messages back (input {len(messages)})" + ) diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index 4ef94f0965b..09f468d1f5c 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -1827,6 +1827,105 @@ async def test_count_input_file_usage_streams_without_building_list(): mock_dict_list.assert_not_called() +@pytest.mark.asyncio +async def test_file_content_block_row_reserves_size_based_floor(): + """A chat row carrying an OpenAI `file` content block must reserve at + least the size-based estimate (serialized bytes / 4). + + token_counter deliberately counts only the block's filename/id fields — + the document payload is opaque to it — so without this floor a row + carrying a large base64 `file_data` counts as a handful of tokens and the + whole batch slides under the TPM limit. Before file blocks were countable + (#33659) such a row RAISED inside token_counter and fell back to the + size-based estimate; the floor restores exactly that conservatism. + """ + 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 = "data:application/pdf;base64," + "A" * 400_000 + file_row_bytes = _json.dumps( + { + "custom_id": "row-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o-mini", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this document."}, + { + "type": "file", + "file": {"filename": "report.pdf", "file_data": blob}, + }, + ], + } + ], + }, + } + ).encode("utf-8") + fake_content = MagicMock() + fake_content.content = file_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(file_row_bytes) // 4 + assert usage.request_count == 1 + assert usage.total_tokens >= size_based_floor, ( + f"file-block row must reserve at least the size-based estimate " + f"({size_based_floor} tokens for {len(file_row_bytes)} bytes), got " + f"{usage.total_tokens} — a large file_data 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-mini", + "messages": [{"role": "user", "content": "Summarize this document."}], + }, + } + ).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 0 < text_usage.total_tokens < len(text_row_bytes) // 4, ( + "plain-text rows must keep the measured token count, not the size floor " + f"— got {text_usage.total_tokens} for a {len(text_row_bytes)}-byte row" + ) + + def _one_row_batch_bytes(model: str) -> bytes: import json as _json diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 38a346e7fb7..e2677363c5d 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -1115,6 +1115,56 @@ def test_reservation_uses_most_expensive_deployment_in_group(): assert estimated == pytest.approx(expected_expensive) +def test_estimate_input_tokens_reserves_max_for_file_content_blocks(): + """A `file` content block's real token cost is the provider's server-side + extraction of the referenced document — token_counter only sees the + filename/id, so counting it would under-reserve arbitrarily large uploads. + Reservation must fall back to the conservative max_input_tokens instead.""" + from litellm.proxy.spend_tracking.budget_reservation import ( + _estimate_input_tokens, + ) + + file_request = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this document."}, + {"type": "file", "file": {"file_id": "file-abc123"}}, + ], + } + ] + } + estimated = _estimate_input_tokens( + request_body=file_request, + route="/chat/completions", + model="gpt-4o", + model_info={"max_input_tokens": 128000}, + ) + assert estimated == 128000 + + # A stray text field on the same body must not be counted in place of the + # file — the messages branch owns the request once `messages` is present. + decoy_request = {**file_request, "prompt": "hi"} + decoy_estimated = _estimate_input_tokens( + request_body=decoy_request, + route="/chat/completions", + model="gpt-4o", + model_info={"max_input_tokens": 128000}, + ) + assert decoy_estimated == 128000 + + # Plain text messages must still be counted, not blanket-reserved. + text_request = {"messages": [{"role": "user", "content": "hi"}]} + counted = _estimate_input_tokens( + request_body=text_request, + route="/chat/completions", + model="gpt-4o", + model_info={"max_input_tokens": 128000}, + ) + assert counted is not None and 0 < counted < 128000 + + @pytest.mark.asyncio async def test_should_clamp_reservation_to_model_ceiling_when_caller_overrequests( spend_counter_state,