diff --git a/docs/my-website/docs/proxy/users.md b/docs/my-website/docs/proxy/users.md index 88a7a0f1e07..0e36e84c208 100644 --- a/docs/my-website/docs/proxy/users.md +++ b/docs/my-website/docs/proxy/users.md @@ -333,6 +333,67 @@ curl 'http://0.0.0.0:4000/key/generate' \ }' ``` +#### **Set multiple budget windows on a key** + +Apply multiple concurrent budget limits at different time scales on the same key — for example, cap a key at **$10/day** AND **$100/month**. + +**When is this useful?** + +A single `budget_duration` window can't prevent a bad day from burning your entire month. Multiple budget windows let you: + +- Block a runaway usage spike within the day while still allowing normal monthly spend. +- Give Claude Code rollouts a daily guardrail (`24h`) and a monthly ceiling (`30d`) so a single heavy session doesn't exhaust the whole month. +- Layer fine-grained hourly limits for bursty workloads on top of a weekly cap. + +:::info + +See [User Budget docs](https://docs.litellm.ai/docs/proxy/users) for more on how budgets work across keys, teams, and users. + +::: + +**Via API** + +Pass `budget_limits` as a list of `{budget_duration, max_budget}` objects: + +```bash +curl 'http://0.0.0.0:4000/key/generate' \ +--header 'Authorization: Bearer ' \ +--header 'Content-Type: application/json' \ +--data-raw '{ + "budget_limits": [ + {"budget_duration": "24h", "max_budget": 10}, + {"budget_duration": "30d", "max_budget": 100} + ] +}' +``` + +Each window is tracked independently and resets on its own schedule: + +| `budget_duration` | Resets | +|---|---| +| `1h` | Every hour | +| `24h` | Daily at midnight UTC | +| `7d` | Every Sunday at midnight UTC | +| `30d` | 1st of every month at midnight UTC | + +**Via Dashboard** + +Open **Virtual Keys → Create Key → Optional Settings → Budget Windows**. + +![Step 1 - open key settings](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/18930ba5-67c0-4031-afc0-57f37b4e59e4/ascreenshot_ef79d8a000bb41cdacf1bd9827732ee8_text_export.jpeg) + +Click **+ Add Budget Window** to add a row, choose the period from the dropdown, and enter the spend cap. + +![Step 2 - add a window](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/5ae8c0b3-2d03-41ad-a63c-47b20c350dfe/ascreenshot_1a7dc6c7d65544f38fd8a65604674f22_text_export.jpeg) + +Add a second row for a different time period (e.g. monthly $100 on top of a daily $10). + +![Step 3 - add second window](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/cbded3a7-1086-4e20-8f0f-de154b76146c/ascreenshot_c51c18752c3b4f8b976d28799b2638b6_text_export.jpeg) + +Each window shows the reset schedule below the input so it's always clear when spend resets. + +![Step 4 - reset hints](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/8754f121-1640-4892-9dd0-fd4a870418bf/ascreenshot_8079eb0df2194e8f99e5258ba4b3c082_text_export.jpeg) + ### ✨ Virtual Key (Model Specific) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql new file mode 100644 index 00000000000..fdd65543a3a --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable: add budget_limits column to LiteLLM_VerificationToken +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB; + +-- AlterTable: add budget_limits column to LiteLLM_TeamTable +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d8d6015ce87..c724d957a4e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -140,6 +140,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) @@ -386,6 +387,7 @@ model LiteLLM_VerificationToken { rotation_interval String? // How often to rotate (e.g., "30d", "90d") last_rotation_at DateTime? // When this key was last rotated key_rotation_at DateTime? // When this key should next be rotated + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index abdba09dd8d..51810c5643f 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -281,7 +281,7 @@ class Timeout(openai.APITimeoutError): # type: ignore return _message -class PermissionDeniedError(openai.PermissionDeniedError): # type:ignore +class PermissionDeniedError(openai.PermissionDeniedError): # type: ignore def __init__( self, message, @@ -847,6 +847,7 @@ class BudgetExceededError(Exception): ): self.current_cost = current_cost self.max_budget = max_budget + self.status_code = 429 message = ( message or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e4702364a3e..cc43e50df07 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -862,6 +862,14 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): models: Optional[List[str]] = None +class BudgetLimitEntry(LiteLLMPydanticObjectBase): + """A single budget window with its own limit and independent reset schedule.""" + + budget_duration: str # e.g. "24h", "7d", "30d" + max_budget: float # max spend in USD for this window + reset_at: Optional[datetime] = None # populated at creation/reset time + + class GenerateRequestBase(LiteLLMPydanticObjectBase): """ Overlapping schema between key and user generate/update requests @@ -881,12 +889,15 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None + budget_limits: Optional[List[BudgetLimitEntry]] = ( + None # multiple concurrent budget windows + ) allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[ - dict - ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[dict] = ( + {} + ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -985,6 +996,7 @@ class GenerateKeyResponse(KeyRequestBase): "permissions", "model_max_budget", "router_settings", + "budget_limits", ] for field in dict_fields: value = values.get(field) @@ -1028,9 +1040,9 @@ class RegenerateKeyRequest(GenerateKeyRequest): spend: Optional[float] = None metadata: Optional[dict] = None new_master_key: Optional[str] = None - grace_period: Optional[ - str - ] = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke + grace_period: Optional[str] = ( + None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke + ) class ResetSpendRequest(LiteLLMPydanticObjectBase): @@ -1540,12 +1552,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[ - AllowedModelRegion - ] = None # require all user requests to use models in this specific region - default_model: Optional[ - str - ] = None # if no equivalent model in allowed region - default all requests to this model + allowed_model_region: Optional[AllowedModelRegion] = ( + None # require all user requests to use models in this specific region + ) + default_model: Optional[str] = ( + None # if no equivalent model in allowed region - default all requests to this model + ) object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="before") @@ -1568,12 +1580,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[ - AllowedModelRegion - ] = None # require all user requests to use models in this specific region - default_model: Optional[ - str - ] = None # if no equivalent model in allowed region - default all requests to this model + allowed_model_region: Optional[AllowedModelRegion] = ( + None # require all user requests to use models in this specific region + ) + default_model: Optional[str] = ( + None # if no equivalent model in allowed region - default all requests to this model + ) object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -1638,6 +1650,9 @@ class TeamBase(LiteLLMPydanticObjectBase): max_budget: Optional[float] = None soft_budget: Optional[float] = None budget_duration: Optional[str] = None + budget_limits: Optional[List[BudgetLimitEntry]] = ( + None # multiple concurrent budget windows + ) models: list = [] blocked: bool = False @@ -1663,15 +1678,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[ - float - ] = None # allow user to set a budget for all team members - team_member_rpm_limit: Optional[ - int - ] = None # allow user to set RPM limit for all team members - team_member_tpm_limit: Optional[ - int - ] = None # allow user to set TPM limit for all team members + team_member_budget: Optional[float] = ( + None # allow user to set a budget for all team members + ) + team_member_rpm_limit: Optional[int] = ( + None # allow user to set RPM limit for all team members + ) + team_member_tpm_limit: Optional[int] = ( + None # allow user to set TPM limit for all team members + ) team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1736,6 +1751,9 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): enforced_file_expires_after: Optional[dict] = None router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None + budget_limits: Optional[List[BudgetLimitEntry]] = ( + None # multiple concurrent budget windows + ) class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): @@ -1768,9 +1786,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[ - Literal["success", "failure", "success_and_failure"] - ] = "success_and_failure" + callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( + "success_and_failure" + ) callback_vars: Dict[str, str] @model_validator(mode="before") @@ -1879,6 +1897,7 @@ class LiteLLM_TeamTable(TeamBase): "model_max_budget", "model_aliases", "router_settings", + "budget_limits", ] if isinstance(values, BaseModel): @@ -2110,9 +2129,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[ - List[FieldDetail] - ] = None # For nested dictionary or Pydantic fields + nested_fields: Optional[List[FieldDetail]] = ( + None # For nested dictionary or Pydantic fields + ) class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2361,6 +2380,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): last_rotation_at: Optional[datetime] = None # When this key was last rotated key_rotation_at: Optional[datetime] = None # When this key should next be rotated router_settings: Optional[dict] = None + budget_limits: Optional[List[dict]] = None # multiple concurrent budget windows model_config = ConfigDict(protected_namespaces=()) @@ -2470,9 +2490,9 @@ class UserAPIKeyAuth( user_max_budget: Optional[float] = None request_route: Optional[str] = None 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 + created_by_user: Optional[Any] = ( + None # Expanded created_by user when expand=user is used + ) end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None # Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery # and forwarded into outbound tokens by guardrails such as MCPJWTSigner. @@ -2611,9 +2631,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[ - Any - ] = None # You might want to replace 'Any' with a more specific type if available + user: Optional[Any] = ( + None # You might want to replace 'Any' with a more specific type if available + ) litellm_budget_table: Optional[LiteLLM_BudgetTable] = None user_email: Optional[str] = None @@ -3764,9 +3784,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[ - float - ] = None # Users max budget within the organization + max_budget_in_organization: Optional[float] = ( + None # Users max budget within the organization + ) class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -4017,9 +4037,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[ - str, ProviderBudgetResponseObject - ] = {} # Dictionary mapping provider names to their budget configurations + providers: Dict[str, ProviderBudgetResponseObject] = ( + {} + ) # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -4181,9 +4201,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[ - str - ] = None # can be either user / team, inferred from the role mapping + object_id_jwt_field: Optional[str] = ( + None # can be either user / team, inferred from the role mapping + ) scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 68bde8434a6..5f05f1513d4 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -8,6 +8,7 @@ Run checks for: 2. If user is in budget 3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget """ + import asyncio import re import time @@ -456,9 +457,9 @@ async def common_checks( # noqa: PLR0915 model=_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=valid_token.team_model_aliases - if valid_token - else None, + team_model_aliases=( + valid_token.team_model_aliases if valid_token else None + ), ): raise ProxyException( message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", @@ -518,6 +519,17 @@ async def common_checks( # noqa: PLR0915 valid_token=valid_token, ) + # 3.1. Multi-window budget check for team + with tracer.trace("litellm.proxy.auth.common_checks.team_multi_budget_check"): + await _team_multi_budget_check(team_object=team_object) + + # 3.2. Multi-window budget check for key + with tracer.trace( + "litellm.proxy.auth.common_checks.virtual_key_multi_budget_check" + ): + if valid_token is not None: + await _virtual_key_multi_budget_check(valid_token=valid_token) + # 3.0.5. If team is over soft budget (alert only, doesn't block) with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"): await _team_soft_budget_check( @@ -2966,6 +2978,43 @@ async def _virtual_key_max_budget_check( ) +async def _virtual_key_multi_budget_check( + valid_token: UserAPIKeyAuth, +): + """ + Raises BudgetExceededError if any budget window in valid_token.budget_limits is exceeded. + + Each window has its own Redis counter keyed by spend:key:{token}:window:{budget_duration}. + Using budget_duration (not list index) keeps counters stable when windows are reordered + or removed during a key update. + + Note: counters are not seeded from DB on Redis cold-start. After a Redis flush, + per-window spend resets to zero within the current window period. This is an acceptable + trade-off: the DB stores reset_at timestamps but not per-window accumulated spend. + """ + if not valid_token.budget_limits: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in valid_token.budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() + counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + ) + if window_spend >= w["max_budget"]: + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], + message=( + f"ExceededBudget: Key over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + ) + + async def _virtual_key_soft_budget_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, @@ -3155,6 +3204,39 @@ async def _team_max_budget_check( ) +async def _team_multi_budget_check( + team_object: Optional[LiteLLM_TeamTable], +): + """ + Raises BudgetExceededError if any budget window in team_object.budget_limits is exceeded. + + Each window has its own Redis counter keyed by spend:team:{team_id}:window:{budget_duration}. + Using budget_duration (not list index) keeps counters stable when windows are reordered + or removed during a team update. + """ + if team_object is None or not team_object.budget_limits: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in team_object.budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() + counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + ) + if window_spend >= w["max_budget"]: + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], + message=( + f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + ) + + async def _team_soft_budget_check( team_object: Optional[LiteLLM_TeamTable], valid_token: Optional[UserAPIKeyAuth], diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 9c306acd2c6..b9b614994c2 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -106,7 +106,7 @@ class UserAPIKeyAuthExceptionHandler: message=e.message, type=ProxyErrorTypes.budget_exceeded, param=None, - code=400, + code=429, ) if isinstance(e, HTTPException): raise ProxyException( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 046c39a9101..1d5b913f9b8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -984,9 +984,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params[ - "allowed_model_region" - ] = _end_user_object.allowed_model_region + end_user_params["allowed_model_region"] = ( + _end_user_object.allowed_model_region + ) if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -1540,9 +1540,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict[ - "end_user_object_permission" - ] = _end_user_object.object_permission + valid_token_dict["end_user_object_permission"] = ( + _end_user_object.object_permission + ) # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bcfaed24398..0f9ca4e37db 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -48,6 +48,9 @@ class ResetBudgetJob: ### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ### await self.reset_budget_for_litellm_budget_table() + ### RESET MULTI-WINDOW BUDGETS ### + await self.reset_budget_windows() + async def reset_budget_for_litellm_team_members( self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): @@ -549,6 +552,115 @@ class ResetBudgetJob: ) verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e) + async def reset_budget_windows(self) -> None: + """ + For keys and teams with budget_limits, reset any individual windows where + reset_at <= now. Only the expired windows are reset; other windows are untouched. + """ + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + from litellm.proxy.proxy_server import spend_counter_cache + + now = datetime.utcnow() + + # --- Keys --- + try: + all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many( + where={"budget_limits": {"not": None}} # type: ignore[arg-type] + ) + for key in all_keys: + raw = key.budget_limits # type: ignore[attr-defined] + if not raw: + continue + windows: list = raw if isinstance(raw, list) else json.loads(raw) + changed = False + for window in windows: + reset_at_str = window.get("reset_at") + if not reset_at_str: + continue + reset_at = datetime.fromisoformat( + reset_at_str.replace("Z", "+00:00") + ).replace(tzinfo=None) + if reset_at <= now: + # Reset this window's counter (keyed by duration, not index) + counter_key = ( + f"spend:key:{key.token}:window:{window['budget_duration']}" + ) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=0.0 + ) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=0.0 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to reset Redis counter %s: %s", + counter_key, + redis_err, + ) + window["reset_at"] = get_budget_reset_time( + budget_duration=window["budget_duration"] + ).isoformat() + changed = True + if changed: + await self.prisma_client.db.litellm_verificationtoken.update( + where={"token": key.token}, + data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] + ) + except Exception as e: + verbose_proxy_logger.exception( + "Failed to reset budget windows for keys: %s", e + ) + + # --- Teams --- + try: + all_teams = await self.prisma_client.db.litellm_teamtable.find_many( + where={"budget_limits": {"not": None}} # type: ignore[arg-type] + ) + for team in all_teams: + raw = team.budget_limits # type: ignore[attr-defined] + if not raw: + continue + windows = raw if isinstance(raw, list) else json.loads(raw) + changed = False + for window in windows: + reset_at_str = window.get("reset_at") + if not reset_at_str: + continue + reset_at = datetime.fromisoformat( + reset_at_str.replace("Z", "+00:00") + ).replace(tzinfo=None) + if reset_at <= now: + counter_key = f"spend:team:{team.team_id}:window:{window['budget_duration']}" + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=0.0 + ) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=0.0 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to reset Redis counter %s: %s", + counter_key, + redis_err, + ) + window["reset_at"] = get_budget_reset_time( + budget_duration=window["budget_duration"] + ).isoformat() + changed = True + if changed: + await self.prisma_client.db.litellm_teamtable.update( + where={"team_id": team.team_id}, + data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] + ) + except Exception as e: + verbose_proxy_logger.exception( + "Failed to reset budget windows for teams: %s", e + ) + @staticmethod async def _reset_budget_common( item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken], @@ -570,14 +682,14 @@ class ResetBudgetJob: from litellm.proxy.proxy_server import spend_counter_cache counter_key = None - if item_type == "key" and hasattr(item, "token") and item.token is not None: - counter_key = f"spend:key:{item.token}" + if item_type == "key" and hasattr(item, "token") and item.token is not None: # type: ignore[union-attr] + counter_key = f"spend:key:{item.token}" # type: ignore[union-attr] elif ( item_type == "team" and hasattr(item, "team_id") - and item.team_id is not None + and item.team_id is not None # type: ignore[union-attr] ): - counter_key = f"spend:team:{item.team_id}" + counter_key = f"spend:team:{item.team_id}" # type: ignore[union-attr] if counter_key is not None: # Always reset in-memory (local fallback) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e105d9b8586..db9765fd084 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -732,9 +732,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 request_type="key", **data_json, table_name="key" ) - response[ - "soft_budget" - ] = data.soft_budget # include the user-input soft budget in the response + response["soft_budget"] = ( + data.soft_budget + ) # include the user-input soft budget in the response response = GenerateKeyResponse(**response) @@ -1559,6 +1559,19 @@ async def prepare_key_update_data( non_default_values["budget_reset_at"] = key_reset_at non_default_values["budget_duration"] = budget_duration + if "budget_limits" in non_default_values and non_default_values["budget_limits"]: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + raw_windows = non_default_values["budget_limits"] + initialized_windows = [] + for window in raw_windows: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + non_default_values["budget_limits"] = json.dumps(initialized_windows) + if "object_permission" in non_default_values: non_default_values = await _handle_update_object_permission( data_json=non_default_values, @@ -2785,6 +2798,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 rotation_interval: Optional[str] = None, router_settings: Optional[dict] = None, access_group_ids: Optional[list] = None, + budget_limits: Optional[list] = None, # multiple concurrent budget windows ): from litellm.proxy.proxy_server import premium_user, prisma_client @@ -2816,6 +2830,18 @@ async def generate_key_helper_fn( # noqa: PLR0915 else: reset_at = get_budget_reset_time(budget_duration=budget_duration) + # Initialize reset_at for each budget window + budget_limits_json: Optional[str] = None + if budget_limits: + initialized_windows = [] + for window in budget_limits: + w = dict(window) if not isinstance(window, dict) else {**window} + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + budget_limits_json = json.dumps(initialized_windows) + aliases_json = json.dumps(aliases) config_json = json.dumps(config) permissions_json = json.dumps(permissions) @@ -2898,6 +2924,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 "organization_id": organization_id, "budget_id": budget_id, "blocked": blocked, + "budget_limits": budget_limits_json, "created_by": created_by, "updated_by": updated_by, "allowed_routes": allowed_routes or [], @@ -3175,10 +3202,10 @@ async def delete_verification_tokens( try: if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"token": {"in": tokens}} + _keys_being_deleted: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"token": {"in": tokens}} + ) ) if len(_keys_being_deleted) == 0: @@ -3378,9 +3405,9 @@ async def _rotate_master_key( # noqa: PLR0915 from litellm.proxy.proxy_server import proxy_config try: - models: Optional[ - List - ] = await prisma_client.db.litellm_proxymodeltable.find_many() + models: Optional[List] = ( + await prisma_client.db.litellm_proxymodeltable.find_many() + ) except Exception: models = None # 2. process model table @@ -4020,11 +4047,11 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[ - BaseModel - ] = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id}, - include={"organization_memberships": True}, + complete_user_info_db_obj: Optional[BaseModel] = ( + await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, + ) ) if complete_user_info_db_obj is None: @@ -4107,10 +4134,10 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Optional[ - List[BaseModel] - ] = await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": complete_user_info.teams}} + teams: Optional[List[BaseModel]] = ( + await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": complete_user_info.teams}} + ) ) if teams is None: return [] diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3643373be65..78d4dcc4438 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -978,6 +978,19 @@ async def new_team( # noqa: PLR0915 budget_duration=complete_team_data.budget_duration, ) + # If budget_limits is set, initialize reset_at for each window + if complete_team_data.budget_limits: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + initialized_windows = [] + for window in complete_team_data.budget_limits: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + complete_team_data.budget_limits = initialized_windows + ## Add Team Member Budget Table members_with_roles: List[Member] = [] if complete_team_data.members_with_roles is not None: @@ -1545,12 +1558,12 @@ async def update_team( # noqa: PLR0915 updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[ - LiteLLM_TeamTable - ] = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - include={"litellm_model_table": True}, # type: ignore + team_row: Optional[LiteLLM_TeamTable] = ( + await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data=updated_kv, + include={"litellm_model_table": True}, # type: ignore + ) ) if team_row is None or team_row.team_id is None: @@ -1593,6 +1606,18 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: reset_at = get_budget_reset_time(budget_duration=data.budget_duration) updated_kv["budget_reset_at"] = reset_at + if data.budget_limits is not None and len(data.budget_limits) > 0: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + initialized_windows = [] + for window in data.budget_limits: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + updated_kv["budget_limits"] = json.dumps(initialized_windows) + async def handle_update_object_permission( data_json: dict, existing_team_row: LiteLLM_TeamTable @@ -2297,13 +2322,13 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } + ) ) if keys_to_delete: @@ -2687,10 +2712,10 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + team_row_base: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) ) if team_row_base is None: raise Exception @@ -2749,10 +2774,10 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": {"in": data.team_ids}} + ) ) if keys_to_delete: @@ -2972,11 +2997,11 @@ async def team_info( ) try: - team_info: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, + team_info: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, + ) ) if team_info is None: raise Exception @@ -3732,9 +3757,7 @@ async def list_team( except Exception as e: team_exception = """Invalid team object for team_id: {}. team_object={}. Error: {} - """.format( - team.team_id, team.model_dump(), str(e) - ) + """.format(team.team_id, team.model_dump(), str(e)) verbose_proxy_logger.exception(team_exception) continue # Sort the responses by team_alias diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d83e3500ebe..21a62f88dfc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -641,9 +641,9 @@ except ImportError: server_root_path = get_server_root_path() _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional[ - "EnterpriseLicenseData" -] = _license_check.airgapped_license_data +premium_user_data: Optional["EnterpriseLicenseData"] = ( + _license_check.airgapped_license_data +) global_max_parallel_request_retries_env: Optional[str] = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" ) @@ -1537,9 +1537,9 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional[ - "ClientSession" -] = None # Global shared session for connection reuse +shared_aiohttp_session: Optional["ClientSession"] = ( + None # Global shared session for connection reuse +) user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1550,13 +1550,13 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[ - RedisCache -] = None # redis cache used for tracking spend, tpm/rpm limits +redis_usage_cache: Optional[RedisCache] = ( + None # redis cache used for tracking spend, tpm/rpm limits +) polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False -native_background_mode: List[ - str -] = [] # Models that should use native provider background mode instead of polling +native_background_mode: List[str] = ( + [] +) # Models that should use native provider background mode instead of polling polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None user_custom_key_generate = None @@ -1783,6 +1783,26 @@ async def increment_spend_counters( increment=response_cost, ) + # Increment per-window budget counters for multi-budget keys + key_obj = await user_api_key_cache.async_get_cache(key=hashed_token) + if key_obj is not None: + key_budget_limits = getattr(key_obj, "budget_limits", None) or ( + key_obj.get("budget_limits") if isinstance(key_obj, dict) else None + ) + if isinstance(key_budget_limits, str): + key_budget_limits = json.loads(key_budget_limits) + if isinstance(key_budget_limits, list): + for window in key_budget_limits: + duration = ( + window["budget_duration"] + if isinstance(window, dict) + else window.budget_duration + ) + await spend_counter_cache.async_increment_cache( + key=f"spend:key:{hashed_token}:window:{duration}", + value=response_cost, + ) + if team_id is not None: await _init_and_increment_spend_counter( counter_key=f"spend:team:{team_id}", @@ -1790,6 +1810,26 @@ async def increment_spend_counters( increment=response_cost, ) + # Increment per-window budget counters for multi-budget teams + team_obj = await user_api_key_cache.async_get_cache(key=f"team_id:{team_id}") + if team_obj is not None: + team_budget_limits = getattr(team_obj, "budget_limits", None) or ( + team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + ) + if isinstance(team_budget_limits, str): + team_budget_limits = json.loads(team_budget_limits) + if isinstance(team_budget_limits, list): + for window in team_budget_limits: + duration = ( + window["budget_duration"] + if isinstance(window, dict) + else window.budget_duration + ) + await spend_counter_cache.async_increment_cache( + key=f"spend:team:{team_id}:window:{duration}", + value=response_cost, + ) + if user_id is not None and team_id is not None: await _init_and_increment_spend_counter( counter_key=f"spend:team_member:{user_id}:{team_id}", @@ -2040,9 +2080,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[ - LiteLLM_TeamTable - ] = await user_api_key_cache.async_get_cache(key=_id) + existing_spend_obj: Optional[LiteLLM_TeamTable] = ( + await user_api_key_cache.async_get_cache(key=_id) + ) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -2163,11 +2203,9 @@ def run_ollama_serve(): with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - verbose_proxy_logger.debug( - f""" + verbose_proxy_logger.debug(f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` - """ - ) + """) def _get_process_rss_mb() -> Optional[float]: @@ -5215,10 +5253,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client + guardrails_in_db: List[Guardrail] = ( + await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client + ) ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -5600,9 +5638,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ[ - "AZURE_API_VERSION" - ] = api_version # set this for azure - litellm can read this from the env + os.environ["AZURE_API_VERSION"] = ( + api_version # set this for azure - litellm can read this from the env + ) if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -5944,9 +5982,9 @@ class ProxyStartupEvent: """ from litellm.secret_managers.main import str_to_bool - _use_redis_transaction_buffer: Optional[ - Union[bool, str] - ] = general_settings.get("use_redis_transaction_buffer", False) + _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( + general_settings.get("use_redis_transaction_buffer", False) + ) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) @@ -12627,9 +12665,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[ - idx - ].field_description = sub_field_info.description + nested_fields[idx].field_description = ( + sub_field_info.description + ) idx += 1 _stored_in_db = None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 46be6b31e1f..6dec754e436 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -140,6 +140,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) @@ -387,6 +388,7 @@ model LiteLLM_VerificationToken { rotation_interval String? // How often to rotate (e.g., "30d", "90d") last_rotation_at DateTime? // When this key was last rotated key_rotation_at DateTime? // When this key should next be rotated + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 845919e9120..be62b21d84a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1898,9 +1898,9 @@ class ProxyLogging: normalized_call_type = CallTypes.aembedding.value if normalized_call_type is not None: litellm_logging_obj.call_type = normalized_call_type - litellm_logging_obj.model_call_details[ - "call_type" - ] = normalized_call_type + litellm_logging_obj.model_call_details["call_type"] = ( + normalized_call_type + ) # Pass-through endpoints are logged via the callback loop's # async_post_call_failure_hook — skip pre_call and failure handlers. if litellm_logging_obj.call_type == CallTypes.pass_through.value: @@ -2524,8 +2524,7 @@ class PrismaClient: required_view = "LiteLLM_VerificationTokenView" expected_views_str = ", ".join(f"'{view}'" for view in expected_views) pg_schema = os.getenv("DATABASE_SCHEMA", "public") - ret = await self.db.query_raw( - f""" + ret = await self.db.query_raw(f""" WITH existing_views AS ( SELECT viewname FROM pg_views @@ -2537,8 +2536,7 @@ class PrismaClient: (SELECT COUNT(*) FROM existing_views) AS view_count, ARRAY_AGG(viewname) AS view_names FROM existing_views - """ - ) + """) expected_total_views = len(expected_views) if ret[0]["view_count"] == expected_total_views: verbose_proxy_logger.info("All necessary views exist!") @@ -2547,8 +2545,7 @@ class PrismaClient: ## check if required view exists ## if ret[0]["view_names"] and required_view not in ret[0]["view_names"]: await self.health_check() # make sure we can connect to db - await self.db.execute_raw( - """ + await self.db.execute_raw(""" CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -2558,8 +2555,7 @@ class PrismaClient: t.rpm_limit AS team_rpm_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; - """ - ) + """) verbose_proxy_logger.info( "LiteLLM_VerificationTokenView Created in DB!" @@ -3142,6 +3138,10 @@ class PrismaClient: hashed_token = self.hash_token(token=token) db_data = self.jsonify_object(data=data) db_data["token"] = hashed_token + # Prisma rejects nullable JSON fields set to None (no default). + # Strip them so the DB stores NULL via the column's nullable constraint. + if db_data.get("budget_limits") is None: + db_data.pop("budget_limits", None) print_verbose( "PrismaClient: Before upsert into litellm_verificationtoken" ) diff --git a/tests/test_litellm/proxy/auth/test_multi_budget_windows.py b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py new file mode 100644 index 00000000000..ed94fca837b --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py @@ -0,0 +1,137 @@ +""" +Unit tests for multi-budget-window enforcement on API keys. +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import _virtual_key_multi_budget_check + + +def _make_valid_token(**kwargs) -> UserAPIKeyAuth: + defaults = dict( + token="sk-test-token", + key_name="test", + spend=0.0, + max_budget=None, + budget_limits=[], + ) + defaults.update(kwargs) + return UserAPIKeyAuth(**defaults) + + +@pytest.mark.asyncio +async def test_no_budget_limits_passes(): + """Keys with empty budget_limits should pass without raising.""" + token = _make_valid_token(budget_limits=[]) + # Should not raise + await _virtual_key_multi_budget_check(valid_token=token) + + +@pytest.mark.asyncio +async def test_under_budget_passes(): + """Key with spend under all windows should pass.""" + token = _make_valid_token( + budget_limits=[ + {"budget_duration": "24h", "max_budget": 10.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": None}, + ] + ) + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=1.0, # well under both windows + ): + await _virtual_key_multi_budget_check(valid_token=token) + + +@pytest.mark.asyncio +async def test_over_first_window_raises(): + """Key exceeding the first (daily) window should raise BudgetExceededError.""" + token = _make_valid_token( + budget_limits=[ + {"budget_duration": "24h", "max_budget": 5.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": None}, + ] + ) + + spend_by_window = [6.0, 6.0] # over daily, under monthly + + call_count = 0 + + async def fake_get_spend(counter_key, fallback_spend): + nonlocal call_count + val = spend_by_window[call_count] + call_count += 1 + return val + + with patch( + "litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_multi_budget_check(valid_token=token) + + err = exc_info.value + assert err.status_code == 429 + assert "24h" in str(err) + assert "Key over" in str(err) + + +@pytest.mark.asyncio +async def test_over_second_window_raises(): + """Key exceeding only the monthly window should raise BudgetExceededError referencing 30d.""" + token = _make_valid_token( + budget_limits=[ + {"budget_duration": "24h", "max_budget": 50.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 5.0, "reset_at": None}, + ] + ) + + spend_by_window = [1.0, 10.0] # under daily, over monthly + + call_count = 0 + + async def fake_get_spend(counter_key, fallback_spend): + nonlocal call_count + val = spend_by_window[call_count] + call_count += 1 + return val + + with patch( + "litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_multi_budget_check(valid_token=token) + + err = exc_info.value + assert err.status_code == 429 + assert "30d" in str(err) + + +@pytest.mark.asyncio +async def test_budget_limit_entry_objects_coerced(): + """BudgetLimitEntry Pydantic objects (not dicts) must be handled without KeyError. + + While budget_limits is normally serialized as List[dict], the auth check must + tolerate BudgetLimitEntry objects in case they arrive without prior serialization. + """ + from litellm.proxy._types import BudgetLimitEntry + + token = _make_valid_token(budget_limits=[]) + # Bypass Pydantic validation to simulate BudgetLimitEntry objects reaching the check + object.__setattr__( + token, + "budget_limits", + [BudgetLimitEntry(budget_duration="24h", max_budget=10.0)], + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=1.0, + ): + # Should not raise TypeError / KeyError — model_dump() coerces the object + await _virtual_key_multi_budget_check(valid_token=token) diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index fcba5774644..aa74bc60aa1 100644 --- a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -25,6 +25,7 @@ const BudgetDurationDropdown: React.FC = ({ placeholder="n/a" allowClear > + @@ -36,6 +37,7 @@ export const getBudgetDurationLabel = (value: string | null | undefined): string if (!value) return "Not set"; const budgetDurationMap: Record = { + "1h": "hourly", "24h": "daily", "7d": "weekly", "30d": "monthly", diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetWindowsEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetWindowsEditor.tsx new file mode 100644 index 00000000000..45687461510 --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetWindowsEditor.tsx @@ -0,0 +1,84 @@ +import { Button, InputNumber, Select } from "antd"; +import React from "react"; + +export interface BudgetWindowEntry { + budget_duration: string; + max_budget: number | null; +} + +export const BUDGET_WINDOW_OPTIONS = [ + { value: "1h", label: "Hourly", resetHint: "Resets every hour" }, + { value: "24h", label: "Daily", resetHint: "Resets daily at midnight UTC" }, + { value: "7d", label: "Weekly", resetHint: "Resets every Sunday at midnight UTC" }, + { value: "30d", label: "Monthly", resetHint: "Resets on the 1st of every month at midnight UTC" }, +]; + +interface BudgetWindowsEditorProps { + value: BudgetWindowEntry[]; + onChange: (v: BudgetWindowEntry[]) => void; +} + +export function BudgetWindowsEditor({ value, onChange }: BudgetWindowsEditorProps) { + const addWindow = () => { + onChange([...value, { budget_duration: "24h", max_budget: null }]); + }; + + const removeWindow = (idx: number) => { + onChange(value.filter((_, i) => i !== idx)); + }; + + const updateWindow = (idx: number, field: keyof BudgetWindowEntry, fieldValue: string | number | null) => { + const updated = value.map((w, i) => (i === idx ? { ...w, [field]: fieldValue } : w)); + onChange(updated); + }; + + return ( +
+ {value.map((window, idx) => { + const hint = BUDGET_WINDOW_OPTIONS.find((o) => o.value === window.budget_duration)?.resetHint; + return ( +
+
+ + + Budget Windows{" "} + + + + + } + > + + +