mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge 1dfe1d968c into 9071ca503e
This commit is contained in:
commit
a140747ca0
5 changed files with 373 additions and 2 deletions
|
|
@ -253,7 +253,10 @@ class InMemoryCache(BaseCache):
|
|||
) -> list[float] | None:
|
||||
results: Final = []
|
||||
for increment in increment_list:
|
||||
result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs)
|
||||
# forward each op's own ttl; allow_ttl_override leaves an already-live ttl untouched
|
||||
result = await self.async_increment(
|
||||
increment["key"], increment["increment_value"], ttl=increment.get("ttl"), **kwargs
|
||||
)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc
|
|||
|
||||
import datetime
|
||||
import enum
|
||||
import math
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints
|
||||
|
|
@ -160,6 +161,126 @@ 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`, used by `TagRateLimitEntry.enabled_for`/`disabled_for`."""
|
||||
|
||||
tag_id: str
|
||||
values: tuple[str, ...]
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_tag_id(self) -> "TagRateLimitScope":
|
||||
if not self.tag_id:
|
||||
raise ValueError("tag_id must be a non-empty string")
|
||||
return self
|
||||
|
||||
@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+deduped so dedup-signature comparisons aren't order-sensitive; __setattr__ works around frozen=True
|
||||
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 the default bucket ttl (period_seconds + 3600); see _PROXY_ModelBasedTagRateLimitsHook._ttl_for
|
||||
key_ttl_seconds: int | None = None
|
||||
# overrides the shared litellm.model_based_tag_rate_limits_max_in_memory_cache_size (200) with a dedicated partition
|
||||
max_in_memory_cache_size: int | None = None
|
||||
# gate this entry on another tag; disabled_for wins over enabled_for; an absent tag never matches either allowlist
|
||||
enabled_for: TagRateLimitScope | None = None
|
||||
disabled_for: TagRateLimitScope | None = None
|
||||
apply_to_key_alias: tuple[str, ...] | None = None
|
||||
apply_to_models: tuple[str, ...] | None = None
|
||||
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_tag_id(self) -> "TagRateLimitEntry":
|
||||
if not self.tag_id:
|
||||
raise ValueError("tag_id must be a non-empty string")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_limit(self) -> "TagRateLimitEntry":
|
||||
# NaN compares False against every ordering operator, silently defeating whichever check gates this limit
|
||||
if math.isnan(self.limit):
|
||||
raise ValueError("limit must not be NaN")
|
||||
if math.isinf(self.limit):
|
||||
raise ValueError("limit must be finite")
|
||||
if self.limit <= 0:
|
||||
raise ValueError("limit must be a positive number")
|
||||
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")
|
||||
# a shorter ttl than period_seconds resets the counter early, letting it exceed the limit
|
||||
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")
|
||||
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":
|
||||
# same reason as TagRateLimitScope._normalize_values
|
||||
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.
|
||||
|
|
@ -214,6 +335,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
|
||||
|
|
|
|||
|
|
@ -45,6 +45,39 @@ async def test_async_increment_delegates_to_locked_sync_path():
|
|||
assert cache.get_cache("counter") == 5
|
||||
|
||||
|
||||
async def test_async_increment_pipeline_applies_each_operations_own_ttl():
|
||||
"""
|
||||
Bugbot finding: async_increment_pipeline dropped each operation's own
|
||||
"ttl" field, so a counter created through it always got the cache's
|
||||
600-second default_ttl regardless of what the caller actually configured
|
||||
(e.g. an hourly or daily rate-limit window) -- the counter (and whatever
|
||||
it was tracking against a limit) silently reset mid-window.
|
||||
"""
|
||||
cache = InMemoryCache(default_ttl=600)
|
||||
await cache.async_increment_pipeline([{"key": "long-window-counter", "increment_value": 1, "ttl": 7200}])
|
||||
|
||||
ttl_remaining = await cache.async_get_ttl("long-window-counter")
|
||||
assert ttl_remaining is not None
|
||||
assert ttl_remaining > time.time() + 600
|
||||
|
||||
|
||||
async def test_async_increment_pipeline_preserves_an_existing_live_ttl_on_later_increments():
|
||||
"""
|
||||
A second increment on the same still-live counter must not reset its
|
||||
remaining ttl back up to the full window -- only the first increment
|
||||
(the one that actually creates the counter) should set it.
|
||||
"""
|
||||
cache = InMemoryCache(default_ttl=600)
|
||||
await cache.async_increment_pipeline([{"key": "counter", "increment_value": 1, "ttl": 7200}])
|
||||
ttl_after_first = cache.ttl_dict["counter"]
|
||||
|
||||
await cache.async_increment_pipeline([{"key": "counter", "increment_value": 1, "ttl": 7200}])
|
||||
ttl_after_second = cache.ttl_dict["counter"]
|
||||
|
||||
assert ttl_after_second == ttl_after_first
|
||||
assert cache.get_cache("counter") == 2
|
||||
|
||||
|
||||
def test_in_memory_openai_obj_cache():
|
||||
from openai import OpenAI
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,10 @@ from litellm.types.router import (
|
|||
GenericLiteLLMParams,
|
||||
LiteLLM_Params,
|
||||
ModelInfo,
|
||||
TagRateLimitEntry,
|
||||
TagRateLimitGroup,
|
||||
TagRateLimits,
|
||||
TagRateLimitScope,
|
||||
)
|
||||
from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams
|
||||
|
||||
|
|
@ -91,7 +95,7 @@ def test_pricing_strings_are_coerced_to_float():
|
|||
|
||||
|
||||
def test_invalid_pricing_is_rejected():
|
||||
with pytest.raises(ValueError, match='validation error for ModelInfo'):
|
||||
with pytest.raises(ValueError, match="validation error for ModelInfo"):
|
||||
ModelInfo(id="x", input_cost_per_token="free")
|
||||
|
||||
|
||||
|
|
@ -146,3 +150,156 @@ def test_aws_session_tags_round_trip_as_sts_shaped_pairs():
|
|||
def test_aws_session_tags_reject_shapes_sts_would_refuse(aws_session_tags):
|
||||
with pytest.raises(ValidationError, match="aws_session_tags"):
|
||||
LiteLLM_Params(model="bedrock/anthropic.claude-opus-5", aws_session_tags=aws_session_tags)
|
||||
|
||||
|
||||
# litellm/types/router.py is imported by plain SDK users, not just the proxy, so a
|
||||
# TagRateLimitEntry validation error should read as a generic config-validation
|
||||
# message and not describe the proxy rate-limit hook's internal admission mechanics.
|
||||
_INTERNAL_ENFORCEMENT_JARGON = (
|
||||
"admission",
|
||||
"tagged request",
|
||||
"tagged traffic",
|
||||
"check-and-increment",
|
||||
"read-only",
|
||||
"atomic",
|
||||
"window rolls over",
|
||||
)
|
||||
|
||||
|
||||
def _assert_message_has_no_internal_jargon(excinfo: pytest.ExceptionInfo) -> None:
|
||||
message = str(excinfo.value).lower()
|
||||
leaked = [term for term in _INTERNAL_ENFORCEMENT_JARGON if term in message]
|
||||
assert not leaked, f"validation message leaked internal enforcement jargon: {leaked}"
|
||||
|
||||
|
||||
def test_limit_infinite_rejected_without_internal_enforcement_jargon():
|
||||
with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo:
|
||||
TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=float("inf"), period_seconds=60)
|
||||
_assert_message_has_no_internal_jargon(excinfo)
|
||||
|
||||
|
||||
def test_limit_non_positive_rejected_without_internal_enforcement_jargon():
|
||||
with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo:
|
||||
TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=0, period_seconds=60)
|
||||
_assert_message_has_no_internal_jargon(excinfo)
|
||||
|
||||
|
||||
def test_key_ttl_seconds_shorter_than_period_rejected_without_internal_enforcement_jargon():
|
||||
with pytest.raises(ValueError, match="validation error for TagRateLimitEntry") as excinfo:
|
||||
TagRateLimitEntry(
|
||||
name="daily",
|
||||
tag_id="end_user_id",
|
||||
limit=500,
|
||||
period_seconds=86400,
|
||||
key_ttl_seconds=60,
|
||||
)
|
||||
_assert_message_has_no_internal_jargon(excinfo)
|
||||
|
||||
|
||||
def test_limit_nan_rejected():
|
||||
"""NaN compares False against every ordering operator, so it would otherwise slip
|
||||
past a `limit <= 0` style guard and silently defeat whichever check gates it."""
|
||||
with pytest.raises(ValueError, match="limit must not be NaN"):
|
||||
TagRateLimitEntry(name="daily", limit=float("nan"), period_seconds=60)
|
||||
|
||||
|
||||
def test_period_seconds_non_positive_rejected():
|
||||
with pytest.raises(ValueError, match="period_seconds must be a positive integer"):
|
||||
TagRateLimitEntry(name="daily", limit=10, period_seconds=0)
|
||||
|
||||
|
||||
def test_key_ttl_seconds_non_positive_rejected():
|
||||
with pytest.raises(ValueError, match="key_ttl_seconds must be a positive integer when set"):
|
||||
TagRateLimitEntry(name="daily", limit=10, period_seconds=60, key_ttl_seconds=0)
|
||||
|
||||
|
||||
def test_max_in_memory_cache_size_non_positive_rejected():
|
||||
with pytest.raises(ValueError, match="max_in_memory_cache_size must be a positive integer"):
|
||||
TagRateLimitEntry(name="daily", limit=10, period_seconds=60, max_in_memory_cache_size=0)
|
||||
|
||||
|
||||
def test_apply_to_key_alias_empty_list_rejected():
|
||||
with pytest.raises(ValueError, match="apply_to_key_alias must be a non-empty list"):
|
||||
TagRateLimitEntry(name="daily", limit=10, period_seconds=60, apply_to_key_alias=())
|
||||
|
||||
|
||||
def test_apply_to_key_alias_sorted_and_deduped():
|
||||
entry = TagRateLimitEntry(name="daily", limit=10, period_seconds=60, apply_to_key_alias=("b", "a", "a"))
|
||||
assert entry.apply_to_key_alias == ("a", "b")
|
||||
|
||||
|
||||
def test_apply_to_models_empty_list_rejected():
|
||||
with pytest.raises(ValueError, match="apply_to_models must be a non-empty list"):
|
||||
TagRateLimitEntry(name="daily", limit=10, period_seconds=60, apply_to_models=())
|
||||
|
||||
|
||||
def test_apply_to_models_sorted_and_deduped():
|
||||
entry = TagRateLimitEntry(name="daily", limit=10, period_seconds=60, apply_to_models=("gpt-4o", "claude", "claude"))
|
||||
assert entry.apply_to_models == ("claude", "gpt-4o")
|
||||
|
||||
|
||||
def test_scope_tag_id_empty_string_rejected():
|
||||
"""An empty tag_id makes identity lookup search for a bare `:` prefix, silently
|
||||
never matching instead of erroring at config load time (Greptile P1 on #39902)."""
|
||||
with pytest.raises(ValueError, match="tag_id must be a non-empty string"):
|
||||
TagRateLimitScope(tag_id="", values=("1032",))
|
||||
|
||||
|
||||
def test_scope_values_empty_list_rejected():
|
||||
with pytest.raises(ValueError, match="values must be a non-empty list"):
|
||||
TagRateLimitScope(tag_id="company_id", values=())
|
||||
|
||||
|
||||
def test_scope_values_sorted_and_deduped():
|
||||
scope = TagRateLimitScope(tag_id="company_id", values=("1032", "1001", "1001"))
|
||||
assert scope.values == ("1001", "1032")
|
||||
|
||||
|
||||
def test_scope_is_frozen():
|
||||
scope = TagRateLimitScope(tag_id="company_id", values=("1032",))
|
||||
with pytest.raises(ValidationError):
|
||||
scope.tag_id = "other_tag"
|
||||
|
||||
|
||||
def test_entry_tag_id_empty_string_rejected():
|
||||
"""Same silent-no-match failure mode as TagRateLimitScope.tag_id (Greptile P1 on
|
||||
#39902); the default is non-empty, but an explicit override could still be empty."""
|
||||
with pytest.raises(ValueError, match="tag_id must be a non-empty string"):
|
||||
TagRateLimitEntry(name="daily", tag_id="", limit=10, period_seconds=60)
|
||||
|
||||
|
||||
def test_entry_accepts_enabled_for_and_disabled_for_scopes():
|
||||
entry = TagRateLimitEntry(
|
||||
name="daily",
|
||||
limit=10,
|
||||
period_seconds=60,
|
||||
enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)),
|
||||
disabled_for=TagRateLimitScope(tag_id="company_id", values=("9999",)),
|
||||
)
|
||||
assert entry.enabled_for.values == ("1032",)
|
||||
assert entry.disabled_for.values == ("9999",)
|
||||
|
||||
|
||||
def test_group_defaults_to_no_limits():
|
||||
assert TagRateLimitGroup().limits == ()
|
||||
|
||||
|
||||
def test_rate_limits_group_holds_multiple_entries():
|
||||
entries = (
|
||||
TagRateLimitEntry(name="daily", limit=10, period_seconds=60),
|
||||
TagRateLimitEntry(name="weekly", limit=100, period_seconds=604800),
|
||||
)
|
||||
group = TagRateLimitGroup(limits=entries)
|
||||
assert group.limits == entries
|
||||
|
||||
|
||||
def test_model_info_tag_rate_limits_defaults_to_none():
|
||||
assert ModelInfo(id="x").tag_rate_limits is None
|
||||
|
||||
|
||||
def test_model_info_accepts_tag_rate_limits():
|
||||
limits = TagRateLimits(
|
||||
token_limits=TagRateLimitGroup(limits=(TagRateLimitEntry(name="daily", limit=10, period_seconds=60),))
|
||||
)
|
||||
info = ModelInfo(id="x", tag_rate_limits=limits)
|
||||
assert info.tag_rate_limits.token_limits.limits[0].name == "daily"
|
||||
|
|
|
|||
55
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
55
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -36665,6 +36665,60 @@ export interface components {
|
|||
/** Tpm Limit */
|
||||
tpm_limit?: number | null;
|
||||
};
|
||||
/** TagRateLimitEntry */
|
||||
TagRateLimitEntry: {
|
||||
/** Apply To Key Alias */
|
||||
apply_to_key_alias?: string[] | null;
|
||||
/** Apply To Models */
|
||||
apply_to_models?: string[] | null;
|
||||
disabled_for?: components["schemas"]["TagRateLimitScope"] | null;
|
||||
enabled_for?: components["schemas"]["TagRateLimitScope"] | null;
|
||||
/** Key Ttl Seconds */
|
||||
key_ttl_seconds?: number | null;
|
||||
/** Limit */
|
||||
limit: number;
|
||||
/** Max In Memory Cache Size */
|
||||
max_in_memory_cache_size?: number | null;
|
||||
/** Name */
|
||||
name: string;
|
||||
/** Period Seconds */
|
||||
period_seconds: number;
|
||||
/**
|
||||
* Scope By Key Hash
|
||||
* @default false
|
||||
*/
|
||||
scope_by_key_hash: boolean;
|
||||
/**
|
||||
* Tag Id
|
||||
* @default end_user_id
|
||||
*/
|
||||
tag_id: string;
|
||||
};
|
||||
/** TagRateLimitGroup */
|
||||
TagRateLimitGroup: {
|
||||
/**
|
||||
* Limits
|
||||
* @default []
|
||||
*/
|
||||
limits: components["schemas"]["TagRateLimitEntry"][];
|
||||
};
|
||||
/**
|
||||
* TagRateLimitScope
|
||||
* @description A gate on a tag other than the entry's own `tag_id`, used by `TagRateLimitEntry.enabled_for`/`disabled_for`.
|
||||
*/
|
||||
TagRateLimitScope: {
|
||||
/** Tag Id */
|
||||
tag_id: string;
|
||||
/** Values */
|
||||
values: string[];
|
||||
};
|
||||
/** TagRateLimits */
|
||||
TagRateLimits: {
|
||||
concurrency_limits?: components["schemas"]["TagRateLimitGroup"] | null;
|
||||
dollar_limits?: components["schemas"]["TagRateLimitGroup"] | null;
|
||||
request_limits?: components["schemas"]["TagRateLimitGroup"] | null;
|
||||
token_limits?: components["schemas"]["TagRateLimitGroup"] | null;
|
||||
};
|
||||
/**
|
||||
* TagSummaryMetrics
|
||||
* @description Summary metrics for a tag
|
||||
|
|
@ -39656,6 +39710,7 @@ export interface components {
|
|||
ptu_effective_from?: string | null;
|
||||
/** Ptu Effective To */
|
||||
ptu_effective_to?: string | null;
|
||||
tag_rate_limits?: components["schemas"]["TagRateLimits"] | null;
|
||||
/** Team Id */
|
||||
team_id?: string | null;
|
||||
/** Team Public Model Name */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue