scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 90

refactor(proxy): type the user window cache fallback and tighten carried window metadata types

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus 2026-09-18 23:50:50 +00:00
parent b730b06691
commit 9024c2784d
4 changed files with 25 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/proxy/auth/auth_checks.py
View file

@ -40,6 +40,7 @@ from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.models.project import LiteLLM_ProjectTable
from litellm.models.team import BudgetLimitEntry
from litellm.proxy._types import (
RBAC_ROLES,
CallInfo,
@ -5599,7 +5600,7 @@ async def _team_multi_budget_check(
async def _user_multi_budget_check(
valid_token: UserAPIKeyAuth | None,
team_object: LiteLLM_TeamTable | None,
general_settings: dict,
general_settings: Mapping[str, object],
):
"""
Raises BudgetExceededError if any budget window in valid_token.user_budget_limits is exceeded.
@ -5617,30 +5618,26 @@ async def _user_multi_budget_check(
from litellm.proxy.proxy_server import get_current_spend
windows: Final[tuple[dict, ...]] = tuple(
window if isinstance(window, dict) else window.model_dump() for window in valid_token.user_budget_limits
)
bind_spend_counter_keys(
frozenset(f"spend:user:{valid_token.user_id}:window:{w['budget_duration']}" for w in windows)
)
windows: Final[tuple[BudgetLimitEntry, ...]] = tuple(valid_token.user_budget_limits)
bind_spend_counter_keys(frozenset(f"spend:user:{valid_token.user_id}:window:{w.budget_duration}" for w in windows))
for w in windows:
counter_key = f"spend:user:{valid_token.user_id}:window:{w['budget_duration']}"
counter_key = f"spend:user:{valid_token.user_id}:window:{w.budget_duration}"
window_spend = await get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
max_budget=w["max_budget"],
max_budget=w.max_budget,
window_entity_type="User",
window_entity_id=valid_token.user_id,
window_duration=str(w["budget_duration"]),
window_duration=str(w.budget_duration),
window_start=get_budget_window_start(w),
)
if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]:
if math.isfinite(w.max_budget) and window_spend >= w.max_budget:
raise litellm.BudgetExceededError(
current_cost=window_spend,
max_budget=w["max_budget"],
max_budget=w.max_budget,
message=(
f"ExceededBudget: User={valid_token.user_id} over {w['budget_duration']} budget. "
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
f"ExceededBudget: User={valid_token.user_id} over {w.budget_duration} budget. "
f"Spend=${window_spend:.4f}, Limit=${w.max_budget:.2f}"
),
entity_type=Litellm_EntityType.USER.value,
entity_id=valid_token.user_id,

25
litellm/proxy/proxy_server.py
View file

@ -3061,15 +3061,11 @@ async def _increment_spend_counters_batched(
)
)
async def _user_window_increment(window: object) -> PendingSpendIncrement | None:
duration = (
window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None)
)
user_window_reset_at = (
window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None)
)
async def _user_window_increment(window: BudgetLimitEntry) -> PendingSpendIncrement | None:
duration: Final = window.budget_duration
user_window_reset_at: Final = window.reset_at
user_window_counter: Final = f"spend:user:{scope_user_id}:window:{duration}"
user_window_start = get_budget_window_start(window)
user_window_start: Final = get_budget_window_start(window)
pending_window: Final = (
await _prepare_window_spend_counter_increment(
counter_key=user_window_counter,
@ -3380,16 +3376,11 @@ async def _enqueue_window_spend_row_update(
)
def _cached_user_budget_limits(cached_user: object) -> tuple[object, ...]:
if cached_user is None:
def _cached_user_budget_limits(cached_user: object) -> tuple[BudgetLimitEntry, ...]:
user: Final = CacheCodec.deserialize(cached_user, LiteLLM_UserTable)
if user is None or user.budget_limits is None:
return ()
raw: Final[object] = (
cached_user.get("budget_limits")
if isinstance(cached_user, dict)
else getattr(cached_user, "budget_limits", None)
)
parsed: Final[object] = json.loads(raw) if isinstance(raw, str) else raw
return tuple(parsed) if isinstance(parsed, list) else ()
return tuple(user.budget_limits)
async def _prepare_window_spend_counter_increment(

2
litellm/proxy/spend_tracking/carried_budget_state.py
View file

@ -66,7 +66,7 @@ def carried_budget_metadata(valid_token: UserAPIKeyAuth) -> Mapping[str, object]
)
def carried_user_budget_limits_metadata(valid_token: UserAPIKeyAuth) -> tuple[dict[str, object], ...] | None:
def carried_user_budget_limits_metadata(valid_token: UserAPIKeyAuth) -> tuple[Mapping[str, object], ...] | None:
if valid_token.user_budget_limits is None:
return None
try:

7
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -13034,11 +13034,14 @@ async def test_team_window_spend_row_is_enqueued():
@pytest.mark.asyncio
async def test_user_window_spend_row_is_enqueued():
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy.proxy_server import increment_spend_counters
reset_at = datetime.now(timezone.utc) + timedelta(days=3)
user_obj = MagicMock()
user_obj.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}]
user_obj = LiteLLM_UserTable(
user_id="user-1",
budget_limits=[{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}],
)
with _window_spend_enqueue_env({"user-1": user_obj}) as queue:
await increment_spend_counters(token=None, team_id=None, user_id="user-1", response_cost=1.5)