mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
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.
This commit is contained in:
parent
92122086ec
commit
4492799fad
1 changed files with 177 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue