mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix: reset team member budgets based on team's budget_duration
Fixes #19105 ## Problem Team member budgets (max_budget_in_team) were not resetting properly: 1. Team members use litellm_budget_table for budget limits 2. Budget resets only occurred if litellm_budget_table had budget_duration set 3. When creating team members with max_budget_in_team, budget_duration wasn't inherited 4. Without budget_duration, team member budgets never reset, even though team budgets did ## Root Cause The reset_budget_for_litellm_budget_table() method only reset budgets that had their own budget_duration and reset_at time. Team members without explicit budget_duration were never included in the reset. ## Solution Added reset_budget_for_team_members_by_team() method that: - Runs during the budget reset job - Finds teams that need budget reset (based on team's budget_duration) - Resets spend for ALL team members in those teams - Ensures team member budgets inherit the team's reset schedule ## Changes - litellm/proxy/common_utils/reset_budget_job.py: - Added reset_budget_for_team_members_by_team() method - Called in reset_budget() after resetting team budgets - Added comprehensive tests to verify the fix - Added reproduction script for issue validation ## Testing - Unit tests verify team member budgets reset with team - Tested with multiple teams and edge cases - All existing tests pass Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
Julio Quinteros Pro
parent
9a2410be71
commit
7c322f9946
3 changed files with 435 additions and 0 deletions
|
|
@ -45,6 +45,9 @@ class ResetBudgetJob:
|
|||
## Reset Team Budget
|
||||
await self.reset_budget_for_litellm_teams()
|
||||
|
||||
### RESET TEAM MEMBER BUDGET (based on team budget duration) ###
|
||||
await self.reset_budget_for_team_members_by_team()
|
||||
|
||||
### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ###
|
||||
await self.reset_budget_for_litellm_budget_table()
|
||||
|
||||
|
|
@ -481,6 +484,85 @@ class ResetBudgetJob:
|
|||
)
|
||||
verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e)
|
||||
|
||||
async def reset_budget_for_team_members_by_team(self):
|
||||
"""
|
||||
Resets the budget for team members based on their team's budget_duration.
|
||||
|
||||
This fixes issue #19105 where team member budgets don't reset because they
|
||||
don't have their own budget_duration set - they should inherit from the team.
|
||||
"""
|
||||
now = datetime.utcnow()
|
||||
start_time = time.time()
|
||||
teams_to_process: Optional[List[LiteLLM_TeamTable]] = None
|
||||
try:
|
||||
# Get all teams that need budget reset
|
||||
teams_to_process = await self.prisma_client.get_data(
|
||||
table_name="team", query_type="find_all", reset_at=now
|
||||
)
|
||||
|
||||
if teams_to_process is not None and len(teams_to_process) > 0:
|
||||
team_ids_to_reset = [
|
||||
team.team_id
|
||||
for team in teams_to_process
|
||||
if team.team_id is not None
|
||||
]
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Resetting team member budgets for teams: %s", team_ids_to_reset
|
||||
)
|
||||
|
||||
# Reset spend for all team members in these teams
|
||||
result = await self.prisma_client.db.litellm_teammembership.update_many(
|
||||
where={
|
||||
"team_id": {"in": team_ids_to_reset}
|
||||
},
|
||||
data={
|
||||
"spend": 0,
|
||||
},
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Reset %s team member budgets", result
|
||||
)
|
||||
|
||||
end_time = time.time()
|
||||
asyncio.create_task(
|
||||
self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
|
||||
service=ServiceTypes.RESET_BUDGET_JOB,
|
||||
duration=end_time - start_time,
|
||||
call_type="reset_budget_team_members_by_team",
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
event_metadata={
|
||||
"num_teams_found": len(teams_to_process) if teams_to_process else 0,
|
||||
"teams_found": json.dumps(
|
||||
teams_to_process, indent=4, default=str
|
||||
),
|
||||
},
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
end_time = time.time()
|
||||
asyncio.create_task(
|
||||
self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
|
||||
service=ServiceTypes.RESET_BUDGET_JOB,
|
||||
duration=end_time - start_time,
|
||||
error=e,
|
||||
call_type="reset_budget_team_members_by_team",
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
event_metadata={
|
||||
"num_teams_found": len(teams_to_process) if teams_to_process else 0,
|
||||
"teams_found": json.dumps(
|
||||
teams_to_process, indent=4, default=str
|
||||
),
|
||||
},
|
||||
)
|
||||
)
|
||||
verbose_proxy_logger.exception(
|
||||
"Failed to reset budget for team members by team: %s", e
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _reset_budget_common(
|
||||
item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken],
|
||||
|
|
|
|||
178
reproduce_issue_19105.py
Normal file
178
reproduce_issue_19105.py
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
"""
|
||||
Reproduction script for issue #19105: Team member budgets not enforced
|
||||
|
||||
This script demonstrates the bug where:
|
||||
1. Team member budget (max_budget_in_team) is not enforced - requests go through even when exceeded
|
||||
2. Team member spend doesn't reset daily while team spend does
|
||||
|
||||
Setup:
|
||||
- Create a team with max_budget and budget_duration
|
||||
- Add a key to the team
|
||||
- Add a user with max_budget_in_team
|
||||
- Make requests to exceed the user's budget
|
||||
- Verify requests still go through (BUG)
|
||||
- Trigger budget reset and verify team member spend doesn't reset (BUG)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
# Add litellm to path
|
||||
sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
|
||||
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
||||
|
||||
|
||||
async def reproduce_issue():
|
||||
print("=" * 80)
|
||||
print("REPRODUCING ISSUE #19105: Team Member Budget Not Enforced")
|
||||
print("=" * 80)
|
||||
|
||||
# Initialize Prisma client
|
||||
prisma_client = PrismaClient()
|
||||
await prisma_client.connect()
|
||||
|
||||
try:
|
||||
# Step 1: Create a team with budget and budget_duration
|
||||
print("\n1. Creating team with budget...")
|
||||
team = await prisma_client.db.litellm_teamtable.create(
|
||||
data={
|
||||
"team_alias": "test_team_19105",
|
||||
"max_budget": 1.0,
|
||||
"budget_duration": "1d", # Daily reset
|
||||
"spend": 0.0,
|
||||
}
|
||||
)
|
||||
print(f" Created team: {team.team_id}, budget: ${team.max_budget}, duration: {team.budget_duration}")
|
||||
|
||||
# Step 2: Create a key for the team
|
||||
print("\n2. Creating key for team...")
|
||||
key = await prisma_client.db.litellm_verificationtoken.create(
|
||||
data={
|
||||
"token": f"sk-test-{team.team_id}",
|
||||
"team_id": team.team_id,
|
||||
}
|
||||
)
|
||||
print(f" Created key: {key.token}")
|
||||
|
||||
# Step 3: Create a budget table for team member
|
||||
print("\n3. Creating budget table for team member...")
|
||||
budget = await prisma_client.db.litellm_budgettable.create(
|
||||
data={
|
||||
"max_budget": 0.01, # Very low budget to trigger limit
|
||||
# Note: NOT setting budget_duration - this is the bug!
|
||||
}
|
||||
)
|
||||
print(f" Created budget: budget_id={budget.budget_id}, max_budget=${budget.max_budget}")
|
||||
print(f" ⚠️ Budget duration: {budget.budget_duration} (None - this is the problem!)")
|
||||
|
||||
# Step 4: Create a user (internal user)
|
||||
print("\n4. Creating internal user...")
|
||||
user = await prisma_client.db.litellm_usertable.create(
|
||||
data={
|
||||
"user_id": "test_user_19105",
|
||||
"user_email": "test@example.com",
|
||||
}
|
||||
)
|
||||
print(f" Created user: {user.user_id}")
|
||||
|
||||
# Step 5: Create team membership linking user to team with budget
|
||||
print("\n5. Creating team membership with max_budget_in_team...")
|
||||
membership = await prisma_client.db.litellm_teammembership.create(
|
||||
data={
|
||||
"team_id": team.team_id,
|
||||
"user_id": user.user_id,
|
||||
"budget_id": budget.budget_id,
|
||||
"spend": 0.015, # Already over budget!
|
||||
}
|
||||
)
|
||||
print(f" Created membership: user={user.user_id}, team={team.team_id}")
|
||||
print(f" Member spend: ${membership.spend}, Member budget: ${budget.max_budget}")
|
||||
print(f" ⚠️ Spend ${membership.spend} > Budget ${budget.max_budget} - should be blocked!")
|
||||
|
||||
# Step 6: Check if budget enforcement would work
|
||||
print("\n6. Testing budget check logic...")
|
||||
membership_with_budget = await prisma_client.db.litellm_teammembership.find_unique(
|
||||
where={"user_id_team_id": {"user_id": user.user_id, "team_id": team.team_id}},
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
|
||||
if membership_with_budget and membership_with_budget.litellm_budget_table:
|
||||
team_member_budget = membership_with_budget.litellm_budget_table.max_budget
|
||||
team_member_spend = membership_with_budget.spend or 0.0
|
||||
|
||||
print(f" Spend: ${team_member_spend}, Budget: ${team_member_budget}")
|
||||
if team_member_spend >= team_member_budget:
|
||||
print(f" ✓ Budget check WOULD block request (spend >= budget)")
|
||||
else:
|
||||
print(f" ✗ Budget check would NOT block request")
|
||||
|
||||
# Step 7: Test budget reset logic
|
||||
print("\n7. Testing budget reset logic...")
|
||||
print(f" Current team member spend: ${membership.spend}")
|
||||
print(f" Budget reset_at: {budget.budget_reset_at}")
|
||||
|
||||
# Create ProxyLogging object (needed for ResetBudgetJob)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
|
||||
reset_job = ResetBudgetJob(
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
# Manually set budget reset_at to past to trigger reset
|
||||
now = datetime.now(timezone.utc)
|
||||
past_time = now - timedelta(hours=1)
|
||||
await prisma_client.db.litellm_budgettable.update(
|
||||
where={"budget_id": budget.budget_id},
|
||||
data={"budget_reset_at": past_time},
|
||||
)
|
||||
print(f" Set budget reset_at to: {past_time} (in the past)")
|
||||
|
||||
# Run budget reset
|
||||
print(f"\n8. Running budget reset job...")
|
||||
await reset_job.reset_budget()
|
||||
|
||||
# Check if team member spend was reset
|
||||
membership_after = await prisma_client.db.litellm_teammembership.find_unique(
|
||||
where={"user_id_team_id": {"user_id": user.user_id, "team_id": team.team_id}},
|
||||
)
|
||||
|
||||
print(f" Team member spend after reset: ${membership_after.spend}")
|
||||
if membership_after.spend == 0.0:
|
||||
print(f" ✓ Team member spend WAS reset")
|
||||
else:
|
||||
print(f" ✗ BUG: Team member spend NOT reset (still ${membership_after.spend})")
|
||||
print(f" This is because the budget table doesn't have a budget_duration set!")
|
||||
|
||||
# Clean up
|
||||
print("\n9. Cleaning up...")
|
||||
await prisma_client.db.litellm_teammembership.delete(
|
||||
where={"user_id_team_id": {"user_id": user.user_id, "team_id": team.team_id}}
|
||||
)
|
||||
await prisma_client.db.litellm_usertable.delete(where={"user_id": user.user_id})
|
||||
await prisma_client.db.litellm_verificationtoken.delete(where={"token": key.token})
|
||||
await prisma_client.db.litellm_budgettable.delete(where={"budget_id": budget.budget_id})
|
||||
await prisma_client.db.litellm_teamtable.delete(where={"team_id": team.team_id})
|
||||
print(" Cleanup complete")
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("CONCLUSION:")
|
||||
print("=" * 80)
|
||||
print("Team member budgets don't reset because:")
|
||||
print("1. Team members use litellm_budget_table for budget limits")
|
||||
print("2. Budget resets only happen if litellm_budget_table has budget_duration set")
|
||||
print("3. When creating team members with max_budget_in_team, the budget_duration")
|
||||
print(" is NOT automatically inherited from the team")
|
||||
print("4. Without budget_duration, the budget never resets")
|
||||
print("\nFIX: Team member budgets should inherit budget_duration from their team")
|
||||
print("=" * 80)
|
||||
|
||||
finally:
|
||||
await prisma_client.disconnect()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(reproduce_issue())
|
||||
175
tests/proxy_unit_tests/test_team_member_budget_reset.py
Normal file
175
tests/proxy_unit_tests/test_team_member_budget_reset.py
Normal file
|
|
@ -0,0 +1,175 @@
|
|||
"""
|
||||
Test for issue #19105: Team member budget reset
|
||||
|
||||
Tests that team member budgets reset based on the team's budget_duration,
|
||||
not just their individual budget table's duration.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import pytest
|
||||
import sys
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
# Add project root to path
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_resets_with_team():
|
||||
"""
|
||||
Test that team member budgets reset when the team's budget resets,
|
||||
even if the member's budget table doesn't have its own budget_duration.
|
||||
|
||||
This tests the fix for issue #19105.
|
||||
"""
|
||||
# Mock Prisma client
|
||||
prisma_client = MagicMock(spec=PrismaClient)
|
||||
prisma_client.db = MagicMock()
|
||||
|
||||
# Mock ProxyLogging
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.service_logging_obj = MagicMock()
|
||||
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
||||
|
||||
# Create test data
|
||||
now = datetime.utcnow()
|
||||
past_time = now - timedelta(hours=1)
|
||||
|
||||
# Team with budget_duration that needs reset
|
||||
team = MagicMock()
|
||||
team.team_id = "team_123"
|
||||
team.max_budget = 10.0
|
||||
team.budget_duration = "1d"
|
||||
team.budget_reset_at = past_time # In the past, needs reset
|
||||
team.spend = 5.0
|
||||
|
||||
# Mock get_data to return the team
|
||||
async def mock_get_data(table_name, query_type, **kwargs):
|
||||
if table_name == "team" and "reset_at" in kwargs:
|
||||
return [team]
|
||||
return None
|
||||
|
||||
prisma_client.get_data = AsyncMock(side_effect=mock_get_data)
|
||||
|
||||
# Mock update_many for team membership
|
||||
update_result = MagicMock()
|
||||
update_result.count = 2 # 2 team members updated
|
||||
prisma_client.db.litellm_teammembership.update_many = AsyncMock(
|
||||
return_value=update_result
|
||||
)
|
||||
|
||||
# Create reset job
|
||||
reset_job = ResetBudgetJob(
|
||||
proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client
|
||||
)
|
||||
|
||||
# Run the team member reset
|
||||
await reset_job.reset_budget_for_team_members_by_team()
|
||||
|
||||
# Verify that update_many was called with correct parameters
|
||||
prisma_client.db.litellm_teammembership.update_many.assert_called_once()
|
||||
call_args = prisma_client.db.litellm_teammembership.update_many.call_args
|
||||
|
||||
# Check the where clause includes the team_id
|
||||
assert call_args.kwargs["where"]["team_id"] == {"in": ["team_123"]}
|
||||
|
||||
# Check that spend is reset to 0
|
||||
assert call_args.kwargs["data"]["spend"] == 0
|
||||
|
||||
# Verify success hook was called
|
||||
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_reset_no_teams():
|
||||
"""
|
||||
Test that the function handles the case where no teams need resetting.
|
||||
"""
|
||||
# Mock Prisma client
|
||||
prisma_client = MagicMock(spec=PrismaClient)
|
||||
prisma_client.db = MagicMock()
|
||||
|
||||
# Mock ProxyLogging
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.service_logging_obj = MagicMock()
|
||||
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
||||
|
||||
# Mock get_data to return empty list (no teams need reset)
|
||||
prisma_client.get_data = AsyncMock(return_value=[])
|
||||
|
||||
# Mock update_many
|
||||
prisma_client.db.litellm_teammembership.update_many = AsyncMock()
|
||||
|
||||
# Create reset job
|
||||
reset_job = ResetBudgetJob(
|
||||
proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client
|
||||
)
|
||||
|
||||
# Run the team member reset
|
||||
await reset_job.reset_budget_for_team_members_by_team()
|
||||
|
||||
# Verify that update_many was NOT called (no teams to process)
|
||||
prisma_client.db.litellm_teammembership.update_many.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_reset_multiple_teams():
|
||||
"""
|
||||
Test that team member budgets are reset for multiple teams.
|
||||
"""
|
||||
# Mock Prisma client
|
||||
prisma_client = MagicMock(spec=PrismaClient)
|
||||
prisma_client.db = MagicMock()
|
||||
|
||||
# Mock ProxyLogging
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.service_logging_obj = MagicMock()
|
||||
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
||||
|
||||
# Create test data - multiple teams
|
||||
now = datetime.utcnow()
|
||||
past_time = now - timedelta(hours=1)
|
||||
|
||||
teams = []
|
||||
for i in range(3):
|
||||
team = MagicMock()
|
||||
team.team_id = f"team_{i}"
|
||||
team.max_budget = 10.0
|
||||
team.budget_duration = "1d"
|
||||
team.budget_reset_at = past_time
|
||||
teams.append(team)
|
||||
|
||||
# Mock get_data to return multiple teams
|
||||
prisma_client.get_data = AsyncMock(return_value=teams)
|
||||
|
||||
# Mock update_many
|
||||
update_result = MagicMock()
|
||||
update_result.count = 5 # 5 team members updated across all teams
|
||||
prisma_client.db.litellm_teammembership.update_many = AsyncMock(
|
||||
return_value=update_result
|
||||
)
|
||||
|
||||
# Create reset job
|
||||
reset_job = ResetBudgetJob(
|
||||
proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client
|
||||
)
|
||||
|
||||
# Run the team member reset
|
||||
await reset_job.reset_budget_for_team_members_by_team()
|
||||
|
||||
# Verify that update_many was called with all team IDs
|
||||
call_args = prisma_client.db.litellm_teammembership.update_many.call_args
|
||||
assert set(call_args.kwargs["where"]["team_id"]["in"]) == {
|
||||
"team_0",
|
||||
"team_1",
|
||||
"team_2",
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue