mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Keep budget_reset_at off the user-settable budget allowlist
LiteLLM_BudgetTable is documented as "user-controllable params" and its model_fields.keys() is used as the allowlist for extracting budget fields from incoming API request bodies (management_helpers/utils.py:88, organization_endpoints.py:112/255/537/549, project_endpoints.py:197/245/632, customer_endpoints.py:598). Request models like NewOrganizationRequest inherit from LiteLLM_BudgetTable, so anything on the base class becomes user-settable — a caller could set budget_reset_at far in the future and evade budget cycling. Move budget_reset_at from the base class to LiteLLM_BudgetTableFull so it appears on API responses without becoming writable, and type LiteLLM_TeamMembership.litellm_budget_table as Union[Full, Base] so Pydantic picks Full when the data has server-managed fields (/team/info reads Prisma rows that include budget_reset_at and created_at) and Base when callers construct with only user-settable fields (existing auth tests and caches).
This commit is contained in:
parent
e5f3e15969
commit
1a0ac9634c
1 changed files with 14 additions and 4 deletions
|
|
@ -1997,7 +1997,12 @@ class TeamRequest(LiteLLMPydanticObjectBase):
|
|||
|
||||
|
||||
class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
||||
"""Represents user-controllable params for a LiteLLM_BudgetTable record"""
|
||||
"""Represents user-controllable params for a LiteLLM_BudgetTable record.
|
||||
|
||||
Budget-write paths use `model_fields.keys()` on this class as an allowlist
|
||||
for user input. Keep server-managed fields (e.g. `budget_reset_at`) on
|
||||
`LiteLLM_BudgetTableFull` so they aren't user-settable.
|
||||
"""
|
||||
|
||||
budget_id: Optional[str] = None
|
||||
soft_budget: Optional[float] = None
|
||||
|
|
@ -2007,7 +2012,6 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
|||
rpm_limit: Optional[int] = None
|
||||
model_max_budget: Optional[dict] = None
|
||||
budget_duration: Optional[str] = None
|
||||
budget_reset_at: Optional[datetime] = None
|
||||
allowed_models: Optional[List[str]] = (
|
||||
None # per-member model scope; empty = inherit team models
|
||||
)
|
||||
|
|
@ -2016,8 +2020,9 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
|||
|
||||
|
||||
class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable):
|
||||
"""Represents all params for a LiteLLM_BudgetTable record"""
|
||||
"""LiteLLM_BudgetTable + server-managed fields returned on API responses."""
|
||||
|
||||
budget_reset_at: Optional[datetime] = None
|
||||
created_at: datetime
|
||||
|
||||
|
||||
|
|
@ -3696,7 +3701,12 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase):
|
|||
budget_id: Optional[str] = None
|
||||
spend: Optional[float] = 0.0
|
||||
total_spend: Optional[float] = 0.0
|
||||
litellm_budget_table: Optional[LiteLLM_BudgetTable]
|
||||
# Union so Pydantic picks Full when data has server-managed fields
|
||||
# (/team/info) and Base when callers/tests construct with only
|
||||
# user-settable fields.
|
||||
litellm_budget_table: Optional[
|
||||
Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable]
|
||||
]
|
||||
|
||||
def safe_get_team_member_rpm_limit(self) -> Optional[int]:
|
||||
if self.litellm_budget_table is not None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue