scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-28 06:23:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 94

feat(proxy): add project-level model_itpm_limit/model_otpm_limit

Browse source 
We already support model_tpm_limit/model_rpm_limit per project, but
that treats input and output tokens as one combined bucket, which
doesn't match how Bedrock Mantle actually meters usage. This adds
model_itpm_limit/model_otpm_limit so a project can cap input and output
tokens separately, same shape as the existing fields.

There's already deployment-level itpm/otpm support in
router_utils/pre_call_checks/io_token_rate_limit_check.py, but that
code only ever deals with one deployment at a time, whereas the proxy
hook (parallel_request_limiter_v3.py) has to juggle key/team/org/project
limits together on the same request. Rather than bolt a second copy of
the reservation logic on, this just adds project-scoped ITPM/OTPM as
two more descriptors into the same reservation path the hook already
uses for TPM/RPM.

If a project sets both the combined TPM limit and the new itpm/otpm
limits on the same model, both get enforced, not one overriding the
other, matching what the deployment-level check already does. Input
and output are reserved separately up front, and if one side reserves
fine but the other is over limit, we roll back the side that already
went through so nothing gets left over-counted. On the way back,
cached prompt tokens get excluded from the ITPM count since that's how
Bedrock bills it, but that's purely for the rate limit math, cost and
usage logging still see the full token count.
This commit is contained in:
Shivi Jain 2026-07-22 22:02:52 +05:30
parent f1f0a0bacd
commit aa130f4e4c
12 changed files with 1019 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260722200000_add_project_itpm_otpm_limit/migration.sql Normal file
View file

@ -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 '{}';

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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")

2
litellm/models/project.py
View file

@ -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

6
litellm/proxy/_types.py
View file

@ -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",

18
litellm/proxy/auth/auth_utils.py
View file

@ -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,

569
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -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)}")

2
litellm/proxy/schema.prisma
View file

@ -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")

2
schema.prisma
View file

@ -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")

52
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -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."""

134
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -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

211
tests/test_litellm/proxy/hooks/test_tpm_concurrent.py
View file

@ -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"])

32
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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;