From 2fa5ac37254afe9ff2b7f0d9e34566fffb693197 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:46:45 +0000 Subject: [PATCH] fix(proxy): count Responses API tokens in the TPM limiter Classify the pre-call reservation by call type instead of the presence of `input`, honor `max_output_tokens`, and normalize `ResponseAPIUsage` into the chat `Usage` shape before reconciliation. --- .../hooks/parallel_request_limiter_v3.py | 102 ++++++++-- .../proxy/hooks/test_tpm_concurrent.py | 186 ++++++++++++++++++ 2 files changed, 273 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index f72492881b3..1027585888b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -43,8 +43,9 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( map_v3_rate_limit_type, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.responses.utils import ResponseAPILoggingUtils 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, @@ -290,6 +291,25 @@ 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 +# Call types whose requests generate no output tokens, so the upfront +# reservation covers input only. +_EMBEDDING_CALL_TYPES: frozenset[str] = frozenset( + { + CallTypes.embedding.value, + CallTypes.aembedding.value, + } +) +# Call types that cap generation with ``max_output_tokens`` rather than +# ``max_tokens``. +_RESPONSES_CALL_TYPES: frozenset[str] = frozenset( + { + CallTypes.responses.value, + CallTypes.aresponses.value, + } +) +# Every usage shape the TPM counters have to read: the chat ``Usage`` object, +# the Responses/realtime ``ResponseAPIUsage`` object, or either as a raw dict. +UsageLike = Usage | ResponseAPIUsage | dict[str, Any] | None # 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" @@ -478,18 +498,49 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return baseline return min(baseline, max(1, min_configured_tpm_limit // _TPM_FLOOR_FRACTION)) + @staticmethod + def _is_embedding_request(data: dict, call_type: str | None) -> bool: + """True for requests that generate no output tokens. + + ``input`` alone can't decide this: embeddings and the Responses API + both use it, and only the former has a zero output budget. Fall back + to the field heuristic (with the Responses-only ``max_output_tokens`` + as a tiebreaker) when the caller has no call type to offer. + """ + if call_type: + return call_type in _EMBEDDING_CALL_TYPES + return data.get("input") is not None and data.get("max_output_tokens") is None + + @staticmethod + def _max_output_tokens_field(call_type: str | None) -> str: + """Name of the field that caps generation for this call type.""" + if call_type in _RESPONSES_CALL_TYPES: + return "max_output_tokens" + return "max_tokens" + + @staticmethod + def _get_explicit_max_output_tokens(data: dict) -> int | None: + """Client-requested output cap, across chat, completions and Responses.""" + for field in ("max_tokens", "max_completion_tokens", "max_output_tokens"): + value = data.get(field) + if value: + return int(value) + return None + def _estimate_tokens_for_request( self, 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``). ``min_configured_tpm_limit`` is the smallest ``tokens_per_unit`` among the TPM-bearing descriptors this request will be charged against. When @@ -499,7 +550,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 and Responses API match (messages, prompt, input_text): case (messages, _, _) if messages: @@ -517,12 +568,12 @@ 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 = self._get_explicit_max_output_tokens(data) - match (explicit_max_tokens, input_text): + match (explicit_max_tokens, self._is_embedding_request(data=data, call_type=call_type)): 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: @@ -2490,12 +2541,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # 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 - ) - is_embedding = data.get("input") is not None + has_explicit_max_tokens = self._get_explicit_max_output_tokens(data) is not None + is_embedding = self._is_embedding_request(data=data, call_type=call_type) if capped_floor < baseline_floor and not has_explicit_max_tokens and not is_embedding: - data["max_tokens"] = capped_floor + data[self._max_output_tokens_field(call_type)] = capped_floor # Floor at 1 token so contentless requests (/responses, # tool-call continuations, empty messages) still flow @@ -2509,6 +2558,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 +2673,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return pipeline_operations + @staticmethod + def _normalize_usage(usage: UsageLike) -> UsageLike: + """Fold Responses-API usage into the chat ``Usage`` shape. + + ``/v1/responses`` (and realtime) report ``input_tokens`` / + ``output_tokens`` instead of ``prompt_tokens`` / ``completion_tokens``, + as a ``ResponseAPIUsage`` object or the equivalent dict. Left as-is + they read as zero tokens and refund the whole reservation. + """ + if not isinstance(usage, ResponseAPIUsage) and not ( + isinstance(usage, dict) and ResponseAPILoggingUtils._is_response_api_usage(usage) + ): + return usage + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + reported_total = usage.total_tokens if isinstance(usage, ResponseAPIUsage) else usage.get("total_tokens") + if not chat_usage.total_tokens and reported_total: + # The transform derives the total from the input/output split; keep + # an aggregate-only total so it can still charge the TPM window. + return chat_usage.model_copy(update={"total_tokens": int(reported_total)}) + return chat_usage + def _get_total_tokens_from_usage( self, usage: Optional[Any], rate_limit_type: Literal["output", "input", "total"] ) -> int: @@ -2635,6 +2706,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ total_tokens = 0 cached_tokens = 0 + usage = self._normalize_usage(usage) if usage: if isinstance(usage, Usage): @@ -2651,7 +2723,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 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": @@ -2671,8 +2742,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return total_tokens - @staticmethod - def _aggregate_only_total_tokens(usage: Union[Usage, dict, None]) -> int: + @classmethod + def _aggregate_only_total_tokens(cls, usage: UsageLike) -> 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,6 +2753,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): uncharged, which is how pass-through traffic slips past a TPM limit it is supposed to share. """ + usage = cls._normalize_usage(usage) if isinstance(usage, Usage): prompt_tokens, completion_tokens, total_tokens = ( usage.prompt_tokens or 0, @@ -3120,7 +3192,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: UsageLike = None if isinstance( response_obj, ( diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index b02f6c15168..8b7ed4a265b 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -23,6 +23,7 @@ import pytest from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + DEFAULT_MAX_TOKENS_ESTIMATE, RATE_LIMIT_DESCRIPTORS_KEY, TPM_RESERVATION_RELEASED_KEY, TPM_RESERVED_MODEL_KEY, @@ -31,6 +32,11 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, ) from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.types.llms.openai import ( + InputTokensDetails, + ResponseAPIUsage, + ResponsesAPIResponse, +) from litellm.types.utils import ModelResponse, Usage @@ -1192,5 +1198,185 @@ async def test_small_tpm_cap_preserves_explicit_max_tokens(rate_limiter): assert data["max_tokens"] == 500 +@pytest.mark.asyncio +async def test_responses_request_reserves_declared_output_budget(rate_limiter): + """ + A /v1/responses request carries its prompt in `input` and its output cap in + `max_output_tokens`. Classifying it as an embedding (any `input` present) + zeroed the output budget and reserved a single token, leaving Responses + traffic effectively unthrottled. + """ + handler, _cache = rate_limiter + + estimate = handler._estimate_tokens_for_request( + data={"model": "gpt-4.1", "input": "abcd" * 4, "max_output_tokens": 4096}, + call_type="aresponses", + ) + + assert estimate == 4 + 4096 + + +@pytest.mark.asyncio +async def test_responses_request_without_output_cap_reserves_floor(rate_limiter): + """A Responses request that omits max_output_tokens still gets the + conservative output floor, same as a chat request.""" + handler, _cache = rate_limiter + + estimate = handler._estimate_tokens_for_request( + data={"model": "gpt-4.1", "input": "abcd" * 4}, + call_type="aresponses", + ) + + assert estimate == 4 + DEFAULT_MAX_TOKENS_ESTIMATE // 4 + + +@pytest.mark.asyncio +async def test_embedding_call_type_reserves_input_only(rate_limiter): + """Embeddings really do have no output budget; call-type classification + must keep that behaviour.""" + handler, _cache = rate_limiter + + estimate = handler._estimate_tokens_for_request( + data={"model": "text-embedding-3-small", "input": "hello world"}, + call_type="aembedding", + ) + + assert estimate == 2 + + +@pytest.mark.asyncio +async def test_small_tpm_cap_caps_responses_output_field(rate_limiter): + """ + The small-TPM hard cap must bound generation with the field the Responses + API actually honours (`max_output_tokens`), and must apply to Responses + requests at all — the old `data.get("input") is not None` embedding test + skipped them. + """ + handler, cache = rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-responses-tpm-cap"), + project_id="proj-responses-tpm-cap", + project_metadata={ + "model_tpm_limit": {"gpt-4.1": 1000}, + }, + ) + + data: dict = {"model": "gpt-4.1", "input": "hello"} + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="aresponses", + ) + + assert data.get("max_output_tokens") == 1000 // 4 + assert "max_tokens" not in data + + +@pytest.mark.parametrize( + "rate_limit_type, expected", + [("input", 123), ("output", 456), ("total", 579)], +) +def test_response_api_usage_read_for_every_rate_limit_type( + rate_limiter, rate_limit_type, expected +): + """`ResponseAPIUsage` reports input_tokens/output_tokens, so the + prompt_tokens/completion_tokens reader saw zero usage for every mode.""" + handler, _cache = rate_limiter + + usage = ResponseAPIUsage(input_tokens=123, output_tokens=456, total_tokens=579) + + assert ( + handler._get_total_tokens_from_usage( + usage=usage, rate_limit_type=rate_limit_type + ) + == expected + ) + + +@pytest.mark.parametrize( + "rate_limit_type, expected", + [("input", 23), ("output", 456), ("total", 479)], +) +def test_response_api_usage_excludes_cached_input_tokens( + rate_limiter, rate_limit_type, expected +): + """Cached input tokens don't count toward provider TPM, so they must be + subtracted from input/total for Responses usage too.""" + handler, _cache = rate_limiter + + usage = ResponseAPIUsage( + input_tokens=123, + output_tokens=456, + total_tokens=579, + input_tokens_details=InputTokensDetails(cached_tokens=100), + ) + + assert ( + handler._get_total_tokens_from_usage( + usage=usage, rate_limit_type=rate_limit_type + ) + == expected + ) + + +@pytest.mark.asyncio +async def test_response_api_usage_reconciled_on_success(rate_limiter): + """ + End-to-end reconciliation for /v1/responses: with 100 tokens reserved and + 579 actually used, the reserved scope must be topped up by +479. Reading + `ResponseAPIUsage` as zero instead refunded the whole reservation, so + Responses traffic never charged the TPM window. + """ + handler, _cache = rate_limiter + + api_key = hash_token("sk-responses-reconcile") + + mock_kwargs = { + "standard_logging_object": { + "metadata": { + "user_api_key_hash": api_key, + TPM_RESERVED_TOKENS_KEY: 100, + TPM_RESERVED_SCOPES_KEY: [["api_key", api_key]], + } + }, + "model": "gpt-4.1", + } + + mock_response = ResponsesAPIResponse( + id="resp_test", + created_at=int(datetime.now().timestamp()), + model="gpt-4.1", + output=[], + usage=ResponseAPIUsage(input_tokens=123, output_tokens=456, total_tokens=579), + ) + + increments: list[dict[str, Any]] = [] + + async def mock_increment(increment_list, **kwargs): + for op in increment_list: + increments.append({"key": op["key"], "increment": op["increment_value"]}) + + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( + mock_increment + ) + + await handler.async_log_success_event( + kwargs=mock_kwargs, + response_obj=mock_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + token_adjustments = [i for i in increments if "tokens" in i["key"]] + + assert any(i["increment"] == 479 for i in token_adjustments), ( + f"Expected a +479 token adjustment (579 actual - 100 reserved) but got: " + f"{token_adjustments}" + ) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"])