Merge 8faa75a9f3 into 58896e2009

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):

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": []})

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."""
 

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,

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):

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