diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103000000_add_soft_budget_to_user_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103000000_add_soft_budget_to_user_table/migration.sql
new file mode 100644
index 00000000000..e7259dc7a5b
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261103000000_add_soft_budget_to_user_table/migration.sql
@@ -0,0 +1,3 @@
+SET LOCAL lock_timeout = '10s';
+
+ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "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 818f2dc4858..600678d744c 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
@@ -310,6 +310,7 @@ model LiteLLM_UserTable {
teams String[] @default([])
user_role String?
max_budget Float?
+ soft_budget Float?
spend Float @default(0.0)
user_email String?
models String[]
diff --git a/litellm/integrations/SlackAlerting/budget_alert_types.py b/litellm/integrations/SlackAlerting/budget_alert_types.py
index d25e24abc11..0d335c026bb 100644
--- a/litellm/integrations/SlackAlerting/budget_alert_types.py
+++ b/litellm/integrations/SlackAlerting/budget_alert_types.py
@@ -31,6 +31,8 @@ class SoftBudgetAlert(BaseBudgetAlertType):
def get_id(self, user_info: CallInfo) -> str:
if user_info.event_group == Litellm_EntityType.TEAM:
return user_info.team_id or "default_id"
+ if user_info.event_group == Litellm_EntityType.USER:
+ return f"user:{user_info.user_id}" if user_info.user_id else "default_id"
return user_info.token or "default_id"
diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py
index dd611947228..fcd660e6b21 100644
--- a/litellm/integrations/SlackAlerting/slack_alerting.py
+++ b/litellm/integrations/SlackAlerting/slack_alerting.py
@@ -78,6 +78,35 @@ else:
Router = Any
+def _budget_alert_email_html(webhook_event: WebhookEvent, email_logo_url: str, email_support_contact: str) -> str:
+ if webhook_event.event == "budget_crossed":
+ return f"""
+
+
Hi {webhook_event.user_id},
+ Your LLM API usage this month has reached your account's monthly budget of ${webhook_event.max_budget}
+ API requests will be rejected until either (a) you increase your monthly budget or (b) your monthly usage resets at the beginning of the next calendar month.
+ If you have any questions, please send an email to {email_support_contact}
+ Best,
+ The LiteLLM team
+ """
+ if webhook_event.event == "soft_budget_crossed":
+ hard_limit: Final = (
+ f" until your budget of ${webhook_event.max_budget} is reached"
+ if webhook_event.max_budget is not None
+ else ""
+ )
+ return f"""
+
+
Hi {webhook_event.user_id},
+ Your LLM API spend of ${webhook_event.spend} has reached your soft budget of ${webhook_event.soft_budget} .
+ This is a warning only. API requests keep going through{hard_limit}.
+ If you have any questions, please send an email to {email_support_contact}
+ Best,
+ The LiteLLM team
+ """
+ return "Alert from LiteLLM Server"
+
+
def _proxy_llm_router() -> Router | None:
from litellm.proxy.proxy_server import llm_router
@@ -1390,28 +1419,13 @@ Model Info:
event_name: Final = webhook_event.event_message
recipient_email: Final = webhook_event.user_email
- user_name: Final = webhook_event.user_id
- max_budget: Final = webhook_event.max_budget
- email_html_content = "Alert from LiteLLM Server"
if recipient_email is None:
verbose_proxy_logger.error("Trying to send email alert to no recipient", extra=webhook_event.dict())
-
- if webhook_event.event == "budget_crossed":
- email_html_content = f"""
-
-
-
Hi {user_name},
-
- Your LLM API usage this month has reached your account's monthly budget of ${max_budget}
-
- API requests will be rejected until either (a) you increase your monthly budget or (b) your monthly usage resets at the beginning of the next calendar month.
-
- If you have any questions, please send an email to {email_support_contact}
-
- Best,
- The LiteLLM team
- """
-
+ email_html_content: Final = _budget_alert_email_html(
+ webhook_event=webhook_event,
+ email_logo_url=email_logo_url,
+ email_support_contact=email_support_contact,
+ )
webhook_event.model_dump_json()
email_event: Final = {
"to": recipient_email,
diff --git a/litellm/models/user.py b/litellm/models/user.py
index 7dcaed9f77f..351d332d0a3 100644
--- a/litellm/models/user.py
+++ b/litellm/models/user.py
@@ -29,6 +29,7 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase):
teams: list[str] = []
user_role: str | None = None
max_budget: float | None = None
+ soft_budget: float | None = None
spend: float = 0.0
user_email: str | None = None
models: list = []
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 5b5c0c5f8d3..30f1b602b97 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1753,6 +1753,7 @@ from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E4
class NewUserResponse(GenerateKeyResponse):
max_budget: float | None = None
+ soft_budget: float | None = None
user_email: str | None = None
user_role: (
Literal[
@@ -3255,6 +3256,7 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase):
user_role: str | None = None
spend: float = 0.0
max_budget: float | None = None
+ soft_budget: float | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
models: list[str] = []
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index af60a315c60..f50e5202027 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -1416,7 +1416,7 @@ async def common_checks(
async def _user_max_budget_check() -> None:
# 4.1 personal budget
- if user_object is None or user_object.max_budget is None:
+ if user_object is None or (user_object.max_budget is None and user_object.soft_budget is None):
return
is_team_key: Final = team_object is not None and team_object.team_id is not None
if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True:
@@ -1430,6 +1430,9 @@ async def common_checks(
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
)
+ _user_soft_budget_check(user_object=user_object, user_spend=user_spend, proxy_logging_obj=proxy_logging_obj)
+ if user_budget is None:
+ return
call_info: Final = CallInfo(
spend=user_spend,
max_budget=user_budget,
@@ -7306,6 +7309,34 @@ async def _team_multi_budget_check(
)
+def _user_soft_budget_check(
+ user_object: LiteLLM_UserTable,
+ user_spend: float,
+ proxy_logging_obj: ProxyLogging,
+) -> None:
+ if user_object.soft_budget is None or user_spend < user_object.soft_budget:
+ return
+ verbose_proxy_logger.debug(
+ "Crossed Soft Budget for user %s, spend %s, soft_budget %s",
+ user_object.user_id,
+ user_spend,
+ user_object.soft_budget,
+ )
+ call_info: Final = CallInfo(
+ spend=user_spend,
+ soft_budget=user_object.soft_budget,
+ user_id=user_object.user_id,
+ user_email=user_object.user_email,
+ event_group=Litellm_EntityType.USER,
+ )
+ asyncio.create_task(
+ proxy_logging_obj.budget_alerts(
+ type="soft_budget",
+ user_info=call_info,
+ )
+ )
+
+
async def _team_soft_budget_check(
team_object: LiteLLM_TeamTable | None,
valid_token: UserAPIKeyAuth | None,
diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py
index d2aac9cae38..caa6fd1b100 100644
--- a/litellm/proxy/management_endpoints/internal_user_endpoints.py
+++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py
@@ -15,6 +15,7 @@ These are members of a Team on LiteLLM
import asyncio
import contextlib
import json
+import math
import traceback
from collections.abc import Awaitable, Mapping, Sequence
from datetime import datetime, timezone
@@ -23,7 +24,7 @@ from typing import Any, Final, Literal, Protocol, cast, overload
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
-from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, with_config
from typing_extensions import ReadOnly, TypedDict
import litellm
@@ -248,10 +249,37 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d
if litellm.internal_user_budget_duration is not None and data_json.get("budget_duration") is None:
data_json["budget_duration"] = litellm.internal_user_budget_duration
+ final_budgets: Final = _UserBudgetRow.model_validate(data_json)
+ _validate_user_soft_budget(soft_budget=final_budgets.soft_budget, max_budget=final_budgets.max_budget)
data_json.pop("teams", None) # handled separately
return data_json
+class _UserBudgetRow(BaseModel):
+ model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True)
+
+ max_budget: float | None = None
+ soft_budget: float | None = None
+
+
+_USER_BUDGET_ROWS: Final = TypeAdapter(tuple[_UserBudgetRow, ...])
+
+
+def _validate_user_soft_budget(soft_budget: float | None, max_budget: float | None) -> None:
+ if soft_budget is None:
+ return
+ if not math.isfinite(soft_budget) or soft_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"soft_budget must be a non-negative finite number. Received: {soft_budget}"},
+ )
+ if max_budget is not None and soft_budget >= max_budget:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"soft_budget ({soft_budget}) must be strictly lower than max_budget ({max_budget})"},
+ )
+
+
async def _check_duplicate_user_field(
field_name: str,
field_value: str | None,
@@ -496,6 +524,7 @@ async def new_user(
- send_invite_email: Optional[bool] - Specify if an invite email should be sent.
- user_role: Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20`
- max_budget: Optional[float] - Specify max budget for a given user.
+ - soft_budget: Optional[float] - Spend threshold that fires a budget alert for the user without blocking requests. Must be lower than max_budget when both are set.
- budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo").
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models). Set to ['no-default-models'] to block all model access. Restricting user to only team-based model access.
- tpm_limit: Optional[int] - Specify tpm limit for a given user (Tokens per minute)
@@ -1202,6 +1231,7 @@ async def user_info_v2(
user_role=user_data.get("user_role"),
spend=user_data.get("spend", 0.0),
max_budget=user_data.get("max_budget"),
+ soft_budget=user_data.get("soft_budget"),
tpm_limit=user_data.get("tpm_limit"),
rpm_limit=user_data.get("rpm_limit"),
models=user_data.get("models") or [],
@@ -1346,7 +1376,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
fields_set: Final = data.fields_set() if hasattr(data, "fields_set") else set()
for k, v in data_json.items():
- if k in ("max_budget", "budget_duration", "tpm_limit", "rpm_limit"):
+ if k in ("max_budget", "soft_budget", "budget_duration", "tpm_limit", "rpm_limit"):
if k in fields_set:
non_default_values[k] = v
elif k == "model_max_budget":
@@ -1380,6 +1410,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
if non_default_values["budget_duration"] is not None
else None
)
+ _validate_user_soft_budget(soft_budget=data.soft_budget, max_budget=None)
if "max_budget" not in non_default_values:
if (
@@ -1558,6 +1589,11 @@ async def _update_single_user_helper(
if existing_user_row is not None:
existing_user_row = LiteLLM_UserTable.model_validate(existing_user_row.model_dump(exclude_none=True))
+ requested: Final = user_request.model_fields_set
+ _validate_user_soft_budget(
+ soft_budget=user_request.soft_budget if "soft_budget" in requested else existing_user_row.soft_budget,
+ max_budget=user_request.max_budget if "max_budget" in requested else existing_user_row.max_budget,
+ )
# Prevent budget self-escalation (GHSA-wvg4-6222-3q4r): non-admin callers
# must not be able to raise their own budget/spend fields.
@@ -1749,6 +1785,7 @@ async def user_update(
- send_invite_email: Optional[bool] - Specify if an invite email should be sent.
- user_role: Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20`
- max_budget: Optional[float] - Specify max budget for a given user.
+ - soft_budget: Optional[float] - Spend threshold that fires a budget alert for the user without blocking requests. Must be lower than max_budget when both are set.
- budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo").
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
- tpm_limit: Optional[int] - Specify tpm limit for a given user (Tokens per minute)
@@ -2005,6 +2042,19 @@ async def bulk_user_update(
non_default_values.pop("user_id", None)
non_default_values.pop("user_email", None)
+ if "soft_budget" in non_default_values or "max_budget" in non_default_values:
+ for user_budget in _USER_BUDGET_ROWS.validate_python(all_users_in_db, from_attributes=True):
+ _validate_user_soft_budget(
+ soft_budget=(
+ data.user_updates.soft_budget
+ if "soft_budget" in non_default_values
+ else user_budget.soft_budget
+ ),
+ max_budget=(
+ data.user_updates.max_budget if "max_budget" in non_default_values else user_budget.max_budget
+ ),
+ )
+
successful_updates = 0
failed_updates: Final = 0
results: Final[list[UserUpdateResult]] = []
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index c8895c01f7b..798497f4470 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -5375,6 +5375,7 @@ async def generate_key_helper_fn(
user_data: Final = {
"max_budget": max_budget,
+ "soft_budget": soft_budget,
"user_email": user_email,
"user_id": user_id,
"user_alias": user_alias,
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index 818f2dc4858..600678d744c 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -310,6 +310,7 @@ model LiteLLM_UserTable {
teams String[] @default([])
user_role String?
max_budget Float?
+ soft_budget Float?
spend Float @default(0.0)
user_email String?
models String[]
diff --git a/litellm/types/proxy/management_endpoints/user_requests.py b/litellm/types/proxy/management_endpoints/user_requests.py
index e53dc4152c3..34328104f1a 100644
--- a/litellm/types/proxy/management_endpoints/user_requests.py
+++ b/litellm/types/proxy/management_endpoints/user_requests.py
@@ -19,6 +19,7 @@ class NewUserRequest(GenerateRequestBase):
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
max_budget: float | None = None
+ soft_budget: float | None = None
user_email: str | None = None
user_alias: str | None = None
user_role: (
@@ -66,6 +67,7 @@ class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with Bulk
| None
) = None
max_budget: float | None = None
+ soft_budget: float | None = None
class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail):
diff --git a/schema.prisma b/schema.prisma
index 818f2dc4858..600678d744c 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -310,6 +310,7 @@ model LiteLLM_UserTable {
teams String[] @default([])
user_role String?
max_budget Float?
+ soft_budget Float?
spend Float @default(0.0)
user_email String?
models String[]
diff --git a/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py b/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py
index f3199d9ebf9..088e22c71d0 100644
--- a/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py
+++ b/tests/unit/integrations/SlackAlerting/test_budget_alert_types.py
@@ -1,3 +1,5 @@
+from typing import Final
+
from litellm.integrations.SlackAlerting.budget_alert_types import (
SoftBudgetAlert,
TokenBudgetAlert,
@@ -6,6 +8,16 @@ from litellm.proxy._types import CallInfo, Litellm_EntityType
class TestSoftBudgetAlert:
+ def test_get_id_for_user_event_is_the_user_id(self) -> None:
+ alert: Final = SoftBudgetAlert()
+ first: Final = CallInfo(spend=1.0, user_id="user-a", event_group=Litellm_EntityType.USER)
+ second: Final = CallInfo(spend=1.0, user_id="user-b", event_group=Litellm_EntityType.USER)
+
+ team: Final = CallInfo(spend=1.0, team_id="user-a", event_group=Litellm_EntityType.TEAM)
+
+ assert alert.get_id(first) != alert.get_id(second)
+ assert alert.get_id(first) != alert.get_id(team)
+
def test_get_id_with_token(self):
"""Test that get_id returns user_info.token when token is provided"""
token_value = "test_token_123"
diff --git a/tests/unit/integrations/SlackAlerting/test_slack_alerting.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py
index aee8a3a692b..f824143c552 100644
--- a/tests/unit/integrations/SlackAlerting/test_slack_alerting.py
+++ b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py
@@ -621,9 +621,8 @@ async def test_daily_report_schedule_cache_calls_declare_their_key_family():
@pytest.fixture
def slack_alerting() -> SlackAlerting:
- return SlackAlerting(
- alerting_threshold=1, internal_usage_cache=DualCache(), alerting=["slack"]
- )
+ return SlackAlerting(alerting_threshold=1, internal_usage_cache=DualCache(), alerting=["slack"])
+
key_info: Final = CallInfo(
token="test_token",
@@ -665,6 +664,7 @@ key_no_max_budget_info: Final = CallInfo(
event_group=Litellm_EntityType.KEY,
)
+
@pytest.mark.parametrize(
"model, optional_params, expected_api_base",
[
@@ -701,11 +701,10 @@ async def test_response_taking_too_long_callback(slack_alerting):
end_time = start_time + timedelta(seconds=301)
kwargs = {"model": "test_model", "messages": "test_messages", "litellm_params": {}}
with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert:
- await slack_alerting.response_taking_too_long_callback(
- kwargs, None, start_time, end_time
- )
+ await slack_alerting.response_taking_too_long_callback(kwargs, None, start_time, end_time)
mock_send_alert.assert_awaited_once()
+
@pytest.mark.asyncio
async def test_alerting_metadata(slack_alerting):
"""
@@ -719,15 +718,13 @@ async def test_alerting_metadata(slack_alerting):
"litellm_params": {"metadata": {"alerting_metadata": {"hello": "world"}}},
}
with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert:
-
## RESPONSE TAKING TOO LONG
- await slack_alerting.response_taking_too_long_callback(
- kwargs, None, start_time, end_time
- )
+ await slack_alerting.response_taking_too_long_callback(kwargs, None, start_time, end_time)
mock_send_alert.assert_awaited_once()
assert "hello" in mock_send_alert.call_args[1]["alerting_metadata"]
+
@pytest.mark.asyncio
async def test_budget_alerts_crossed(slack_alerting):
user_max_budget = 100
@@ -744,6 +741,7 @@ async def test_budget_alerts_crossed(slack_alerting):
)
mock_send_alert.assert_awaited_once()
+
@pytest.mark.asyncio
async def test_budget_alerts_crossed_again(slack_alerting):
user_max_budget = 100
@@ -771,6 +769,7 @@ async def test_budget_alerts_crossed_again(slack_alerting):
)
mock_send_alert.assert_not_awaited()
+
@pytest.mark.asyncio
async def test_daily_reports_unit_test(slack_alerting):
with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert:
@@ -790,9 +789,7 @@ async def test_daily_reports_unit_test(slack_alerting):
updated_at=litellm.utils.get_utc_datetime(),
)
- updated_val = await slack_alerting.async_update_daily_reports(
- deployment_metrics=deployment_metrics
- )
+ updated_val = await slack_alerting.async_update_daily_reports(deployment_metrics=deployment_metrics)
assert updated_val == 1
@@ -800,6 +797,7 @@ async def test_daily_reports_unit_test(slack_alerting):
mock_send_alert.assert_awaited_once()
+
@pytest.mark.asyncio
async def test_send_daily_reports_ignores_zero_values():
router = MagicMock()
@@ -807,9 +805,7 @@ async def test_send_daily_reports_ignores_zero_values():
slack_alerting = SlackAlerting(internal_usage_cache=MagicMock())
# model1:failed=None, model2:failed=0, model3:failed=10, model1:latency=0; model2:latency=0; model3:latency=None
- slack_alerting.internal_usage_cache.async_batch_get_cache = AsyncMock(
- return_value=[None, 0, 10, 0, 0, None]
- )
+ slack_alerting.internal_usage_cache.async_batch_get_cache = AsyncMock(return_value=[None, 0, 10, 0, 0, None])
slack_alerting.internal_usage_cache.async_set_cache_pipeline = AsyncMock()
router.registry.return_value.deployment_dict.side_effect = lambda x: {"litellm_params": {"model": x}}
@@ -828,15 +824,14 @@ async def test_send_daily_reports_ignores_zero_values():
assert result == True
+
@pytest.mark.asyncio
async def test_send_daily_reports_all_zero_or_none():
router = MagicMock()
router.registry.return_value.names.return_value = ["model1", "model2", "model3"]
slack_alerting = SlackAlerting(internal_usage_cache=MagicMock())
- slack_alerting.internal_usage_cache.async_batch_get_cache = AsyncMock(
- return_value=[None, 0, None, 0, None, 0]
- )
+ slack_alerting.internal_usage_cache.async_batch_get_cache = AsyncMock(return_value=[None, 0, None, 0, None, 0])
with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert:
result = await slack_alerting.send_daily_reports(router)
@@ -846,6 +841,7 @@ async def test_send_daily_reports_all_zero_or_none():
assert result == False
+
@pytest.mark.parametrize(
"alerting_type",
[
@@ -883,6 +879,7 @@ async def test_send_token_budget_crossed_alerts(alerting_type):
)
mock_send_alert.assert_awaited_once()
+
@pytest.mark.parametrize(
"alerting_type",
[
@@ -898,9 +895,7 @@ async def test_send_token_budget_crossed_alerts(alerting_type):
async def test_webhook_alerting(alerting_type):
slack_alerting = SlackAlerting(alerting=["webhook"])
- with patch.object(
- slack_alerting, "send_webhook_alert", new=AsyncMock()
- ) as mock_send_alert:
+ with patch.object(slack_alerting, "send_webhook_alert", new=AsyncMock()) as mock_send_alert:
user_info = {
"token": "sk-test-mock-token-606",
"spend": 1,
@@ -921,6 +916,7 @@ async def test_webhook_alerting(alerting_type):
)
mock_send_alert.assert_awaited_once()
+
@pytest.mark.parametrize(
"model, api_base, llm_provider, vertex_project, vertex_location",
[
@@ -937,9 +933,7 @@ async def test_webhook_alerting(alerting_type):
)
@pytest.mark.parametrize("error_code", [500, 408, 400])
@pytest.mark.asyncio
-async def test_outage_alerting_called(
- model, api_base, llm_provider, vertex_project, vertex_location, error_code
-):
+async def test_outage_alerting_called(model, api_base, llm_provider, vertex_project, vertex_location, error_code):
"""
If call fails, outage alert is called
@@ -960,9 +954,7 @@ async def test_outage_alerting_called(
)
elif error_code == 408:
print("RAISING 408 ERROR CODE")
- error_to_raise = litellm.Timeout(
- message="A timeout occurred", model=model, llm_provider=llm_provider
- )
+ error_to_raise = litellm.Timeout(message="A timeout occurred", model=model, llm_provider=llm_provider)
elif error_code == 500:
print("RAISING 500 ERROR CODE")
error_to_raise = litellm.ServiceUnavailableError(
@@ -996,9 +988,7 @@ async def test_outage_alerting_called(
)
slack_alerting.update_values(llm_router=router)
- with patch.object(
- slack_alerting, "outage_alerts", new=AsyncMock()
- ) as mock_outage_alert:
+ with patch.object(slack_alerting, "outage_alerts", new=AsyncMock()) as mock_outage_alert:
try:
await router.acompletion(
model=model,
@@ -1026,6 +1016,7 @@ async def test_outage_alerting_called(
else:
mock_send_alert.assert_not_called()
+
@pytest.mark.parametrize(
"model, api_base, llm_provider, vertex_project, vertex_location",
[
@@ -1050,9 +1041,7 @@ async def test_region_outage_alerting_called(
If multiple calls fail, outage alert is sent
"""
- slack_alerting = SlackAlerting(
- alerting=["webhook"], alert_types=[AlertType.region_outage_alerts]
- )
+ slack_alerting = SlackAlerting(alerting=["webhook"], alert_types=[AlertType.region_outage_alerts])
litellm.callbacks = [slack_alerting]
@@ -1067,9 +1056,7 @@ async def test_region_outage_alerting_called(
)
elif error_code == 408:
print("RAISING 408 ERROR CODE")
- error_to_raise = litellm.Timeout(
- message="A timeout occurred", model=model, llm_provider=llm_provider
- )
+ error_to_raise = litellm.Timeout(message="A timeout occurred", model=model, llm_provider=llm_provider)
elif error_code == 500:
print("RAISING 500 ERROR CODE")
error_to_raise = litellm.ServiceUnavailableError(
@@ -1122,13 +1109,15 @@ async def test_region_outage_alerting_called(
else:
deployment_id = "2"
await slack_alerting.region_outage_alerts(
- exception=error_to_raise, deployment_id=deployment_id # type: ignore
+ exception=error_to_raise,
+ deployment_id=deployment_id, # type: ignore
)
if model == "gemini-3.8-flash" and (error_code == 500 or error_code == 408):
mock_send_alert.assert_called_once()
else:
mock_send_alert.assert_not_called()
+
@pytest.mark.asyncio
async def test_print_alerting_payload_warning():
"""
@@ -1158,9 +1147,7 @@ async def test_print_alerting_payload_warning():
test_payload = {"text": "Test alert message"}
# Send an alert
- with patch.object(
- slack_alerting.async_http_handler, "post", new=AsyncMock()
- ) as mock_post:
+ with patch.object(slack_alerting.async_http_handler, "post", new=AsyncMock()) as mock_post:
await send_to_webhook(
slackAlertingInstance=slack_alerting,
item={
@@ -1180,6 +1167,7 @@ async def test_print_alerting_payload_warning():
verbose_proxy_logger.removeHandler(handler)
log_stream.close()
+
@pytest.mark.asyncio
async def test_soft_budget_alerts():
"""
@@ -1223,6 +1211,7 @@ async def test_soft_budget_alerts():
)
assert alert_message == expected_message
+
@pytest.mark.parametrize(
"entity_info",
[
@@ -1265,3 +1254,47 @@ async def test_soft_budget_alerts_webhook(entity_info):
assert logged_webhook_event.user_email == entity_info.user_email
assert logged_webhook_event.key_alias == entity_info.key_alias
assert logged_webhook_event.event_group == entity_info.event_group
+
+
+def test_soft_budget_crossed_email_names_spend_and_threshold() -> None:
+ from litellm.integrations.SlackAlerting.slack_alerting import _budget_alert_email_html
+ from litellm.proxy._types import Litellm_EntityType, WebhookEvent
+
+ event: Final = WebhookEvent(
+ spend=52.5,
+ soft_budget=50,
+ max_budget=100,
+ user_id="user-a",
+ user_email="a@example.com",
+ event="soft_budget_crossed",
+ event_group=Litellm_EntityType.USER,
+ event_message="Soft budget crossed",
+ )
+ html: Final = _budget_alert_email_html(
+ webhook_event=event, email_logo_url="https://logo", email_support_contact="support@example.com"
+ )
+
+ assert "$52.5" in html
+ assert "soft budget of $50" in html
+ assert "until your budget of $100" in html
+ assert "will be rejected" not in html
+
+
+def test_budget_crossed_email_keeps_the_max_budget_body() -> None:
+ from litellm.integrations.SlackAlerting.slack_alerting import _budget_alert_email_html
+ from litellm.proxy._types import Litellm_EntityType, WebhookEvent
+
+ event: Final = WebhookEvent(
+ spend=100,
+ max_budget=100,
+ user_id="user-a",
+ event="budget_crossed",
+ event_group=Litellm_EntityType.USER,
+ event_message="Budget crossed",
+ )
+ html: Final = _budget_alert_email_html(
+ webhook_event=event, email_logo_url="https://logo", email_support_contact="support@example.com"
+ )
+
+ assert "monthly budget of $100" in html
+ assert "will be rejected" in html
diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py
index 9d4fe67614e..f66cce7026c 100644
--- a/tests/unit/proxy/auth/test_auth_checks.py
+++ b/tests/unit/proxy/auth/test_auth_checks.py
@@ -12,7 +12,7 @@ load_dotenv()
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from copy import deepcopy
from types import SimpleNamespace
-from typing import Final, Literal
+from typing import Final, Literal, cast
from unittest.mock import AsyncMock, MagicMock, patch
import pytest, litellm
@@ -44,6 +44,7 @@ from litellm.proxy.auth.auth_checks import (
is_model_cost_zero,
virtual_key_soft_budget_check,
_team_soft_budget_check,
+ _user_soft_budget_check,
common_checks,
)
from litellm.proxy.utils import InternalUsageCache, ProxyLogging
@@ -160,12 +161,15 @@ async def test_team_require_trace_id_accepts_supported_trace_carriers(
@pytest.mark.parametrize("route", ["/v1/traces", "/v1/logs", "/v1/traces/query"])
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_trace_telemetry_routes(route: str) -> None:
- assert await _run_team_trace_id_check(
- route=route,
- method="POST",
- request_body={},
- team_metadata={"require_trace_id": True},
- ) is True
+ assert (
+ await _run_team_trace_id_check(
+ route=route,
+ method="POST",
+ request_body={},
+ team_metadata={"require_trace_id": True},
+ )
+ is True
+ )
@pytest.mark.asyncio
@@ -197,11 +201,14 @@ async def test_team_require_trace_id_rejects_overridden_body_metadata_on_chat()
async def test_team_require_trace_id_accepts_effective_metadata_on_responses_route(
request_body: dict[str, object],
) -> None:
- assert await _run_team_trace_id_check(
- route="/v1/responses",
- request_body=request_body,
- team_metadata={"require_trace_id": True},
- ) is True
+ assert (
+ await _run_team_trace_id_check(
+ route="/v1/responses",
+ request_body=request_body,
+ team_metadata={"require_trace_id": True},
+ )
+ is True
+ )
@pytest.mark.parametrize(
@@ -240,12 +247,15 @@ async def test_team_require_trace_id_rejects_fallback_carriers_when_litellm_trac
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_mcp_with_empty_body() -> None:
- assert await _run_team_trace_id_check(
- route="/mcp/",
- request_body={},
- headers={"x-litellm-trace-id": "trace-mcp-123"},
- team_metadata={"require_trace_id": True},
- ) is True
+ assert (
+ await _run_team_trace_id_check(
+ route="/mcp/",
+ request_body={},
+ headers={"x-litellm-trace-id": "trace-mcp-123"},
+ team_metadata={"require_trace_id": True},
+ )
+ is True
+ )
@pytest.mark.parametrize(
@@ -278,12 +288,15 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_pass_through_
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_pass_through_route() -> None:
- assert await _run_team_trace_id_check(
- route="/anthropic/v1/messages",
- request_body={},
- headers={"x-litellm-trace-id": "trace-provider-123"},
- team_metadata={"require_trace_id": True},
- ) is True
+ assert (
+ await _run_team_trace_id_check(
+ route="/anthropic/v1/messages",
+ request_body={},
+ headers={"x-litellm-trace-id": "trace-provider-123"},
+ team_metadata={"require_trace_id": True},
+ )
+ is True
+ )
@pytest.mark.parametrize(
@@ -329,7 +342,9 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_registered_pa
"jsonrpc": "2.0",
"id": "1",
"method": "message/send",
- "params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
+ "params": {
+ "message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}
+ },
},
),
],
@@ -370,7 +385,9 @@ async def test_team_require_trace_id_rejects_traceparent_on_mcp_and_agent_routes
"jsonrpc": "2.0",
"id": "1",
"method": "message/send",
- "params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
+ "params": {
+ "message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}
+ },
"metadata": {"trace_id": "trace-body-123"},
},
),
@@ -463,7 +480,7 @@ def test_lens_management_does_not_require_an_inference_user_parameter(route: str
async def test_get_end_user_object(customer_spend, customer_budget):
"""
Scenario 1: normal - get_end_user_object returns the cached user
- Scenario 2: user over budget - NOTE: budget enforcement now happens in
+ Scenario 2: user over budget - NOTE: budget enforcement now happens in
common_checks() via _check_end_user_budget(), not in get_end_user_object()
This test verifies that get_end_user_object correctly retrieves the end user
@@ -591,7 +608,7 @@ async def test_can_key_call_model(model, expect_to_work):
if expect_to_work:
await can_key_call_model(**args)
else:
- with pytest.raises(Exception, match='is not available for this API key') as e:
+ with pytest.raises(Exception, match="is not available for this API key") as e:
await can_key_call_model(**args)
print(e)
@@ -908,16 +925,12 @@ async def test_is_valid_fallback_model():
)
try:
- await is_valid_fallback_model(
- model="gpt-3.5-turbo", llm_router=router, user_model=None
- )
+ await is_valid_fallback_model(model="gpt-3.5-turbo", llm_router=router, user_model=None)
except Exception as e:
pytest.fail(f"Expected is_valid_fallback_model to work, got exception: {e}")
with pytest.raises(Exception, match="Invalid") as exc_info:
- await is_valid_fallback_model(
- model="gpt-4o", llm_router=router, user_model=None
- )
+ await is_valid_fallback_model(model="gpt-4o", llm_router=router, user_model=None)
e = exc_info.value
assert "Invalid" in str(e)
@@ -931,9 +944,7 @@ async def test_is_valid_fallback_model():
],
)
@pytest.mark.asyncio
-async def test_virtual_key_max_budget_check(
- token_spend, max_budget, expect_budget_error
-):
+async def test_virtual_key_max_budget_check(token_spend, max_budget, expect_budget_error):
"""
Test if virtual key budget checks work as expected:
1. Triggers budget alert for all cases
@@ -1048,9 +1059,7 @@ async def test_can_team_access_model(model, team_models, expect_to_work):
team_model_aliases=None,
)
if not expect_to_work:
- pytest.fail(
- f"Expected model access check to fail for model={model}, team_models={team_models}"
- )
+ pytest.fail(f"Expected model access check to fail for model={model}, team_models={team_models}")
except Exception as e:
if expect_to_work:
pytest.fail(
@@ -1103,9 +1112,9 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
await asyncio.sleep(0.1) # Allow time for the alert task to complete
- assert (
- alert_triggered == expect_alert
- ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
+ assert alert_triggered == expect_alert, (
+ f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
+ )
@pytest.mark.parametrize(
@@ -1158,9 +1167,7 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
],
)
@pytest.mark.asyncio
-async def test_team_soft_budget_check(
- spend, soft_budget, expect_alert, metadata, expected_alert_emails
-):
+async def test_team_soft_budget_check(spend, soft_budget, expect_alert, metadata, expected_alert_emails):
"""
Test cases for _team_soft_budget_check:
1. Spend over soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured)
@@ -1208,9 +1215,9 @@ async def test_team_soft_budget_check(
await asyncio.sleep(0.1) # Allow time for the alert task to complete
- assert (
- alert_triggered == expect_alert
- ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
+ assert alert_triggered == expect_alert, (
+ f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
+ )
if expect_alert:
assert captured_call_info is not None
@@ -1222,10 +1229,7 @@ async def test_team_soft_budget_check(
if expected_alert_emails is not None:
assert captured_call_info.alert_emails == expected_alert_emails
else:
- assert (
- captured_call_info.alert_emails is None
- or captured_call_info.alert_emails == []
- )
+ assert captured_call_info.alert_emails is None or captured_call_info.alert_emails == []
@pytest.mark.asyncio
@@ -1356,9 +1360,7 @@ async def test_get_fuzzy_user_object():
# Test 5: Only email provided (no SSO ID)
mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user)
- result = await _get_fuzzy_user_object(
- prisma_client=mock_prisma, user_email="test@example.com"
- )
+ result = await _get_fuzzy_user_object(prisma_client=mock_prisma, user_email="test@example.com")
assert result == test_user
mock_prisma.db.litellm_usertable.find_first.assert_called_with(
where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}},
@@ -1367,9 +1369,7 @@ async def test_get_fuzzy_user_object():
# Test 6: Only SSO ID provided (no email)
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=test_user)
- result = await _get_fuzzy_user_object(
- prisma_client=mock_prisma, sso_user_id="sso_123"
- )
+ result = await _get_fuzzy_user_object(prisma_client=mock_prisma, sso_user_id="sso_123")
assert result == test_user
mock_prisma.db.litellm_usertable.find_unique.assert_called_with(
where={"sso_user_id": "sso_123"}, include={"organization_memberships": True}
@@ -1416,7 +1416,7 @@ async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work)
llm_router=router,
)
else:
- with pytest.raises(Exception, match='is not available for this API key') as e:
+ with pytest.raises(Exception, match="is not available for this API key") as e:
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
@@ -1502,9 +1502,7 @@ async def test_delete_cache_access_object():
],
)
@pytest.mark.asyncio
-async def test_get_resources_from_access_groups(
- resource_field, access_group_data, expected
-):
+async def test_get_resources_from_access_groups(resource_field, access_group_data, expected):
"""Test _get_resources_from_access_groups returns correct resource list from access groups."""
from unittest.mock import AsyncMock, MagicMock, patch
@@ -2094,7 +2092,9 @@ def test_is_model_cost_zero_judges_an_alias_chain_by_the_deployment_its_entry_ro
model_group_alias={"chain-entry": "local-free", "local-free": "paid-gpt"},
)
expected: Final = {"chain-entry": True, "local-free": False, "paid-gpt": False}
- order: Final = ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry")
+ order: Final = (
+ ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry")
+ )
verdicts: Final = {name: is_model_cost_zero(model=name, llm_router=router) for name in order}
@@ -2337,10 +2337,7 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis
monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
client = _WedgedPublishRedisClient()
cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
- hashed_tokens: Final = tuple(
- hashlib.sha256(f"token-{index}".encode()).hexdigest()
- for index in range(600)
- )
+ hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600))
for hashed_token in hashed_tokens:
cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token})
seeded: Final = tuple(
@@ -2400,10 +2397,7 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_
monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
client = _WedgedPublishRedisClient()
- hashed_tokens: Final = tuple(
- hashlib.sha256(f"token-{index}".encode()).hexdigest()
- for index in range(600)
- )
+ hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600))
redis = _AuthRedisCache({token: {"token": token} for token in hashed_tokens})
cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer
@@ -2473,3 +2467,46 @@ async def test_bulk_eviction_broadcasts_while_the_usage_cache_redis_is_wedged(
assert usage_redis.finished
assert {json.loads(message)["cache_key"] for message in client.published} == set(hashed_tokens)
+
+
+@pytest.mark.parametrize(
+ "spend, soft_budget, expect_alert",
+ [
+ (100, 50, True),
+ (50, 50, True),
+ (25, 50, False),
+ (100, None, False),
+ ],
+)
+@pytest.mark.asyncio
+async def test_user_soft_budget_check(spend: float, soft_budget: float | None, expect_alert: bool) -> None:
+ from litellm.models.user import LiteLLM_UserTable
+
+ alerts: list[tuple[str, CallInfo]] = []
+
+ class MockProxyLogging:
+ async def budget_alerts(self, type: str, user_info: CallInfo) -> None:
+ alerts.append((type, user_info))
+
+ user_object: Final = LiteLLM_UserTable(
+ user_id="soft-budget-user",
+ user_email="soft@example.com",
+ max_budget=200,
+ soft_budget=soft_budget,
+ spend=spend,
+ )
+
+ proxy_logging_obj: Final = cast(ProxyLogging, MockProxyLogging()) # cast-ok: test double for budget_alerts
+ _user_soft_budget_check(user_object=user_object, user_spend=spend, proxy_logging_obj=proxy_logging_obj)
+ await asyncio.sleep(0.1)
+
+ assert (len(alerts) == 1) == expect_alert
+ if expect_alert:
+ alert_type, call_info = alerts[0]
+ assert alert_type == "soft_budget"
+ assert call_info.user_id == "soft-budget-user"
+ assert call_info.user_email == "soft@example.com"
+ assert call_info.soft_budget == soft_budget
+ assert call_info.max_budget is None
+ assert call_info.spend == spend
+ assert call_info.event_group == Litellm_EntityType.USER
diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
index d8aa0239e2d..62a35a2e0a5 100644
--- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
+++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
@@ -1672,9 +1672,7 @@ async def test_new_user_clears_recent_missing_user_lookup(mocker: MockerFixture)
)
mocker.patch(
"litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn",
- new=mocker.AsyncMock(
- return_value={"user_id": user_id, "token": "sk-sso-created-user", "expires": None}
- ),
+ new=mocker.AsyncMock(return_value={"user_id": user_id, "token": "sk-sso-created-user", "expires": None}),
)
try:
@@ -2301,7 +2299,11 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo
"user_alias": "Spruce",
}
- replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0}
+ replacement: Final = {
+ "model-spruce": {"budget_limit": 0, "time_period": "1d"},
+ "model-birch": 5.0,
+ "model-cedar": 0,
+ }
request: Final = UpdateUserRequest(
user_id="user-spruce",
model_max_budget=replacement,
@@ -2343,10 +2345,14 @@ async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: M
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
- mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
+ mocker.patch(
+ "litellm.proxy.proxy_server.prisma_client", prisma_client
+ ) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
- mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
+ mocker.patch(
+ "litellm.proxy.proxy_server.user_api_key_cache", cache
+ ) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@@ -2382,10 +2388,14 @@ async def test_user_status_update_refreshes_cached_user(
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
- mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
+ mocker.patch(
+ "litellm.proxy.proxy_server.prisma_client", prisma_client
+ ) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
- mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
+ mocker.patch(
+ "litellm.proxy.proxy_server.user_api_key_cache", cache
+ ) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@@ -2419,10 +2429,14 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
prisma_client: Final = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user])
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
- mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
+ mocker.patch(
+ "litellm.proxy.proxy_server.prisma_client", prisma_client
+ ) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
- mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
+ mocker.patch(
+ "litellm.proxy.proxy_server.user_api_key_cache", cache
+ ) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@@ -2454,7 +2468,9 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
@pytest.mark.asyncio
@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"])
-async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> None:
+async def test_user_max_budget_update_evicts_cached_user_on_every_worker(
+ mocker: MockerFixture, all_users: bool
+) -> None:
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper, bulk_user_update
@@ -2467,10 +2483,14 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker:
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
- mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
+ mocker.patch(
+ "litellm.proxy.proxy_server.prisma_client", prisma_client
+ ) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
- mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
+ mocker.patch(
+ "litellm.proxy.proxy_server.user_api_key_cache", cache
+ ) # test-quality-ok: exercise a real isolated cache
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@@ -2875,24 +2895,12 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker):
mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[])
# Mock all delete_many calls
- mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(
- return_value=[]
- )
- mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(
- return_value=0
- )
- mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(
- return_value=1
- )
- mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(
- return_value=0
- )
- mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(
- return_value=0
- )
- mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(
- return_value=1
- )
+ mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[])
+ mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0)
+ mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=1)
+ mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0)
+ mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
+ mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
_serve_locked_user_delete(mock_prisma_client, mock_user_row)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
@@ -2985,9 +2993,15 @@ async def test_delete_user_evicts_jwt_key_mapping_cache_of_its_keys(mocker):
mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
_serve_locked_user_delete(mock_prisma_client, user_row)
- mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: substitute the database dependency
- mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
- mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time
+ mocker.patch(
+ "litellm.proxy.proxy_server.prisma_client", mock_prisma_client
+ ) # test-quality-ok: substitute the database dependency
+ mocker.patch(
+ "litellm.proxy.proxy_server.user_api_key_cache", cache
+ ) # test-quality-ok: exercise a real isolated cache
+ mocker.patch(
+ "litellm.proxy.proxy_server.proxy_logging_obj", None
+ ) # test-quality-ok: delete_user reads it off proxy_server at call time
await delete_user(
data=DeleteUserRequest(user_ids=["jwt-user"]),
@@ -3129,6 +3143,7 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker: MockerFixture
user_role="internal_user",
spend=42.5,
max_budget=100.0,
+ soft_budget=80.0,
tpm_limit=100000,
rpm_limit=1000,
models=["gpt-4"],
@@ -3160,9 +3175,7 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker: MockerFixture
return None
mock_prisma_client: Final = _PrismaClient(
- db=_Database(
- litellm_usertable=_UserTable(find_unique=mocker.AsyncMock(side_effect=mock_find_unique))
- )
+ db=_Database(litellm_usertable=_UserTable(find_unique=mocker.AsyncMock(side_effect=mock_find_unique)))
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
@@ -3184,6 +3197,7 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker: MockerFixture
assert response.user_role == "internal_user"
assert response.spend == 42.5
assert response.max_budget == 100.0
+ assert response.soft_budget == 80.0
assert response.tpm_limit == 100000
assert response.rpm_limit == 1000
assert response.models == ["gpt-4"]
@@ -3552,6 +3566,7 @@ async def test_user_info_v2_response_shape(mocker):
"user_role",
"spend",
"max_budget",
+ "soft_budget",
"tpm_limit",
"rpm_limit",
"models",
@@ -4597,7 +4612,9 @@ async def test_user_update_broadcasts_the_entitlement_invalidation_to_other_work
_object_permission_mocks(mocker)
cache: Final = mocker.MagicMock()
cache.async_delete_cache = mocker.AsyncMock()
- mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: substitute the cache dependency
+ mocker.patch(
+ "litellm.proxy.proxy_server.user_api_key_cache", cache
+ ) # test-quality-ok: substitute the cache dependency
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@@ -5144,12 +5161,18 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non
prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
_serve_locked_user_delete(prisma_client, deleted)
- mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
+ mocker.patch(
+ "litellm.proxy.proxy_server.prisma_client", prisma_client
+ ) # test-quality-ok: substitute the database dependency
cache: Final = UserApiKeyCache()
for row in (deleted, survivor):
await cache.async_set_cache(key=row.user_id, value=row, model_type=LiteLLM_UserTable)
- mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
- mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time
+ mocker.patch(
+ "litellm.proxy.proxy_server.user_api_key_cache", cache
+ ) # test-quality-ok: exercise a real isolated cache
+ mocker.patch(
+ "litellm.proxy.proxy_server.proxy_logging_obj", None
+ ) # test-quality-ok: delete_user reads it off proxy_server at call time
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new_callable=mocker.AsyncMock,
@@ -5255,7 +5278,9 @@ async def test_ui_view_users_answers_a_db_outage_as_503_no_db_connection_not_as_
assert raised.value.code == "503"
assert raised.value.type == ProxyErrorTypes.no_db_connection
assert isinstance(raised.value.__cause__, httpx.ConnectError)
- outage_logs: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()]
+ outage_logs: Final = [
+ r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()
+ ]
assert outage_logs == ["Database unavailable during user search: ConnectError"]
@@ -5413,7 +5438,9 @@ async def test_user_update_without_password_revokes_nothing(_admin_prisma, mocke
@pytest.mark.asyncio
-async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(_admin_prisma: MagicMock) -> None:
+async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(
+ _admin_prisma: MagicMock,
+) -> None:
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
with pytest.raises(ProxyException) as exc_info:
@@ -5425,3 +5452,93 @@ async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_
assert (exc_info.value.code, exc_info.value.type) == ("400", "invalid_request_error")
assert "Invalid budget_duration" in exc_info.value.message
assert "Authentication Error" not in exc_info.value.message
+
+
+def test_update_internal_user_params_sets_and_clears_soft_budget() -> None:
+ from litellm.proxy._types import UpdateUserRequest
+ from litellm.proxy.management_endpoints.internal_user_endpoints import (
+ _update_internal_user_params,
+ )
+
+ set_request: Final = UpdateUserRequest(user_id="u1", soft_budget=25.5)
+ set_values: dict[str, object] = _update_internal_user_params(
+ data_json=set_request.model_dump(exclude_unset=True), data=set_request
+ )
+ assert set_values["soft_budget"] == 25.5
+
+ clear_request: Final = UpdateUserRequest(user_id="u1", soft_budget=None)
+ clear_values: dict[str, object] = _update_internal_user_params(
+ data_json=clear_request.model_dump(exclude_unset=True), data=clear_request
+ )
+ assert "soft_budget" in clear_values and clear_values["soft_budget"] is None
+
+ untouched_request: Final = UpdateUserRequest(user_id="u1", max_budget=10)
+ untouched_values: dict[str, object] = _update_internal_user_params(
+ data_json=untouched_request.model_dump(exclude_unset=True), data=untouched_request
+ )
+ assert "soft_budget" not in untouched_values
+
+
+@pytest.mark.parametrize("bad_soft_budget", [-1, float("inf"), float("nan")])
+def test_update_internal_user_params_rejects_invalid_soft_budget(bad_soft_budget: float) -> None:
+ from fastapi import HTTPException
+
+ from litellm.proxy._types import UpdateUserRequest
+ from litellm.proxy.management_endpoints.internal_user_endpoints import (
+ _update_internal_user_params,
+ )
+
+ request: Final = UpdateUserRequest(user_id="u1", soft_budget=bad_soft_budget)
+ with pytest.raises(HTTPException) as exc_info:
+ _update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request)
+ assert exc_info.value.status_code == 400
+
+
+def test_update_internal_new_user_params_keeps_soft_budget_below_max_budget() -> None:
+ from fastapi import HTTPException
+
+ from litellm.proxy._types import NewUserRequest
+ from litellm.proxy.management_endpoints.internal_user_endpoints import (
+ _update_internal_new_user_params,
+ )
+
+ ok: Final = NewUserRequest(user_id="u1", max_budget=100, soft_budget=50)
+ assert _update_internal_new_user_params(ok.model_dump(), ok)["soft_budget"] == 50
+
+ soft_only: Final = NewUserRequest(user_id="u1", soft_budget=50)
+ assert _update_internal_new_user_params(soft_only.model_dump(), soft_only)["soft_budget"] == 50
+
+ too_high: Final = NewUserRequest(user_id="u1", max_budget=100, soft_budget=100)
+ with pytest.raises(HTTPException) as exc_info:
+ _update_internal_new_user_params(too_high.model_dump(), too_high)
+ assert exc_info.value.status_code == 400
+
+
+def test_update_internal_new_user_params_checks_soft_budget_against_defaulted_max_budget(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ import litellm
+ from litellm.proxy.management_endpoints.internal_user_endpoints import _update_internal_new_user_params
+
+ monkeypatch.setattr(litellm, "default_internal_user_params", {"max_budget": 10})
+ request: Final = NewUserRequest(user_id="u1", soft_budget=50)
+
+ with pytest.raises(HTTPException) as exc_info:
+ _update_internal_new_user_params(request.model_dump(), request)
+ assert exc_info.value.status_code == 400
+
+
+def test_validate_user_soft_budget_uses_existing_max_budget_on_update() -> None:
+ from fastapi import HTTPException
+
+ from litellm.proxy.management_endpoints.internal_user_endpoints import (
+ _validate_user_soft_budget,
+ )
+
+ _validate_user_soft_budget(soft_budget=None, max_budget=10)
+ _validate_user_soft_budget(soft_budget=5, max_budget=None)
+ _validate_user_soft_budget(soft_budget=5, max_budget=10)
+ with pytest.raises(HTTPException):
+ _validate_user_soft_budget(soft_budget=10, max_budget=10)
+ with pytest.raises(HTTPException):
+ _validate_user_soft_budget(soft_budget=20, max_budget=10)
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.test.tsx
index f797d6e4f36..aaff3502db8 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.test.tsx
@@ -14,6 +14,7 @@ vi.mock("./user_edit_view", () => ({
UserEditView: ({ onSubmit, onCancel }: { onSubmit: (values: any) => void; onCancel: () => void }) => (
{t("users.userInfo.details.maxBudget")}
-- {userData.max_budget !== null && userData.max_budget !== undefined - ? `$${formatNumberWithCommas(userData.max_budget, 4)}` - : unlimitedLabel} -
+{formatBudgetValue(userData.max_budget, unlimitedLabel)}
+{t("users.userInfo.details.softBudget")}
+{formatBudgetValue(userData.soft_budget, t("users.userInfo.details.notSet"))}