scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 152

Organization Budget resets

Browse source
This commit is contained in:
yuneng-jiang 2026-01-30 13:07:58 -08:00
parent f7e1a22947
commit 8397d4b8f8
2 changed files with 281 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

100
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -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

194
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -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