diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260722200000_add_project_itpm_otpm_limit/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260722200000_add_project_itpm_otpm_limit/migration.sql new file mode 100644 index 00000000000..df3c351d335 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260722200000_add_project_itpm_otpm_limit/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable: Add ITPM/OTPM fields to LiteLLM_ProjectTable +ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN IF NOT EXISTS "model_itpm_limit" JSONB NOT NULL DEFAULT '{}'; +ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN IF NOT EXISTS "model_otpm_limit" JSONB NOT NULL DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 23a9c086c73..d0baa6e0110 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -169,6 +169,8 @@ model LiteLLM_ProjectTable { model_spend Json @default("{}") model_rpm_limit Json @default("{}") model_tpm_limit Json @default("{}") + model_itpm_limit Json @default("{}") + model_otpm_limit Json @default("{}") blocked Boolean @default(false) object_permission_id String? created_at DateTime @default(now()) @map("created_at") diff --git a/litellm/models/project.py b/litellm/models/project.py index 083c7ee3cc5..beaed4af5f1 100644 --- a/litellm/models/project.py +++ b/litellm/models/project.py @@ -27,6 +27,8 @@ class LiteLLM_ProjectTable(LiteLLMPydanticObjectBase): model_spend: Optional[dict] = None model_rpm_limit: Optional[dict] = None model_tpm_limit: Optional[dict] = None + model_itpm_limit: dict | None = None + model_otpm_limit: dict | None = None blocked: bool = False object_permission_id: Optional[str] = None created_by: Optional[str] = None diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7df725bf965..34b6574363e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2818,6 +2818,8 @@ class NewProjectRequest(LiteLLM_BudgetTable): models: List[str] = [] model_rpm_limit: Optional[dict] = None model_tpm_limit: Optional[dict] = None + model_itpm_limit: dict | None = None + model_otpm_limit: dict | None = None blocked: bool = False object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -2850,6 +2852,8 @@ class UpdateProjectRequest(LiteLLM_BudgetTable): models: Optional[List[str]] = None model_rpm_limit: Optional[dict] = None model_tpm_limit: Optional[dict] = None + model_itpm_limit: dict | None = None + model_otpm_limit: dict | None = None blocked: Optional[bool] = None budget_id: Optional[str] = None object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -3978,6 +3982,8 @@ class PassThroughEndpointLoggingTypedDict(TypedDict): LiteLLM_ManagementEndpoint_MetadataFields = [ "model_rpm_limit", "model_tpm_limit", + "model_itpm_limit", + "model_otpm_limit", "mcp_rpm_limit", "tag_rpm_limit", "rpm_limit_type", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index ecb37e67c14..73f6f912afc 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -990,7 +990,7 @@ def get_key_model_tpm_limit( def get_model_rate_limit_from_metadata( user_api_key_dict: UserAPIKeyAuth, metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], - rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit", "model_itpm_limit", "model_otpm_limit"], ) -> Optional[Dict[str, int]]: if getattr(user_api_key_dict, metadata_accessor_key): return getattr(user_api_key_dict, metadata_accessor_key).get(rate_limit_key) @@ -1077,6 +1077,22 @@ def get_project_model_tpm_limit( return None +def get_project_model_itpm_limit( + user_api_key_dict: UserAPIKeyAuth, +) -> dict[str, int] | None: + if user_api_key_dict.project_metadata: + return user_api_key_dict.project_metadata.get("model_itpm_limit") + return None + + +def get_project_model_otpm_limit( + user_api_key_dict: UserAPIKeyAuth, +) -> dict[str, int] | None: + if user_api_key_dict.project_metadata: + return user_api_key_dict.project_metadata.get("model_otpm_limit") + return None + + def custom_auth_common_checks_warning( *, custom_auth_configured: bool, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 22ea9fe176a..b2f3191db55 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -307,6 +307,19 @@ TPM_RESERVED_SCOPES_KEY = "_litellm_tpm_reserved_scopes" # (e.g. async_log_failure_event firing after async_post_call_failure_hook) # does not double-refund. TPM_RESERVATION_RELEASED_KEY = "_litellm_tpm_reservation_released" +# Project-scoped ITPM/OTPM reservation stash, mirroring the TPM_RESERVED_* +# keys above but tracked separately since ITPM and OTPM are reserved from +# different estimates (input tokens vs. max_tokens) and reconciled against +# different actual-usage fields (prompt_tokens vs. completion_tokens). +ITPM_RESERVED_TOKENS_KEY = "_litellm_itpm_reserved_tokens" +ITPM_RESERVED_SCOPES_KEY = "_litellm_itpm_reserved_scopes" +OTPM_RESERVED_TOKENS_KEY = "_litellm_otpm_reserved_tokens" +OTPM_RESERVED_SCOPES_KEY = "_litellm_otpm_reserved_scopes" +# Descriptor "key" values for project-scoped ITPM/OTPM. Distinct from +# "model_per_project" (the combined-TPM descriptor) so both can be enforced +# on the same project+model simultaneously without colliding on cache keys. +PROJECT_ITPM_DESCRIPTOR_KEY = "model_per_project_itpm" +PROJECT_OTPM_DESCRIPTOR_KEY = "model_per_project_otpm" RATE_LIMIT_DESCRIPTORS_KEY = "_litellm_rate_limit_descriptors" # Pre-call RateLimitResponse stashed here so streaming success logging can # mirror ``x-ratelimit-*`` headers into the SLP. Streaming exits @@ -333,6 +346,10 @@ _LITELLM_STASH_KEYS: Tuple[str, ...] = ( TPM_RESERVED_MODEL_KEY, TPM_RESERVED_SCOPES_KEY, TPM_RESERVATION_RELEASED_KEY, + ITPM_RESERVED_TOKENS_KEY, + ITPM_RESERVED_SCOPES_KEY, + OTPM_RESERVED_TOKENS_KEY, + OTPM_RESERVED_SCOPES_KEY, RATE_LIMIT_DESCRIPTORS_KEY, RATE_LIMIT_RESPONSE_KEY, MAX_PARALLEL_SLOT_ACQUIRED_KEY, @@ -418,6 +435,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) + # project_id:model pairs already warned about model_itpm_limit/model_otpm_limit + # configured alongside model_tpm_limit, so the warning is logged once per pair + # rather than per request. + self._project_io_token_conflict_warned: set[str] = set() + # When disabled, TPM is enforced post-call from actual usage (pre-v1.82 # behavior) instead of reserving an estimated budget upfront, shedding # the extra per-request Redis Lua round-trip and the global-lock @@ -491,6 +513,35 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Supports chat (messages), completions (prompt), and embeddings (input). + ``min_configured_tpm_limit`` is the smallest ``tokens_per_unit`` among + the TPM-bearing descriptors this request will be charged against. When + provided, the no-``max_tokens`` output-budget floor is capped at a + fraction of that limit so small TPM caps remain usable. Omit to + preserve the unconstrained floor. + """ + estimated_input_tokens, max_tokens_estimate = self._estimate_input_and_output_tokens( + data=data, + min_configured_tpm_limit=min_configured_tpm_limit, + ) + total_estimated = estimated_input_tokens + max_tokens_estimate + + verbose_proxy_logger.debug( + f"TPM reservation estimate: input={estimated_input_tokens}, " + f"max_tokens={max_tokens_estimate}, total={total_estimated}" + ) + + return total_estimated + + def _estimate_input_and_output_tokens( + self, + data: dict, + min_configured_tpm_limit: int | None = None, + ) -> tuple[int, int]: + """ + Estimate input tokens and output (max_tokens) budget separately, so + callers needing independent ITPM/OTPM reservations (rather than one + combined TPM reservation) can use each half on its own. + ``min_configured_tpm_limit`` is the smallest ``tokens_per_unit`` among the TPM-bearing descriptors this request will be charged against. When provided, the no-``max_tokens`` output-budget floor is capped at a @@ -542,15 +593,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): output_floor = self._no_max_tokens_output_floor(min_configured_tpm_limit) max_tokens_estimate = max(estimated_input_tokens, output_floor) - total_estimated = estimated_input_tokens + max_tokens_estimate - - verbose_proxy_logger.debug( - f"TPM reservation estimate: input={estimated_input_tokens}, " - f"max_tokens={max_tokens_estimate} (explicit={explicit_max_tokens is not None}), " - f"total={total_estimated}" - ) - - return total_estimated + return estimated_input_tokens, max_tokens_estimate def _is_redis_cluster(self) -> bool: """ @@ -1561,9 +1604,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): TPM-only descriptor/increment list and delegates the all-or-nothing atomicity (Lua on Redis, asyncio-locked DualCache otherwise) to the shared primitive. + + Excludes project ITPM/OTPM descriptors -- those are reserved + separately (different estimate per bucket) via ``reserve_io_tokens``. """ tpm_descriptors: List[RateLimitDescriptor] = [ - d for d in descriptors if (d.get("rate_limit") or {}).get("tokens_per_unit") is not None + d + for d in descriptors + if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) + and (d.get("rate_limit") or {}).get("tokens_per_unit") is not None ] if not tpm_descriptors: return RateLimitResponse(overall_code="OK", statuses=[]) @@ -1577,6 +1626,100 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=parent_otel_span, ) + async def _refund_reserved_tokens( + self, + scopes: list[tuple[str, str]], + amount: int, + parent_otel_span: Span | None = None, + ) -> None: + """ + Directly decrement previously-reserved token counters for ``scopes`` + by ``amount``. Used to roll back a reservation that already + succeeded once a *different* bucket in the same request turns out to + be over its limit (e.g. ITPM reserved fine, OTPM then hits its + limit -- the ITPM reservation must not be left inflated). + """ + if amount <= 0 or not scopes: + return + pipeline_operations = [ + RedisPipelineIncrementOperation( + key=self.create_rate_limit_keys(scope_key, scope_value, "tokens"), + increment_value=-amount, + ttl=self.window_size, + ) + for scope_key, scope_value in scopes + ] + await self.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations, + parent_otel_span=parent_otel_span, + ) + + async def reserve_io_tokens( + self, + descriptors: list[RateLimitDescriptor], + estimated_input_tokens: int, + estimated_output_tokens: int, + parent_otel_span: Span | None = None, + ) -> tuple[RateLimitResponse, int, int]: + """ + Reserve ``estimated_input_tokens`` against project ITPM descriptors + and ``estimated_output_tokens`` against project OTPM descriptors. + + ITPM and OTPM are reserved from different-sized estimates, so unlike + same-size TPM descriptors they can't share a single + ``atomic_check_and_increment_by_n`` call -- each bucket gets its own + all-or-nothing atomic call. If the OTPM reservation is over limit + after ITPM already succeeded, the ITPM reservation this call made is + rolled back before returning, so a partial reservation never leaks. + + Returns ``(response, itpm_reserved, otpm_reserved)`` -- the latter two + are the amounts actually reserved (0 if that bucket wasn't + configured, or if the reservation failed), for the caller to stash + for post-call reconciliation. + """ + itpm_descriptors = [d for d in descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY] + otpm_descriptors = [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 + + statuses: list[RateLimitStatus] = [] + itpm_reserved = 0 + + if itpm_descriptors: + itpm_response = await self.atomic_check_and_increment_by_n( + descriptors=itpm_descriptors, + increments=[{"tokens": estimated_input_tokens} for _ in itpm_descriptors], + parent_otel_span=parent_otel_span, + ) + if itpm_response["overall_code"] == "OVER_LIMIT": + return itpm_response, 0, 0 + itpm_reserved = estimated_input_tokens + statuses.extend(itpm_response["statuses"]) + + if otpm_descriptors: + otpm_response = await self.atomic_check_and_increment_by_n( + descriptors=otpm_descriptors, + increments=[{"tokens": estimated_output_tokens} for _ in otpm_descriptors], + parent_otel_span=parent_otel_span, + ) + if otpm_response["overall_code"] == "OVER_LIMIT": + if itpm_reserved > 0: + await self._refund_reserved_tokens( + scopes=[(d["key"], d["value"]) for d in itpm_descriptors], + amount=itpm_reserved, + parent_otel_span=parent_otel_span, + ) + return otpm_response, 0, 0 + statuses.extend(otpm_response["statuses"]) + return ( + RateLimitResponse(overall_code="OK", statuses=statuses), + itpm_reserved, + estimated_output_tokens, + ) + + return RateLimitResponse(overall_code="OK", statuses=statuses), itpm_reserved, 0 + def create_organization_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, requested_model: Optional[str] = None ) -> List[RateLimitDescriptor]: @@ -2275,6 +2418,86 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + def _warn_project_io_token_and_tpm_coexist_once( + self, + user_api_key_dict: UserAPIKeyAuth, + requested_model: str, + ) -> None: + """Log once per project:model pair when model_itpm_limit/model_otpm_limit + is configured alongside model_tpm_limit. Both are enforced (mirrors the + deployment-level itpm/otpm + tpm/rpm coexistence behavior in + ``model_rate_limit_check.py``); this is purely a heads-up warning. + """ + tpm_limit_for_project_model = ( + get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_tpm_limit") or {} + ) + if requested_model not in tpm_limit_for_project_model: + return + warn_key = f"{user_api_key_dict.project_id}:{requested_model}" + if warn_key in self._project_io_token_conflict_warned: + return + self._project_io_token_conflict_warned.add(warn_key) + verbose_proxy_logger.warning( + f"Project '{user_api_key_dict.project_id}' configures model_itpm_limit/model_otpm_limit " + f"alongside model_tpm_limit for model '{requested_model}'; both limit types are enforced" + ) + + def _add_project_io_token_rate_limit_descriptors_from_metadata( + self, + user_api_key_dict: UserAPIKeyAuth, + requested_model: str | None, + descriptors: list[RateLimitDescriptor], + ) -> None: + """Add project-scoped ITPM/OTPM descriptors from project_metadata. + + Enforced independently of, and alongside, the combined ``model_per_project`` + TPM descriptor above -- these give Bedrock Mantle-style separate input/output + token quotas at the project level. + """ + if requested_model is None or user_api_key_dict.project_id is None: + return + + itpm_limit_for_project_model = ( + get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_itpm_limit") or {} + ) + otpm_limit_for_project_model = ( + get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_otpm_limit") or {} + ) + + model_itpm_limit = itpm_limit_for_project_model.get(requested_model) + model_otpm_limit = otpm_limit_for_project_model.get(requested_model) + + if model_itpm_limit is None and model_otpm_limit is None: + return + + self._warn_project_io_token_and_tpm_coexist_once(user_api_key_dict, requested_model) + + descriptor_value = f"{user_api_key_dict.project_id}:{requested_model}" + if model_itpm_limit is not None: + descriptors.append( + RateLimitDescriptor( + key=PROJECT_ITPM_DESCRIPTOR_KEY, + value=descriptor_value, + rate_limit={ + "requests_per_unit": None, + "tokens_per_unit": model_itpm_limit, + "window_size": self.window_size, + }, + ) + ) + if model_otpm_limit is not None: + descriptors.append( + RateLimitDescriptor( + key=PROJECT_OTPM_DESCRIPTOR_KEY, + value=descriptor_value, + rate_limit={ + "requests_per_unit": None, + "tokens_per_unit": model_otpm_limit, + "window_size": self.window_size, + }, + ) + ) + def _handle_rate_limit_error( self, response: RateLimitResponse, @@ -2319,6 +2542,125 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): llm_provider=llm_provider, ) + async def _reserve_project_io_tokens_or_raise( + self, + descriptors: list[RateLimitDescriptor], + data: dict, + requested_model: str | None, + user_api_key_dict: UserAPIKeyAuth, + tpm_reservation_scopes: list[tuple[str, str]], + tpm_reservation_amount: int, + ) -> None: + """ + Reserve project-scoped ITPM/OTPM tokens (Bedrock Mantle-style + separate input/output token buckets), independently of -- and, when + both are configured, in addition to -- the combined-TPM reservation + the caller already made. Raises (via ``_handle_rate_limit_error``) on + an over-limit reservation, first rolling back the combined-TPM + reservation named by ``tpm_reservation_scopes``/``tpm_reservation_amount`` + if one was made, so a partial reservation never leaks. + """ + io_token_descriptors = [ + 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 = [ + int(v) + for d in io_token_descriptors + if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY + for v in [(d.get("rate_limit") or {}).get("tokens_per_unit")] + if v is not None + ] + min_configured_otpm_limit = min(configured_otpm_limits) if configured_otpm_limits else None + + estimated_input_tokens, estimated_output_tokens = self._estimate_input_and_output_tokens( + data=data, + min_configured_tpm_limit=min_configured_otpm_limit, + ) + estimated_input_tokens = max(estimated_input_tokens, 1) + estimated_output_tokens = max(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 + # combined-TPM floor cap in the caller. Only tightens + # data["max_tokens"]; never loosens a cap already set. + capped_output_floor = self._no_max_tokens_output_floor(min_configured_otpm_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 + if capped_output_floor < baseline_floor and not has_explicit_max_tokens and not is_embedding: + existing_cap = data.get("max_tokens") + if existing_cap is None or capped_output_floor < existing_cap: + data["max_tokens"] = capped_output_floor + + io_response, itpm_reserved, otpm_reserved = await self.reserve_io_tokens( + descriptors=io_token_descriptors, + estimated_input_tokens=estimated_input_tokens, + estimated_output_tokens=estimated_output_tokens, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + + if io_response["overall_code"] == "OVER_LIMIT": + # A combined-TPM reservation may have already succeeded above for + # this same request; refund it too, or its counter stays inflated + # until the window's TTL expires. + await self._refund_reserved_tokens( + scopes=tpm_reservation_scopes, + amount=tpm_reservation_amount, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + acquisition = self._get_parallel_slot_acquisition(kwargs=data) + if acquisition is not None: + await self._release_parallel_request_slots( + acquisition=acquisition, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + self._clear_parallel_slot_marker(data) + self._handle_rate_limit_error( + response=io_response, + descriptors=descriptors, + requested_model=requested_model, + ) + return + + if itpm_reserved > 0: + itpm_scopes = [ + (d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY + ] + self._stash_value_in_metadata_channels(data=data, key=ITPM_RESERVED_TOKENS_KEY, value=itpm_reserved) + self._stash_value_in_metadata_channels( + data=data, + key=ITPM_RESERVED_SCOPES_KEY, + value=[[k, v] for k, v in itpm_scopes], + ) + if otpm_reserved > 0: + otpm_scopes = [ + (d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_OTPM_DESCRIPTOR_KEY + ] + self._stash_value_in_metadata_channels(data=data, key=OTPM_RESERVED_TOKENS_KEY, value=otpm_reserved) + self._stash_value_in_metadata_channels( + data=data, + key=OTPM_RESERVED_SCOPES_KEY, + value=[[k, v] for k, v in otpm_scopes], + ) + + stored_response = data.get("litellm_proxy_rate_limit_response") + if isinstance(stored_response, dict): + stored_response.setdefault("statuses", []).extend(io_response["statuses"]) + elif io_response["statuses"]: + data["litellm_proxy_rate_limit_response"] = io_response + self._stash_value_in_metadata_channels( + data=data, + key=RATE_LIMIT_RESPONSE_KEY, + value=io_response, + ) + + verbose_proxy_logger.debug( + f"ITPM/OTPM tokens reserved: itpm={itpm_reserved}, otpm={otpm_reserved} for model {requested_model}" + ) + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -2396,6 +2738,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, descriptors=descriptors, ) + self._add_project_io_token_rate_limit_descriptors_from_metadata( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model, + descriptors=descriptors, + ) # Org Level Rate Limits descriptors.extend(self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model)) @@ -2464,11 +2811,21 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): configured_tpm_limits = [ int(v) for d in descriptors + if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) for v in [(d.get("rate_limit") or {}).get("tokens_per_unit")] if v is not None ] has_tpm_limits = bool(configured_tpm_limits) + # Populated on a successful combined-TPM reservation below, so the + # project ITPM/OTPM block further down can roll it back if a + # different bucket in the same request subsequently hits its + # 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: list[tuple[str, str]] = [] + tpm_reservation_amount = 0 + if has_tpm_limits and self.tpm_reservation_enabled: min_configured_tpm_limit = min(configured_tpm_limits) @@ -2535,8 +2892,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_scopes: List[Tuple[str, str]] = [ (d["key"], d["value"]) for d in descriptors - if (d.get("rate_limit") or {}).get("tokens_per_unit") is not None + if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) + and (d.get("rate_limit") or {}).get("tokens_per_unit") is not None ] + tpm_reservation_scopes = reserved_scopes + tpm_reservation_amount = estimated_tokens self._stash_reservation_in_data( data=data, estimated_tokens=estimated_tokens, @@ -2564,6 +2924,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug(f"TPM tokens reserved: {estimated_tokens} for model {requested_model}") + await self._reserve_project_io_tokens_or_raise( + descriptors=descriptors, + data=data, + requested_model=requested_model, + user_api_key_dict=user_api_key_dict, + tpm_reservation_scopes=tpm_reservation_scopes, + tpm_reservation_amount=tpm_reservation_amount, + ) + # Defense-in-depth: scrub any stash key that escaped onto data # top-level (stale cache hit, router pass, test fixture) before the # body is forwarded to the provider. @@ -2880,6 +3249,150 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): scopes.add((entry[0], entry[1])) return scopes + @classmethod + def _get_reserved_itpm_tokens_from_kwargs( + cls, + kwargs: Any, + standard_logging_metadata: dict[str, Any] | None = None, + ) -> int: + candidate = cls._lookup_stashed_value(kwargs, standard_logging_metadata, ITPM_RESERVED_TOKENS_KEY) + try: + return int(candidate or 0) + except (TypeError, ValueError): + return 0 + + @classmethod + def _get_reserved_otpm_tokens_from_kwargs( + cls, + kwargs: Any, + standard_logging_metadata: dict[str, Any] | None = None, + ) -> int: + candidate = cls._lookup_stashed_value(kwargs, standard_logging_metadata, OTPM_RESERVED_TOKENS_KEY) + try: + return int(candidate or 0) + except (TypeError, ValueError): + return 0 + + @classmethod + def _get_reserved_itpm_scopes_from_kwargs( + cls, + kwargs: Any, + standard_logging_metadata: dict[str, Any] | None = None, + ) -> set[tuple[str, str]]: + return cls._narrow_reserved_scopes( + cls._lookup_stashed_value(kwargs, standard_logging_metadata, ITPM_RESERVED_SCOPES_KEY) + ) + + @classmethod + def _get_reserved_otpm_scopes_from_kwargs( + cls, + kwargs: Any, + standard_logging_metadata: dict[str, Any] | None = None, + ) -> set[tuple[str, str]]: + return cls._narrow_reserved_scopes( + cls._lookup_stashed_value(kwargs, standard_logging_metadata, OTPM_RESERVED_SCOPES_KEY) + ) + + @staticmethod + def _narrow_reserved_scopes(candidate: Any) -> set[tuple[str, str]]: + if not isinstance(candidate, list): + return set() + scopes: set[tuple[str, str]] = set() + for entry in candidate: + if ( + isinstance(entry, (list, tuple)) + and len(entry) == 2 + and isinstance(entry[0], str) + and isinstance(entry[1], str) + ): + scopes.add((entry[0], entry[1])) + return scopes + + def _resolve_io_token_reconcile_usage( + self, + response_obj: Any, + ) -> tuple[int, int, bool]: + """ + Resolve ``(billable_input_tokens, completion_tokens, usage_resolved)`` + for ITPM/OTPM reconciliation. Cache-read tokens are excluded from + billable input -- Bedrock Mantle doesn't count them toward ITPM -- + but they're untouched everywhere else (cost/usage logging still sees + the full prompt token count). + """ + usage: Any | None = None + if isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse, BaseLiteLLMOpenAIResponseObject), + ): + usage = getattr(response_obj, "usage", None) + + 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, dict): + prompt_tokens = usage.get("prompt_tokens", 0) or 0 + completion_tokens = usage.get("completion_tokens", 0) or 0 + prompt_details = usage.get("prompt_tokens_details") or {} + cached_tokens = (prompt_details.get("cached_tokens", 0) or 0) if isinstance(prompt_details, dict) else 0 + else: + return 0, 0, False + + return max(0, prompt_tokens - cached_tokens), completion_tokens, True + + def _build_io_token_reservation_ops( + self, + kwargs: Any, + response_obj: Any, + ) -> list["RedisPipelineIncrementOperation"]: + """ + Reconcile project ITPM/OTPM reservations to actual usage on success: + ITPM to billable input tokens, OTPM to actual completion tokens. + Reuses ``_build_reservation_aware_tpm_ops``'s delta pattern -- ITPM/OTPM + are stored in the same ":tokens" cache bucket as combined TPM, just + under distinct scope keys, so the reservation-aware increment math is + identical; only the usage fields being reconciled against differ. + """ + standard_logging_object = kwargs.get("standard_logging_object") or {} + standard_logging_metadata = standard_logging_object.get("metadata") or {} + + itpm_reserved = self._get_reserved_itpm_tokens_from_kwargs(kwargs, standard_logging_metadata) + otpm_reserved = self._get_reserved_otpm_tokens_from_kwargs(kwargs, standard_logging_metadata) + 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: + # Usage missing (e.g. a response the SDK couldn't parse) -- keep + # the reservation rather than guess; it self-expires at the + # window TTL instead of drifting the counter on a bad guess. + return [] + + ops: list[RedisPipelineIncrementOperation] = [] + if itpm_reserved > 0: + itpm_scopes = self._get_reserved_itpm_scopes_from_kwargs(kwargs, standard_logging_metadata) + ops.extend( + self._build_reservation_aware_tpm_ops( + targets=list(itpm_scopes), + reserved_scopes=itpm_scopes, + actual_tokens=billable_input, + reserved_tokens=itpm_reserved, + ) + ) + if otpm_reserved > 0: + otpm_scopes = self._get_reserved_otpm_scopes_from_kwargs(kwargs, standard_logging_metadata) + ops.extend( + self._build_reservation_aware_tpm_ops( + targets=list(otpm_scopes), + reserved_scopes=otpm_scopes, + actual_tokens=completion_tokens, + reserved_tokens=otpm_reserved, + ) + ) + return ops + @classmethod def _is_reservation_released( cls, @@ -3175,6 +3688,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): response_obj=response_obj, rate_limit_type=rate_limit_type, ) + pipeline_operations.extend(self._build_io_token_reservation_ops(kwargs=kwargs, response_obj=response_obj)) if pipeline_operations: await self.async_increment_tokens_with_ttl_preservation( @@ -3358,12 +3872,43 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + # Refund project ITPM/OTPM reservations the same way -- full + # refund, since a failed call has no billable usage to reconcile + # against. + itpm_reserved = ( + 0 if already_released else self._get_reserved_itpm_tokens_from_kwargs(kwargs, standard_logging_metadata) + ) + if itpm_reserved > 0: + itpm_scopes = self._get_reserved_itpm_scopes_from_kwargs(kwargs, standard_logging_metadata) + pipeline_operations.extend( + self._build_reservation_aware_tpm_ops( + targets=list(itpm_scopes), + reserved_scopes=itpm_scopes, + actual_tokens=0, + reserved_tokens=itpm_reserved, + ) + ) + + otpm_reserved = ( + 0 if already_released else self._get_reserved_otpm_tokens_from_kwargs(kwargs, standard_logging_metadata) + ) + if otpm_reserved > 0: + otpm_scopes = self._get_reserved_otpm_scopes_from_kwargs(kwargs, standard_logging_metadata) + pipeline_operations.extend( + self._build_reservation_aware_tpm_ops( + targets=list(otpm_scopes), + reserved_scopes=otpm_scopes, + actual_tokens=0, + reserved_tokens=otpm_reserved, + ) + ) + if pipeline_operations: await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( increment_list=pipeline_operations, litellm_parent_otel_span=litellm_parent_otel_span, ) - if reserved_tokens > 0: + if reserved_tokens > 0 or itpm_reserved > 0 or otpm_reserved > 0: self._mark_reservation_released(kwargs) except Exception as e: verbose_proxy_logger.exception(f"Error in rate limit failure event: {str(e)}") diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 23a9c086c73..d0baa6e0110 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -169,6 +169,8 @@ model LiteLLM_ProjectTable { model_spend Json @default("{}") model_rpm_limit Json @default("{}") model_tpm_limit Json @default("{}") + model_itpm_limit Json @default("{}") + model_otpm_limit Json @default("{}") blocked Boolean @default(false) object_permission_id String? created_at DateTime @default(now()) @map("created_at") diff --git a/schema.prisma b/schema.prisma index 23a9c086c73..d0baa6e0110 100644 --- a/schema.prisma +++ b/schema.prisma @@ -169,6 +169,8 @@ model LiteLLM_ProjectTable { model_spend Json @default("{}") model_rpm_limit Json @default("{}") model_tpm_limit Json @default("{}") + model_itpm_limit Json @default("{}") + model_otpm_limit Json @default("{}") blocked Boolean @default(false) object_permission_id String? created_at DateTime @default(now()) @map("created_at") diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 9f24c662581..a65b68f9a2e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -22,6 +22,8 @@ from litellm.proxy.auth.auth_utils import ( get_key_model_tpm_limit, get_key_tag_rpm_limit, get_model_from_request, + get_project_model_itpm_limit, + get_project_model_otpm_limit, get_project_model_rpm_limit, get_project_model_tpm_limit, get_request_route_template, @@ -1428,6 +1430,56 @@ class TestGetProjectModelTpmLimit: assert result is None +class TestGetProjectModelItpmLimit: + """Tests for get_project_model_itpm_limit function.""" + + def test_returns_project_metadata_itpm_limit(self): + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + project_metadata={"model_itpm_limit": {"bedrock_mantle/claude-opus": 20000000}}, + ) + result = get_project_model_itpm_limit(user_api_key_dict) + assert result == {"bedrock_mantle/claude-opus": 20000000} + + def test_returns_none_when_no_project_metadata(self): + user_api_key_dict = UserAPIKeyAuth(api_key="sk-123") + result = get_project_model_itpm_limit(user_api_key_dict) + assert result is None + + def test_returns_none_when_project_metadata_missing_key(self): + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + project_metadata={"other_key": "value"}, + ) + result = get_project_model_itpm_limit(user_api_key_dict) + assert result is None + + +class TestGetProjectModelOtpmLimit: + """Tests for get_project_model_otpm_limit function.""" + + def test_returns_project_metadata_otpm_limit(self): + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + project_metadata={"model_otpm_limit": {"bedrock_mantle/claude-opus": 4000000}}, + ) + result = get_project_model_otpm_limit(user_api_key_dict) + assert result == {"bedrock_mantle/claude-opus": 4000000} + + def test_returns_none_when_no_project_metadata(self): + user_api_key_dict = UserAPIKeyAuth(api_key="sk-123") + result = get_project_model_otpm_limit(user_api_key_dict) + assert result is None + + def test_returns_none_when_project_metadata_missing_key(self): + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + project_metadata={"other_key": "value"}, + ) + result = get_project_model_otpm_limit(user_api_key_dict) + assert result is None + + class TestCheckCompleteCredentials: """Tests for the api_key validation in check_complete_credentials.""" 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 c76e1a60afd..b5ec62474b7 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 @@ -3101,6 +3101,140 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3(): ), f"model_per_project should not be added for unrelated model, got: {descriptor_keys}" +@pytest.mark.asyncio +async def test_project_model_itpm_otpm_limits_enforced_v3(): + """ + Regression test: project-level model_itpm_limit/model_otpm_limit (Bedrock + Mantle-style separate input/output token quotas) must produce their own + rate limit descriptors, distinct from the combined model_per_project TPM + descriptor. + """ + _api_key = hash_token("sk-project-io-test") + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + captured_descriptors = [] + + async def mock_should_rate_limit(descriptors, **kwargs): + captured_descriptors.extend(descriptors) + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + project_id="proj-mantle", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 20000000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 4000000}, + }, + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "bedrock_mantle/claude-opus"}, + call_type="", + ) + + descriptor_keys = [d["key"] for d in captured_descriptors] + assert "model_per_project_itpm" in descriptor_keys, descriptor_keys + assert "model_per_project_otpm" in descriptor_keys, descriptor_keys + assert "model_per_project" not in descriptor_keys, descriptor_keys + + itpm_descriptor = next(d for d in captured_descriptors if d["key"] == "model_per_project_itpm") + otpm_descriptor = next(d for d in captured_descriptors if d["key"] == "model_per_project_otpm") + assert itpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" + assert itpm_descriptor["rate_limit"]["tokens_per_unit"] == 20000000 + assert otpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" + assert otpm_descriptor["rate_limit"]["tokens_per_unit"] == 4000000 + + +@pytest.mark.asyncio +async def test_project_model_itpm_otpm_limits_not_triggered_for_other_model_v3(): + """model_itpm_limit/model_otpm_limit should not fire for a model not in project_metadata.""" + _api_key = hash_token("sk-project-io-test-2") + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + captured_descriptors = [] + + async def mock_should_rate_limit(descriptors, **kwargs): + captured_descriptors.extend(descriptors) + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + project_id="proj-mantle", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 20000000}, + }, + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + descriptor_keys = [d["key"] for d in captured_descriptors] + assert "model_per_project_itpm" not in descriptor_keys, descriptor_keys + assert "model_per_project_otpm" not in descriptor_keys, descriptor_keys + + +@pytest.mark.asyncio +async def test_project_model_itpm_and_tpm_limits_coexist_v3(): + """ + When a project configures model_itpm_limit/model_otpm_limit alongside + model_tpm_limit for the same model, both limit types are enforced + simultaneously (mirrors the deployment-level itpm/otpm + tpm/rpm + coexistence behavior) -- the combined TPM descriptor is not suppressed. + """ + _api_key = hash_token("sk-project-io-test-3") + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + captured_descriptors = [] + + async def mock_should_rate_limit(descriptors, **kwargs): + captured_descriptors.extend(descriptors) + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + project_id="proj-mantle", + project_metadata={ + "model_tpm_limit": {"bedrock_mantle/claude-opus": 1000}, + "model_itpm_limit": {"bedrock_mantle/claude-opus": 20000000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 4000000}, + }, + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "bedrock_mantle/claude-opus"}, + call_type="", + ) + + descriptor_keys = [d["key"] for d in captured_descriptors] + assert "model_per_project" in descriptor_keys, descriptor_keys + assert "model_per_project_itpm" in descriptor_keys, descriptor_keys + assert "model_per_project_otpm" in descriptor_keys, descriptor_keys + assert ("proj-mantle:bedrock_mantle/claude-opus") in parallel_request_handler._project_io_token_conflict_warned + + @pytest.mark.asyncio async def test_pre_call_hook_does_not_leak_internal_stash_to_request_body(): """Regression for #27001: stash keys must stay in metadata, never on diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index b02f6c15168..cbc19c9447a 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -23,6 +23,10 @@ import pytest from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + ITPM_RESERVED_SCOPES_KEY, + ITPM_RESERVED_TOKENS_KEY, + OTPM_RESERVED_SCOPES_KEY, + OTPM_RESERVED_TOKENS_KEY, RATE_LIMIT_DESCRIPTORS_KEY, TPM_RESERVATION_RELEASED_KEY, TPM_RESERVED_MODEL_KEY, @@ -31,7 +35,7 @@ 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.utils import ModelResponse, Usage +from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage @pytest.fixture @@ -1192,5 +1196,210 @@ async def test_small_tpm_cap_preserves_explicit_max_tokens(rate_limiter): assert data["max_tokens"] == 500 +@pytest.mark.asyncio +async def test_project_otpm_reservation_prevents_concurrent_bypass(rate_limiter): + """ + Bedrock Mantle-style OTPM: with a 100 OTPM limit and 5 concurrent + requests each reserving 50+ output tokens, upfront reservation must + reject the late arrivals -- not let all 5 through. Exercises the + in-memory fallback in ``atomic_check_and_increment_by_n`` for the + project-scoped ITPM/OTPM descriptors specifically. + """ + handler, cache = rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-otpm-bypass"), + project_id="proj-mantle-bypass", + project_metadata={ + "model_otpm_limit": {"bedrock_mantle/claude-opus": 100}, + }, + ) + + request_data = { + "model": "bedrock_mantle/claude-opus", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 50, + } + + async def make_request(request_id: int) -> Dict[str, Any]: + data = request_data.copy() + try: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + return {"request_id": request_id, "success": True} + except Exception as e: + return { + "request_id": request_id, + "success": False, + "status_code": getattr(e, "status_code", None), + } + + results = await asyncio.gather(*[make_request(i) for i in range(5)]) + + successful = [r for r in results if r["success"]] + rate_limited = [r for r in results if not r["success"] and r.get("status_code") == 429] + + assert len(rate_limited) > 0, ( + f"Expected some OTPM-rate-limited requests but all {len(successful)} succeeded." + ) + + +@pytest.mark.asyncio +async def test_project_otpm_over_limit_rolls_back_itpm_reservation(rate_limiter): + """ + When ITPM reserves fine but OTPM is then over limit, the ITPM + reservation this same pre-call already made must be rolled back -- + otherwise it leaks until the window's TTL, silently shrinking the ITPM + budget for every other request in that minute. + """ + handler, cache = rate_limiter + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-otpm-rollback"), + project_id="proj-mantle-rollback", + project_metadata={ + "model_itpm_limit": {"bedrock_mantle/claude-opus": 1000000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 10}, + }, + ) + + itpm_counter_key = handler.create_rate_limit_keys( + key="model_per_project_itpm", + value="proj-mantle-rollback:bedrock_mantle/claude-opus", + rate_limit_type="tokens", + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 500, # blows past the 10-token OTPM limit + } + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + assert getattr(exc_info.value, "status_code", None) == 429 + + cached_value = await cache.async_get_cache(key=itpm_counter_key, local_only=True) + assert int(cached_value or 0) == 0, ( + f"ITPM reservation leaked after OTPM rejection: counter={cached_value}" + ) + + +@pytest.mark.asyncio +async def test_project_itpm_reconciled_on_success_excludes_cached_tokens(rate_limiter): + """ + On success, ITPM reconciles to billable input tokens (prompt_tokens + minus cached_tokens) -- not raw prompt_tokens. Cached prompt-read tokens + are free under Bedrock Mantle and must not count against the ITPM quota, + even though they still appear in usage/cost logging elsewhere. + """ + handler, _cache = rate_limiter + + itpm_scope = ("model_per_project_itpm", "proj-mantle:model") + otpm_scope = ("model_per_project_otpm", "proj-mantle:model") + + mock_kwargs = { + "standard_logging_object": { + "metadata": { + ITPM_RESERVED_TOKENS_KEY: 100, + ITPM_RESERVED_SCOPES_KEY: [list(itpm_scope)], + OTPM_RESERVED_TOKENS_KEY: 60, + OTPM_RESERVED_SCOPES_KEY: [list(otpm_scope)], + } + }, + } + + mock_response = ModelResponse( + id="test", + object="chat.completion", + created=int(datetime.now().timestamp()), + model="bedrock_mantle/claude-opus", + usage=Usage( + prompt_tokens=80, + completion_tokens=40, + total_tokens=120, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=30), + ), + choices=[], + ) + + increments = [] + + 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(), + ) + + itpm_adjustments = [i for i in increments if "model_per_project_itpm" in i["key"]] + otpm_adjustments = [i for i in increments if "model_per_project_otpm" in i["key"]] + + # billable_input = 80 - 30 cached = 50; delta = 50 - 100 reserved = -50 + assert any(i["increment"] == -50 for i in itpm_adjustments), ( + f"Expected a -50 ITPM adjustment (50 billable - 100 reserved), got: {itpm_adjustments}" + ) + # delta = 40 actual completion - 60 reserved = -20 + assert any(i["increment"] == -20 for i in otpm_adjustments), ( + f"Expected a -20 OTPM adjustment (40 actual - 60 reserved), got: {otpm_adjustments}" + ) + + +@pytest.mark.asyncio +async def test_project_itpm_otpm_released_on_failure(rate_limiter): + """On failure, the full ITPM and OTPM reservations must be refunded.""" + handler, _cache = rate_limiter + + itpm_scope = ("model_per_project_itpm", "proj-mantle:model") + otpm_scope = ("model_per_project_otpm", "proj-mantle:model") + + mock_kwargs = { + "standard_logging_object": { + "metadata": { + ITPM_RESERVED_TOKENS_KEY: 100, + ITPM_RESERVED_SCOPES_KEY: [list(itpm_scope)], + OTPM_RESERVED_TOKENS_KEY: 60, + OTPM_RESERVED_SCOPES_KEY: [list(otpm_scope)], + } + }, + } + + increments = [] + + 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_failure_event( + kwargs=mock_kwargs, + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + itpm_releases = [i for i in increments if "model_per_project_itpm" in i["key"]] + otpm_releases = [i for i in increments if "model_per_project_otpm" in i["key"]] + + assert any(i["increment"] == -100 for i in itpm_releases), itpm_releases + assert any(i["increment"] == -60 for i in otpm_releases), otpm_releases + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 203057f23f1..a0ff232013a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25925,6 +25925,14 @@ export interface components { metadata?: { [key: string]: unknown; } | null; + /** Model Itpm Limit */ + model_itpm_limit?: { + [key: string]: unknown; + } | null; + /** Model Otpm Limit */ + model_otpm_limit?: { + [key: string]: unknown; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: unknown; @@ -28012,10 +28020,18 @@ export interface components { metadata?: { [key: string]: unknown; } | null; + /** Model Itpm Limit */ + model_itpm_limit?: { + [key: string]: unknown; + } | null; /** Model Max Budget */ model_max_budget?: { [key: string]: unknown; } | null; + /** Model Otpm Limit */ + model_otpm_limit?: { + [key: string]: unknown; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: unknown; @@ -28073,6 +28089,14 @@ export interface components { metadata?: { [key: string]: unknown; } | null; + /** Model Itpm Limit */ + model_itpm_limit?: { + [key: string]: unknown; + } | null; + /** Model Otpm Limit */ + model_otpm_limit?: { + [key: string]: unknown; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: unknown; @@ -32239,10 +32263,18 @@ export interface components { metadata?: { [key: string]: unknown; } | null; + /** Model Itpm Limit */ + model_itpm_limit?: { + [key: string]: unknown; + } | null; /** Model Max Budget */ model_max_budget?: { [key: string]: unknown; } | null; + /** Model Otpm Limit */ + model_otpm_limit?: { + [key: string]: unknown; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: unknown;