From 2b23295f82d29cac7eb97474ea0d53b992c3f511 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Fri, 14 Aug 2026 23:31:18 +0530 Subject: [PATCH] fix(proxy): reconcile project quota reservations --- basedpyright-code-budget.json | 6 +- litellm/llms/custom_httpx/llm_http_handler.py | 26 +- litellm/proxy/hooks/batch_rate_limiter.py | 50 +- .../hooks/parallel_request_limiter_v3.py | 589 +++++++++--------- litellm/responses/streaming_iterator.py | 59 +- .../proxy/hooks/test_tpm_concurrent.py | 10 +- type-discipline-budget.json | 8 +- 7 files changed, 407 insertions(+), 341 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 521b4315e6e..67b0575a0cb 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5707 }, "reportMissingTypeArgument": { - "limit": 15642 + "limit": 15641 }, "reportMissingTypeStubs": { "limit": 40 @@ -108,7 +108,7 @@ "limit": 39237 }, "reportUnknownParameterType": { - "limit": 19969 + "limit": 19968 }, "reportUnknownVariableType": { "limit": 30881 @@ -132,7 +132,7 @@ "limit": 27 }, "reportUnusedClass": { - "limit": 23 + "limit": 22 }, "reportUnusedFunction": { "limit": 139 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e4c3a29b956..8627d1797a5 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2,7 +2,7 @@ import asyncio import json import os import ssl -from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from functools import lru_cache from types import ModuleType @@ -69,6 +69,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, + ProjectQuotaCallback, ResponsesAPIStreamingIterator, ResponsesWebSocketStreaming, SyncResponsesAPIStreamingIterator, @@ -252,7 +253,7 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool: return False -def _collect_ws_project_quota_callbacks() -> list: +def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]: """Duck-type discover proxy hooks exposing per-frame project ITPM/OTPM enforcement, so the Responses WebSocket loop can charge every ``response.create`` frame, not just the connection's first one. @@ -261,19 +262,16 @@ def _collect_ws_project_quota_callbacks() -> list: proxy hook directly) to avoid a layering violation (SDK importing from the proxy layer). """ - try: - import litellm as _litellm + import litellm as _litellm - return [ - cb for cb in _litellm.callbacks if callable(getattr(cb, "enforce_project_io_token_quota_for_frame", None)) - ] - except Exception as exc: # noqa: BLE001 - discovery must not block the connection - verbose_logger.warning( - "Responses WebSocket: failed to collect project quota callbacks — " - "per-frame ITPM/OTPM enforcement will be skipped. Error: %s", - exc, - ) - return [] + callbacks: Final = cast( # cast-ok: callback registry is inspected before protocol use + Sequence[object], _litellm.callbacks + ) + return tuple( + cast(ProjectQuotaCallback, callback) # cast-ok: required callback method is callable + for callback in callbacks + if callable(getattr(callback, "enforce_project_io_token_quota_for_frame", None)) + ) class BaseLLMHTTPHandler: diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 3091b0f6973..efef246a7a6 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -18,11 +18,12 @@ Quick summary: """ import json -from collections.abc import Iterable +from collections.abc import Iterable, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn from fastapi import HTTPException -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -77,6 +78,9 @@ else: RateLimitDescriptor = dict[str, Any] +_BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object]) + + class BatchFileUsage(BaseModel): """ Internal model for batch file usage tracking, used for batch rate limiting @@ -214,7 +218,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): tpm_limit_type=None, model_has_failures=False, ) - self.parallel_request_limiter._add_project_io_token_rate_limit_descriptors_from_metadata( + self.parallel_request_limiter.add_project_io_token_rate_limit_descriptors_from_metadata( user_api_key_dict=user_api_key_dict, requested_model=self._get_batch_routing_model(data), descriptors=descriptors, @@ -314,7 +318,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): def _estimate_entry_output_tokens( self, - entry: dict, + entry: Mapping[str, object], min_configured_otpm_limit: int | None, ) -> int: """Conservative per-row output-token estimate for the project OTPM reservation. @@ -325,16 +329,21 @@ class _PROXY_BatchRateLimiter(CustomLogger): that omits ``max_tokens`` can't be used to bypass OTPM the way an unbounded streaming request could. """ - body: Final = entry.get("body", {}) or {} + 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 = body.get("max_tokens", body.get("max_completion_tokens")) + explicit_cap: Final = body.get("max_tokens", body.get("max_completion_tokens")) if explicit_cap is not None: try: return max(0, int(explicit_cap)) except (TypeError, ValueError): pass - return self.parallel_request_limiter._no_max_tokens_output_floor(min_configured_otpm_limit) + return self.parallel_request_limiter.no_max_tokens_output_floor(min_configured_otpm_limit) @staticmethod def _has_applicable_batch_rate_limits( @@ -446,7 +455,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): f"Limit resets at: {reset_time_formatted}" ) else: # tokens - batch_token_count = ( + batch_token_count: Final = ( batch_usage.output_tokens if descriptor.get("key") == PROJECT_OTPM_DESCRIPTOR_KEY else batch_usage.total_tokens @@ -496,8 +505,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): data=data, ) - increments: Final[list[dict[Literal["requests", "tokens"], int]]] = [ - { + increments: Final = [ # mutable-ok: atomic limiter API requires mutable increment records + { # mutable-ok: atomic limiter API requires mutable increment records "requests": batch_usage.request_count, "tokens": batch_usage.output_tokens if d.get("key") == PROJECT_OTPM_DESCRIPTOR_KEY @@ -530,7 +539,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", user_api_key_dict: UserAPIKeyAuth | None = None, data: dict | None = None, - descriptors: list["RateLimitDescriptor"] | None = None, + descriptors: Sequence["RateLimitDescriptor"] | None = None, ) -> BatchFileUsage: """ Count number of requests and tokens in a batch input file. @@ -545,13 +554,14 @@ class _PROXY_BatchRateLimiter(CustomLogger): Returns: BatchFileUsage with total_tokens, output_tokens, and request_count """ - otpm_limits: Final = [ + otpm_limits: Final = tuple( int(v) - for d in (descriptors or []) + for d in (descriptors or ()) if d.get("key") == PROJECT_OTPM_DESCRIPTOR_KEY - for v in [(d.get("rate_limit") or {}).get("tokens_per_unit")] + for rate_limit in (d.get("rate_limit"),) + for v in (rate_limit.get("tokens_per_unit") if rate_limit is not None else None,) if v is not None - ] + ) min_configured_otpm_limit: Final = min(otpm_limits) if otpm_limits else None try: # Check if this is a managed file (base64 encoded unified file ID) @@ -605,7 +615,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): # the counter can't measure is estimated, not hard-rejected. models: Final[set] = set() total_tokens = 0 - output_tokens = 0 + output_tokens = 0 # rebind-ok: accumulated per JSONL row in the loop below request_count = 0 for raw_line in _iter_batch_input_lines(file_content_bytes): request_count += 1 @@ -613,9 +623,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): entry = json.loads(raw_line) except Exception: total_tokens += _estimate_batch_entry_tokens(raw_line) - output_tokens += self.parallel_request_limiter._no_max_tokens_output_floor( - min_configured_otpm_limit - ) + output_tokens += self.parallel_request_limiter.no_max_tokens_output_floor(min_configured_otpm_limit) continue if isinstance(entry, dict): model = (entry.get("body") or {}).get("model") @@ -623,9 +631,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): models.add(model) output_tokens += self._estimate_entry_output_tokens(entry, min_configured_otpm_limit) else: - output_tokens += self.parallel_request_limiter._no_max_tokens_output_floor( - min_configured_otpm_limit - ) + output_tokens += self.parallel_request_limiter.no_max_tokens_output_floor(min_configured_otpm_limit) try: total_tokens += _count_entry_tokens(entry) except Exception: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 7a97718247e..f858fd2af98 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -22,7 +22,7 @@ from typing import ( TypedDict, ) -from typing_extensions import NotRequired +from typing_extensions import NotRequired, ReadOnly from litellm import DualCache from litellm._logging import verbose_proxy_logger @@ -69,6 +69,7 @@ else: Span = Any InternalUsageCache = Any + BATCH_RATE_LIMITER_SCRIPT: Final = """ local results = {} local now = tonumber(ARGV[1]) @@ -413,7 +414,7 @@ class RateLimitStatus(TypedDict): class RateLimitResponse(TypedDict): overall_code: str statuses: list[RateLimitStatus] - reservation_windows: NotRequired[frozenset[tuple[str, str, Literal["redis", "local"]]]] + reservation_windows: NotRequired[ReadOnly[frozenset[tuple[str, str, Literal["redis", "local"]]]]] class ReservationAwareIncrementOperation(RedisPipelineIncrementOperation): @@ -452,6 +453,7 @@ class AtomicCounterMeta(TypedDict): class AtomicCounterState(TypedDict): window_expired: bool current: int + window_start: ReadOnly[str] DescriptorAtomicGroup: TypeAlias = tuple[list[str], list[int], list[AtomicCounterMeta]] @@ -639,7 +641,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return self._time_provider() @staticmethod - def _no_max_tokens_output_floor( + def no_max_tokens_output_floor( min_configured_tpm_limit: int | None, ) -> int: """Output-budget floor used when the request omits max_tokens. @@ -670,7 +672,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data: object, call_type: str | None, ) -> Mapping[str, object] | None: - contents = data.get("contents") if isinstance(data, dict) else None + contents: Final = data.get("contents") if isinstance(data, dict) else None if ( not isinstance(data, dict) or call_type not in GOOGLE_GENAI_NATIVE_CALL_TYPES @@ -679,7 +681,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return None from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter - config = data.get("config") if "config" in data else data.get("generationConfig") + config: Final = data.get("config") if "config" in data else data.get("generationConfig") return GoogleGenAIAdapter().translate_generate_content_to_completion( model=data.get("model") if isinstance(data.get("model"), str) else "", contents=contents, @@ -696,15 +698,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not isinstance(data, dict): return None if call_type in GOOGLE_GENAI_NATIVE_CALL_TYPES: - config = data.get("config") if "config" in data else data.get("generationConfig") - values = tuple( - int(config[field]) + config: Final = data.get("config") if "config" in data else data.get("generationConfig") + google_cap_values: Final = tuple( + int(raw_value) for field in ("maxOutputTokens", "max_output_tokens") - if isinstance(config, dict) and isinstance(config.get(field), (int, float, str)) + if isinstance(config, dict) + for raw_value in (config.get(field),) + if isinstance(raw_value, (int, float, str)) ) - return max(values, default=None) + return max(google_cap_values, default=None) if call_type in RESPONSES_API_CALL_TYPES: - value = data.get("max_output_tokens") + value: Final = data.get("max_output_tokens") if value is None: return None if not isinstance(value, (int, float, str)): @@ -712,13 +716,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return max(RESPONSES_API_MIN_OUTPUT_TOKENS, int(value)) if call_type in EMBEDDING_API_CALL_TYPES: return None - fields = ( + fields: Final = ( ("max_tokens", "max_completion_tokens") if call_type else ("max_tokens", "max_completion_tokens", "max_output_tokens") ) - values = tuple(int(data[field]) for field in fields if isinstance(data.get(field), (int, float, str))) - return max(values, default=None) + output_cap_values: Final = tuple( + int(raw_value) + for field in fields + for raw_value in (data.get(field),) + if isinstance(raw_value, (int, float, str)) + ) + return max(output_cap_values, default=None) @classmethod def _has_explicit_output_cap(cls, data: object, call_type: str | None) -> bool: @@ -733,18 +742,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _get_output_candidate_count(data: object, call_type: str | None = None) -> int: if not isinstance(data, dict): return 1 - config = ( + config: Final = ( (data.get("config") if "config" in data else data.get("generationConfig")) if call_type in GOOGLE_GENAI_NATIVE_CALL_TYPES else None ) - candidate_values = ( + candidate_values: Final = ( data.get("n"), data.get("best_of"), config.get("candidateCount") if isinstance(config, dict) else None, config.get("candidate_count") if isinstance(config, dict) else None, ) - candidate_count = 1 + candidate_count = 1 # rebind-ok: running maximum across candidate-count aliases for value in candidate_values: try: candidate_count = max(candidate_count, int(value or 1)) @@ -774,31 +783,34 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(data, dict): return - capped_floor = _PROXY_MaxParallelRequestsHandler_v3._no_max_tokens_output_floor(min_configured_limit) - if call_type in RESPONSES_API_CALL_TYPES: - capped_floor = max(capped_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) - baseline_floor = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION - is_embedding = _PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type) + base_capped_floor: Final = _PROXY_MaxParallelRequestsHandler_v3.no_max_tokens_output_floor(min_configured_limit) + capped_floor: Final = ( + max(base_capped_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) + if call_type in RESPONSES_API_CALL_TYPES + else base_capped_floor + ) + baseline_floor: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION + is_embedding: Final = _PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type) if ( capped_floor >= baseline_floor or _PROXY_MaxParallelRequestsHandler_v3._has_explicit_output_cap(data, call_type) or is_embedding ): return - effective_cap = max(capped_floor, configured_output_tokens or 0) + effective_cap: Final = max(capped_floor, configured_output_tokens or 0) if call_type in GOOGLE_GENAI_NATIVE_CALL_TYPES: - config_field = "config" if "config" in data or "generationConfig" not in data else "generationConfig" - config = data.get(config_field) + config_field: Final = "config" if "config" in data or "generationConfig" not in data else "generationConfig" + config: Final = data.get(config_field) if config is None or isinstance(config, dict): - data[config_field] = { # mutable-ok: downstream native routing requires a mutable request config + data[config_field] = { # rebind-ok: routed request needs cap # mutable-ok: downstream needs dict **(config or {}), # mutable-ok: downstream native routing requires a mutable request config "maxOutputTokens": effective_cap, } return - cap_field = "max_output_tokens" if call_type in RESPONSES_API_CALL_TYPES else "max_tokens" - existing_cap = data.get(cap_field) + cap_field: Final = "max_output_tokens" if call_type in RESPONSES_API_CALL_TYPES else "max_tokens" + existing_cap: Final = data.get(cap_field) if existing_cap is None or effective_cap < existing_cap: - data[cap_field] = effective_cap + data[cap_field] = effective_cap # rebind-ok: downstream routing requires the bounded output cap def _estimate_tokens_for_request( self, @@ -878,69 +890,57 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return 0, 0 translated_data: Final = self._translate_google_genai_native_request(data, call_type) estimable_data: Final = translated_data if translated_data is not None else data - messages = estimable_data.get("messages") - prompt = estimable_data.get("prompt") - input_text = estimable_data.get("input") + selected_fields: Final[tuple[object | None, object | None, object | None]] = ( + (None, None, estimable_data.get("input")) + if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES + else (None, estimable_data.get("prompt"), None) + if call_type in TEXT_COMPLETION_API_CALL_TYPES + else (estimable_data.get("messages"), None, None) + if call_type + else ( + estimable_data.get("messages"), + estimable_data.get("prompt"), + estimable_data.get("input"), + ) + ) + messages, prompt, input_text = selected_fields - if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES: - messages = None - prompt = None - elif call_type in TEXT_COMPLETION_API_CALL_TYPES: - messages = None - input_text = None - elif call_type: - prompt = None - input_text = None - - match (messages, prompt, input_text): - case (selected_messages, _, _) if selected_messages: - total_chars = len(get_str_from_messages(selected_messages)) - case (_, str() as selected_prompt, _): - total_chars = len(selected_prompt) - case (_, list() as selected_prompt, _): - total_chars = sum(len(str(item)) for item in selected_prompt) - case (_, _, str() as selected_input): - total_chars = len(selected_input) - case (_, _, list() as selected_input): - total_chars = sum(len(str(item)) for item in selected_input) - case _: - total_chars = 0 + total_chars: Final = ( + len(get_str_from_messages(messages)) + if isinstance(messages, list) and messages + else len(prompt) + if isinstance(prompt, str) + else sum(len(str(item)) for item in prompt) + if isinstance(prompt, list) + else len(input_text) + if isinstance(input_text, str) + else sum(len(str(item)) for item in input_text) + if isinstance(input_text, list) + else 0 + ) estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type) is_embedding: Final = self._is_embedding_request(data, call_type) - match (explicit_max_tokens, is_embedding): - case (_, True): - max_tokens_estimate = 0 - case (mt, _) if mt is not None: - max_tokens_estimate = mt - case _ if total_chars == 0 and configured_output_tokens is None: - # Fully contentless request (no messages, prompt, or input). - # Don't apply the conservative output-budget floor here — it - # would over-reserve and could push small TPM limits into a - # false 429. The caller floors at 1 so backpressure still - # applies once the counter is at limit. - max_tokens_estimate = 0 - case _: - # No max_tokens specified — reserve at least the input size with a - # conservative floor so a stream of small concurrent requests can't - # collectively bypass the limit. Cap the floor by a fraction of - # the smallest TPM limit this request will be charged against, - # so a small per-tenant TPM cap can't be tripped by the floor - # alone. - output_floor = self._no_max_tokens_output_floor(min_configured_tpm_limit) - if call_type in RESPONSES_API_CALL_TYPES: - output_floor = max(output_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) - max_tokens_estimate = ( - configured_output_tokens - if configured_output_tokens is not None - else max(estimated_input_tokens, output_floor) - ) + base_output_floor: Final = self.no_max_tokens_output_floor(min_configured_tpm_limit) + output_floor: Final = ( + max(base_output_floor, RESPONSES_API_MIN_OUTPUT_TOKENS) + if call_type in RESPONSES_API_CALL_TYPES + else base_output_floor + ) + max_tokens_estimate: Final = ( + 0 + if is_embedding or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None) + else explicit_max_tokens + if explicit_max_tokens is not None + else configured_output_tokens + if configured_output_tokens is not None + else max(estimated_input_tokens, output_floor) + ) - max_tokens_estimate *= self._get_output_candidate_count(data, call_type) - return estimated_input_tokens, max_tokens_estimate + return estimated_input_tokens, max_tokens_estimate * self._get_output_candidate_count(data, call_type) def _is_redis_cluster(self) -> bool: """ @@ -1755,7 +1755,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): mid-loop, refund applied increments and fall back to in-memory. """ if not descriptor_groups: - return RateLimitResponse(overall_code="OK", statuses=[]) + return RateLimitResponse( + overall_code="OK", + statuses=[], # mutable-ok: response contract requires a status list + ) applied: Final[list[list[AtomicCounterMeta]]] = [] statuses: Final[list[RateLimitStatus]] = [] raw: list[CacheCounterValue] @@ -1946,7 +1949,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ], ) descriptor_state.append( - { + { # mutable-ok: local atomic-counter state is updated during pass two "window_expired": window_expired, "current": current_counter, "window_start": str(now_int if window_expired else int(window_start)), @@ -2096,21 +2099,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): configured, or if the reservation failed), for the caller to stash for post-call reconciliation. """ - itpm_descriptors = [ # mutable-ok: atomic limiter API requires lists + itpm_descriptors: Final = [ # mutable-ok: atomic limiter API requires lists d for d in descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY ] - otpm_descriptors = [ # mutable-ok: atomic limiter API requires lists + otpm_descriptors: Final = [ # mutable-ok: atomic limiter API requires lists d for d in descriptors if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY ] if not itpm_descriptors and not otpm_descriptors: return RateLimitResponse(overall_code="OK", statuses=[]), 0, 0 # mutable-ok: response contract uses a list - itpm_response: RateLimitResponse | None = None - itpm_reserved = 0 - - if itpm_descriptors: - itpm_response = await self.atomic_check_and_increment_by_n( + itpm_response: Final = ( + await self.atomic_check_and_increment_by_n( descriptors=itpm_descriptors, increments=[ # mutable-ok: atomic limiter API requires mutable increment records {"tokens": estimated_input_tokens} # mutable-ok: atomic limiter increment record @@ -2118,12 +2118,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ], parent_otel_span=parent_otel_span, ) - if itpm_response["overall_code"] == "OVER_LIMIT": - return itpm_response, 0, 0 - itpm_reserved = estimated_input_tokens + if itpm_descriptors + else None + ) + if itpm_response is not None and itpm_response["overall_code"] == "OVER_LIMIT": + return itpm_response, 0, 0 + itpm_reserved: Final = estimated_input_tokens if itpm_response is not None else 0 if otpm_descriptors: - otpm_response = await self.atomic_check_and_increment_by_n( + otpm_response: Final = await self.atomic_check_and_increment_by_n( descriptors=otpm_descriptors, increments=[ # mutable-ok: atomic limiter API requires mutable increment records {"tokens": estimated_output_tokens} # mutable-ok: atomic limiter increment record @@ -2142,7 +2145,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=parent_otel_span, ) return otpm_response, 0, 0 - statuses = ( + statuses: Final = ( [ # mutable-ok: response contract uses a list *itpm_response["statuses"], *otpm_response["statuses"], @@ -2172,7 +2175,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def enforce_project_io_token_quota_for_frame( self, - user_api_key_dict: UserAPIKeyAuth, + user_api_key_dict: UserAPIKeyAuth | None, requested_model: str | None, estimated_input_tokens: int, estimated_output_tokens: int, @@ -2188,8 +2191,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): like the batch rate limiter -- this charges the estimate immediately and never refunds it. """ - descriptors: Final[list[RateLimitDescriptor]] = [] - self._add_project_io_token_rate_limit_descriptors_from_metadata( + if user_api_key_dict is None: + return + descriptors: Final[list[RateLimitDescriptor]] = [] # mutable-ok: descriptor helper appends in place + self.add_project_io_token_rate_limit_descriptors_from_metadata( user_api_key_dict=user_api_key_dict, requested_model=requested_model, descriptors=descriptors, @@ -2911,7 +2916,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - def _add_project_io_token_rate_limit_descriptors_from_metadata( + def add_project_io_token_rate_limit_descriptors_from_metadata( self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None, @@ -2926,22 +2931,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if requested_model is None or user_api_key_dict.project_id is None: return - itpm_limit_for_project_model = ( + itpm_limit_for_project_model: Final = ( get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_itpm_limit") or {} # mutable-ok: metadata helper returns an optional mapping ) - otpm_limit_for_project_model = ( + otpm_limit_for_project_model: Final = ( get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_otpm_limit") or {} # mutable-ok: metadata helper returns an optional mapping ) - model_itpm_limit = itpm_limit_for_project_model.get(requested_model) - model_otpm_limit = otpm_limit_for_project_model.get(requested_model) + model_itpm_limit: Final = itpm_limit_for_project_model.get(requested_model) + model_otpm_limit: Final = otpm_limit_for_project_model.get(requested_model) if model_itpm_limit is None and model_otpm_limit is None: return - descriptor_value = f"{user_api_key_dict.project_id}:{requested_model}" + descriptor_value: Final = f"{user_api_key_dict.project_id}:{requested_model}" if model_itpm_limit is not None: descriptors.append( RateLimitDescriptor( @@ -3026,10 +3031,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(block, dict): return DEFAULT_AUDIO_TOKEN_ESTIMATE - input_audio = block.get("input_audio") - b64_data = input_audio.get("data") if isinstance(input_audio, dict) else None + input_audio: Final = block.get("input_audio") + b64_data: Final = input_audio.get("data") if isinstance(input_audio, dict) else None if b64_data and isinstance(b64_data, str): - decoded_bytes = len(b64_data) * 3 // 4 + decoded_bytes: Final = len(b64_data) * 3 // 4 return max(decoded_bytes // _AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE) return DEFAULT_AUDIO_TOKEN_ESTIMATE @@ -3042,17 +3047,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(messages, list): return 0 - total = 0 - for message in messages: - content = message.get("content") if isinstance(message, dict) else None - if not isinstance(content, list): - continue - total += sum( - cls._estimate_audio_block_tokens(block) - for block in content - if isinstance(block, dict) and block.get("type") == "input_audio" - ) - return total + return sum( + cls._estimate_audio_block_tokens(block) + for message in messages + if isinstance(message, dict) + for content in (message.get("content"),) + if isinstance(content, list) + for block in content + if isinstance(block, dict) and block.get("type") == "input_audio" + ) @staticmethod def _strip_audio_content_blocks(messages: object) -> object: @@ -3066,7 +3069,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(messages, list): return messages - sanitized = [] # mutable-ok: token_counter requires a list of message dicts + sanitized: Final[list[object]] = [] # mutable-ok: token_counter requires a list of message dicts for message in messages: if not isinstance(message, dict): sanitized.append(message) @@ -3127,7 +3130,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): @classmethod def _contains_image_content(cls, value: object) -> bool: if isinstance(value, dict): - media_type = value.get("media_type") or value.get("mime_type") + media_type: Final = value.get("media_type") or value.get("mime_type") return ( value.get("type") in ("image", "image_url", "input_image") or (isinstance(media_type, str) and media_type.startswith("image/")) @@ -3159,9 +3162,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): @staticmethod def _rerank_input_to_text(data: Mapping[str, object]) -> str: - documents = data.get("documents") - document_items: Sequence[object] = documents if isinstance(documents, list) else () # pyright: ignore[reportUnknownVariableType] # rerank documents are validated runtime JSON - input_parts: tuple[object, ...] = ( # pyright: ignore[reportUnknownVariableType] # list narrowing preserves unknown JSON element types + documents: Final = data.get("documents") + document_items: Final[Sequence[object]] = documents if isinstance(documents, list) else () # pyright: ignore[reportUnknownVariableType] # rerank documents are validated runtime JSON + input_parts: Final[tuple[object, ...]] = ( # pyright: ignore[reportUnknownVariableType] # list narrowing preserves unknown JSON element types data.get("query"), *document_items, ) @@ -3199,39 +3202,45 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not isinstance(data, dict): return 0 - selected_text = None - countable_tools = data.get("tools") - countable_tool_choice = data.get("tool_choice") - if call_type in RESPONSES_API_CALL_TYPES: - messages = self._responses_input_to_chat_messages(data) - elif (translated_request := self._translate_google_genai_native_request(data, call_type)) is not None: - messages = translated_request.get("messages") - countable_tools = translated_request.get("tools") - countable_tool_choice = translated_request.get("tool_choice") - elif self._is_embedding_request(data, call_type): - messages = None - selected_text = data.get("input") - pretokenized_input_tokens = self._count_pretokenized_embedding_input(selected_text) - if pretokenized_input_tokens is not None: - return pretokenized_input_tokens - elif call_type in RERANK_API_CALL_TYPES: - messages = None - selected_text = self._rerank_input_to_text(data) # pyright: ignore[reportUnknownArgumentType] # proxy request bodies are runtime-validated JSON - elif call_type in TEXT_COMPLETION_API_CALL_TYPES: - messages = None - selected_text = data.get("prompt") - else: - messages = data.get("messages") - if messages is None: - selected_text = data.get("prompt") - if messages is None and selected_text is None: - selected_text = data.get("input") + is_responses_request: Final = call_type in RESPONSES_API_CALL_TYPES + translated_request: Final = ( + None if is_responses_request else self._translate_google_genai_native_request(data, call_type) + ) + is_embedding_request: Final = self._is_embedding_request(data, call_type) + embedding_text: Final = data.get("input") if is_embedding_request else None + pretokenized_input_tokens: Final = ( + self._count_pretokenized_embedding_input(embedding_text) if is_embedding_request else None + ) + if pretokenized_input_tokens is not None: + return pretokenized_input_tokens - audio_token_estimate = self._estimate_audio_content_tokens(messages) - countable_messages = self._strip_audio_content_blocks(messages) if audio_token_estimate > 0 else messages + prompt: Final = data.get("prompt") + fallback_text: Final = prompt if prompt is not None else data.get("input") + selected_inputs: Final[tuple[object | None, object | None, object | None, object | None]] = ( + (self._responses_input_to_chat_messages(data), None, data.get("tools"), data.get("tool_choice")) + if is_responses_request + else ( + translated_request.get("messages"), + None, + translated_request.get("tools"), + translated_request.get("tool_choice"), + ) + if translated_request is not None + else (None, embedding_text, data.get("tools"), data.get("tool_choice")) + if is_embedding_request + else (None, self._rerank_input_to_text(data), data.get("tools"), data.get("tool_choice")) + if call_type in RERANK_API_CALL_TYPES + else (None, prompt, data.get("tools"), data.get("tool_choice")) + if call_type in TEXT_COMPLETION_API_CALL_TYPES + else (data.get("messages"), fallback_text, data.get("tools"), data.get("tool_choice")) + ) + messages, selected_text, countable_tools, countable_tool_choice = selected_inputs + + audio_token_estimate: Final = self._estimate_audio_content_tokens(messages) + countable_messages: Final = self._strip_audio_content_blocks(messages) if audio_token_estimate > 0 else messages try: - estimate = max( + estimate: Final = max( 0, int( token_counter( @@ -3245,7 +3254,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ), ) return estimate + audio_token_estimate - except Exception: # noqa: BLE001 - any tokenizer/model-resolution/transform failure degrades to the cheap estimate, never a 500 + except Exception: # noqa: BLE001 # tokenizer failures degrade to the cheap estimate if call_type in RERANK_API_CALL_TYPES and isinstance(selected_text, str): return max(0, len(selected_text) // DEFAULT_CHARS_PER_TOKEN) estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data, call_type=call_type) @@ -3272,14 +3281,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(data, dict): return - stash = claim_request_stash_for_data(data) - io_token_descriptors = [ # mutable-ok: reservation API requires descriptor lists + stash: Final = claim_request_stash_for_data(data) + io_token_descriptors: Final = [ # mutable-ok: reservation API requires descriptor lists d for d in descriptors if d["key"] in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) ] if not io_token_descriptors: return - configured_otpm_limits = [ # mutable-ok: min calculation materializes validated limits + configured_otpm_limits: Final = [ # mutable-ok: min calculation materializes validated limits int(v) for d in io_token_descriptors if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY @@ -3290,8 +3299,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ] if v is not None ] - min_configured_otpm_limit = min(configured_otpm_limits) if configured_otpm_limits else None - configured_itpm_limits = [ # mutable-ok: min calculation materializes validated limits + min_configured_otpm_limit: Final = min(configured_otpm_limits) if configured_otpm_limits else None + configured_itpm_limits: Final = [ # mutable-ok: min calculation materializes validated limits int(v) for d in io_token_descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY @@ -3302,14 +3311,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ] if v is not None ] - min_configured_itpm_limit = min(configured_itpm_limits) if configured_itpm_limits else None + min_configured_itpm_limit: Final = min(configured_itpm_limits) if configured_itpm_limits else None - _, estimated_output_tokens = self._estimate_input_and_output_tokens( + _, raw_estimated_output_tokens = self._estimate_input_and_output_tokens( data=data, min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, ) - estimated_input_tokens = ( + raw_estimated_input_tokens: Final = ( min_configured_itpm_limit if min_configured_itpm_limit is not None and ( @@ -3319,9 +3328,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) else self._estimate_precise_input_tokens(data=data, model=requested_model, call_type=call_type) ) - estimated_input_tokens = max(estimated_input_tokens, 1) - if not self._has_explicit_output_cap(data, call_type): - estimated_output_tokens = max(estimated_output_tokens, 1) + estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1) + estimated_output_tokens: Final = ( + raw_estimated_output_tokens + if self._has_explicit_output_cap(data, call_type) + else max(raw_estimated_output_tokens, 1) + ) # Hard-cap generation length so an unbounded response can't overshoot # the OTPM budget before post-call reconciliation runs, mirroring the @@ -3353,7 +3365,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.reservation_released = True - acquisition = stash.parallel_slot + acquisition: Final = stash.parallel_slot if acquisition is not None: await self._release_parallel_request_slots( acquisition=acquisition, @@ -3367,9 +3379,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if itpm_reserved > 0: - itpm_scopes = [ # mutable-ok: request stash freezes the collected scopes + itpm_scopes: Final = tuple( (d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY - ] + ) stash.itpm_reserved_tokens = itpm_reserved stash.itpm_reserved_scopes = frozenset(itpm_scopes) stash.itpm_reserved_window_identities = frozenset( @@ -3378,9 +3390,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if "model_per_project_itpm" in counter_key ) if otpm_reserved > 0: - otpm_scopes = [ # mutable-ok: request stash freezes the collected scopes + otpm_scopes: Final = tuple( (d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY - ] + ) stash.otpm_reserved_tokens = otpm_reserved stash.otpm_reserved_scopes = frozenset(otpm_scopes) stash.otpm_reserved_window_identities = frozenset( @@ -3473,7 +3485,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, descriptors=descriptors, ) - self._add_project_io_token_rate_limit_descriptors_from_metadata( + self.add_project_io_token_rate_limit_descriptors_from_metadata( user_api_key_dict=user_api_key_dict, requested_model=requested_model, descriptors=descriptors, @@ -3546,8 +3558,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # limit. Stays empty/0 whenever no combined-TPM reservation was # made (or it was over limit, in which case execution never # reaches the ITPM/OTPM block -- `_handle_rate_limit_error` raises). - tpm_reservation_scopes: Sequence[tuple[str, str]] = () - tpm_reservation_amount = 0 + tpm_reservation_scopes: Sequence[tuple[str, str]] = () # rebind-ok: set after successful reservation + tpm_reservation_amount = 0 # rebind-ok: set after successful reservation if has_tpm_limits and self.tpm_reservation_enabled: min_configured_tpm_limit: Final = min(configured_tpm_limits) @@ -3633,8 +3645,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) is not None ) - tpm_reservation_scopes = tuple(stash.reserved_scopes) - tpm_reservation_amount = estimated_tokens + tpm_reservation_scopes = tuple( # rebind-ok: record successful reservation scopes + stash.reserved_scopes + ) + tpm_reservation_amount = estimated_tokens # rebind-ok: record successful reservation amount # Merge TPM statuses into the stored rate-limit response # so x-ratelimit-{key}-remaining-tokens / -limit-tokens @@ -3648,7 +3662,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug( "TPM tokens reserved: %s for model %s", estimated_tokens, requested_model ) - await self._reserve_project_io_tokens_or_raise( descriptors=descriptors, data=data, @@ -3734,7 +3747,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return total_tokens @staticmethod - def _aggregate_only_total_tokens(usage: Usage | dict | None) -> int: + def _aggregate_only_total_tokens(usage: Usage | ResponseAPIUsage | Mapping[str, object] | 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 @@ -3744,24 +3757,43 @@ 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, + if usage is None: + return 0 + token_counts: Final = ( + (usage.prompt_tokens or 0, usage.completion_tokens or 0, usage.total_tokens or 0) + if isinstance(usage, Usage) + else (usage.input_tokens or 0, usage.output_tokens or 0, usage.total_tokens or 0) + if isinstance(usage, ResponseAPIUsage) + else ( + usage.get("prompt_tokens") or usage.get("input_tokens") or 0, + usage.get("completion_tokens") or usage.get("output_tokens") or 0, usage.get("total_tokens") or 0, ) - else: - return 0 - if prompt_tokens or completion_tokens: + ) + prompt_tokens, completion_tokens, total_tokens = token_counts + if prompt_tokens or completion_tokens or not isinstance(total_tokens, int): return 0 return total_tokens + @staticmethod + def _response_usage( + response_obj: object, + ) -> Usage | ResponseAPIUsage | Mapping[str, object] | None: + if isinstance(response_obj, (Usage, ResponseAPIUsage)): + return response_obj + if isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse, BaseLiteLLMOpenAIResponseObject), + ): + usage: Final = getattr(response_obj, "usage", None) + return usage if isinstance(usage, (Usage, ResponseAPIUsage, dict)) else None + if isinstance(response_obj, dict): + nested_usage: Final = response_obj.get("usage") + if isinstance(nested_usage, (Usage, ResponseAPIUsage, dict)): + return nested_usage + return response_obj + return None + async def _execute_token_increment_script( self, pipeline_operations: list["RedisPipelineIncrementOperation"], @@ -3884,15 +3916,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if self.window_guarded_token_increment_script is not None: try: await self.window_guarded_token_increment_script( - keys=[window_key, operation["key"]], - args=[ + keys=[ # mutable-ok: Redis script interface requires a key list + window_key, + operation["key"], + ], + args=[ # mutable-ok: Redis script interface requires an argument list expected_window_start, operation["increment_value"], operation["ttl"] or 0, ], ) continue - except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the plain increment fallback, never a 500 + except Exception as e: # noqa: BLE001 # Redis failures use the plain increment fallback verbose_proxy_logger.warning( "Window-guarded token adjustment failed for %s: %s", operation["key"], @@ -3978,16 +4013,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not isinstance(response_obj, RerankResponse) or response_obj.meta is None: return None - rerank_tokens = response_obj.meta.get("tokens") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads + rerank_tokens: Final = response_obj.meta.get("tokens") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads if rerank_tokens is not None: - input_tokens = rerank_tokens.get("input_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload - output_tokens = rerank_tokens.get("output_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload + input_tokens: Final = rerank_tokens.get("input_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload + output_tokens: Final = rerank_tokens.get("output_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # token fields are typed integers despite the generic get overload if input_tokens or output_tokens: return max(0, input_tokens), max(0, output_tokens), True - billed_units = response_obj.meta.get("billed_units") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads + billed_units: Final = response_obj.meta.get("billed_units") # pyright: ignore[reportUnknownMemberType] # TypedDict's optional generic metadata widens get overloads if billed_units is not None: - total_tokens = billed_units.get("total_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # billed total is a typed integer despite the generic get overload + total_tokens: Final = billed_units.get("total_tokens") or 0 # pyright: ignore[reportUnknownMemberType] # billed total is a typed integer despite the generic get overload if total_tokens: return max(0, total_tokens), 0, True return None @@ -4003,68 +4038,57 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): but they're untouched everywhere else (cost/usage logging still sees the full prompt token count). """ - rerank_usage = self._resolve_rerank_token_usage(response_obj) + rerank_usage: Final = self._resolve_rerank_token_usage(response_obj) if rerank_usage is not None: return rerank_usage - usage: object | None = None - if isinstance(response_obj, (Usage, ResponseAPIUsage)): - usage = response_obj - elif isinstance( - response_obj, - (ModelResponse, EmbeddingResponse, TextCompletionResponse, BaseLiteLLMOpenAIResponseObject), - ): - usage = getattr(response_obj, "usage", None) - elif isinstance(response_obj, dict): - usage = response_obj.get("usage") - if usage is None and any( - key in response_obj - for key in ( - "prompt_tokens", - "completion_tokens", - "input_tokens", - "output_tokens", - ) - ): - usage = response_obj + usage: Final = self._response_usage(response_obj) if isinstance(usage, Usage): - prompt_tokens = usage.prompt_tokens or 0 - completion_tokens = usage.completion_tokens or 0 - cached_tokens = 0 - if usage.prompt_tokens_details is not None: - cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 - elif isinstance(usage, ResponseAPIUsage): - # Responses API usage uses input_tokens/output_tokens instead of - # prompt_tokens/completion_tokens. - prompt_tokens = usage.input_tokens or 0 - completion_tokens = usage.output_tokens or 0 - cached_tokens = 0 - if usage.input_tokens_details is not None: - cached_tokens = usage.input_tokens_details.cached_tokens or 0 - elif isinstance(usage, dict): - prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens") or 0 - completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens") or 0 - prompt_details = ( - usage.get("prompt_tokens_details") - or usage.get("input_tokens_details") - or {} # mutable-ok: usage details are optional mappings + prompt_tokens: Final = usage.prompt_tokens or 0 + completion_tokens: Final = usage.completion_tokens or 0 + cached_tokens: Final = ( + getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 + if usage.prompt_tokens_details is not None + else 0 ) - cached_tokens = ( - (prompt_details.get("cached_tokens", 0) or 0) if isinstance(prompt_details, dict) else 0 - ) or (usage.get("cache_read_input_tokens") or 0) - else: - return 0, 0, False + if prompt_tokens == 0 and completion_tokens == 0: + return 0, 0, False + return max(0, prompt_tokens - cached_tokens), completion_tokens, True - if prompt_tokens == 0 and completion_tokens == 0: - return 0, 0, False - return max(0, prompt_tokens - cached_tokens), completion_tokens, True + if isinstance(usage, ResponseAPIUsage): + response_input_tokens: Final = usage.input_tokens or 0 + response_output_tokens: Final = usage.output_tokens or 0 + response_cached_tokens: Final = ( + usage.input_tokens_details.cached_tokens or 0 if usage.input_tokens_details is not None else 0 + ) + if response_input_tokens == 0 and response_output_tokens == 0: + return 0, 0, False + return max(0, response_input_tokens - response_cached_tokens), response_output_tokens, True + + if isinstance(usage, Mapping): + raw_prompt_tokens: Final = usage.get("prompt_tokens") or usage.get("input_tokens") or 0 + raw_completion_tokens: Final = usage.get("completion_tokens") or usage.get("output_tokens") or 0 + mapped_prompt_tokens: Final = raw_prompt_tokens if isinstance(raw_prompt_tokens, int) else 0 + mapped_completion_tokens: Final = raw_completion_tokens if isinstance(raw_completion_tokens, int) else 0 + prompt_details: Final = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") + raw_cached_tokens: Final = ( + (prompt_details.get("cached_tokens", 0) if isinstance(prompt_details, dict) else 0) + or usage.get("cache_read_input_tokens") + or 0 + ) + mapped_cached_tokens: Final = raw_cached_tokens if isinstance(raw_cached_tokens, int) else 0 + if mapped_prompt_tokens == 0 and mapped_completion_tokens == 0: + return 0, 0, False + return max(0, mapped_prompt_tokens - mapped_cached_tokens), mapped_completion_tokens, True + + return 0, 0, False def _build_io_token_reservation_ops( self, kwargs: object, response_obj: object, - ) -> list[RedisPipelineIncrementOperation] | tuple[ReservationAwareIncrementOperation, ...]: + ) -> Sequence[RedisPipelineIncrementOperation]: """ Reconcile project ITPM/OTPM reservations to actual usage on success: ITPM to billable input tokens, OTPM to actual completion tokens. @@ -4075,25 +4099,33 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(kwargs, dict): return () - stash = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) if stash is None: return () - itpm_reserved = stash.itpm_reserved_tokens - otpm_reserved = stash.otpm_reserved_tokens + itpm_reserved: Final = stash.itpm_reserved_tokens + otpm_reserved: Final = stash.otpm_reserved_tokens if itpm_reserved <= 0 and otpm_reserved <= 0: return () - billable_input, completion_tokens, usage_resolved = self._resolve_io_token_reconcile_usage(response_obj) - if not usage_resolved: - billable_input, completion_tokens, usage_resolved = self._resolve_io_token_reconcile_usage( - kwargs.get("combined_usage_object") - ) - if not usage_resolved: - if not stash.reservation_released: - return () - billable_input = itpm_reserved - completion_tokens = otpm_reserved + response_usage: Final = self._resolve_io_token_reconcile_usage(response_obj) + combined_usage: Final = self._resolve_io_token_reconcile_usage(kwargs.get("combined_usage_object")) + aggregate_total: Final = self._aggregate_only_total_tokens( + self._response_usage(response_obj) + ) or self._aggregate_only_total_tokens(self._response_usage(kwargs.get("combined_usage_object"))) + + if not response_usage[2] and not combined_usage[2] and aggregate_total <= 0 and not stash.reservation_released: + return () + resolved_usage: Final = ( + response_usage + if response_usage[2] + else combined_usage + if combined_usage[2] + else (aggregate_total, aggregate_total, True) + if aggregate_total > 0 + else (itpm_reserved, otpm_reserved, False) + ) + billable_input, completion_tokens, _ = resolved_usage if stash.reservation_released or ( not stash.itpm_reserved_window_identities and not stash.otpm_reserved_window_identities @@ -4110,24 +4142,28 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens=0 if stash.reservation_released else otpm_reserved, ) - itpm_ops: Sequence[ReservationAwareIncrementOperation] = () - if itpm_reserved > 0: - itpm_ops = self._build_project_reservation_ops( + itpm_ops: Final[Sequence[ReservationAwareIncrementOperation]] = ( + self._build_project_reservation_ops( targets=tuple(stash.itpm_reserved_scopes), reserved_scopes=frozenset() if stash.reservation_released else stash.itpm_reserved_scopes, actual_tokens=billable_input, reserved_tokens=itpm_reserved, reservation_window_identities=stash.itpm_reserved_window_identities, ) - otpm_ops: Sequence[ReservationAwareIncrementOperation] = () - if otpm_reserved > 0: - otpm_ops = self._build_project_reservation_ops( + if itpm_reserved > 0 + else () + ) + otpm_ops: Final[Sequence[ReservationAwareIncrementOperation]] = ( + self._build_project_reservation_ops( targets=tuple(stash.otpm_reserved_scopes), reserved_scopes=frozenset() if stash.reservation_released else stash.otpm_reserved_scopes, actual_tokens=completion_tokens, reserved_tokens=otpm_reserved, reservation_window_identities=stash.otpm_reserved_window_identities, ) + if otpm_reserved > 0 + else () + ) return tuple((*itpm_ops, *otpm_ops)) def _collect_tpm_scope_targets( @@ -4522,14 +4558,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # already released it (proxy-level rejection that also bubbles up # here as an LLM-error callback). max_parallel_requests is its # own counter and is always decremented per call. - if stash is None or stash.reservation_released: - reserved_tokens = 0 - itpm_reserved = 0 - otpm_reserved = 0 - else: - reserved_tokens = stash.reserved_tokens - itpm_reserved = stash.itpm_reserved_tokens - otpm_reserved = stash.otpm_reserved_tokens + reserved_tokens, itpm_reserved, otpm_reserved = ( + (0, 0, 0) + if stash is None or stash.reservation_released + else (stash.reserved_tokens, stash.itpm_reserved_tokens, stash.otpm_reserved_tokens) + ) if stash is not None and reserved_tokens > 0: verbose_proxy_logger.debug("Releasing reserved TPM tokens on failure: %s", reserved_tokens) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index af561c9092b..e678fba2852 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -50,6 +50,16 @@ if TYPE_CHECKING: ) +class ProjectQuotaCallback(Protocol): + async def enforce_project_io_token_quota_for_frame( + self, + user_api_key_dict: UserAPIKeyAuth | None, + requested_model: str | None, + estimated_input_tokens: int, + estimated_output_tokens: int, + ) -> None: ... + + @lru_cache(maxsize=1) def _get_openai_response_types(): from litellm.types.llms import openai as openai_types @@ -1345,11 +1355,19 @@ def _extract_frame_quota_estimate_inputs(msg_obj: Mapping[str, object]) -> tuple """ nested: Final = msg_obj.get("response") params: Final[Mapping[str, object]] = ( - nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"} + nested + if _is_json_object(nested) and nested + else MappingProxyType( # mutable-ok: immediately frozen filtered frame + {k: v for k, v in msg_obj.items() if k != "type"} + ) ) - text_parts: list[str] = [] # mutable-ok: local accumulator built in one pass, not shared - - def _collect_text(value: object) -> None: + text_parts: Final[list[str]] = [] # mutable-ok: local accumulator built in one pass, not shared + pending: Final[list[object]] = [ # mutable-ok: explicit worklist avoids recursion + params.get("input"), + params.get("instructions"), + ] + while pending: + value = pending.pop() if isinstance(value, str): text_parts.append(value) elif _is_json_array(value): @@ -1357,20 +1375,17 @@ def _extract_frame_quota_estimate_inputs(msg_obj: Mapping[str, object]) -> tuple if isinstance(item, str): text_parts.append(item) elif _is_json_object(item): - _collect_text(item.get("content")) - _collect_text(item.get("text")) - - _collect_text(params.get("input")) - _collect_text(params.get("instructions")) + pending.append(item.get("content")) + pending.append(item.get("text")) total_chars: Final = sum(len(part) for part in text_parts) estimated_input_tokens: Final = max(1, total_chars // _FRAME_CHARS_PER_TOKEN_ESTIMATE) if total_chars else 0 - max_output_tokens = params.get("max_output_tokens") + max_output_tokens: Final = params.get("max_output_tokens") return estimated_input_tokens, max_output_tokens if isinstance(max_output_tokens, int) else None async def _enforce_frame_project_quota( - quota_callbacks: Sequence[Any], + quota_callbacks: Sequence[ProjectQuotaCallback], user_api_key_dict: UserAPIKeyAuth | None, model: str | None, raw_message: str, @@ -1387,7 +1402,7 @@ async def _enforce_frame_project_quota( if not _is_json_object(msg_obj) or msg_obj.get("type") != "response.create": return estimated_input_tokens, explicit_max_output_tokens = _extract_frame_quota_estimate_inputs(msg_obj) - estimated_output_tokens = ( + estimated_output_tokens: Final = ( explicit_max_output_tokens if explicit_max_output_tokens is not None else _FRAME_NO_MAX_OUTPUT_TOKENS_FLOOR ) for callback in quota_callbacks: @@ -1433,7 +1448,7 @@ class ResponsesWebSocketStreaming: first_message: str | None = None, guardrail_callbacks: list[Any] | None = None, output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, - quota_callbacks: list[Any] | None = None, + quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, authorized_model: str | None = None, ): self.websocket = websocket @@ -1446,7 +1461,7 @@ class ResponsesWebSocketStreaming: self.first_message = first_message self.guardrail_callbacks: list[Any] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] - self.quota_callbacks: list[Any] = quota_callbacks or [] + self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model @@ -1870,9 +1885,17 @@ class ResponsesWebSocketStreaming: except RateLimitError as e: try: await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": "rate_limit_exceeded", "message": str(e)}}) + json.dumps( # mutable-ok: WebSocket wire payload requires JSON objects + { # mutable-ok: WebSocket wire payload requires JSON objects + "type": "error", + "error": { # mutable-ok: nested WebSocket error object + "type": "rate_limit_exceeded", + "message": str(e), + }, + } + ) ) - except Exception: # noqa: BLE001, S110 - best-effort notification, client may already be gone + except Exception: # noqa: BLE001, S110 # client may already be gone pass return False return True @@ -1969,7 +1992,7 @@ class ManagedResponsesWebSocketHandler: timeout: float | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, - quota_callbacks: list[Any] | None = None, + quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, **kwargs: object, ) -> None: self.websocket = websocket @@ -1986,7 +2009,7 @@ class ManagedResponsesWebSocketHandler: self.custom_llm_provider = custom_llm_provider self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message - self.quota_callbacks: list[Any] = quota_callbacks or [] + self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 66fc84ab1e0..e55185cfa67 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -3582,18 +3582,24 @@ async def test_streaming_combined_usage_reconciles_project_io_reservations( assert [operation["increment_value"] for operation in otpm_adjustments] == [-45] -def test_aggregate_only_combined_usage_keeps_project_io_reservations(rate_limiter): +def test_aggregate_only_combined_usage_reconciles_project_io_reservations(rate_limiter): handler, _cache = rate_limiter stash = get_or_create_request_stash() stash.itpm_reserved_tokens = 100 stash.itpm_reserved_scopes = frozenset( {(PROJECT_ITPM_DESCRIPTOR_KEY, "project:model")} ) + stash.otpm_reserved_tokens = 80 + stash.otpm_reserved_scopes = frozenset( + {(PROJECT_OTPM_DESCRIPTOR_KEY, "project:model")} + ) kwargs = { "combined_usage_object": Usage(total_tokens=55), } - assert handler._build_io_token_reservation_ops(kwargs, object()) == () + operations = handler._build_io_token_reservation_ops(kwargs, object()) + + assert [operation["increment_value"] for operation in operations] == [-45, -25] def test_raw_split_usage_dict_reconciles_project_io_tokens(rate_limiter): diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 894d99c92e0..21dc0bc2791 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22941 + "limit": 22936 }, "LIT002": { - "limit": 27139 + "limit": 27133 }, "LIT003": { "limit": 269 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16716 + "limit": 16701 }, "LIT011": { - "limit": 5596 + "limit": 5595 }, "LIT012": { "limit": 4519