From 4492799fadd76ffe53a1b47905a4e0db4787ff6d Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 21:53:39 -0400 Subject: [PATCH] feat(types): add TagRateLimitEntry/TagRateLimitScope config types Introduces the tag-scoped rate limit config schema (TagRateLimitEntry, TagRateLimitScope, TagRateLimitGroup, TagRateLimits) and wires it onto ModelInfo.tag_rate_limits, giving tag-based rate limiting hooks a config shape to validate and consume. --- litellm/types/router.py | 177 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) diff --git a/litellm/types/router.py b/litellm/types/router.py index 4f4df1a8d2e..8beb83787fd 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -4,6 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc import datetime import enum +import math from collections.abc import Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints @@ -155,6 +156,180 @@ def _as_utc(value: datetime.datetime | None) -> datetime.datetime | None: return value.astimezone(datetime.timezone.utc) +class TagRateLimitScope(BaseModel): + """ + A gate on a tag OTHER than the entry's own `tag_id` -- e.g. scoping an + entry to `tag_id: company_id, values: ["1032"]` so it only applies to + requests tagged as belonging to company 1032, independent of whichever + tag the entry itself keys its bucket by. See `TagRateLimitEntry.enabled_for`/ + `disabled_for`, which are the only two fields that construct this. + """ + + tag_id: str + values: tuple[str, ...] + + model_config = ConfigDict(frozen=True) + + @model_validator(mode="after") + def _validate_values(self) -> "TagRateLimitScope": + if not self.values: + raise ValueError("values must be a non-empty list of strings") + return self + + @model_validator(mode="after") + def _normalize_values(self) -> "TagRateLimitScope": + # Sorted and deduplicated: `values` is only ever used for membership + # tests (see _entry_applies), never order-dependent, but is also + # folded verbatim into the dedup signature two deployments' entries + # are compared by (see _scope_signature) -- an unsorted tuple would + # make config-order alone, not policy, decide whether two entries + # dedup to one shared bucket or wrongly split into two. + # object.__setattr__ bypasses this frozen model's own assignment + # guard -- returning a replacement instance from an "after" validator + # is silently ignored when constructing via __init__ (only takes + # effect via model_validate), so mutating in place is the only way + # this normalization reliably applies regardless of construction path. + object.__setattr__(self, "values", tuple(sorted(set(self.values)))) # mutable-ok: frozen before escaping + return self + + +class TagRateLimitEntry(BaseModel): + name: str + tag_id: str = "end_user_id" + limit: float + period_seconds: int + scope_by_key_hash: bool = False + # Overrides this entry's bucket/reservation key TTL (Redis, and the + # in-memory fallback when Redis isn't configured). Defaults to + # period_seconds + 3600 when unset -- see _PROXY_ModelBasedTagRateLimitsHook._ttl_for. + # A high-cardinality tag_id can keep many keys alive at once; lowering + # this lets an operator shed them sooner without shortening + # period_seconds itself. + key_ttl_seconds: int | None = None + # Overrides the size of the dedicated in-memory cache partition this + # entry's own keys live in, when Redis isn't configured (or as a local + # fast-path cache when it is). Unset means this entry shares the hook's + # single default partition, sized by + # litellm.model_based_tag_rate_limits_max_in_memory_cache_size (200 if that's also + # unset). A high-cardinality tag_id can churn past that shared cap and + # evict another entry's active counters; setting this gives the entry + # its own dedicated partition instead. + max_in_memory_cache_size: int | None = None + # Gate this entry on a tag -- often a SECOND, independent tag (e.g. + # `enabled_for: {tag_id: company_id, values: ["1032"]}` to scope an + # override to one company's traffic), but `tag_id` can equally be set to + # this same entry's own `tag_id` to scope by a subset of its own + # resolved identity instead, without a second tag at all. + # `disabled_for` is checked first (deny overrides allow) when both are + # set. An absent gate tag never satisfies `enabled_for` (an allowlist + # gate requires an explicit match) but never triggers `disabled_for` + # either (nothing to match against a denylist). + enabled_for: TagRateLimitScope | None = None + disabled_for: TagRateLimitScope | None = None + # Restrict this entry to requests authenticated with one of these virtual + # keys' own `key_alias`. Unset (the default) means the entry applies to + # every request regardless of which key made it. A key with no alias set + # never satisfies this allowlist, same "absent gate never matches an + # allowlist" precedent as `enabled_for`. + apply_to_key_alias: tuple[str, ...] | None = None + # Restrict this entry to requests whose caller-facing `model` matches one + # of these names. Unset (the default) means the entry applies to every + # model. A request with no `model` field never satisfies this allowlist, + # same "absent gate never matches an allowlist" precedent as + # `apply_to_key_alias`. + apply_to_models: tuple[str, ...] | None = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="after") + def _validate_limit(self) -> "TagRateLimitEntry": + # NaN compares False against every ordering operator, so a NaN limit + # makes the atomic requests/concurrency check-and-increment (which + # rejects when the new value exceeds the limit) never reject -- + # admitting indefinitely -- while the read-only tokens/dollars check + # (which admits when the current value is under the limit) never + # admits, rejecting every tagged request. Either outcome silently + # defeats the entry; reject it at config load time instead. + if math.isnan(self.limit): + raise ValueError("limit must not be NaN") + if math.isinf(self.limit): + raise ValueError( + "limit must be finite -- positive infinity makes admission never reject (current + increment " + "> limit is always false), negative infinity makes it always reject every tagged request" + ) + if self.limit <= 0: + raise ValueError( + "limit must be a positive number -- zero or negative makes the atomic requests/concurrency " + "check (current + increment > limit) reject every admission and the read-only tokens/dollars " + "check (current < limit) never admit, silently blocking all matching traffic instead of the " + "likely intended config" + ) + return self + + @model_validator(mode="after") + def _validate_period_seconds(self) -> "TagRateLimitEntry": + if self.period_seconds <= 0: + raise ValueError("period_seconds must be a positive integer") + return self + + @model_validator(mode="after") + def _validate_key_ttl_seconds(self) -> "TagRateLimitEntry": + if self.key_ttl_seconds is not None and self.key_ttl_seconds <= 0: + raise ValueError("key_ttl_seconds must be a positive integer when set") + if self.key_ttl_seconds is not None and self.key_ttl_seconds < self.period_seconds: + raise ValueError( + "key_ttl_seconds must be at least period_seconds when set -- a shorter TTL expires the " + "counter before its window rolls over, letting tagged traffic reset to zero and exceed the limit" + ) + return self + + @model_validator(mode="after") + def _validate_max_in_memory_cache_size(self) -> "TagRateLimitEntry": + if self.max_in_memory_cache_size is not None and self.max_in_memory_cache_size <= 0: + raise ValueError("max_in_memory_cache_size must be a positive integer when set") + return self + + @model_validator(mode="after") + def _validate_apply_to_key_alias(self) -> "TagRateLimitEntry": + if self.apply_to_key_alias is not None and not self.apply_to_key_alias: + raise ValueError("apply_to_key_alias must be a non-empty list of strings when set") + return self + + @model_validator(mode="after") + def _normalize_apply_to_key_alias(self) -> "TagRateLimitEntry": + # Sorted and deduplicated for the same reason as + # TagRateLimitScope._normalize_values: only ever used for membership + # tests, but also folded verbatim into the dedup signature, where an + # unsorted tuple would make config-order alone decide whether two + # deployments' entries dedup to one shared bucket. + if self.apply_to_key_alias is not None: + self.apply_to_key_alias = tuple(sorted(set(self.apply_to_key_alias))) # mutable-ok: frozen before escaping + return self + + @model_validator(mode="after") + def _validate_apply_to_models(self) -> "TagRateLimitEntry": + if self.apply_to_models is not None and not self.apply_to_models: + raise ValueError("apply_to_models must be a non-empty list of strings when set") + return self + + @model_validator(mode="after") + def _normalize_apply_to_models(self) -> "TagRateLimitEntry": + if self.apply_to_models is not None: + self.apply_to_models = tuple(sorted(set(self.apply_to_models))) # mutable-ok: frozen before escaping + return self + + +class TagRateLimitGroup(BaseModel): + limits: tuple[TagRateLimitEntry, ...] = () + + +class TagRateLimits(BaseModel): + token_limits: TagRateLimitGroup | None = None + request_limits: TagRateLimitGroup | None = None + dollar_limits: TagRateLimitGroup | None = None + concurrency_limits: TagRateLimitGroup | None = None + + class ModelInfo(MirroredPricingParams): id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. @@ -209,6 +384,8 @@ class ModelInfo(MirroredPricingParams): # in the spend log row's metadata. Set it on every deployment of the group. internal_router_model: bool | None = None + tag_rate_limits: TagRateLimits | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided