diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 46381eebc0c..adcb467e34d 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -42,6 +42,7 @@ from litellm.caching.redis_batch import ( ) from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, @@ -767,6 +768,20 @@ def _call_id_from_callback_kwargs(kwargs: object) -> str | None: return call_id if isinstance(call_id, str) else None +_RESPONSES_WS_TURN_EVENTS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def _unmodeled_response_usage(kwargs: Mapping[str, object], response_obj: object) -> Usage | None: + combined_usage: Final = kwargs.get("combined_usage_object") + if isinstance(combined_usage, Usage): + return combined_usage + if kwargs.get("call_type") != CallTypes.aresponses_websocket.value or not isinstance(response_obj, list): + return None + return ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( + _RESPONSES_WS_TURN_EVENTS.validate_python(response_obj) + ) + + def _parse_output_cap_value(raw_value: object) -> int | None: if isinstance(raw_value, bool) or not isinstance(raw_value, (int, float, str)): return None @@ -4793,7 +4808,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): response_usage: Final = self._resolve_io_token_reconcile_usage(response_obj) if response_usage[2]: return response_usage - return self._resolve_io_token_reconcile_usage(kwargs.get("combined_usage_object")) + return self._resolve_io_token_reconcile_usage(_unmodeled_response_usage(kwargs, response_obj)) def _build_websocket_turn_io_ops( self, @@ -5036,7 +5051,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Get total tokens from response. Responses LiteLLM does not model # (e.g. pass-through, whose usage is reported by the upstream rather # than parsed out of the body) carry their usage in - # ``combined_usage_object`` instead, and would otherwise never charge + # ``combined_usage_object`` instead, and a native Responses WebSocket + # turn in its raw event list; either would otherwise never charge # the TPM window. _usage: Usage | dict | None = None if isinstance( @@ -5050,9 +5066,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ): _usage = getattr(response_obj, "usage", None) else: - _combined_usage: Final = kwargs.get("combined_usage_object") - if isinstance(_combined_usage, Usage): - _usage = _combined_usage + _usage = _unmodeled_response_usage(kwargs, response_obj) total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type) if total_tokens == 0: total_tokens = self._aggregate_only_total_tokens(usage=_usage) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index faf7dd7283e..49d063aa0df 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -8063,6 +8063,45 @@ async def test_websocket_turn_failure_counts_its_usage_toward_tpm_v3(monkeypatch assert get_request_stash().reservation_released is False +@pytest.mark.asyncio +async def test_websocket_turn_success_counts_its_usage_toward_tpm_v3(monkeypatch): + """ + A successful native Responses WebSocket turn is logged with the raw event list + as its response, not a modeled response object. Its usage must still reach the + key's TPM window right away, while the socket stays open. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-ws-successful-turn") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") + reserved = await _admit_with_tpm_reservation( + handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value + ) + turn_kwargs = { + key: value + for key, value in _websocket_turn_kwargs("turn-1", _api_key, Usage()).items() + if key != "combined_usage_object" + } + + await handler.async_log_success_event( + kwargs={**turn_kwargs, "call_type": CallTypes.aresponses_websocket.value}, + response_obj=[ + {"type": "response.created", "response": {"id": "resp_turn_1"}}, + { + "type": "response.completed", + "response": {"id": "resp_turn_1", "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}}, + }, + ], + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved + 17 + assert get_request_stash().reservation_released is False + + @pytest.mark.asyncio async def test_failure_event_without_a_stash_charges_nothing_outside_a_websocket_connection_v3(monkeypatch): """ diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 51c8de77d0c..699b7e1cfb1 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -3768,13 +3768,14 @@ async def test_websocket_turn_usage_counts_toward_project_itpm_and_otpm(rate_lim assert await counters() == (stash.itpm_reserved_tokens, stash.otpm_reserved_tokens) await handler.async_log_success_event( - kwargs=turn_kwargs("turn-1"), - response_obj=ResponsesAPIResponse( - id="resp_1", - created_at=int(datetime.now().timestamp()), - output=[], - usage=ResponseAPIUsage(input_tokens=30, output_tokens=7, total_tokens=37), - ), + kwargs={**turn_kwargs("turn-1"), "call_type": CallTypes.aresponses_websocket.value}, + response_obj=[ + {"type": "response.created", "response": {"id": "resp_1"}}, + { + "type": "response.completed", + "response": {"id": "resp_1", "usage": {"input_tokens": 30, "output_tokens": 7, "total_tokens": 37}}, + }, + ], start_time=datetime.now(), end_time=datetime.now(), )