diff --git a/litellm/models/team.py b/litellm/models/team.py index 8edf10703b1..5e7dab9f958 100644 --- a/litellm/models/team.py +++ b/litellm/models/team.py @@ -10,7 +10,7 @@ import json from datetime import datetime from typing import Final, Literal, Optional -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -86,6 +86,7 @@ class TeamBase(LiteLLMPydanticObjectBase): class LiteLLM_TeamTable(TeamBase): team_id: str spend: float | None = None + max_parallel_requests: int | None = None budget_duration: str | None = None budget_reset_at: datetime | None = None @@ -122,6 +123,10 @@ class LiteLLM_TeamTable(TeamBase): if isinstance(values.get("members_with_roles"), dict) and not values["members_with_roles"]: values["members_with_roles"] = [] + spend = values.get("spend") + if spend is not None and isinstance(spend, (int, float)) and spend < 0: + values["spend"] = 0.0 + for field in dict_fields: value = values.get(field) if value is not None and isinstance(value, str): diff --git a/litellm/models/user.py b/litellm/models/user.py index 7dcaed9f77f..1a1808db66b 100644 --- a/litellm/models/user.py +++ b/litellm/models/user.py @@ -7,7 +7,7 @@ Canonical definition for ``litellm_usertable``. Re-exported from from datetime import datetime -from pydantic import ConfigDict, Field, model_validator +from pydantic import ConfigDict, Field, field_validator, model_validator from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.models.organization_membership import ( @@ -49,10 +49,12 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) + @model_validator(mode="before") @classmethod def set_model_info(cls, values): - if values.get("spend") is None: + spend = values.get("spend") + if spend is None or (isinstance(spend, (int, float)) and spend < 0): values.update({"spend": 0.0}) if values.get("models") is None: values.update({"models": []}) diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index 89048c56a9f..35a1ea72e51 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -7,13 +7,22 @@ Canonical definition for ``litellm_verificationtoken``. Re-exported from from datetime import datetime -from pydantic import ConfigDict +from pydantic import ConfigDict, model_validator, field_validator from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.types.llms.base import LiteLLMPydanticObjectBase class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): + @model_validator(mode="before") + @classmethod + def clamp_negative_spend(cls, values): + if isinstance(values, dict): + spend = values.get("spend") + if spend is not None and isinstance(spend, (int, float)) and spend < 0: + values["spend"] = 0.0 + return values + token: str | None = None key_name: str | None = None key_alias: str | None = None @@ -66,6 +75,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) + class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): """Audit record for deleted keys; mirrors the token plus deletion metadata.""" diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 9c433af487e..26354b1e67d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1037,14 +1037,17 @@ class ResetBudgetJob: db_span("reset_spend_rows", "LiteLLM_VerificationToken"), spend_reset_unit_of_work(self._new_batch) as uow, ): + tokens: list[str] = [] for k in updated_keys: if k.row.token is None: continue + tokens.append(k.row.token) uow.keys.queue_spend_reset( token=k.row.token, budget_reset_at=k.row.budget_reset_at, spend_decrement=k.spend_decrement, ) + uow.keys.queue_clamp_negative_spend(tokens) async def _write_user_reset_updates(self, updated_users: Sequence[_RowReset[LiteLLM_UserTable]]) -> None: """ @@ -1064,12 +1067,17 @@ class ResetBudgetJob: db_span("reset_spend_rows", "LiteLLM_UserTable"), spend_reset_unit_of_work(self._new_batch) as uow, ): + user_ids: list[str] = [] for u in updated_users: + if u.row.user_id is None: + continue + user_ids.append(u.row.user_id) uow.users.queue_spend_reset( user_id=u.row.user_id, budget_reset_at=u.row.budget_reset_at, spend_decrement=u.spend_decrement, ) + uow.users.queue_clamp_negative_spend(user_ids) async def _write_team_reset_updates(self, updated_teams: Sequence[_RowReset[LiteLLM_TeamTable]]) -> None: """ @@ -1089,12 +1097,17 @@ class ResetBudgetJob: db_span("reset_spend_rows", "LiteLLM_TeamTable"), spend_reset_unit_of_work(self._new_batch) as uow, ): + team_ids: list[str] = [] for t in updated_teams: + if t.row.team_id is None: + continue + team_ids.append(t.row.team_id) uow.teams.queue_spend_reset( team_id=t.row.team_id, budget_reset_at=t.row.budget_reset_at, spend_decrement=t.spend_decrement, ) + uow.teams.queue_clamp_negative_spend(team_ids) def _emit_phase_failure( self, diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index b78f7e5a9b7..7d9bbb63307 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -39,6 +39,14 @@ class KeySpendResetWrites: data=_spend_reset_data(budget_reset_at, spend_decrement), ) + def queue_clamp_negative_spend(self, tokens: Sequence[str]) -> None: + if not tokens: + return + self.table.update_many( + where={"token": {"in": list(tokens)}, "spend": {"lt": 0}}, + data={"spend": 0}, + ) + @dataclass(frozen=True, slots=True) class UserSpendResetWrites: @@ -50,6 +58,14 @@ class UserSpendResetWrites: data=_spend_reset_data(budget_reset_at, spend_decrement), ) + def queue_clamp_negative_spend(self, user_ids: Sequence[str]) -> None: + if not user_ids: + return + self.table.update_many( + where={"user_id": {"in": list(user_ids)}, "spend": {"lt": 0}}, + data={"spend": 0}, + ) + @dataclass(frozen=True, slots=True) class TeamSpendResetWrites: @@ -61,6 +77,14 @@ class TeamSpendResetWrites: data=_spend_reset_data(budget_reset_at, spend_decrement), ) + def queue_clamp_negative_spend(self, team_ids: Sequence[str]) -> None: + if not team_ids: + return + self.table.update_many( + where={"team_id": {"in": list(team_ids)}, "spend": {"lt": 0}}, + data={"spend": 0}, + ) + @runtime_checkable class SpendResetWrites(Protocol): diff --git a/tests/unit/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index 049e5275db6..9b3df6d645f 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -214,3 +214,44 @@ async def test_membership_total_spend_repair_is_queued_on_the_cascade_batch_ahea assert repair[0] == "execute_raw" assert repair[2] == {"args": (["budget-1", "budget-2"],)} assert reset == ("litellm_teammembership.update_many", {"budget_id": "budget-1"}, {"spend": 0}) + + +async def test_queue_clamp_negative_spend_across_tables(): + batch = FakeBatch() + + async with spend_reset_unit_of_work(lambda: batch) as uow: + uow.keys.queue_clamp_negative_spend(["tok-1", "tok-2"]) + uow.users.queue_clamp_negative_spend(["user-1"]) + uow.teams.queue_clamp_negative_spend(["team-1", "team-2"]) + # Empty calls should be no-ops + uow.keys.queue_clamp_negative_spend([]) + uow.users.queue_clamp_negative_spend([]) + uow.teams.queue_clamp_negative_spend([]) + + assert batch.commit_count == 1 + assert batch.calls == [ + ("litellm_verificationtoken.update_many", {"token": {"in": ["tok-1", "tok-2"]}, "spend": {"lt": 0}}, {"spend": 0}), + ("litellm_usertable.update_many", {"user_id": {"in": ["user-1"]}, "spend": {"lt": 0}}, {"spend": 0}), + ("litellm_teamtable.update_many", {"team_id": {"in": ["team-1", "team-2"]}, "spend": {"lt": 0}}, {"spend": 0}), + ] + + +def test_models_clamp_negative_fractional_spend(): + from litellm.models.team import LiteLLM_TeamTable + from litellm.models.user import LiteLLM_UserTable + from litellm.models.verification_token import LiteLLM_VerificationToken + + # Negative epsilon spend from budget reset floating-point decrement is clamped to 0.0 + team = LiteLLM_TeamTable(team_id="team-neg", spend=-2.7e-17) + assert team.spend == 0.0 + + user = LiteLLM_UserTable(user_id="user-neg", spend=-1.5e-16) + assert user.spend == 0.0 + + token = LiteLLM_VerificationToken(token="tok-neg", spend=-3.1e-17) + assert token.spend == 0.0 + + # Normal non-negative spend remains unchanged + assert LiteLLM_TeamTable(team_id="team-pos", spend=10.5).spend == 10.5 + assert LiteLLM_UserTable(user_id="user-pos", spend=5.25).spend == 5.25 + assert LiteLLM_VerificationToken(token="tok-pos", spend=1.5).spend == 1.5