diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index cf99c5cd9fa..590a81e7338 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2897,6 +2897,7 @@ class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): spend: float = 0.0 allowed_model_region: Optional[AllowedModelRegion] = None default_model: Optional[str] = None + budget_id: Optional[str] = None litellm_budget_table: Optional[LiteLLM_BudgetTable] = None object_permission_id: Optional[str] = None object_permission: Optional[LiteLLM_ObjectPermissionTable] = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 68bde8434a6..98ed24773bf 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -918,9 +918,26 @@ async def _apply_default_budget_to_end_user( if default_budget is not None: # Apply default budget to end user object end_user_obj.litellm_budget_table = default_budget - verbose_proxy_logger.debug( - f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}" - ) + + # Backfill budget_id to DB for existing users that were created without it + if end_user_obj.budget_id is None: + try: + await prisma_client.db.litellm_endusertable.update( + where={"user_id": end_user_obj.user_id}, + data={"budget_id": litellm.max_end_user_budget_id}, + ) + end_user_obj.budget_id = litellm.max_end_user_budget_id + verbose_proxy_logger.debug( + f"Persisted default budget_id {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}" + ) + except Exception: + verbose_proxy_logger.debug( + f"Failed to persist default budget_id for end user {end_user_obj.user_id}" + ) + else: + verbose_proxy_logger.debug( + f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}" + ) return end_user_obj diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ec98cfd4d1e..32730877918 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4679,14 +4679,17 @@ class ProxyUpdateSpend: ) in end_user_list_transactions.items(): if litellm.max_end_user_budget is not None: pass + _create_data: dict = { + "user_id": end_user_id, + "spend": response_cost, + "blocked": False, + } + if litellm.max_end_user_budget_id is not None: + _create_data["budget_id"] = litellm.max_end_user_budget_id batcher.litellm_endusertable.upsert( where={"user_id": end_user_id}, data={ - "create": { - "user_id": end_user_id, - "spend": response_cost, - "blocked": False, - }, + "create": _create_data, "update": {"spend": {"increment": response_cost}}, }, )