From 81324aef21d41a52309daf3deae5aea0c62586f3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 26 Jul 2026 22:17:12 +0000 Subject: [PATCH] fix(proxy): count /v1/responses tokens in the TPM limiter (#34728) --- .../hooks/parallel_request_limiter_v3.py | 137 +++++++++--------- .../hooks/test_parallel_request_limiter_v3.py | 105 ++++++++++++++ 2 files changed, 177 insertions(+), 65 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index f72492881b3..42afb98b149 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -44,7 +44,10 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.types.caching import RedisPipelineIncrementOperation -from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject +from litellm.types.llms.openai import ( + BaseLiteLLMOpenAIResponseObject, + ResponseAPIUsage, +) from litellm.types.utils import ( CallTypes, EmbeddingResponse, @@ -483,13 +486,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data: dict, model: Optional[str] = None, min_configured_tpm_limit: Optional[int] = None, + call_type: str | None = None, ) -> int: """ Estimate total tokens this request will consume so we can reserve them upfront (input + output budget): estimated = input_tokens + max_tokens. - Supports chat (messages), completions (prompt), and embeddings (input). + Supports chat (messages), completions (prompt), embeddings (input), and + the Responses API (input + ``max_output_tokens``). Embeddings and + Responses both carry their prompt in ``input``, so ``call_type`` is used + to tell them apart: embeddings generate no output tokens, whereas a + Responses request budgets output like a chat completion. ``min_configured_tpm_limit`` is the smallest ``tokens_per_unit`` among the TPM-bearing descriptors this request will be charged against. When @@ -499,7 +507,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ messages = data.get("messages") prompt = data.get("prompt") - input_text = data.get("input") # embeddings + input_text = data.get("input") # embeddings / responses match (messages, prompt, input_text): case (messages, _, _) if messages: @@ -517,12 +525,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): estimated_input_tokens = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 - explicit_max_tokens = data.get("max_tokens") or data.get("max_completion_tokens") + explicit_max_tokens = ( + data.get("max_tokens") or data.get("max_completion_tokens") or data.get("max_output_tokens") + ) - match (explicit_max_tokens, input_text): + is_embedding = call_type in (CallTypes.embedding.value, CallTypes.aembedding.value) + + match (explicit_max_tokens, is_embedding): case (mt, _) if mt is not None: max_tokens_estimate = int(mt) - case (_, embeddings_input) if embeddings_input: + case (_, True): # Embeddings have no output tokens max_tokens_estimate = 0 case _ if total_chars == 0: @@ -2483,19 +2495,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit = min(configured_tpm_limits) # When the configured TPM cap is small enough to constrain the - # no-max_tokens floor, also hard-cap the model output via - # data["max_tokens"] so concurrent unbounded generations can't - # spend past the limit before post-call reconciliation runs. - # Skip when the request already sets max_tokens or has no + # no-max_tokens floor, also hard-cap the model output so + # concurrent unbounded generations can't spend past the limit + # before post-call reconciliation runs. The Responses API caps + # output via max_output_tokens rather than max_tokens. + # Skip when the request already sets an output budget or has no # generation budget at all (embeddings). capped_floor = self._no_max_tokens_output_floor(min_configured_tpm_limit) baseline_floor = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION has_explicit_max_tokens = ( - data.get("max_tokens") is not None or data.get("max_completion_tokens") is not None + data.get("max_tokens") is not None + or data.get("max_completion_tokens") is not None + or data.get("max_output_tokens") is not None ) - is_embedding = data.get("input") is not None + is_responses = call_type in (CallTypes.responses.value, CallTypes.aresponses.value) + is_embedding = call_type in (CallTypes.embedding.value, CallTypes.aembedding.value) if capped_floor < baseline_floor and not has_explicit_max_tokens and not is_embedding: - data["max_tokens"] = capped_floor + data["max_output_tokens" if is_responses else "max_tokens"] = capped_floor # Floor at 1 token so contentless requests (/responses, # tool-call continuations, empty messages) still flow @@ -2509,6 +2525,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data=data, model=requested_model, min_configured_tpm_limit=min_configured_tpm_limit, + call_type=call_type, ), 1, ) @@ -2623,6 +2640,35 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return pipeline_operations + @staticmethod + def _normalize_usage_counts(usage: Any) -> tuple[int, int, int, int]: + """Return ``(input, output, total, cached)`` for any supported usage shape. + + Handles the legacy chat ``Usage`` object (prompt / completion split), + the Responses API ``ResponseAPIUsage`` (input / output split), and the + dict form the Responses API sometimes carries. Unknown shapes read as + all zeros. ``cached`` is the cached input-token count providers like AWS + Bedrock don't charge against input / total rate limits. + """ + if isinstance(usage, Usage): + details = usage.prompt_tokens_details + cached = (getattr(details, "cached_tokens", 0) or 0) if details is not None else 0 + return (usage.prompt_tokens or 0, usage.completion_tokens or 0, usage.total_tokens or 0, cached) + if isinstance(usage, ResponseAPIUsage): + details = usage.input_tokens_details + cached = (details.cached_tokens or 0) if details is not None else 0 + return (usage.input_tokens or 0, usage.output_tokens or 0, usage.total_tokens or 0, cached) + if isinstance(usage, dict): + prompt_details = usage.get("prompt_tokens_details") or {} + cached = prompt_details.get("cached_tokens", 0) or 0 if isinstance(prompt_details, dict) else 0 + return ( + usage.get("prompt_tokens") or 0, + usage.get("completion_tokens") or 0, + usage.get("total_tokens") or 0, + cached, + ) + return (0, 0, 0, 0) + def _get_total_tokens_from_usage( self, usage: Optional[Any], rate_limit_type: Literal["output", "input", "total"] ) -> int: @@ -2633,46 +2679,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): because providers like AWS Bedrock don't count cached tokens toward rate limits. This aligns LiteLLM's TPM calculation with provider behavior. """ - total_tokens = 0 - cached_tokens = 0 + if not usage: + return 0 - if usage: - if isinstance(usage, Usage): - if rate_limit_type == "output": - total_tokens = usage.completion_tokens or 0 - elif rate_limit_type == "input": - total_tokens = usage.prompt_tokens or 0 - elif rate_limit_type == "total": - total_tokens = usage.total_tokens or 0 + input_tokens, output_tokens, total_tokens, cached_tokens = self._normalize_usage_counts(usage) + selected = {"input": input_tokens, "output": output_tokens, "total": total_tokens}[rate_limit_type] - # Get cached tokens to exclude from input/total - if rate_limit_type in ("input", "total"): - if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None: - cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 + # Providers don't count cached tokens toward input / total rate limits. + if rate_limit_type in ("input", "total"): + return max(0, selected - cached_tokens) + return selected - elif isinstance(usage, dict): - # Responses API usage comes as a dict - if rate_limit_type == "output": - total_tokens = usage.get("completion_tokens", 0) or 0 - elif rate_limit_type == "input": - total_tokens = usage.get("prompt_tokens", 0) or 0 - elif rate_limit_type == "total": - total_tokens = usage.get("total_tokens", 0) or 0 - - # Get cached tokens from dict - if rate_limit_type in ("input", "total"): - prompt_details = usage.get("prompt_tokens_details") or {} - if isinstance(prompt_details, dict): - cached_tokens = prompt_details.get("cached_tokens", 0) or 0 - - # Subtract cached tokens for input/total (providers don't count them) - if cached_tokens > 0: - total_tokens = max(0, total_tokens - cached_tokens) - - return total_tokens - - @staticmethod - def _aggregate_only_total_tokens(usage: Union[Usage, dict, None]) -> int: + def _aggregate_only_total_tokens(self, usage: Union[Usage, ResponseAPIUsage, dict, None]) -> int: """Total for usage that carries no input/output split, else 0. A source that can only report one number for the whole request (a @@ -2682,21 +2700,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): uncharged, which is how pass-through traffic slips past a TPM limit it is supposed to share. """ - if isinstance(usage, Usage): - prompt_tokens, completion_tokens, total_tokens = ( - usage.prompt_tokens or 0, - usage.completion_tokens or 0, - usage.total_tokens or 0, - ) - elif isinstance(usage, dict): - prompt_tokens, completion_tokens, total_tokens = ( - usage.get("prompt_tokens") or 0, - usage.get("completion_tokens") or 0, - usage.get("total_tokens") or 0, - ) - else: + if not usage: return 0 - if prompt_tokens or completion_tokens: + input_tokens, output_tokens, total_tokens, _ = self._normalize_usage_counts(usage) + if input_tokens or output_tokens: return 0 return total_tokens @@ -3120,7 +3127,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # than parsed out of the body) carry their usage in # ``combined_usage_object`` instead, and would otherwise never charge # the TPM window. - _usage: Union[Usage, dict, None] = None + _usage: Union[Usage, ResponseAPIUsage, dict, None] = None if isinstance( response_obj, ( diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 9337050b61c..62a3a20b6ec 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3009,6 +3009,111 @@ class TestGetTotalTokensFromUsageCacheExclusion: assert result == 0, f"Expected 0 for None usage, got {result}" +class TestResponsesAPITokenCounting: + """ + Regression tests for issue #34728. + + The TPM limiter undercounted /v1/responses in both phases: pre-call + estimation classified a Responses request (non-empty ``input``) like an + embedding, reserving zero output and ignoring ``max_output_tokens``; and + post-call reconciliation returned zero for a real ``ResponseAPIUsage`` + object because it is neither a legacy ``Usage`` nor a dict. Together they + could fully refund an already too-small reservation, under-enforcing TPM + limits for Responses traffic. + """ + + @pytest.fixture + def handler(self): + local_cache = DualCache() + return _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + + def test_estimate_responses_reserves_max_output_tokens(self, handler): + """input + max_output_tokens should be reserved, not treated as embeddings.""" + data = {"model": "gpt-4.1", "input": "hello", "max_output_tokens": 4096} + + estimate = handler._estimate_tokens_for_request(data=data, call_type="aresponses") + + # input("hello") ~ 1 token + 4096 output budget + assert estimate == 4097, f"Expected 4097 (1 input + 4096 output), got {estimate}" + + def test_estimate_responses_without_max_output_tokens_uses_output_floor(self, handler): + """A Responses request without an output budget still reserves an output floor.""" + data = {"model": "gpt-4.1", "input": "hello"} + + estimate = handler._estimate_tokens_for_request(data=data, call_type="responses") + + # Unlike embeddings, Responses generates output, so the floor applies. + # input(1) + floor(DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION = 1024) + assert estimate == 1025, f"Expected 1025 (1 input + 1024 output floor), got {estimate}" + + def test_estimate_embeddings_still_reserves_zero_output(self, handler): + """Embeddings carry ``input`` too but must keep a zero output budget.""" + data = {"model": "text-embedding-3-small", "input": "hello"} + + estimate = handler._estimate_tokens_for_request(data=data, call_type="aembedding") + + assert estimate == 1, f"Expected 1 (input only, no output), got {estimate}" + + def test_reconcile_response_api_usage_split(self, handler): + """A real ResponseAPIUsage object must reconcile input/output/total, not zero.""" + from litellm.types.llms.openai import ResponseAPIUsage + + usage = ResponseAPIUsage(input_tokens=123, output_tokens=456, total_tokens=579) + + assert handler._get_total_tokens_from_usage(usage, "input") == 123 + assert handler._get_total_tokens_from_usage(usage, "output") == 456 + assert handler._get_total_tokens_from_usage(usage, "total") == 579 + + def test_reconcile_response_api_usage_excludes_cached_tokens(self, handler): + """Cached input tokens are excluded from input/total but not output.""" + from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage + + usage = ResponseAPIUsage( + input_tokens=123, + output_tokens=456, + total_tokens=579, + input_tokens_details=InputTokensDetails(cached_tokens=23), + ) + + assert handler._get_total_tokens_from_usage(usage, "input") == 100 + assert handler._get_total_tokens_from_usage(usage, "total") == 556 + assert handler._get_total_tokens_from_usage(usage, "output") == 456 + + def test_aggregate_only_total_tokens_ignores_split_response_usage(self, handler): + """ResponseAPIUsage with an input/output split is not an aggregate-only source.""" + from litellm.types.llms.openai import ResponseAPIUsage + + usage = ResponseAPIUsage(input_tokens=123, output_tokens=456, total_tokens=579) + + assert handler._aggregate_only_total_tokens(usage) == 0 + + def test_success_event_reconciles_responses_usage_nonzero(self, handler): + """End-to-end: a ResponsesAPIResponse must charge the TPM window nonzero.""" + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + response_obj = ResponsesAPIResponse( + id="resp_123", + created_at=0, + model="gpt-4.1", + object="response", + output=[], + usage=ResponseAPIUsage(input_tokens=123, output_tokens=456, total_tokens=579), + ) + + ops = handler._build_success_event_pipeline_operations( + kwargs={}, + response_obj=response_obj, + rate_limit_type="total", + ) + + # No descriptors configured, so no ops, but the token count that would + # be charged must be nonzero. Assert on the reconciliation directly. + assert handler._get_total_tokens_from_usage(response_obj.usage, "total") == 579 + assert isinstance(ops, list) + + @pytest.mark.asyncio async def test_project_model_rate_limits_enforced_v3(): """