mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
feat(budget): throttle keys after spend limit instead of revoking access (#31300)
Add an opt-in mode so a key that exceeds its own max_budget is throttled to a globally configured percentage of its TPM/RPM instead of being blocked entirely. A new litellm_settings global, budget_exceeded_throttle_percentage, sets the fraction (e.g. 0.1 = 10%). A per-key throttle_on_budget_exceeded flag (stored in key metadata via the existing management-endpoint metadata routing) opts the key in. When both are set and the key is over budget, the budget check records the percentage on a request-scoped budget_throttle_pct instead of raising, and the rate limiter scales the key's configured TPM/RPM by it. Keys without the flag keep hard-blocking; team/user/org budgets are unaffected. The throttle is recomputed from the key's original limits on every request and the decision is cleared before the auth object is cached, so it never compounds across requests. Both the budget read-time check and the budget reservation path honor the opt-in, and both the v3 and legacy rate limiters apply the scaling. Enabling throttle_on_budget_exceeded is proxy-admin only. It converts an admin-imposed hard budget block into a soft throttle that keeps spending past max_budget, so a non-admin must not be able to self-opt-in and bypass their own spend cap. Both /key/generate and /key/update reject a non-admin setting it to true (update only gates the transition to enabled, so a non-admin can still edit other fields and turn the flag off). This matches the feature being wholly proxy-admin operated: the global percentage is admin-only too. A key that opts in but has no TPM or RPM limit has nothing to scale, so it stays hard-blocked rather than serving unlimited requests past its budget (fail-safe). The global budget_exceeded_throttle_percentage is configurable from the admin UI (Settings -> General Settings), persisted through litellm_settings so it survives a restart, not only from config.yaml. Resolves LIT-3894. Scope for LIT-3893.
This commit is contained in:
Yassin Kortam
GitHub
parent
6041d37414
commit
68f997dd09
22 changed files with 998 additions and 22 deletions
|
|
@ -379,6 +379,7 @@ budget_duration: Optional[str] = (
|
|||
None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
|
||||
)
|
||||
default_soft_budget: float = DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0
|
||||
budget_exceeded_throttle_percentage: Optional[float] = None
|
||||
forward_traceparent_to_llm_provider: bool = False
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1504,6 +1504,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [
|
|||
"public_model_groups_links",
|
||||
"cost_discount_config",
|
||||
"cost_margin_config",
|
||||
"budget_exceeded_throttle_percentage",
|
||||
]
|
||||
SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"]
|
||||
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))
|
||||
|
|
|
|||
|
|
@ -1070,6 +1070,7 @@ class KeyRequestBase(GenerateRequestBase):
|
|||
budget_id: Optional[str] = None
|
||||
tags: Optional[List[str]] = None
|
||||
disable_global_guardrails: Optional[bool] = None
|
||||
throttle_on_budget_exceeded: Optional[bool] = None
|
||||
enforced_params: Optional[List[str]] = None
|
||||
allowed_routes: Optional[list] = []
|
||||
allowed_passthrough_routes: Optional[list] = None
|
||||
|
|
@ -2469,6 +2470,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
request_route: Optional[str] = None
|
||||
is_session_token: bool = False
|
||||
budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True)
|
||||
budget_throttle_pct: Optional[float] = Field(default=None, exclude=True)
|
||||
user: Optional[Any] = None # Expanded user object when expand=user is used
|
||||
created_by_user: Optional[Any] = None # Expanded created_by user when expand=user is used
|
||||
end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
|
||||
|
|
@ -3859,6 +3861,7 @@ LiteLLM_ManagementEndpoint_MetadataFields = [
|
|||
"allowed_vector_store_indexes",
|
||||
"enforced_batch_output_expires_after",
|
||||
"enforced_file_expires_after",
|
||||
"throttle_on_budget_exceeded",
|
||||
]
|
||||
|
||||
LiteLLM_ManagementEndpoint_MetadataFields_Premium = [
|
||||
|
|
|
|||
|
|
@ -60,6 +60,10 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.budget_throttle import (
|
||||
budget_throttle_percentage,
|
||||
should_throttle_budget_exceeded,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
|
|
@ -2496,6 +2500,7 @@ def _copy_user_api_key_auth_for_cache(
|
|||
) -> UserAPIKeyAuth:
|
||||
copied_key_obj = user_api_key_obj.model_copy()
|
||||
copied_key_obj.budget_reservation = None
|
||||
copied_key_obj.budget_throttle_pct = None
|
||||
copied_key_obj.parent_otel_span = None
|
||||
copied_key_obj.request_route = None
|
||||
return copied_key_obj
|
||||
|
|
@ -3428,6 +3433,24 @@ async def is_valid_fallback_model(
|
|||
return True
|
||||
|
||||
|
||||
def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool:
|
||||
"""
|
||||
Throttle an over-budget key instead of blocking it, when the key opted in
|
||||
via `throttle_on_budget_exceeded` and a global percentage is configured.
|
||||
|
||||
Records the percentage on the request-scoped `budget_throttle_pct` so the
|
||||
rate limiter scales the key's TPM/RPM down to it; the persistent limits are
|
||||
left untouched so the throttle never compounds across requests. Returns True
|
||||
when the key was throttled (caller skips raising), False when it should still
|
||||
be hard-blocked.
|
||||
"""
|
||||
pct = budget_throttle_percentage()
|
||||
if pct is None or not should_throttle_budget_exceeded(valid_token):
|
||||
return False
|
||||
valid_token.budget_throttle_pct = pct
|
||||
return True
|
||||
|
||||
|
||||
async def _virtual_key_max_budget_check(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
|
|
@ -3488,6 +3511,8 @@ async def _virtual_key_max_budget_check(
|
|||
# so a NaN max_budget would silently disable enforcement. Treat a
|
||||
# non-finite max_budget as "no configured limit" rather than as a bypass.
|
||||
if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget:
|
||||
if _apply_budget_exceeded_throttle(valid_token):
|
||||
return
|
||||
# name the key in the error so operators don't have to reverse-map
|
||||
# spend back to a key; key_name is the masked form (last 4 chars)
|
||||
key_label = valid_token.key_alias or "key"
|
||||
|
|
|
|||
56
litellm/proxy/auth/budget_throttle.py
Normal file
56
litellm/proxy/auth/budget_throttle.py
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
"""
|
||||
Throttle a key after it exceeds its own ``max_budget`` instead of blocking it.
|
||||
|
||||
When a key opts in via ``throttle_on_budget_exceeded`` and a global
|
||||
``budget_exceeded_throttle_percentage`` is configured, an over-budget key keeps
|
||||
serving requests but at a reduced TPM/RPM (the configured percentage of its
|
||||
configured limits). The decision (over budget + opted in) is made once during
|
||||
auth; the scaling is recomputed from the key's original limits on every request
|
||||
so it never compounds across requests.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
def budget_throttle_percentage() -> Optional[float]:
|
||||
"""
|
||||
The global throttle percentage, or None when throttling is disabled /
|
||||
misconfigured (in which case an over-budget key is hard-blocked, the safe
|
||||
default).
|
||||
"""
|
||||
pct = litellm.budget_exceeded_throttle_percentage
|
||||
if not isinstance(pct, (int, float)) or isinstance(pct, bool):
|
||||
return None
|
||||
if not 0 < pct <= 1:
|
||||
return None
|
||||
return float(pct)
|
||||
|
||||
|
||||
def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool:
|
||||
"""
|
||||
True when a key that exceeded its own ``max_budget`` should be throttled
|
||||
rather than blocked: it opted in, a valid global percentage is set, and the
|
||||
key has a TPM or RPM limit to scale down. A key with neither limit has
|
||||
nothing to throttle, so it stays hard-blocked (the safe default) rather than
|
||||
serving unlimited requests past its budget.
|
||||
"""
|
||||
if (valid_token.metadata or {}).get("throttle_on_budget_exceeded") is not True:
|
||||
return False
|
||||
if valid_token.tpm_limit is None and valid_token.rpm_limit is None:
|
||||
return False
|
||||
return budget_throttle_percentage() is not None
|
||||
|
||||
|
||||
def throttled_limit(limit: Optional[int], pct: Optional[float]) -> Optional[int]:
|
||||
"""
|
||||
Scale a TPM/RPM limit to ``pct`` of its value, keeping a trickle of at least
|
||||
1 so a throttled key is slowed rather than fully locked out. An unset limit
|
||||
or unset percentage leaves the limit unchanged.
|
||||
"""
|
||||
if limit is None or pct is None:
|
||||
return limit
|
||||
return max(1, math.floor(limit * pct))
|
||||
|
|
@ -2004,6 +2004,11 @@ async def _user_api_key_auth_builder(
|
|||
|
||||
valid_token_dict = valid_token.model_dump(exclude_none=True)
|
||||
valid_token_dict.pop("token", None)
|
||||
# budget_throttle_pct is excluded from model_dump (it must not leak
|
||||
# into serialized responses), so carry the request-scoped decision
|
||||
# forward by hand to the auth object the rate limiter receives.
|
||||
if valid_token.budget_throttle_pct is not None:
|
||||
valid_token_dict["budget_throttle_pct"] = valid_token.budget_throttle_pct
|
||||
|
||||
if _end_user_object is not None:
|
||||
valid_token_dict.update(end_user_params)
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
get_key_model_rpm_limit,
|
||||
get_key_model_tpm_limit,
|
||||
)
|
||||
from litellm.proxy.auth.budget_throttle import throttled_limit
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
|
||||
|
||||
|
|
@ -248,10 +249,11 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
|
|||
if data is None:
|
||||
data = {}
|
||||
global_max_parallel_requests = data.get("metadata", {}).get("global_max_parallel_requests", None)
|
||||
tpm_limit = getattr(user_api_key_dict, "tpm_limit", sys.maxsize)
|
||||
throttle_pct = getattr(user_api_key_dict, "budget_throttle_pct", None)
|
||||
tpm_limit = throttled_limit(getattr(user_api_key_dict, "tpm_limit", sys.maxsize), throttle_pct)
|
||||
if tpm_limit is None:
|
||||
tpm_limit = sys.maxsize
|
||||
rpm_limit = getattr(user_api_key_dict, "rpm_limit", sys.maxsize)
|
||||
rpm_limit = throttled_limit(getattr(user_api_key_dict, "rpm_limit", sys.maxsize), throttle_pct)
|
||||
if rpm_limit is None:
|
||||
rpm_limit = sys.maxsize
|
||||
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_utils import get_model_rate_limit_from_metadata
|
||||
from litellm.proxy.auth.budget_throttle import throttled_limit
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import (
|
||||
ProxyRateLimitError,
|
||||
map_v3_rate_limit_type,
|
||||
|
|
@ -1549,18 +1550,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
or user_api_key_dict.tpm_limit is not None
|
||||
or user_api_key_dict.max_parallel_requests is not None
|
||||
):
|
||||
throttle_pct = user_api_key_dict.budget_throttle_pct
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="api_key",
|
||||
value=user_api_key_dict.api_key,
|
||||
rate_limit={
|
||||
"requests_per_unit": self._get_enforced_limit(
|
||||
limit_value=user_api_key_dict.rpm_limit,
|
||||
limit_value=throttled_limit(user_api_key_dict.rpm_limit, throttle_pct),
|
||||
limit_type=rpm_limit_type,
|
||||
model_has_failures=model_has_failures,
|
||||
),
|
||||
"tokens_per_unit": self._get_enforced_limit(
|
||||
limit_value=user_api_key_dict.tpm_limit,
|
||||
limit_value=throttled_limit(user_api_key_dict.tpm_limit, throttle_pct),
|
||||
limit_type=tpm_limit_type,
|
||||
model_has_failures=model_has_failures,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -758,6 +758,12 @@ async def _common_key_generation_helper(
|
|||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."},
|
||||
)
|
||||
|
||||
if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None:
|
||||
await validate_team_id_used_in_service_account_request(
|
||||
team_id=data.team_id,
|
||||
|
|
@ -1483,6 +1489,7 @@ async def generate_key_fn(
|
|||
- guardrails: Optional[List[str]] - List of active guardrails for the key
|
||||
- policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
|
||||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
|
||||
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
|
||||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
|
|
@ -2317,6 +2324,16 @@ async def _validate_update_key_data(
|
|||
or "budget_limits" in data.model_fields_set
|
||||
)
|
||||
|
||||
_existing_metadata = getattr(existing_key_row, "metadata", None)
|
||||
_existing_throttle = (
|
||||
_existing_metadata.get("throttle_on_budget_exceeded") if isinstance(_existing_metadata, dict) else None
|
||||
)
|
||||
if data.throttle_on_budget_exceeded is True and _existing_throttle is not True and not _is_proxy_admin:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."},
|
||||
)
|
||||
|
||||
# Personal-key bypass: the caller both created the key AND still owns it
|
||||
# (user_id == caller). Checking only created_by would let a demoted admin
|
||||
# who originally created a key for another user continue editing it without
|
||||
|
|
@ -2507,6 +2524,7 @@ async def update_key_fn(
|
|||
- guardrails: Optional[List[str]] - List of active guardrails for the key
|
||||
- policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
|
||||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
|
||||
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
|
||||
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
|
||||
- blocked: Optional[bool] - Whether the key is blocked
|
||||
- aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases)
|
||||
|
|
|
|||
|
|
@ -14285,6 +14285,9 @@ async def update_config_general_settings(
|
|||
detail={"error": CommonProxyErrors.not_allowed_access.value},
|
||||
)
|
||||
|
||||
if data.field_name in _GENERAL_SETTINGS_UI_LITELLM_FIELDS:
|
||||
return await _persist_general_settings_ui_litellm_field(data.field_name, data.field_value, user_api_key_dict)
|
||||
|
||||
if data.field_name not in ConfigGeneralSettings.model_fields:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
|
|
@ -14550,6 +14553,55 @@ async def get_config_general_settings(
|
|||
)
|
||||
|
||||
|
||||
_GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, dict[str, str]] = {
|
||||
"budget_exceeded_throttle_percentage": {
|
||||
"type": "Float",
|
||||
"description": (
|
||||
"Fraction (0, 1] of a key's configured TPM/RPM that an over-budget key with "
|
||||
"'Throttle on budget exceeded' enabled keeps serving at. Leave empty to hard-block "
|
||||
"over-budget keys."
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> Optional[float]:
|
||||
if value is None or value == "":
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or not (0 < float(value) <= 1):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"{field_name} must be a number in (0, 1] or empty"},
|
||||
)
|
||||
return float(value)
|
||||
|
||||
|
||||
async def _persist_general_settings_ui_litellm_field(
|
||||
field_name: str, value: Any, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> dict:
|
||||
validated = _validate_general_settings_ui_litellm_value(field_name, value)
|
||||
config = await proxy_config.get_config()
|
||||
before_value = config.get("litellm_settings", {}).get(field_name)
|
||||
setattr(litellm, field_name, validated)
|
||||
if "litellm_settings" not in config:
|
||||
config["litellm_settings"] = {}
|
||||
config["litellm_settings"][field_name] = validated
|
||||
await proxy_config.save_config(new_config=config)
|
||||
asyncio.create_task(create_config_audit_log(field_name, "updated", before_value, validated, user_api_key_dict))
|
||||
return {"message": f"Field {field_name} updated", "status": "success"}
|
||||
|
||||
|
||||
async def _reset_general_settings_ui_litellm_field(field_name: str, user_api_key_dict: UserAPIKeyAuth) -> dict:
|
||||
config = await proxy_config.get_config()
|
||||
before_value = config.get("litellm_settings", {}).get(field_name)
|
||||
setattr(litellm, field_name, None)
|
||||
if "litellm_settings" in config:
|
||||
config["litellm_settings"].pop(field_name, None)
|
||||
await proxy_config.save_config(new_config=config)
|
||||
asyncio.create_task(create_config_audit_log(field_name, "deleted", before_value, None, user_api_key_dict))
|
||||
return {"message": f"Field {field_name} reset", "status": "success"}
|
||||
|
||||
|
||||
@router.get(
|
||||
"/config/list",
|
||||
tags=["config.yaml"],
|
||||
|
|
@ -14703,6 +14755,35 @@ async def get_config_list(
|
|||
)
|
||||
return_val.append(_response_obj)
|
||||
|
||||
db_litellm_settings_row = await ConfigRepository(prisma_client).table.find_first(
|
||||
where={"param_name": "litellm_settings"}
|
||||
)
|
||||
db_litellm_settings: dict = (
|
||||
dict(db_litellm_settings_row.param_value)
|
||||
if db_litellm_settings_row is not None and db_litellm_settings_row.param_value is not None
|
||||
else {}
|
||||
)
|
||||
for litellm_field_name, spec in _GENERAL_SETTINGS_UI_LITELLM_FIELDS.items():
|
||||
current_value: Optional[float] = getattr(litellm, litellm_field_name, None)
|
||||
stored_in_db_litellm: Optional[bool]
|
||||
if litellm_field_name in db_litellm_settings:
|
||||
stored_in_db_litellm = True
|
||||
elif current_value is not None:
|
||||
stored_in_db_litellm = False
|
||||
else:
|
||||
stored_in_db_litellm = None
|
||||
return_val.append(
|
||||
ConfigList(
|
||||
field_name=litellm_field_name,
|
||||
field_type=spec["type"],
|
||||
field_description=spec["description"],
|
||||
field_value=current_value,
|
||||
stored_in_db=stored_in_db_litellm,
|
||||
field_default_value=None,
|
||||
nested_fields=None,
|
||||
)
|
||||
)
|
||||
|
||||
return return_val
|
||||
|
||||
|
||||
|
|
@ -14743,6 +14824,9 @@ async def delete_config_general_settings(
|
|||
},
|
||||
)
|
||||
|
||||
if data.field_name in _GENERAL_SETTINGS_UI_LITELLM_FIELDS:
|
||||
return await _reset_general_settings_ui_litellm_field(data.field_name, user_api_key_dict)
|
||||
|
||||
if data.field_name not in ConfigGeneralSettings.model_fields:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_utils import get_model_from_request
|
||||
from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
|
@ -54,6 +55,61 @@ def get_reserved_counter_keys(budget_reservation: Optional[dict]) -> set:
|
|||
}
|
||||
|
||||
|
||||
def _key_reservation_should_release_for_throttle(counter_key: str, valid_token: Optional[UserAPIKeyAuth]) -> bool:
|
||||
"""
|
||||
Whether an over-budget key's own ``max_budget`` reservation should be
|
||||
released rather than blocked, because the key opted into throttling: the
|
||||
rate limiter slows it instead. Only the key's own ``max_budget`` counter is
|
||||
exempt; team/user/window counters still enforce normally, and under-budget
|
||||
requests never reach this branch so their concurrent-overspend protection is
|
||||
untouched.
|
||||
"""
|
||||
if valid_token is None:
|
||||
return False
|
||||
return counter_key == f"spend:key:{valid_token.token}" and should_throttle_budget_exceeded(valid_token)
|
||||
|
||||
|
||||
async def _apply_over_budget_reservation_policy(
|
||||
counter: _BudgetCounter,
|
||||
valid_token: Optional[UserAPIKeyAuth],
|
||||
entry: dict[str, Any],
|
||||
applied_entries: list[dict[str, Any]],
|
||||
reservation_cost: float,
|
||||
current_spend: float,
|
||||
) -> float:
|
||||
"""
|
||||
Decide what to do when a counter is over budget, and return the reservation
|
||||
cost to carry into the next counter. Three outcomes: an over-budget key that
|
||||
opted into throttling releases its own reservation (the rate limiter slows
|
||||
it) and keeps the cost; a partially-remaining budget resizes the reservation
|
||||
down to what is left; anything else hard-blocks by raising.
|
||||
"""
|
||||
if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token):
|
||||
await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost)
|
||||
applied_entries.remove(entry)
|
||||
return reservation_cost
|
||||
|
||||
remaining_before_reservation = counter.max_budget - (current_spend - reservation_cost)
|
||||
if remaining_before_reservation > 1e-12:
|
||||
await _resize_applied_reservation(
|
||||
entries=applied_entries,
|
||||
current_reserved_cost=reservation_cost,
|
||||
new_reserved_cost=remaining_before_reservation,
|
||||
)
|
||||
return remaining_before_reservation
|
||||
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=current_spend,
|
||||
max_budget=counter.max_budget,
|
||||
message=(
|
||||
"Budget has been exceeded! "
|
||||
f"{counter.entity_type}={counter.entity_id} "
|
||||
f"Current cost: {current_spend}, "
|
||||
f"Max budget: {counter.max_budget}"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def reserve_budget_for_request(
|
||||
request_body: dict,
|
||||
route: str,
|
||||
|
|
@ -130,25 +186,15 @@ async def reserve_budget_for_request(
|
|||
cached_spend = await _get_current_counter_value(counter=counter)
|
||||
current_spend = cached_spend + reservation_cost
|
||||
if current_spend > counter.max_budget:
|
||||
remaining_before_reservation = counter.max_budget - (current_spend - reservation_cost)
|
||||
if remaining_before_reservation > 1e-12:
|
||||
await _resize_applied_reservation(
|
||||
entries=applied_entries,
|
||||
current_reserved_cost=reservation_cost,
|
||||
new_reserved_cost=remaining_before_reservation,
|
||||
)
|
||||
reservation_cost = remaining_before_reservation
|
||||
continue
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=current_spend,
|
||||
max_budget=counter.max_budget,
|
||||
message=(
|
||||
"Budget has been exceeded! "
|
||||
f"{counter.entity_type}={counter.entity_id} "
|
||||
f"Current cost: {current_spend}, "
|
||||
f"Max budget: {counter.max_budget}"
|
||||
),
|
||||
reservation_cost = await _apply_over_budget_reservation_policy(
|
||||
counter=counter,
|
||||
valid_token=valid_token,
|
||||
entry=entry,
|
||||
applied_entries=applied_entries,
|
||||
reservation_cost=reservation_cost,
|
||||
current_spend=current_spend,
|
||||
)
|
||||
continue
|
||||
except Exception:
|
||||
await _release_applied_entries_best_effort(
|
||||
entries=applied_entries,
|
||||
|
|
|
|||
|
|
@ -2635,6 +2635,170 @@ async def test_virtual_key_budget_check_fallback_no_counter():
|
|||
assert exc_info.value.current_cost == 15.0
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Throttle-on-budget-exceeded tests (LIT-3894): an over-budget key that
|
||||
# opted in is throttled to a global % of its TPM/RPM instead of blocked.
|
||||
# =====================================================================
|
||||
|
||||
|
||||
def _over_budget_token(**overrides) -> UserAPIKeyAuth:
|
||||
base = dict(
|
||||
token="throttle-token",
|
||||
spend=20.0,
|
||||
max_budget=10.0,
|
||||
user_id="test-user",
|
||||
)
|
||||
base.update(overrides)
|
||||
return UserAPIKeyAuth(**base)
|
||||
|
||||
|
||||
def _patched_spend(value: float):
|
||||
async def mock_get_current_spend(
|
||||
counter_key, fallback_spend, max_budget=None, **kwargs
|
||||
):
|
||||
return value
|
||||
|
||||
return patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend)
|
||||
|
||||
|
||||
def _budget_logging_obj():
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
|
||||
proxy_logging_obj.budget_alerts = AsyncMock()
|
||||
return proxy_logging_obj
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"limit, pct, expected",
|
||||
[
|
||||
(1000, 0.1, 100),
|
||||
(100, 0.1, 10),
|
||||
(1, 0.1, 1), # floor would be 0; trickle of 1 keeps the key alive
|
||||
(None, 0.1, None),
|
||||
(50, 0.5, 25),
|
||||
(1000, None, 1000), # no percentage -> limit unchanged
|
||||
],
|
||||
)
|
||||
def test_throttled_limit(limit, pct, expected):
|
||||
from litellm.proxy.auth.budget_throttle import throttled_limit
|
||||
|
||||
assert throttled_limit(limit, pct) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_exceeded_throttles_instead_of_blocking(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1)
|
||||
valid_token = _over_budget_token(
|
||||
tpm_limit=1000,
|
||||
rpm_limit=100,
|
||||
metadata={"throttle_on_budget_exceeded": True},
|
||||
)
|
||||
|
||||
with _patched_spend(20.0):
|
||||
await _virtual_key_max_budget_check(
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=_budget_logging_obj(),
|
||||
)
|
||||
|
||||
# persistent limits are untouched (so the throttle never compounds); the
|
||||
# request-scoped percentage is what the rate limiter scales by
|
||||
assert valid_token.budget_throttle_pct == 0.1
|
||||
assert valid_token.tpm_limit == 1000
|
||||
assert valid_token.rpm_limit == 100
|
||||
# the request-scoped decision must not leak into serialized responses
|
||||
assert "budget_throttle_pct" not in valid_token.model_dump()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_throttle_decision_cleared_before_caching():
|
||||
"""The request-scoped throttle decision must not persist into the key cache,
|
||||
otherwise it would re-apply (and compound) on every subsequent request."""
|
||||
from litellm.proxy.auth.auth_checks import _copy_user_api_key_auth_for_cache
|
||||
|
||||
valid_token = _over_budget_token(
|
||||
tpm_limit=1000, rpm_limit=100, metadata={"throttle_on_budget_exceeded": True}
|
||||
)
|
||||
valid_token.budget_throttle_pct = 0.1
|
||||
|
||||
cached = _copy_user_api_key_auth_for_cache(user_api_key_obj=valid_token)
|
||||
|
||||
assert cached.budget_throttle_pct is None
|
||||
assert cached.tpm_limit == 1000
|
||||
assert cached.rpm_limit == 100
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_exceeded_throttle_no_configured_limits(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1)
|
||||
valid_token = _over_budget_token(metadata={"throttle_on_budget_exceeded": True})
|
||||
assert valid_token.tpm_limit is None
|
||||
assert valid_token.rpm_limit is None
|
||||
|
||||
with _patched_spend(20.0):
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await _virtual_key_max_budget_check(
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=_budget_logging_obj(),
|
||||
)
|
||||
|
||||
assert valid_token.budget_throttle_pct is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_exceeded_not_opted_in_still_blocks(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1)
|
||||
valid_token = _over_budget_token(tpm_limit=1000, rpm_limit=100)
|
||||
|
||||
with _patched_spend(20.0):
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await _virtual_key_max_budget_check(
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=_budget_logging_obj(),
|
||||
)
|
||||
|
||||
assert valid_token.budget_throttle_pct is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pct", [None, 0, 1.5, -0.1, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_exceeded_invalid_percentage_blocks(monkeypatch, pct):
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", pct)
|
||||
valid_token = _over_budget_token(
|
||||
tpm_limit=1000,
|
||||
rpm_limit=100,
|
||||
metadata={"throttle_on_budget_exceeded": True},
|
||||
)
|
||||
|
||||
with _patched_spend(20.0):
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await _virtual_key_max_budget_check(
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=_budget_logging_obj(),
|
||||
)
|
||||
|
||||
assert valid_token.budget_throttle_pct is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_under_budget_does_not_throttle(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1)
|
||||
valid_token = _over_budget_token(
|
||||
max_budget=100.0,
|
||||
tpm_limit=1000,
|
||||
rpm_limit=100,
|
||||
metadata={"throttle_on_budget_exceeded": True},
|
||||
)
|
||||
|
||||
with _patched_spend(5.0):
|
||||
await _virtual_key_max_budget_check(
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=_budget_logging_obj(),
|
||||
)
|
||||
|
||||
assert valid_token.budget_throttle_pct is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_budget_check_reads_from_spend_counter():
|
||||
"""Team budget check should use get_current_spend when counter exists."""
|
||||
|
|
|
|||
|
|
@ -48,6 +48,42 @@ def time_controller(monkeypatch):
|
|||
return controller
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"throttle_pct, expected_rpm, expected_tpm",
|
||||
[
|
||||
(None, 100, 1000), # no throttle -> configured limits
|
||||
(0.1, 10, 100), # 10% of configured
|
||||
(0.5, 50, 500),
|
||||
],
|
||||
)
|
||||
def test_api_key_descriptor_applies_budget_throttle(
|
||||
throttle_pct, expected_rpm, expected_tpm
|
||||
):
|
||||
"""The api_key rate-limit descriptor scales the key's configured TPM/RPM by
|
||||
the request-scoped budget_throttle_pct, leaving the configured limits intact."""
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache())
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-throttle"),
|
||||
rpm_limit=100,
|
||||
tpm_limit=1000,
|
||||
budget_throttle_pct=throttle_pct,
|
||||
)
|
||||
|
||||
descriptors = handler._create_rate_limit_descriptors(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data={},
|
||||
rpm_limit_type=None,
|
||||
tpm_limit_type=None,
|
||||
model_has_failures=False,
|
||||
)
|
||||
|
||||
api_key_descriptor = next(d for d in descriptors if d["key"] == "api_key")
|
||||
assert api_key_descriptor["rate_limit"]["requests_per_unit"] == expected_rpm
|
||||
assert api_key_descriptor["rate_limit"]["tokens_per_unit"] == expected_tpm
|
||||
|
||||
|
||||
@pytest.mark.flaky(reruns=3)
|
||||
@pytest.mark.asyncio
|
||||
async def test_sliding_window_rate_limit_v3(monkeypatch, time_controller):
|
||||
|
|
|
|||
|
|
@ -1495,6 +1495,57 @@ async def test_generate_service_account_works_with_team_id():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_key_throttle_rejected_for_non_admin():
|
||||
"""Security regression: a non-admin creating a key must not be able to set
|
||||
throttle_on_budget_exceeded=true, which would let the new key keep spending
|
||||
past an admin-imposed per-key budget ceiling instead of hard-blocking. The
|
||||
/key/update gate does not cover generate, so generate needs its own admin
|
||||
check. Only the enable value is gated, so this must 403."""
|
||||
mock_prisma_client = AsyncMock()
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(throttle_on_budget_exceeded=True),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-alice",
|
||||
user_id="alice",
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
team_table=None,
|
||||
)
|
||||
assert int(getattr(exc.value, "status_code", 0)) == 403
|
||||
assert "Only proxy admins can enable throttle_on_budget_exceeded" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_key_throttle_allowed_for_admin():
|
||||
"""A proxy admin may create a key with throttle_on_budget_exceeded=true; the
|
||||
generate admin gate must let the admin through to key creation."""
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
|
||||
patch("litellm.proxy.proxy_server.llm_router", None),
|
||||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn"
|
||||
) as mock_generate_key,
|
||||
):
|
||||
mock_generate_key.return_value = {
|
||||
"key": "sk-test-key",
|
||||
"expires": None,
|
||||
"user_id": "admin",
|
||||
"team_id": None,
|
||||
}
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(throttle_on_budget_exceeded=True),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"),
|
||||
litellm_changed_by=None,
|
||||
team_table=None,
|
||||
)
|
||||
assert mock_generate_key.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_service_account_requires_team_id():
|
||||
data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"})
|
||||
|
|
@ -9577,6 +9628,165 @@ async def test_update_key_non_budget_fields_allowed_for_internal_user(monkeypatc
|
|||
assert result is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_throttle_on_budget_exceeded_rejected_for_internal_user(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Security regression: throttle_on_budget_exceeded turns an admin-imposed
|
||||
hard budget block into a soft throttle that keeps spending past max_budget,
|
||||
so it is a budget-enforcement change. A non-admin key owner (same setup that
|
||||
is allowed to change non-budget fields via the caller_is_creator shortcut)
|
||||
must NOT be able to self-opt-in to it; it has to route through the admin-only
|
||||
_check_key_admin_access and return 403. Without treating the flag as a budget
|
||||
change this update would succeed, letting the owner bypass their own cap."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
update_key_fn,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_proxy_logging_obj = MagicMock()
|
||||
|
||||
test_hashed_token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
|
||||
|
||||
# Owner of the key (created_by == user_id) so caller_is_creator is True.
|
||||
# This is exactly the setup that is allowed to change non-budget fields;
|
||||
# the throttle flag must still be rejected.
|
||||
mock_existing_key = MagicMock()
|
||||
mock_existing_key.token = test_hashed_token
|
||||
mock_existing_key.user_id = "internal_user"
|
||||
mock_existing_key.created_by = "internal_user"
|
||||
mock_existing_key.team_id = None
|
||||
mock_existing_key.project_id = None
|
||||
mock_existing_key.max_budget = 10.0
|
||||
mock_existing_key.key_alias = None
|
||||
mock_existing_key.models = []
|
||||
mock_existing_key.metadata = {}
|
||||
mock_existing_key.model_dump.return_value = {
|
||||
"token": test_hashed_token,
|
||||
"user_id": "internal_user",
|
||||
"team_id": None,
|
||||
"max_budget": 10.0,
|
||||
}
|
||||
|
||||
mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=mock_existing_key)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.query_params = {}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-internal",
|
||||
user_id="internal_user",
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_key_fn(
|
||||
request=mock_request,
|
||||
data=UpdateKeyRequest(
|
||||
key=test_hashed_token,
|
||||
throttle_on_budget_exceeded=True,
|
||||
),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert str(exc.value.code) == "403"
|
||||
assert "Only proxy admins can enable throttle_on_budget_exceeded" in str(exc.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_throttle_unchanged_allows_non_budget_edit_for_internal_user(
|
||||
monkeypatch,
|
||||
):
|
||||
"""A non-admin owner editing a non-budget field must not be blocked just
|
||||
because the UI resends throttle_on_budget_exceeded unchanged (the edit form
|
||||
always includes it). Only the transition to enabled is admin-gated, so an
|
||||
unchanged False here leaves the key owner's non-budget edit working."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
update_key_fn,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_proxy_logging_obj = MagicMock()
|
||||
|
||||
test_hashed_token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
|
||||
|
||||
mock_existing_key = MagicMock()
|
||||
mock_existing_key.token = test_hashed_token
|
||||
mock_existing_key.user_id = "internal_user"
|
||||
mock_existing_key.created_by = "internal_user"
|
||||
mock_existing_key.team_id = None
|
||||
mock_existing_key.project_id = None
|
||||
mock_existing_key.max_budget = 10.0
|
||||
mock_existing_key.key_alias = None
|
||||
mock_existing_key.models = []
|
||||
mock_existing_key.metadata = {"throttle_on_budget_exceeded": False}
|
||||
mock_existing_key.model_dump.return_value = {
|
||||
"token": test_hashed_token,
|
||||
"user_id": "internal_user",
|
||||
"team_id": None,
|
||||
"max_budget": 10.0,
|
||||
}
|
||||
|
||||
mock_updated_key = MagicMock()
|
||||
mock_updated_key.token = test_hashed_token
|
||||
mock_updated_key.key_alias = "my-alias"
|
||||
|
||||
mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
|
||||
mock_prisma_client.update_data = AsyncMock(return_value=mock_updated_key)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=mock_existing_key)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
monkeypatch.setattr("litellm.store_audit_logs", False)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", lambda token: test_hashed_token)
|
||||
|
||||
async def _noop(**kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
|
||||
_noop,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._enforce_unique_key_alias",
|
||||
_noop,
|
||||
)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.query_params = {}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-internal",
|
||||
user_id="internal_user",
|
||||
)
|
||||
|
||||
result = await update_key_fn(
|
||||
request=mock_request,
|
||||
data=UpdateKeyRequest(
|
||||
key=test_hashed_token,
|
||||
key_alias="my-alias",
|
||||
throttle_on_budget_exceeded=False,
|
||||
),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_non_budget_rejects_cross_user_modification(monkeypatch):
|
||||
"""Regression: previously _check_key_admin_access was gated on
|
||||
|
|
|
|||
|
|
@ -58,6 +58,101 @@ def _request_body() -> dict:
|
|||
}
|
||||
|
||||
|
||||
async def _reserve(valid_token, cost, key_cache, proxy_logging_obj):
|
||||
with patch(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
|
||||
return_value=cost,
|
||||
):
|
||||
return await reserve_budget_for_request(
|
||||
request_body=_request_body(),
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
valid_token=valid_token,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reservation_still_protects_under_budget_throttled_key(
|
||||
spend_counter_state, monkeypatch
|
||||
):
|
||||
"""An opted-in key that is still under budget keeps its reservation counter,
|
||||
so concurrent requests can't collectively overshoot max_budget."""
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1)
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="key-throttle-under",
|
||||
spend=0.0,
|
||||
max_budget=1.0,
|
||||
metadata={"throttle_on_budget_exceeded": True},
|
||||
)
|
||||
|
||||
reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj)
|
||||
|
||||
assert reservation is not None
|
||||
assert (
|
||||
counter_cache.in_memory_cache.get_cache(key="spend:key:key-throttle-under")
|
||||
== 0.6
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reservation_does_not_block_over_budget_throttled_key(
|
||||
spend_counter_state, monkeypatch
|
||||
):
|
||||
"""Once an opted-in key is over budget the reservation path must not raise;
|
||||
the rate limiter throttles it instead."""
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1)
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="key-throttle-over",
|
||||
spend=0.0,
|
||||
max_budget=1.0,
|
||||
tpm_limit=1000,
|
||||
rpm_limit=100,
|
||||
metadata={"throttle_on_budget_exceeded": True},
|
||||
)
|
||||
|
||||
# first reservation lands under budget (counter -> 0.6)
|
||||
await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj)
|
||||
|
||||
# any further request is over budget (0.6 + 0.6 > 1.0): the opted-in key is
|
||||
# released and allowed through (None), not blocked, and its over-budget
|
||||
# increment is released so the counter is not permanently inflated
|
||||
result = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj)
|
||||
assert result is None
|
||||
assert (
|
||||
counter_cache.in_memory_cache.get_cache(key="spend:key:key-throttle-over")
|
||||
== 0.6
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reservation_blocks_over_budget_non_throttled_key(
|
||||
spend_counter_state, monkeypatch
|
||||
):
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1)
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="key-no-optin-over",
|
||||
spend=0.0,
|
||||
max_budget=1.0,
|
||||
)
|
||||
|
||||
await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj)
|
||||
await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) # counter -> 1.0
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj)
|
||||
|
||||
|
||||
def test_should_not_serialize_budget_reservation_on_user_api_key_auth():
|
||||
auth = UserAPIKeyAuth(
|
||||
token="key-budget-runtime-state",
|
||||
|
|
|
|||
|
|
@ -8644,6 +8644,149 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch):
|
|||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch):
|
||||
"""The throttle fraction is a litellm_settings scalar surfaced on the General
|
||||
Settings table as a Float field so it sits with the other global limits; it
|
||||
must appear in /config/list reading its live litellm.<attr> value."""
|
||||
import types
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_config_table = MagicMock()
|
||||
mock_config_table.find_first = AsyncMock(return_value=None)
|
||||
mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table)
|
||||
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.15)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
try:
|
||||
client = TestClient(app)
|
||||
resp = client.get("/config/list", params={"config_type": "general_settings"})
|
||||
assert resp.status_code == 200, resp.text
|
||||
fields = {item["field_name"]: item for item in resp.json()}
|
||||
assert "budget_exceeded_throttle_percentage" in fields
|
||||
assert fields["budget_exceeded_throttle_percentage"]["field_type"] == "Float"
|
||||
assert fields["budget_exceeded_throttle_percentage"]["field_value"] == 0.15
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_config_field_throttle_persists_to_litellm_settings(monkeypatch):
|
||||
"""Editing the throttle Float row on the General Settings table routes to
|
||||
litellm_settings (not general_settings): it sets litellm.<attr> live and
|
||||
persists under litellm_settings so the runtime read is unchanged."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import (
|
||||
ConfigFieldUpdate,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.proxy_server import update_config_general_settings
|
||||
|
||||
saved: dict = {}
|
||||
|
||||
async def fake_get_config():
|
||||
return {"litellm_settings": {}}
|
||||
|
||||
async def fake_save_config(new_config=None):
|
||||
saved.update(new_config or {})
|
||||
|
||||
monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config)
|
||||
monkeypatch.setattr(ps.proxy_config, "save_config", fake_save_config)
|
||||
monkeypatch.setattr(ps, "prisma_client", MagicMock())
|
||||
monkeypatch.setattr(litellm, "store_audit_logs", False)
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", None)
|
||||
|
||||
admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
await update_config_general_settings(
|
||||
data=ConfigFieldUpdate(
|
||||
field_name="budget_exceeded_throttle_percentage",
|
||||
field_value=0.1,
|
||||
config_type="general_settings",
|
||||
),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
|
||||
assert litellm.budget_exceeded_throttle_percentage == 0.1
|
||||
assert saved["litellm_settings"]["budget_exceeded_throttle_percentage"] == 0.1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_value", [0, -0.1, 1.5, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_config_field_throttle_rejects_invalid(monkeypatch, bad_value):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import (
|
||||
ConfigFieldUpdate,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.proxy_server import update_config_general_settings
|
||||
|
||||
async def fake_get_config():
|
||||
return {"litellm_settings": {}}
|
||||
|
||||
monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config)
|
||||
monkeypatch.setattr(ps, "prisma_client", MagicMock())
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", None)
|
||||
|
||||
admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await update_config_general_settings(
|
||||
data=ConfigFieldUpdate(
|
||||
field_name="budget_exceeded_throttle_percentage",
|
||||
field_value=bad_value,
|
||||
config_type="general_settings",
|
||||
),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
assert exc.value.status_code == 400
|
||||
assert litellm.budget_exceeded_throttle_percentage is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_config_field_throttle_rejected_for_non_admin(monkeypatch):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import (
|
||||
ConfigFieldUpdate,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.proxy_server import update_config_general_settings
|
||||
|
||||
monkeypatch.setattr(ps, "prisma_client", MagicMock())
|
||||
monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", None)
|
||||
|
||||
non_admin = UserAPIKeyAuth(api_key="k", user_id="u", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
with pytest.raises(HTTPException):
|
||||
await update_config_general_settings(
|
||||
data=ConfigFieldUpdate(
|
||||
field_name="budget_exceeded_throttle_percentage",
|
||||
field_value=0.1,
|
||||
config_type="general_settings",
|
||||
),
|
||||
user_api_key_dict=non_admin,
|
||||
)
|
||||
assert litellm.budget_exceeded_throttle_percentage is None
|
||||
|
||||
|
||||
def test_preserve_redacted_plugin_keys_keeps_stored_credential():
|
||||
"""A redacted or blank plugin_key on update must not overwrite the real key."""
|
||||
from litellm.proxy.proxy_server import _preserve_redacted_plugin_keys
|
||||
|
|
|
|||
|
|
@ -161,6 +161,14 @@ const GeneralSettings: React.FC<GeneralSettingsPageProps> = ({ accessToken, user
|
|||
checked={value.field_value === true || value.field_value === "true"}
|
||||
onChange={(checked) => handleInputChange(value.field_name, checked)}
|
||||
/>
|
||||
) : value.field_type == "Float" ? (
|
||||
<InputNumber
|
||||
min={0}
|
||||
max={1}
|
||||
step={0.05}
|
||||
value={value.field_value}
|
||||
onChange={(newValue) => handleInputChange(value.field_name, newValue)}
|
||||
/>
|
||||
) : null}
|
||||
</TableCell>
|
||||
<TableCell>
|
||||
|
|
|
|||
|
|
@ -1163,6 +1163,21 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
form={form}
|
||||
showDetailedDescriptions={true}
|
||||
/>
|
||||
<Form.Item
|
||||
className="mt-4"
|
||||
label={
|
||||
<span>
|
||||
Throttle on budget exceeded{" "}
|
||||
<Tooltip title="When this key exceeds its max budget, throttle its TPM/RPM to the globally configured percentage instead of blocking access entirely. Requires budget_exceeded_throttle_percentage in litellm_settings and a TPM/RPM limit on the key.">
|
||||
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name="throttle_on_budget_exceeded"
|
||||
valuePropName="checked"
|
||||
>
|
||||
<Switch checkedChildren="Yes" unCheckedChildren="No" />
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
|
|
|
|||
|
|
@ -335,6 +335,36 @@ describe("KeyEditView", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("should initialize and submit throttle_on_budget_exceeded from key metadata", async () => {
|
||||
const onSubmitMock = vi.fn().mockResolvedValue(undefined);
|
||||
const keyDataWithThrottle = {
|
||||
...MOCK_KEY_DATA,
|
||||
metadata: { ...MOCK_KEY_DATA.metadata, throttle_on_budget_exceeded: true },
|
||||
};
|
||||
|
||||
renderWithProviders(
|
||||
<KeyEditView
|
||||
keyData={keyDataWithThrottle}
|
||||
onCancel={() => {}}
|
||||
onSubmit={onSubmitMock}
|
||||
accessToken={"test-token"}
|
||||
userID={"test-user"}
|
||||
userRole={"admin"}
|
||||
premiumUser={false}
|
||||
/>,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Throttle on budget exceeded")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(onSubmitMock).toHaveBeenCalledWith(expect.objectContaining({ throttle_on_budget_exceeded: true }));
|
||||
});
|
||||
});
|
||||
|
||||
it("should disable models field when management routes are selected", async () => {
|
||||
const keyDataWithManagementRoutes = {
|
||||
...MOCK_KEY_DATA,
|
||||
|
|
|
|||
|
|
@ -181,6 +181,7 @@ export function KeyEditView({
|
|||
metadata: formatMetadataForDisplay(stripTagsFromMetadata(keyData.metadata)),
|
||||
guardrails: keyData.metadata?.guardrails,
|
||||
disable_global_guardrails: keyData.metadata?.disable_global_guardrails || false,
|
||||
throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false,
|
||||
prompts: keyData.metadata?.prompts,
|
||||
tags: keyData.metadata?.tags,
|
||||
vector_stores: keyData.object_permission?.vector_stores || [],
|
||||
|
|
@ -222,6 +223,7 @@ export function KeyEditView({
|
|||
accessGroups: keyData.object_permission?.mcp_access_groups || [],
|
||||
},
|
||||
mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {},
|
||||
throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false,
|
||||
logging_settings: extractLoggingSettings(keyData.metadata),
|
||||
disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks)
|
||||
? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks)
|
||||
|
|
@ -512,6 +514,21 @@ export function KeyEditView({
|
|||
|
||||
<RateLimitTypeFormItem type="rpm" name="rpm_limit_type" showDetailedDescriptions={false} />
|
||||
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Throttle on budget exceeded{" "}
|
||||
<Tooltip title="When this key exceeds its max budget, throttle its TPM/RPM to the globally configured percentage instead of blocking access entirely. Requires budget_exceeded_throttle_percentage in litellm_settings and a TPM/RPM limit on the key.">
|
||||
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
}
|
||||
name="throttle_on_budget_exceeded"
|
||||
valuePropName="checked"
|
||||
>
|
||||
<Switch checkedChildren="Yes" unCheckedChildren="No" />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item label="Max Parallel Requests" name="max_parallel_requests">
|
||||
<NumericalInput min={0} />
|
||||
</Form.Item>
|
||||
|
|
|
|||
|
|
@ -542,6 +542,9 @@ export default function KeyInfoView({
|
|||
<div className="mt-2">
|
||||
<Text>TPM: {currentKeyData.tpm_limit !== null ? currentKeyData.tpm_limit : "Unlimited"}</Text>
|
||||
<Text>RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"}</Text>
|
||||
{Boolean(currentKeyData.metadata?.throttle_on_budget_exceeded) && (
|
||||
<Text>Throttle on budget exceeded: Yes</Text>
|
||||
)}
|
||||
</div>
|
||||
</Card>
|
||||
|
||||
|
|
|
|||
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -6502,6 +6502,7 @@ export interface paths {
|
|||
* - guardrails: Optional[List[str]] - List of active guardrails for the key
|
||||
* - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
|
||||
* - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
|
||||
* - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
|
||||
* - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
* - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||||
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||||
|
|
@ -6905,6 +6906,7 @@ export interface paths {
|
|||
* - guardrails: Optional[List[str]] - List of active guardrails for the key
|
||||
* - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
|
||||
* - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
|
||||
* - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
|
||||
* - prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
|
||||
* - blocked: Optional[bool] - Whether the key is blocked
|
||||
* - aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases)
|
||||
|
|
@ -23707,6 +23709,8 @@ export interface components {
|
|||
tags?: string[] | null;
|
||||
/** Team Id */
|
||||
team_id?: string | null;
|
||||
/** Throttle On Budget Exceeded */
|
||||
throttle_on_budget_exceeded?: boolean | null;
|
||||
/** Tpm Limit */
|
||||
tpm_limit?: number | null;
|
||||
/** Tpm Limit Type */
|
||||
|
|
@ -23847,6 +23851,8 @@ export interface components {
|
|||
tags?: string[] | null;
|
||||
/** Team Id */
|
||||
team_id?: string | null;
|
||||
/** Throttle On Budget Exceeded */
|
||||
throttle_on_budget_exceeded?: boolean | null;
|
||||
/** Token */
|
||||
token?: string | null;
|
||||
/** Token Id */
|
||||
|
|
@ -28186,6 +28192,8 @@ export interface components {
|
|||
team_id?: string | null;
|
||||
/** Teams */
|
||||
teams?: unknown[] | null;
|
||||
/** Throttle On Budget Exceeded */
|
||||
throttle_on_budget_exceeded?: boolean | null;
|
||||
/** Token */
|
||||
token?: string | null;
|
||||
/** Token Id */
|
||||
|
|
@ -29813,6 +29821,8 @@ export interface components {
|
|||
tags?: string[] | null;
|
||||
/** Team Id */
|
||||
team_id?: string | null;
|
||||
/** Throttle On Budget Exceeded */
|
||||
throttle_on_budget_exceeded?: boolean | null;
|
||||
/** Tpm Limit */
|
||||
tpm_limit?: number | null;
|
||||
/** Tpm Limit Type */
|
||||
|
|
@ -31760,6 +31770,8 @@ export interface components {
|
|||
temp_budget_expiry?: string | null;
|
||||
/** Temp Budget Increase */
|
||||
temp_budget_increase?: number | null;
|
||||
/** Throttle On Budget Exceeded */
|
||||
throttle_on_budget_exceeded?: boolean | null;
|
||||
/** Tpm Limit */
|
||||
tpm_limit?: number | null;
|
||||
/** Tpm Limit Type */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue