This commit is contained in:
jesus-berri 2026-11-05 15:38:45 +00:00 • committed by GitHub
commit ada579fc40
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 599 additions and 187 deletions

View file

@ -0,0 +1,3 @@
SET LOCAL lock_timeout = '10s';
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "soft_budget" DOUBLE PRECISION;

View file

@ -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[]

View file

@ -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"

View file

@ -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"""
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
<p> Hi {webhook_event.user_id}, <br/>
Your LLM API usage this month has reached your account's <b> monthly budget of ${webhook_event.max_budget} </b> <br /> <br />
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. <br /> <br />
If you have any questions, please send an email to {email_support_contact} <br /> <br />
Best, <br />
The LiteLLM team <br />
"""
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"""
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
<p> Hi {webhook_event.user_id}, <br/>
Your LLM API spend of <b> ${webhook_event.spend} </b> has reached your <b> soft budget of ${webhook_event.soft_budget} </b>. <br /> <br />
This is a warning only. API requests keep going through{hard_limit}. <br /> <br />
If you have any questions, please send an email to {email_support_contact} <br /> <br />
Best, <br />
The LiteLLM team <br />
"""
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"""
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
<p> Hi {user_name}, <br/>
Your LLM API usage this month has reached your account's <b> monthly budget of ${max_budget} </b> <br /> <br />
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. <br /> <br />
If you have any questions, please send an email to {email_support_contact} <br /> <br />
Best, <br />
The LiteLLM team <br />
"""
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,

View file

@ -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 = []

View file

@ -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] = []

View file

@ -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,

View file

@ -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]] = []

View file

@ -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,

View file

@ -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[]

View file

@ -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):

View file

@ -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[]

View file

