diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index fb600cee26b..52beed14fc4 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( LiteLLM_BudgetTableFull, LiteLLM_EndUserTable, + LiteLLM_OrganizationTable, LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken, @@ -71,16 +72,19 @@ class ResetBudgetJob: async def reset_budget_for_litellm_budget_table(self): """ - Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired + Resets the budget for all LiteLLM End-Users (Customers), Organizations, and Team Members if their budget has expired The corresponding Budget duration is also updated. """ now = datetime.now(timezone.utc) start_time = time.time() endusers_to_reset: Optional[List[LiteLLM_EndUserTable]] = None + organizations_to_reset: Optional[List[LiteLLM_OrganizationTable]] = None budgets_to_reset: Optional[List[LiteLLM_BudgetTableFull]] = None updated_endusers: List[LiteLLM_EndUserTable] = [] + updated_organizations: List[LiteLLM_OrganizationTable] = [] failed_endusers = [] + failed_organizations = [] try: budgets_to_reset = await self.prisma_client.get_data( table_name="budget", query_type="find_all", reset_at=now @@ -98,14 +102,21 @@ class ResetBudgetJob: table_name="budget", ) + budget_id_list = [ + budget.budget_id + for budget in budgets_to_reset + if budget.budget_id is not None + ] + endusers_to_reset = await self.prisma_client.get_data( table_name="enduser", query_type="find_all", - budget_id_list=[ - budget.budget_id - for budget in budgets_to_reset - if budget.budget_id is not None - ], + budget_id_list=budget_id_list, + ) + + # Query organizations with the same budget_ids + organizations_to_reset = await self.prisma_client.db.litellm_organizationtable.find_many( + where={"budget_id": {"in": budget_id_list}} ) await self.reset_budget_for_litellm_team_members( @@ -146,11 +157,43 @@ class ResetBudgetJob: table_name="enduser", ) + if organizations_to_reset is not None and len(organizations_to_reset) > 0: + # Update organizations directly using prisma client update_many + # since "organization" is not supported in update_data + # Similar to how team members are reset + try: + await self.prisma_client.db.litellm_organizationtable.update_many( + where={"budget_id": {"in": budget_id_list}}, + data={"spend": 0.0}, + ) + # Mark all as updated since update_many succeeded + updated_organizations = organizations_to_reset + verbose_proxy_logger.debug( + "Updated organizations %s", + json.dumps(organizations_to_reset, indent=4, default=str), + ) + except Exception as e: + # If batch update fails, track all organizations as failed + for org in organizations_to_reset: + failed_organizations.append( + {"organization": org, "error": str(e)} + ) + verbose_proxy_logger.exception( + "Failed to reset budget for organizations: %s", e + ) + end_time = time.time() - if len(failed_endusers) > 0: # If any endusers failed to reset - raise Exception( - f"Failed to reset {len(failed_endusers)} endusers: {json.dumps(failed_endusers, default=str)}" - ) + if len(failed_endusers) > 0 or len(failed_organizations) > 0: + error_messages = [] + if len(failed_endusers) > 0: + error_messages.append( + f"Failed to reset {len(failed_endusers)} endusers: {json.dumps(failed_endusers, default=str)}" + ) + if len(failed_organizations) > 0: + error_messages.append( + f"Failed to reset {len(failed_organizations)} organizations: {json.dumps(failed_organizations, default=str)}" + ) + raise Exception("; ".join(error_messages)) asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( @@ -180,6 +223,20 @@ class ResetBudgetJob: "endusers_failed": json.dumps( failed_endusers, indent=4, default=str ), + "num_organizations_found": ( + len(organizations_to_reset) if organizations_to_reset else 0 + ), + "organizations_found": json.dumps( + organizations_to_reset, indent=4, default=str + ), + "num_organizations_updated": len(updated_organizations), + "organizations_updated": json.dumps( + updated_organizations, indent=4, default=str + ), + "num_organizations_failed": len(failed_organizations), + "organizations_failed": json.dumps( + failed_organizations, indent=4, default=str + ), }, ) ) @@ -206,10 +263,18 @@ class ResetBudgetJob: "endusers_found": json.dumps( endusers_to_reset, indent=4, default=str ), + "num_organizations_found": ( + len(organizations_to_reset) if organizations_to_reset else 0 + ), + "organizations_found": json.dumps( + organizations_to_reset, indent=4, default=str + ), }, ) ) - verbose_proxy_logger.exception("Failed to reset budget for endusers: %s", e) + verbose_proxy_logger.exception( + "Failed to reset budget for endusers and organizations: %s", e + ) async def reset_budget_for_litellm_keys(self): """ @@ -541,6 +606,19 @@ class ResetBudgetJob: raise e return enduser + @staticmethod + async def _reset_budget_for_organization( + organization: LiteLLM_OrganizationTable, + ) -> Optional[LiteLLM_OrganizationTable]: + try: + organization.spend = 0.0 + except Exception as e: + verbose_proxy_logger.exception( + "Error resetting budget for organization: %s. Item: %s", e, organization + ) + raise e + return organization + @staticmethod async def _reset_budget_reset_at_date( budget: LiteLLM_BudgetTableFull, current_time: datetime diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index a059a3adcb1..9a1f709f7fb 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -25,9 +25,44 @@ class MockLiteLLMTeamMembership: return {"count": 1} +class MockLiteLLMOrganizationTable: + def __init__(self, organizations_data: List[Any]): + self.organizations_data = organizations_data + + async def find_many(self, where: Dict[str, Any]) -> List[Any]: + # Mock find_many for organizations - filter by budget_id + if "budget_id" in where and "in" in where["budget_id"]: + budget_id_list = where["budget_id"]["in"] + return [ + org + for org in self.organizations_data + if hasattr(org, "budget_id") and org.budget_id in budget_id_list + ] + return [] + + async def update_many( + self, where: Dict[str, Any], data: Dict[str, Any] + ) -> Dict[str, Any]: + # Mock the update_many method for organizations + # Update spend in the organizations_data + if "budget_id" in where and "in" in where["budget_id"]: + budget_id_list = where["budget_id"]["in"] + count = 0 + for org in self.organizations_data: + if hasattr(org, "budget_id") and org.budget_id in budget_id_list: + if "spend" in data: + org.spend = data["spend"] + count += 1 + return {"count": count} + return {"count": 0} + + class MockDB: - def __init__(self): + def __init__(self, organizations_data: List[Any] = None): self.litellm_teammembership = MockLiteLLMTeamMembership() + self.litellm_organizationtable = MockLiteLLMOrganizationTable( + organizations_data or [] + ) class MockPrismaClient: @@ -38,6 +73,7 @@ class MockPrismaClient: "team": [], "budget": [], "enduser": [], + "organization": [], } self.updated_data: Dict[str, List[Any]] = { "key": [], @@ -45,8 +81,15 @@ class MockPrismaClient: "team": [], "budget": [], "enduser": [], + "organization": [], } - self.db = MockDB() + self.db = MockDB(organizations_data=[]) + + def _update_db_organizations(self): + """Update the MockDB with current organization data""" + self.db.litellm_organizationtable = MockLiteLLMOrganizationTable( + self.data.get("organization", []) + ) async def get_data(self, table_name, query_type, **kwargs): data = self.data.get(table_name, []) @@ -225,8 +268,21 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): }, ) + test_organization = type( + "LiteLLM_OrganizationTable", + (), + { + "organization_id": "test-org-1", + "organization_alias": "Test Org", + "budget_id": "test-budget-1", + "spend": 150.0, + }, + ) + mock_prisma_client.data["budget"] = [test_budget] mock_prisma_client.data["enduser"] = [test_enduser] + mock_prisma_client.data["organization"] = [test_organization] + mock_prisma_client._update_db_organizations() # Run the test asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -238,6 +294,8 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): updated_budget = mock_prisma_client.updated_data["budget"][0] assert updated_enduser.spend == 0.0 assert updated_budget.budget_reset_at > now + # Verify organization spend was reset + assert test_organization.spend == 0.0 def test_reset_budget_all(reset_budget_job, mock_prisma_client): @@ -299,11 +357,24 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): }, ) + test_organization = type( + "LiteLLM_OrganizationTable", + (), + { + "organization_id": "test-org-1", + "organization_alias": "Test Org", + "budget_id": "test-budget-1", + "spend": 150.0, + }, + ) + mock_prisma_client.data["key"] = [test_key] mock_prisma_client.data["user"] = [test_user] mock_prisma_client.data["team"] = [test_team] mock_prisma_client.data["budget"] = [test_budget] mock_prisma_client.data["enduser"] = [test_enduser] + mock_prisma_client.data["organization"] = [test_organization] + mock_prisma_client._update_db_organizations() # Run the test asyncio.run(reset_budget_job.reset_budget()) @@ -320,3 +391,122 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): assert mock_prisma_client.updated_data["user"][0].spend == 0.0 assert mock_prisma_client.updated_data["team"][0].spend == 0.0 assert mock_prisma_client.updated_data["enduser"][0].spend == 0.0 + # Verify organization spend was reset + assert test_organization.spend == 0.0 + + +def test_reset_budget_for_organization(reset_budget_job, mock_prisma_client): + """ + Test that organizations with expired budgets have their spend reset to 0. + This tests the new organization budget reset logic. + """ + # Setup test data + now = datetime.now(timezone.utc) + test_budget = type( + "LiteLLM_BudgetTable", + (), + { + "max_budget": 1000.0, + "budget_duration": "1d", + "budget_reset_at": now, + "budget_id": "test-budget-org-1", + "created_at": now - timedelta(days=2), + }, + ) + + test_organization_1 = type( + "LiteLLM_OrganizationTable", + (), + { + "organization_id": "test-org-1", + "organization_alias": "Test Org 1", + "budget_id": "test-budget-org-1", + "spend": 250.0, + }, + ) + + test_organization_2 = type( + "LiteLLM_OrganizationTable", + (), + { + "organization_id": "test-org-2", + "organization_alias": "Test Org 2", + "budget_id": "test-budget-org-1", + "spend": 350.0, + }, + ) + + mock_prisma_client.data["budget"] = [test_budget] + mock_prisma_client.data["organization"] = [ + test_organization_1, + test_organization_2, + ] + mock_prisma_client._update_db_organizations() + + # Run the test + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + # Verify results + assert len(mock_prisma_client.updated_data["budget"]) == 1 + updated_budget = mock_prisma_client.updated_data["budget"][0] + assert updated_budget.budget_reset_at > now + + # Verify both organizations had their spend reset to 0 + assert test_organization_1.spend == 0.0 + assert test_organization_2.spend == 0.0 + + +def test_reset_budget_for_organization_failure(reset_budget_job, mock_prisma_client): + """ + Test that organization reset failures are properly handled and tracked. + This tests the error handling logic for organization budget resets. + The function catches exceptions and logs them, but doesn't re-raise them. + """ + # Setup test data + now = datetime.now(timezone.utc) + test_budget = type( + "LiteLLM_BudgetTable", + (), + { + "max_budget": 1000.0, + "budget_duration": "1d", + "budget_reset_at": now, + "budget_id": "test-budget-org-fail", + "created_at": now - timedelta(days=2), + }, + ) + + test_organization = type( + "LiteLLM_OrganizationTable", + (), + { + "organization_id": "test-org-fail", + "organization_alias": "Test Org Fail", + "budget_id": "test-budget-org-fail", + "spend": 500.0, + }, + ) + + # Create a mock that will raise an exception on update_many + class FailingMockLiteLLMOrganizationTable(MockLiteLLMOrganizationTable): + async def update_many( + self, where: Dict[str, Any], data: Dict[str, Any] + ) -> Dict[str, Any]: + raise Exception("Database connection failed") + + mock_prisma_client.data["budget"] = [test_budget] + mock_prisma_client.data["organization"] = [test_organization] + mock_prisma_client.db.litellm_organizationtable = ( + FailingMockLiteLLMOrganizationTable([test_organization]) + ) + + # Run the test - the function catches exceptions internally and logs them + # It doesn't re-raise them, so the function completes normally + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + # Verify that the organization spend was NOT reset due to the failure + # The spend should still be 500.0 because the update failed + assert test_organization.spend == 500.0 + + # Verify that the budget was still processed (it should be updated) + assert len(mock_prisma_client.updated_data["budget"]) == 1