From a6dc447470f41baea1383b30fcee4c10cf6bdc04 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Sun, 16 Aug 2026 16:24:23 +0530 Subject: [PATCH] fix(proxy): stop Responses batch rows from bypassing project OTPM Embeddings rows were identified by body shape (has `input`, no `messages`/`prompt`), which also matches a `/v1/responses` batch row and reserved zero output tokens for it -- letting a project caller run large Responses generations against a quota-limited model without consuming OTPM. Classify embeddings by the row's own `url` instead, and read `max_output_tokens` as a Responses output cap alongside `max_tokens`/`max_completion_tokens`. Co-authored-by: Cursor --- litellm/proxy/hooks/batch_rate_limiter.py | 30 ++++++- .../proxy/hooks/test_batch_file_validation.py | 80 +++++++++++++++++++ 2 files changed, 106 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 3b53d792ce6..4052a9b04fa 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -395,18 +395,40 @@ class _PROXY_BatchRateLimiter(CustomLogger): Batch completion never reconciles actual usage back into the rate limiter, so this pre-call estimate is the only OTPM enforcement a batch gets. Mirrors the real-time no-``max_tokens`` floor so a row - that omits ``max_tokens`` can't be used to bypass OTPM the way an + that omits an output cap can't be used to bypass OTPM the way an unbounded streaming request could. + + Embeddings rows are identified by the row's own ``url`` (the OpenAI + batch schema puts the target route there, e.g. ``/v1/embeddings``), + never by body shape: a `/v1/responses` row also carries `body.input` + with no `messages`/`prompt`, so guessing from body shape alone would + misclassify a token-generating Responses row as a zero-output + embeddings row and let it skip the OTPM reservation entirely. """ + url: Final = entry.get("url") + if isinstance(url, str) and "embeddings" in url: + return 0 # embeddings: no output tokens raw_body: Final = entry.get("body") body: Final[Mapping[str, object]] = ( MappingProxyType(_BATCH_BODY_ADAPTER.validate_python(raw_body)) if isinstance(raw_body, Mapping) else MappingProxyType({}) # mutable-ok: immediately frozen empty fallback ) - if body.get("input") is not None and body.get("messages") is None and body.get("prompt") is None: - return 0 # embeddings: no output tokens - explicit_cap: Final = body.get("max_tokens", body.get("max_completion_tokens")) + # `max_tokens`/`max_completion_tokens` cap chat completions; `/v1/responses` + # rows cap output with `max_output_tokens` instead -- omitting it here + # would fall through to the floor estimate for every capped Responses row. + explicit_cap: Final = next( + ( + v + for v in ( + body.get("max_tokens"), + body.get("max_completion_tokens"), + body.get("max_output_tokens"), + ) + if v is not None + ), + None, + ) if explicit_cap is not None: try: return max(0, int(explicit_cap)) 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 ee37cf7d302..e5211973ec3 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -1980,3 +1980,83 @@ async def test_count_input_file_usage_collects_models_after_malformed_line(): ) assert exc.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# VERIA-Low regression: Responses batch rows must not bypass project OTPM +# --------------------------------------------------------------------------- + + +def _output_estimator(): + """A `_PROXY_BatchRateLimiter` whose output-token floor is observable: + the no-`max_tokens` floor mock returns a distinctive sentinel so tests can + tell "floor was used" apart from "an explicit cap was read".""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + limiter = MagicMock() + limiter.no_max_tokens_output_floor.return_value = 999 + return _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=limiter, + ) + + +def test_estimate_entry_output_tokens_zero_for_embeddings_url(): + """A real `/v1/embeddings` row reserves zero output tokens.""" + rate_limiter = _output_estimator() + entry = { + "url": "/v1/embeddings", + "body": {"model": "text-embedding-3-small", "input": "hello world"}, + } + + assert rate_limiter._estimate_entry_output_tokens(entry, None) == 0 + + +def test_estimate_entry_output_tokens_does_not_zero_responses_row_with_input(): + """Pre-fix: a `/v1/responses` row carries `body.input` with no `messages`/ + `prompt`, so the old body-shape heuristic misclassified it as embeddings + and reserved zero output tokens -- a project caller could submit large + Responses generations against a quota-limited model without consuming + OTPM. The row's own `url` (not body shape) must decide this.""" + rate_limiter = _output_estimator() + entry = { + "url": "/v1/responses", + "body": {"model": "gpt-4o", "input": "write me an essay"}, + } + + # No explicit cap on the row, so it must fall back to the no-max-tokens + # floor -- never straight to zero. + assert rate_limiter._estimate_entry_output_tokens(entry, None) == 999 + rate_limiter.parallel_request_limiter.no_max_tokens_output_floor.assert_called_once_with(None) + + +def test_estimate_entry_output_tokens_uses_max_output_tokens_for_responses(): + """`/v1/responses` caps output with `max_output_tokens`, not `max_tokens`/ + `max_completion_tokens`. Pre-fix this field was never inspected, so a + capped Responses row still fell through to the (possibly larger) floor + estimate instead of the caller's own declared cap.""" + rate_limiter = _output_estimator() + entry = { + "url": "/v1/responses", + "body": {"model": "gpt-4o", "input": "hi", "max_output_tokens": 123}, + } + + assert rate_limiter._estimate_entry_output_tokens(entry, None) == 123 + rate_limiter.parallel_request_limiter.no_max_tokens_output_floor.assert_not_called() + + +def test_estimate_entry_output_tokens_prefers_max_tokens_over_max_output_tokens(): + """When a row somehow carries both fields, the chat-style cap wins first -- + `max_output_tokens` is only consulted once the chat-style caps are absent.""" + rate_limiter = _output_estimator() + entry = { + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o", + "messages": [], + "max_tokens": 50, + "max_output_tokens": 500, + }, + } + + assert rate_limiter._estimate_entry_output_tokens(entry, None) == 50