@ -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"

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -14,6 +14,7 @@ vi.mock("./user_edit_view", () => ({
UserEditView: ({ onSubmit, onCancel }: { onSubmit: (values: any) => void; onCancel: () => void }) => (
<div data-testid="user-edit-view">
<button onClick={() => onSubmit({ user_role: "admin", max_budget: 100 })}>Submit</button>
<button onClick={() => onSubmit({ max_budget: 100, soft_budget: 20 })}>Submit with soft budget</button>
<button onClick={() => onSubmit({ max_budget: 100, budget_duration: null })}>Submit with unset duration</button>
<button onClick={onCancel}>Cancel</button>
</div>
@ -205,6 +206,20 @@ describe("BulkEditUserModal", () => {
});
});
it("forwards a bulk soft budget to the update payload", async () => {
const user = userEvent.setup();
renderWithProviders(<BulkEditUserModal {...defaultProps} />);
await user.click(screen.getByRole("button", { name: "Submit with soft budget" }));
await waitFor(() => {
expect(mockUserBulkUpdateUserCall).toHaveBeenCalledWith("test-token", { max_budget: 100, soft_budget: 20 }, [
"user1",
"user2",
]);
});
});
it("omits an unset reset period from the bulk update payload", async () => {
const user = userEvent.setup();
renderWithProviders(<BulkEditUserModal {...defaultProps} />);

View file

@ -78,6 +78,7 @@ const BulkEditUserModal: React.FC<BulkEditUserModalProps> = ({
teams: [],
models: [],
max_budget: null,
soft_budget: null,
spend: 0,
metadata: {},
created_at: null,
@ -110,6 +111,10 @@ const BulkEditUserModal: React.FC<BulkEditUserModalProps> = ({
updatePayload.max_budget = formValues.max_budget;
}
if (formValues.soft_budget !== null && formValues.soft_budget !== undefined) {
updatePayload.soft_budget = formValues.soft_budget;
}
if (formValues.models && formValues.models.length > 0) {
updatePayload.models = formValues.models;
}

View file

@ -322,6 +322,34 @@ describe("UserEditView", () => {
expect(callArgs.metadata).toEqual(MOCK_USER_DATA.user_info.metadata);
});
it("sends soft_budget only when it changed", async () => {
const onSubmitMock = vi.fn();
renderWithProviders(
<UserEditView
{...defaultProps}
userData={{ ...MOCK_USER_DATA, user_info: { ...MOCK_USER_DATA.user_info, soft_budget: 40 } }}
onSubmit={onSubmitMock}
/>,
);
const softBudgetInput = await screen.findByLabelText(/soft budget/i);
expect(softBudgetInput).toHaveValue(40);
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(onSubmitMock).toHaveBeenCalled());
expect(onSubmitMock.mock.calls[0][0]).not.toHaveProperty("soft_budget");
fireEvent.change(softBudgetInput, { target: { value: "55.5" } });
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(onSubmitMock).toHaveBeenCalledTimes(2));
expect(onSubmitMock.mock.calls[1][0].soft_budget).toBe(55.5);
await userEvent.clear(softBudgetInput);
await userEvent.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(onSubmitMock).toHaveBeenCalledTimes(3));
expect(onSubmitMock.mock.calls[2][0].soft_budget).toBeNull();
});
it("uses No change for an unset reset period in bulk edit", async () => {
const onSubmit = vi.fn();
const userDataWithUnsetResetPeriod = {

View file

@ -64,6 +64,7 @@ const userEditShape = (t: TFunction) => ({
user_alias: z.string().nullish(),
user_role: z.string().nullish(),
models: z.array(z.string()),
soft_budget: z.union([z.string(), z.number()]).nullish(),
budget_duration: z.string().nullish(),
metadata: z.string().nullish(),
tpm_limit: rateLimitField(t),
@ -117,6 +118,7 @@ const toFormValues = (
user_role: userData.user_info?.user_role,
models: userData.user_info?.models || [],
max_budget: isUnlimited ? "" : maxBudget,
soft_budget: userData.user_info?.soft_budget ?? "",
budget_duration: userData.user_info?.budget_duration,
metadata: userData.user_info?.metadata ? JSON.stringify(userData.user_info.metadata, null, 2) : undefined,
...(canEditMcpPermissions ? buildMcpFieldValues(objectPermission) : {}),
@ -199,8 +201,15 @@ export function UserEditView({
const modelBudgets = modelMaxBudgetUpdate(modelMaxBudget, userData.user_info?.model_max_budget);
const tpmLimit = changedLimit(tpmLimitInput, userData.user_info?.tpm_limit);
const rpmLimit = changedLimit(rpmLimitInput, userData.user_info?.rpm_limit);
const { soft_budget: softBudgetValue, ...formValuesWithoutSoftBudget } = formValues;
const softBudgetInput =
softBudgetValue === "" || softBudgetValue === undefined || softBudgetValue === null
? null
: Number(softBudgetValue);
const softBudget = changedLimit(softBudgetInput, userData.user_info?.soft_budget);
const payload = {
...formValues,
...formValuesWithoutSoftBudget,
...(softBudget !== undefined && { soft_budget: softBudget }),
...("metadata" in values ? { metadata: metadata.value } : {}),
...(modelBudgets !== undefined && { model_max_budget: modelBudgets }),
...(tpmLimit !== undefined && { tpm_limit: tpmLimit }),
@ -319,6 +328,25 @@ export function UserEditView({
)}
</FormField>
<FormField
control={form.control}
name="soft_budget"
label={labelWithHint(t("users.editView.fields.softBudget"), t("users.editView.fields.softBudgetHint"))}
>
{({ ref, value, onChange, ...control }) => (
<Input
{...control}
ref={ref}
type="number"
step={0.01}
value={value ?? ""}
onChange={(event) => onChange(event.target.value)}
onWheel={(event) => event.currentTarget.blur()}
placeholder={t("users.editView.fields.softBudgetPlaceholder")}
/>
)}
</FormField>
<FormField control={form.control} name="budget_duration" label={t("users.editView.fields.resetBudget")}>
{({ id, value, onChange }) => (
<BudgetDurationDropdown

View file

@ -51,6 +51,9 @@ import { AutoRouterUsageView } from "@/app/(dashboard)/cost-optimization/_compon
import { useActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange";
import { BudgetPeriodTooltip } from "@/components/shared/BudgetPeriodTooltip";
const formatBudgetValue = (value: number | null | undefined, fallback: string, precision = 4): string =>
value !== null && value !== undefined ? `$${formatNumberWithCommas(value, precision)}` : fallback;
interface UserInfoViewProps {
userId: string;
onClose: () => void;
@ -345,6 +348,7 @@ export default function UserInfoView({
user_alias: formValues.user_alias ?? userData.user_alias,
models: formValues.models ?? userData.models,
max_budget: formValues.max_budget === undefined ? userData.max_budget : formValues.max_budget,
soft_budget: formValues.soft_budget === undefined ? userData.soft_budget : formValues.soft_budget,
tpm_limit: formValues.tpm_limit === undefined ? userData.tpm_limit : formValues.tpm_limit,
rpm_limit: formValues.rpm_limit === undefined ? userData.rpm_limit : formValues.rpm_limit,
budget_duration:
@ -400,8 +404,7 @@ export default function UserInfoView({
const unknownTimestamp = t("users.userInfo.details.unknown");
const unlimitedLabel = t("users.userInfo.unlimited");
const budgetLimit =
userData.max_budget !== null ? `$${formatNumberWithCommas(userData.max_budget, 2)}` : unlimitedLabel;
const budgetLimit = formatBudgetValue(userData.max_budget, unlimitedLabel, 2);
const teamPickerPlaceholder = t("users.userInfo.addTeamDialog.teamPlaceholder");
// Build a legacy-compatible shape for UserEditView
@ -413,6 +416,7 @@ export default function UserInfoView({
user_role: userData.user_role,
models: userData.models,
max_budget: userData.max_budget,
soft_budget: userData.soft_budget,
tpm_limit: userData.tpm_limit,
rpm_limit: userData.rpm_limit,
budget_duration: userData.budget_duration,
@ -698,11 +702,12 @@ export default function UserInfoView({
<div>
<p className="font-medium">{t("users.userInfo.details.maxBudget")}</p>
<p>
{userData.max_budget !== null && userData.max_budget !== undefined
? `$${formatNumberWithCommas(userData.max_budget, 4)}`
: unlimitedLabel}
</p>
<p>{formatBudgetValue(userData.max_budget, unlimitedLabel)}</p>
</div>
<div>
<p className="font-medium">{t("users.userInfo.details.softBudget")}</p>
<p>{formatBudgetValue(userData.soft_budget, t("users.userInfo.details.notSet"))}</p>
</div>
<div>

View file

@ -1155,6 +1155,7 @@ export interface UserInfoV2Response {
user_role: string | null;
spend: number;
max_budget: number | null;
soft_budget?: number | null;
tpm_limit?: number | null;
rpm_limit?: number | null;
models: string[];

View file

@ -18601,6 +18601,7 @@
"created": "Created",
"lastUpdated": "Last Updated",
"maxBudget": "Max Budget",
"softBudget": "Soft Budget",
"budgetReset": "Budget Reset",
"metadata": "Metadata",
"mcpPermissions": "MCP Permissions",
@ -18654,6 +18655,9 @@
"maxBudget": "Max Budget (USD)",
"unlimitedBudget": "Unlimited Budget",
"maxBudgetPlaceholder": "Enter a numerical value",
"softBudget": "Soft Budget (USD)",
"softBudgetHint": "Sends a budget alert when the user's spend reaches this amount. Requests are not blocked. Must be lower than the max budget when both are set.",
"softBudgetPlaceholder": "Optional",
"resetBudget": "Reset Budget",
"noChange": "No change",
"tpmLimit": "TPM Limit",

View file

@ -19522,6 +19522,7 @@ export interface paths {
* - 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)
@ -19661,6 +19662,7 @@ export interface paths {
* - 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)
@ -28935,6 +28937,8 @@ export interface components {
rpm_limit?: number | null;
/** Send Invite Email */
send_invite_email?: boolean | null;
/** Soft Budget */
soft_budget?: number | null;
/**
* Spend
* @default 0
@ -37307,6 +37311,8 @@ export interface components {
policies: string[];
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
soft_budget?: number | null;
/**
* Spend
* @default 0
@ -37404,6 +37410,8 @@ export interface components {
policies: string[];
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
soft_budget?: number | null;
/**
* Spend
* @default 0
@ -40713,6 +40721,8 @@ export interface components {
rpm_limit?: number | null;
/** Send Invite Email */
send_invite_email?: boolean | null;
/** Soft Budget */
soft_budget?: number | null;
/**
* Spend
* @default 0
@ -40887,6 +40897,8 @@ export interface components {
rpm_limit?: number | null;
/** Rpm Limit Type */
rpm_limit_type?: ("guaranteed_throughput" | "best_effort_throughput" | "dynamic") | null;
/** Soft Budget */
soft_budget?: number | null;
/**
* Spend
* @default 0
@ -50738,6 +50750,8 @@ export interface components {
prompts?: string[] | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
soft_budget?: number | null;
/** Spend */
spend?: number | null;
/** Tag Rpm Limit */
@ -50844,6 +50858,8 @@ export interface components {
prompts?: string[] | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
soft_budget?: number | null;
/** Spend */
spend?: number | null;
/** Tag Rpm Limit */
@ -51630,6 +51646,8 @@ export interface components {
object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
soft_budget?: number | null;
/**
* Spend
* @default 0