scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Use end user budget instead of key budget when creating new team

Browse source
This commit is contained in:
yuneng-jiang 2025-10-29 15:51:23 -07:00
parent 06449df9d3
commit 9a22b890e7
2 changed files with 162 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -688,13 +688,13 @@ async def new_team( # noqa: PLR0915
if (
data.max_budget is not None
and user_api_key_dict.max_budget is not None
and data.max_budget > user_api_key_dict.max_budget
and user_api_key_dict.end_user_max_budget is not None
and data.max_budget > user_api_key_dict.end_user_max_budget
):
raise HTTPException(
status_code=400,
detail={
"error": f"max budget higher than user max. User max budget={user_api_key_dict.max_budget}. User role={user_api_key_dict.user_role}"
"error": f"max budget higher than user max. User max budget={user_api_key_dict.end_user_max_budget}. User role={user_api_key_dict.user_role}"
},
)

159
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -1793,3 +1793,162 @@ async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_a
mock_db_client.db.litellm_teammembership.delete_many.assert_awaited_with(
where={"team_id": test_team_id, "user_id": test_user_id}
)
@pytest.mark.asyncio
async def test_new_team_max_budget_exceeds_user_max_budget():
"""
Test that /team/new raises ProxyException when max_budget exceeds user's end_user_max_budget.
This validates the budget enforcement logic where non-admin users cannot create teams
with budgets higher than their personal maximum budget limit.
"""
from fastapi import Request
from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import new_team
# Create non-admin user with end_user_max_budget set to 100.0
non_admin_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="non-admin-user-123",
end_user_max_budget=100.0,
)
# Create team request with max_budget (200.0) exceeding user's limit (100.0)
team_request = NewTeamRequest(
team_alias="high-budget-team",
max_budget=200.0, # Exceeds user's end_user_max_budget
)
dummy_request = MagicMock(spec=Request)
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.proxy_server._license_check"
) as mock_license, patch(
"litellm.proxy.proxy_server.user_api_key_cache"
) as mock_cache, patch(
"litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"
), patch(
"litellm.proxy.proxy_server.create_audit_log_for_update"
) as mock_audit:
# Setup basic mocks
mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_license.is_team_count_over_limit.return_value = False
mock_prisma.get_data = AsyncMock(return_value=None)
mock_audit.return_value = AsyncMock(return_value=MagicMock())
# The budget validation happens BEFORE trying to add members,
# so we shouldn't need to mock the user table operations
# Should raise ProxyException (HTTPException gets converted by handle_exception_on_proxy)
with pytest.raises(ProxyException) as exc_info:
await new_team(
data=team_request,
http_request=dummy_request,
user_api_key_dict=non_admin_user,
)
# Verify exception details
# ProxyException stores status_code in 'code' attribute
assert exc_info.value.code == '400'
assert "max budget higher than user max" in str(exc_info.value.message)
assert "100.0" in str(exc_info.value.message) # User's max budget should be mentioned
assert LitellmUserRoles.INTERNAL_USER.value in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_new_team_max_budget_within_user_limit():
"""
Test that /team/new succeeds when max_budget is within user's end_user_max_budget.
This ensures that users can create teams with budgets at or below their personal limit.
"""
from fastapi import Request
from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import new_team
# Create non-admin user with end_user_max_budget set to 100.0
non_admin_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="non-admin-user-456",
end_user_max_budget=100.0,
models=[], # Empty models list to bypass model validation
)
# Create team request with max_budget (50.0) within user's limit (100.0)
team_request = NewTeamRequest(
team_alias="within-budget-team",
max_budget=50.0, # Within user's end_user_max_budget
)
dummy_request = MagicMock(spec=Request)
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.proxy_server.user_api_key_cache"
) as mock_cache, patch(
"litellm.proxy.proxy_server._license_check"
) as mock_license, patch(
"litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"
), patch(
"litellm.proxy.proxy_server.create_audit_log_for_update"
) as mock_audit:
# Setup mocks
mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
mock_license.is_team_count_over_limit.return_value = False
mock_prisma.jsonify_team_object = lambda db_data: db_data
mock_prisma.get_data = AsyncMock(return_value=None)
mock_prisma.update_data = AsyncMock()
mock_audit.return_value = AsyncMock(return_value=MagicMock())
# Mock team creation
mock_created_team = MagicMock()
mock_created_team.team_id = "team-within-budget-789"
mock_created_team.team_alias = "within-budget-team"
mock_created_team.max_budget = 50.0
mock_created_team.members_with_roles = []
mock_created_team.metadata = None
mock_created_team.model_dump.return_value = {
"team_id": "team-within-budget-789",
"team_alias": "within-budget-team",
"max_budget": 50.0,
"members_with_roles": [],
}
mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team)
mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team)
# Mock model table
mock_prisma.db.litellm_modeltable = MagicMock()
mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123"))
# Mock user table operations for adding the creator as a member
mock_user = MagicMock()
mock_user.user_id = "non-admin-user-456"
mock_user.model_dump.return_value = {"user_id": "non-admin-user-456", "teams": ["team-within-budget-789"]}
mock_prisma.db.litellm_usertable = MagicMock()
mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user)
mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user)
# Mock team membership table
mock_membership = MagicMock()
mock_membership.model_dump.return_value = {
"team_id": "team-within-budget-789",
"user_id": "non-admin-user-456",
"budget_id": None,
}
mock_prisma.db.litellm_teammembership = MagicMock()
mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership)
# Should NOT raise an exception
result = await new_team(
data=team_request,
http_request=dummy_request,
user_api_key_dict=non_admin_user,
)
# Verify the team was created successfully
assert result is not None
assert result["team_id"] == "team-within-budget-789"
assert result["max_budget"] == 50.0