diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31-py3-none-any.whl new file mode 100644 index 0000000000..90b36bd78a Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31.tar.gz new file mode 100644 index 0000000000..6460723547 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql new file mode 100644 index 0000000000..a64f1de342 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "soft_budget" DOUBLE PRECISION; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index dc49036cb1..e713e6ca87 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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? diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index d43b591686..0e72cd9081 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -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==", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 131a6caab0..5b697249a4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 63db2d72fe..8e903180ac 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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 diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b118400b62..a6a573836b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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? diff --git a/pyproject.toml b/pyproject.toml index 40ca400401..768a41d2e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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} diff --git a/requirements.txt b/requirements.txt index aca27bf284..1fb8a22cc9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/schema.prisma b/schema.prisma index 03b910db85..240e0dfea4 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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? diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 467ee3661d..7e351e2916 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -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(): """