tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-13 02:22:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #20530 from BerriAI/litellm_ui_team_soft_budget

Browse Source 
[Feature] Add soft_budget to Team Table + Create/Update Endpoints
This commit is contained in:
yuneng-jiang
2026-02-05 15:52:17 -08:00
committed by GitHub
parent bff5f8d7b0 45c6382667
commit 523a36ed53
12 changed files with 343 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31-py3-none-any.whl
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31.tar.gz
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql
+3
View File
@@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "soft_budget" DOUBLE PRECISION;
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+1
View File
@@ -113,6 +113,7 @@ model LiteLLM_TeamTable {
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
soft_budget Float?
spend Float @default(0.0)
models String[]
max_parallel_requests Int?
litellm-proxy-extras/pyproject.toml
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-proxy-extras"
version = "0.4.30"
version = "0.4.31"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
authors = ["BerriAI"]
readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.4.30"
version = "0.4.31"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==",
litellm/proxy/_types.py
+2
View File
@@ -1488,6 +1488,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
# Budget fields
max_budget: Optional[float] = None
soft_budget: Optional[float] = None
budget_duration: Optional[str] = None
models: list = []
@@ -1559,6 +1560,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
max_budget: Optional[float] = None
soft_budget: Optional[float] = None
models: Optional[list] = None
blocked: Optional[bool] = None
budget_duration: Optional[str] = None
litellm/proxy/management_endpoints/team_endpoints.py
+46
View File
@@ -685,6 +685,7 @@ async def new_team( # noqa: PLR0915
- rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement.
- tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement.
- max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
- soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set, soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
- budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
- models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
- blocked: bool - Flag indicating if the team is blocked or not - will stop all calls from keys with this team_id.
@@ -760,6 +761,22 @@ async def new_team( # noqa: PLR0915
status_code=400,
detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"}
)
if data.soft_budget is not None and data.soft_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"}
)
if data.soft_budget is not None:
if data.max_budget is not None:
# If max_budget is set, soft_budget must be strictly lower than max_budget
if data.soft_budget >= data.max_budget:
raise HTTPException(
status_code=400,
detail={
"error": f"soft_budget ({data.soft_budget}) must be strictly lower than max_budget ({data.max_budget})"
}
)
# Check if license is over limit
total_teams = await prisma_client.db.litellm_teamtable.count()
@@ -1226,6 +1243,7 @@ async def update_team( # noqa: PLR0915
- tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit
- rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit
- max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
- soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
- budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
- models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
- prompts: Optional[List[str]] - List of prompts that the team is allowed to use.
@@ -1302,6 +1320,11 @@ async def update_team( # noqa: PLR0915
status_code=400,
detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"}
)
if data.soft_budget is not None and data.soft_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"}
)
existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": data.team_id}
@@ -1312,6 +1335,29 @@ async def update_team( # noqa: PLR0915
status_code=404,
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
if data.soft_budget is not None:
max_budget_to_check = data.max_budget if data.max_budget is not None else existing_team_row.max_budget
if max_budget_to_check is not None:
if data.soft_budget >= max_budget_to_check:
raise HTTPException(
status_code=400,
detail={
"error": f"soft_budget ({data.soft_budget}) must be strictly lower than max_budget ({max_budget_to_check})"
}
)
if data.max_budget is not None:
existing_soft_budget = getattr(existing_team_row, 'soft_budget', None)
soft_budget_to_check = data.soft_budget if data.soft_budget is not None else existing_soft_budget
if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)):
if data.max_budget <= soft_budget_to_check:
raise HTTPException(
status_code=400,
detail={
"error": f"max_budget ({data.max_budget}) must be strictly greater than soft_budget ({soft_budget_to_check})"
}
)
if (
data.organization_id is not None and len(data.organization_id) > 0
litellm/proxy/schema.prisma
+1
View File
@@ -113,6 +113,7 @@ model LiteLLM_TeamTable {
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
soft_budget Float?
spend Float @default(0.0)
models String[]
max_parallel_requests Int?
pyproject.toml
+1 -1
View File
@@ -61,7 +61,7 @@ boto3 = { version = "1.40.76", optional = true }
redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"}
a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"}
litellm-proxy-extras = {version = "0.4.30", optional = true}
litellm-proxy-extras = {version = "0.4.31", optional = true}
rich = {version = "13.7.1", optional = true}
litellm-enterprise = {version = "0.1.27", optional = true}
diskcache = {version = "^5.6.1", optional = true}
requirements.txt
+1 -1
View File
@@ -50,7 +50,7 @@ sentry_sdk==2.21.0 # for sentry error handling
detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
cryptography==44.0.1
tzdata==2025.1 # IANA time zone database
litellm-proxy-extras==0.4.30 # for proxy extras - e.g. prisma migrations
litellm-proxy-extras==0.4.31 # for proxy extras - e.g. prisma migrations
llm-sandbox==0.3.31 # for skill execution in sandbox
### LITELLM PACKAGE DEPENDENCIES
python-dotenv==1.0.1 # for env
schema.prisma
+1
View File
@@ -113,6 +113,7 @@ model LiteLLM_TeamTable {
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
soft_budget Float?
spend Float @default(0.0)
models String[]
max_parallel_requests Int?
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+285
View File
@@ -4931,6 +4931,291 @@ async def test_update_team_negative_team_member_budget():
assert request.team_member_budget == -15.0
# Parametrized tests for soft_budget in create endpoint
@pytest.mark.parametrize(
"soft_budget,max_budget,should_succeed,expected_soft_budget,expected_max_budget,error_message",
[
# Test 1: Soft budget only - success + soft budget set
(50.0, None, True, 50.0, None, None),
# Test 2: Soft budget with higher max budget, success with both set
(50.0, 100.0, True, 50.0, 100.0, None),
# Test 3: Soft budget with lower max budget, fail
(100.0, 50.0, False, None, None, "soft_budget (100.0) must be strictly lower than max_budget (50.0)"),
# Test 4: Soft budget equal to max budget, fail
(100.0, 100.0, False, None, None, "soft_budget (100.0) must be strictly lower than max_budget (100.0)"),
],
)
@pytest.mark.asyncio
async def test_new_team_soft_budget_validation(
soft_budget, max_budget, should_succeed, expected_soft_budget, expected_max_budget, error_message
):
"""
Test soft_budget validation in /team/new endpoint.
Covers:
- Soft budget only - success + soft budget set
- Soft budget with higher max budget, success with both set
- Soft budget with lower max budget, fail
"""
from fastapi import Request
from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth
from litellm.proxy.management_endpoints.team_endpoints import new_team
# Create admin user to bypass user budget checks
admin_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
models=[],
)
# Create team request with soft_budget and optionally max_budget
team_request = NewTeamRequest(
team_alias="test-soft-budget-team",
soft_budget=soft_budget,
max_budget=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", new=AsyncMock()
) 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 user cache
from litellm.proxy._types import LiteLLM_UserTable
mock_user_obj = LiteLLM_UserTable(
user_id="admin-user",
max_budget=None, # Admin has no budget limit
)
mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj)
# Mock team creation
mock_created_team = MagicMock()
mock_created_team.team_id = "test-team-123"
mock_created_team.team_alias = "test-soft-budget-team"
mock_created_team.soft_budget = expected_soft_budget
mock_created_team.max_budget = expected_max_budget
mock_created_team.members_with_roles = []
mock_created_team.metadata = None
mock_created_team.model_dump.return_value = {
"team_id": "test-team-123",
"team_alias": "test-soft-budget-team",
"soft_budget": expected_soft_budget,
"max_budget": expected_max_budget,
"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
mock_user = MagicMock()
mock_user.user_id = "admin-user"
mock_user.model_dump.return_value = {"user_id": "admin-user", "teams": ["test-team-123"]}
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": "test-team-123",
"user_id": "admin-user",
"budget_id": None,
}
mock_prisma.db.litellm_teammembership = MagicMock()
mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership)
if should_succeed:
# Should NOT raise an exception
result = await new_team(
data=team_request,
http_request=dummy_request,
user_api_key_dict=admin_user,
)
# Verify the team was created successfully with correct values
assert result is not None
assert result["team_id"] == "test-team-123"
if expected_soft_budget is not None:
assert result["soft_budget"] == expected_soft_budget
if expected_max_budget is not None:
assert result["max_budget"] == expected_max_budget
else:
# Should raise ProxyException
with pytest.raises(ProxyException) as exc_info:
await new_team(
data=team_request,
http_request=dummy_request,
user_api_key_dict=admin_user,
)
# Verify exception details
assert exc_info.value.code == '400'
if error_message:
assert error_message in str(exc_info.value.message)
# Parametrized tests for soft_budget in update endpoint
@pytest.mark.parametrize(
"existing_soft_budget,existing_max_budget,update_soft_budget,update_max_budget,should_succeed,expected_soft_budget,expected_max_budget,error_message",
[
# Test 1: Soft budget only (no previous max_budget) - success with soft budget set
(None, None, 50.0, None, True, 50.0, None, None),
# Test 2: Soft budget with max budget - success if soft budget is strictly lower than max budget
(None, None, 50.0, 100.0, True, 50.0, 100.0, None),
# Test 3: Soft budget with max budget - fail if soft budget >= max budget
(None, None, 100.0, 50.0, False, None, None, "soft_budget (100.0) must be strictly lower than max_budget (50.0)"),
# Test 4: Only max budget with existing soft_budget, success with max_budget strictly greater
(50.0, None, None, 100.0, True, 50.0, 100.0, None),
# Test 5: Only max budget with existing soft_budget, fail if max_budget <= soft_budget
(50.0, None, None, 50.0, False, None, None, "max_budget (50.0) must be strictly greater than soft_budget (50.0)"),
# Test 6: Update both soft_budget and max_budget - success if soft < max
(30.0, 100.0, 40.0, 80.0, True, 40.0, 80.0, None),
# Test 7: Update both soft_budget and max_budget - fail if soft >= max
(30.0, 100.0, 80.0, 40.0, False, None, None, "soft_budget (80.0) must be strictly lower than max_budget (40.0)"),
],
)
@pytest.mark.asyncio
async def test_update_team_soft_budget_validation(
existing_soft_budget, existing_max_budget, update_soft_budget, update_max_budget,
should_succeed, expected_soft_budget, expected_max_budget, error_message
):
"""
Test soft_budget validation in /team/update endpoint.
Covers:
- Soft budget only (no previous max_budget) - success with soft budget set
- Soft budget with max budget - success if soft budget is strictly lower than max budget, fail otherwise
- Only max budget with existing soft_budget, success with max_budget strictly greater, fail otherwise
"""
from fastapi import Request
from litellm.proxy._types import (
LiteLLM_UserTable,
ProxyException,
UpdateTeamRequest,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.team_endpoints import update_team
# Create admin user to bypass user budget checks
admin_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
models=[],
)
# Create update request
update_request = UpdateTeamRequest(
team_id="test-team-123",
soft_budget=update_soft_budget,
max_budget=update_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.litellm_proxy_admin_name", "admin"
), patch(
"litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()
) as mock_audit:
# Mock existing team with existing budgets
mock_existing_team = MagicMock()
mock_existing_team.team_id = "test-team-123"
mock_existing_team.organization_id = None
mock_existing_team.soft_budget = existing_soft_budget
mock_existing_team.max_budget = existing_max_budget
mock_existing_team.model_dump.return_value = {
"team_id": "test-team-123",
"organization_id": None,
"soft_budget": existing_soft_budget,
"max_budget": existing_max_budget,
}
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team)
# Mock user cache
mock_user_obj = LiteLLM_UserTable(
user_id="admin-user",
max_budget=None, # Admin has no budget limit
)
mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj)
# Mock updated team - preserve existing values if not being updated
final_soft_budget = update_soft_budget if update_soft_budget is not None else existing_soft_budget
final_max_budget = update_max_budget if update_max_budget is not None else existing_max_budget
mock_updated_team = MagicMock()
mock_updated_team.team_id = "test-team-123"
mock_updated_team.organization_id = None
mock_updated_team.soft_budget = final_soft_budget
mock_updated_team.max_budget = final_max_budget
mock_updated_team.model_dump.return_value = {
"team_id": "test-team-123",
"organization_id": None,
"soft_budget": final_soft_budget,
"max_budget": final_max_budget,
}
mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team)
mock_prisma.jsonify_team_object = lambda db_data: db_data
mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object
if should_succeed:
# Should NOT raise an exception
result = await update_team(
data=update_request,
http_request=dummy_request,
user_api_key_dict=admin_user,
)
# Verify the team was updated successfully with correct values
assert result is not None
assert result["data"].team_id == "test-team-123"
# Verify soft_budget matches expected value (or final computed value if expected is None)
if expected_soft_budget is not None:
assert result["data"].soft_budget == expected_soft_budget
else:
assert result["data"].soft_budget == final_soft_budget
# Verify max_budget matches expected value (or final computed value if expected is None)
if expected_max_budget is not None:
assert result["data"].max_budget == expected_max_budget
else:
assert result["data"].max_budget == final_max_budget
else:
# Should raise ProxyException
with pytest.raises(ProxyException) as exc_info:
await update_team(
data=update_request,
http_request=dummy_request,
user_api_key_dict=admin_user,
)
# Verify exception details
assert exc_info.value.code == '400'
if error_message:
assert error_message in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_new_team_positive_budgets_accepted():
"""