diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e4c1c46ca1d..a6215ac873e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5aa4ae12b3..2416569c3f3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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( diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index 43cdeec622b..6f1b1040264 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -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: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index eb712064b43..d98a958e02f 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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)