diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index 7a77898b16..3616325383 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -47,6 +47,15 @@ from litellm.secret_managers.main import get_secret_bool from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL +def _parse_email_list(raw) -> List[str]: + """Parse emails from a list or comma-separated string.""" + if isinstance(raw, list): + return [e.strip() for e in raw if isinstance(e, str) and e.strip()] + elif isinstance(raw, str): + return [e.strip() for e in raw.split(",") if e.strip()] + return [] + + class BaseEmailLogger(CustomLogger): DEFAULT_LITELLM_EMAIL = "notifications@alerts.litellm.ai" DEFAULT_SUPPORT_EMAIL = "support@berri.ai" @@ -312,17 +321,22 @@ class BaseEmailLogger(CustomLogger): ) pass - async def send_max_budget_alert_email(self, event: WebhookEvent): + async def send_max_budget_alert_email( + self, + event: WebhookEvent, + threshold_pct: Optional[int] = None, + recipient_emails: Optional[List[str]] = None, + ): """ - Send email to user when max budget alert threshold is reached - """ - email_params = await self._get_email_params( - email_event=EmailEvent.max_budget_alert, - user_id=event.user_id, - user_email=event.user_email, - event_message=event.event_message, - ) + Send email to user when max budget alert threshold is reached. + Args: + event: The webhook event with spend/budget info + threshold_pct: Override percentage for multi-threshold alerts (e.g. 50, 75, 100). + When None, uses EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE (old behavior). + recipient_emails: Override recipient list for multi-threshold alerts. + When None, resolves single owner email via _get_email_params (old behavior). + """ verbose_proxy_logger.debug( f"send_max_budget_alert_email_event: {json.dumps(event.model_dump(exclude_none=True), indent=4, default=str)}" ) @@ -334,30 +348,65 @@ class BaseEmailLogger(CustomLogger): ) # Calculate percentage and alert threshold - percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100) + percentage = threshold_pct if threshold_pct is not None else int( + EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100 + ) + threshold_fraction = percentage / 100.0 alert_threshold_str = ( - f"${event.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE:.2f}" + f"${event.max_budget * threshold_fraction:.2f}" if event.max_budget is not None else "N/A" ) - email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format( - email_logo_url=email_params.logo_url, - recipient_email=email_params.recipient_email, - percentage=percentage, - spend=spend_str, - max_budget=max_budget_str, - alert_threshold=alert_threshold_str, - base_url=email_params.base_url, - email_support_contact=email_params.support_contact, - ) - await self.send_email( - from_email=self.DEFAULT_LITELLM_EMAIL, - to_email=[email_params.recipient_email], - subject=email_params.subject, - html_body=email_html_content, - ) - pass + if recipient_emails is not None: + # Multi-threshold path: batch send with generic key-based greeting + email_params = await self._get_email_params( + email_event=EmailEvent.max_budget_alert, + user_id=event.user_id, + user_email=event.user_email or recipient_emails[0], + event_message=event.event_message, + ) + greeting = event.user_email or event.key_alias or event.token or "" + email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format( + email_logo_url=email_params.logo_url, + recipient_email=greeting, + percentage=percentage, + spend=spend_str, + max_budget=max_budget_str, + alert_threshold=alert_threshold_str, + base_url=email_params.base_url, + email_support_contact=email_params.support_contact, + ) + await self.send_email( + from_email=self.DEFAULT_LITELLM_EMAIL, + to_email=recipient_emails, + subject=email_params.subject, + html_body=email_html_content, + ) + else: + # Old path: single recipient resolved from user_id/user_email + email_params = await self._get_email_params( + email_event=EmailEvent.max_budget_alert, + user_id=event.user_id, + user_email=event.user_email, + event_message=event.event_message, + ) + email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format( + email_logo_url=email_params.logo_url, + recipient_email=email_params.recipient_email, + percentage=percentage, + spend=spend_str, + max_budget=max_budget_str, + alert_threshold=alert_threshold_str, + base_url=email_params.base_url, + email_support_contact=email_params.support_contact, + ) + await self.send_email( + from_email=self.DEFAULT_LITELLM_EMAIL, + to_email=[email_params.recipient_email], + subject=email_params.subject, + html_body=email_html_content, + ) async def budget_alerts( self, @@ -469,6 +518,13 @@ class BaseEmailLogger(CustomLogger): # For max_budget_alert, check if we've already sent an alert if type == "max_budget_alert": if user_info.max_budget is not None and user_info.spend is not None: + if user_info.max_budget_alert_emails: + # New path: multi-threshold alerts + await self._handle_multi_threshold_max_budget_alert( + user_info=user_info, _cache=_cache + ) + return + alert_threshold = ( user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE ) @@ -527,6 +583,77 @@ class BaseEmailLogger(CustomLogger): ) return + async def _handle_multi_threshold_max_budget_alert( + self, + user_info: CallInfo, + _cache: DualCache, + ): + """ + Loop over configured thresholds in max_budget_alert_emails, + check cache per threshold, and send to configured recipients. + """ + for threshold_str, raw_emails in user_info.max_budget_alert_emails.items(): + try: + threshold_pct = int(threshold_str) + except (ValueError, TypeError): + continue + + threshold_amount = user_info.max_budget * (threshold_pct / 100.0) + if user_info.spend < threshold_amount: + continue + + _id = user_info.token or user_info.user_id or "default_id" + _cache_key = ( + f"email_budget_alerts:max_budget_alert:{threshold_pct}:{_id}" + ) + + result = await _cache.async_get_cache(key=_cache_key) + if result is not None: + continue + + # Parse emails + auto-include owner + emails = _parse_email_list(raw_emails) + if user_info.user_email: + emails.append(user_info.user_email) + recipient_emails = list(set(emails)) if emails else None + + event_message = f"Max Budget Alert - {threshold_pct}% of Maximum Budget Reached" + webhook_event = WebhookEvent( + event="max_budget_alert", + event_message=event_message, + spend=user_info.spend, + max_budget=user_info.max_budget, + soft_budget=user_info.soft_budget, + token=user_info.token, + customer_id=user_info.customer_id, + user_id=user_info.user_id, + team_id=user_info.team_id, + team_alias=user_info.team_alias, + organization_id=user_info.organization_id, + user_email=user_info.user_email, + key_alias=user_info.key_alias, + projected_exceeded_date=user_info.projected_exceeded_date, + projected_spend=user_info.projected_spend, + event_group=user_info.event_group, + ) + + try: + await self.send_max_budget_alert_email( + webhook_event, + threshold_pct=threshold_pct, + recipient_emails=recipient_emails, + ) + await _cache.async_set_cache( + key=_cache_key, + value="SENT", + ttl=EMAIL_BUDGET_ALERT_TTL, + ) + except Exception as e: + verbose_proxy_logger.error( + f"Error sending multi-threshold max budget alert email for {threshold_pct}%: {e}", + exc_info=True, + ) + async def _get_email_params( self, email_event: EmailEvent, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0bbee56d5e..d3c8324870 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3095,6 +3095,10 @@ class CallInfo(LiteLLMPydanticObjectBase): default=None, description="Additional email addresses to send alerts to (e.g., from team metadata)", ) + max_budget_alert_emails: Optional[Dict[str, List[str]]] = Field( + default=None, + description="Map of threshold percentage to email recipients (e.g., {'50': ['a@co.com'], '75': ['a@co.com', 'b@co.com']})", + ) class WebhookEvent(CallInfo): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 56958a88f6..fbc24ab17c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2963,22 +2963,13 @@ async def _virtual_key_max_budget_alert_check( and valid_token.spend is not None and valid_token.spend > 0 ): - alert_threshold = ( - valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + owner_email = user_obj.user_email if user_obj else None + alert_email_config = (valid_token.metadata or {}).get( + "max_budget_alert_emails" ) - # Only alert if we've crossed the threshold but haven't exceeded max_budget yet - if ( - valid_token.spend >= alert_threshold - and valid_token.spend < valid_token.max_budget - ): - verbose_proxy_logger.debug( - "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", - valid_token.token, - valid_token.spend, - valid_token.max_budget, - alert_threshold, - ) + if isinstance(alert_email_config, dict) and alert_email_config: + # New path: pass the map through, let the email handler decide what to fire call_info = CallInfo( token=valid_token.token, spend=valid_token.spend, @@ -2988,17 +2979,55 @@ async def _virtual_key_max_budget_alert_check( team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=valid_token.org_id, - user_email=user_obj.user_email if user_obj else None, + user_email=owner_email, key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, + max_budget_alert_emails=alert_email_config, ) - asyncio.create_task( proxy_logging_obj.budget_alerts( type="max_budget_alert", user_info=call_info, ) ) + else: + # Old path: existing single 80% threshold — completely unchanged + alert_threshold = ( + valid_token.max_budget + * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + ) + + if ( + valid_token.spend >= alert_threshold + and valid_token.spend < valid_token.max_budget + ): + verbose_proxy_logger.debug( + "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", + valid_token.token, + valid_token.spend, + valid_token.max_budget, + alert_threshold, + ) + call_info = CallInfo( + token=valid_token.token, + spend=valid_token.spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=owner_email, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="max_budget_alert", + user_info=call_info, + ) + ) async def _check_team_member_budget( diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py index 744195dfb6..af5e234140 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -878,4 +878,200 @@ async def test_budget_alerts_max_budget_alert_crossed( cache_call_args = mock_cache.async_set_cache.call_args[1] assert cache_call_args["key"] == "email_budget_alerts:max_budget_alert:test_user" assert cache_call_args["value"] == "SENT" - assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL \ No newline at end of file + assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL + + +@pytest.mark.asyncio +async def test_multi_threshold_sends_crossed_thresholds( + base_email_logger, mock_send_email +): + """Test that multi-threshold path sends emails for all crossed thresholds""" + user_info = CallInfo( + token="hashed_key_1", + user_id="test_user", + user_email="owner@co.com", + spend=80.0, + max_budget=100.0, + event_group=Litellm_EntityType.KEY, + max_budget_alert_emails={ + "50": ["finance@co.com"], + "75": ["finance@co.com", "bu_lead@co.com"], + "100": ["cto@co.com"], + }, + ) + + mock_cache = mock.AsyncMock() + mock_cache.async_get_cache = mock.AsyncMock(return_value=None) + mock_cache.async_set_cache = mock.AsyncMock() + base_email_logger.internal_usage_cache = mock_cache + + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) + + # spend=80 crosses 50% ($50) and 75% ($75), but not 100% ($100) + assert mock_send_email.call_count == 2 + + # Check cache keys include threshold percentage + cache_keys = [ + c[1]["key"] for c in mock_cache.async_set_cache.call_args_list + ] + assert "email_budget_alerts:max_budget_alert:50:hashed_key_1" in cache_keys + assert "email_budget_alerts:max_budget_alert:75:hashed_key_1" in cache_keys + + +@pytest.mark.asyncio +async def test_multi_threshold_dedup_cache_prevents_resend( + base_email_logger, mock_send_email +): + """Test that cached thresholds are not re-sent""" + user_info = CallInfo( + token="hashed_key_1", + user_id="test_user", + user_email="owner@co.com", + spend=80.0, + max_budget=100.0, + event_group=Litellm_EntityType.KEY, + max_budget_alert_emails={ + "50": ["finance@co.com"], + "75": ["finance@co.com"], + }, + ) + + # Simulate 50% already sent (cached), 75% not yet sent + async def cache_get(key): + if "50:" in key: + return "SENT" + return None + + mock_cache = mock.AsyncMock() + mock_cache.async_get_cache = mock.AsyncMock(side_effect=cache_get) + mock_cache.async_set_cache = mock.AsyncMock() + base_email_logger.internal_usage_cache = mock_cache + + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) + + # Only 75% should fire + assert mock_send_email.call_count == 1 + cache_key = mock_cache.async_set_cache.call_args[1]["key"] + assert "75:" in cache_key + + +@pytest.mark.asyncio +async def test_multi_threshold_owner_email_auto_included( + base_email_logger, mock_send_email +): + """Test that the owner email is auto-appended and deduplicated""" + user_info = CallInfo( + token="hashed_key_1", + user_id="test_user", + user_email="owner@co.com", + spend=60.0, + max_budget=100.0, + event_group=Litellm_EntityType.KEY, + max_budget_alert_emails={ + "50": ["finance@co.com", "owner@co.com"], # owner already in list + }, + ) + + mock_cache = mock.AsyncMock() + mock_cache.async_get_cache = mock.AsyncMock(return_value=None) + mock_cache.async_set_cache = mock.AsyncMock() + base_email_logger.internal_usage_cache = mock_cache + + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) + + mock_send_email.assert_called_once() + to_emails = mock_send_email.call_args[1]["to_email"] + # owner@co.com should appear exactly once (deduplicated) + assert sorted(to_emails) == ["finance@co.com", "owner@co.com"] + + +@pytest.mark.asyncio +async def test_multi_threshold_malformed_keys_skipped( + base_email_logger, mock_send_email +): + """Test that non-numeric threshold keys are skipped""" + user_info = CallInfo( + token="hashed_key_1", + user_id="test_user", + user_email="owner@co.com", + spend=60.0, + max_budget=100.0, + event_group=Litellm_EntityType.KEY, + max_budget_alert_emails={ + "fifty": ["finance@co.com"], # invalid + "50": ["finance@co.com"], # valid, crossed + }, + ) + + mock_cache = mock.AsyncMock() + mock_cache.async_get_cache = mock.AsyncMock(return_value=None) + mock_cache.async_set_cache = mock.AsyncMock() + base_email_logger.internal_usage_cache = mock_cache + + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) + + # Only the valid "50" threshold should fire + assert mock_send_email.call_count == 1 + + +@pytest.mark.asyncio +async def test_multi_threshold_empty_emails_only_owner( + base_email_logger, mock_send_email +): + """Test that empty email list for a threshold sends only to owner""" + user_info = CallInfo( + token="hashed_key_1", + user_id="test_user", + user_email="owner@co.com", + spend=60.0, + max_budget=100.0, + event_group=Litellm_EntityType.KEY, + max_budget_alert_emails={ + "50": [], # empty list + }, + ) + + mock_cache = mock.AsyncMock() + mock_cache.async_get_cache = mock.AsyncMock(return_value=None) + mock_cache.async_set_cache = mock.AsyncMock() + base_email_logger.internal_usage_cache = mock_cache + + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) + + mock_send_email.assert_called_once() + to_emails = mock_send_email.call_args[1]["to_email"] + assert to_emails == ["owner@co.com"] + + +@pytest.mark.asyncio +async def test_no_map_preserves_old_single_threshold( + base_email_logger, mock_send_email +): + """Test that without max_budget_alert_emails, the old 80% single-threshold path works""" + user_info = CallInfo( + user_id="test_user", + user_email="test@example.com", + spend=165.0, + max_budget=200.0, + event_group=Litellm_EntityType.USER, + ) + + mock_cache = mock.AsyncMock() + mock_cache.async_get_cache = mock.AsyncMock(return_value=None) + mock_cache.async_set_cache = mock.AsyncMock() + base_email_logger.internal_usage_cache = mock_cache + + with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}): + await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info) + + mock_send_email.assert_called_once() + call_args = mock_send_email.call_args[1] + assert call_args["to_email"] == ["test@example.com"] + # Old path cache key has no threshold percentage + cache_key = mock_cache.async_set_cache.call_args[1]["key"] + assert cache_key == "email_budget_alerts:max_budget_alert:test_user" \ No newline at end of file diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index bd659ed518..d9ab8e06d0 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1544,6 +1544,114 @@ async def test_virtual_key_max_budget_alert_check_scenarios( ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}" +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_with_multi_threshold_map(): + """Test that max_budget_alert_emails map from metadata is attached to CallInfo on the new path""" + alert_triggered = False + captured_call_info = None + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal alert_triggered, captured_call_info + alert_triggered = True + captured_call_info = user_info + + alert_config = { + "50": ["finance@co.com"], + "75": ["finance@co.com", "bu_lead@co.com"], + } + valid_token = UserAPIKeyAuth( + token="test-token", + spend=30.0, + max_budget=100.0, + user_id="test-user", + key_alias="test-key", + metadata={"max_budget_alert_emails": alert_config}, + ) + user_obj = LiteLLM_UserTable( + user_id="test-user", + user_email="owner@co.com", + max_budget=None, + ) + + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=MockProxyLogging(), + user_obj=user_obj, + ) + await asyncio.sleep(0.1) + + assert alert_triggered is True + assert captured_call_info is not None + assert captured_call_info.max_budget_alert_emails == alert_config + assert captured_call_info.user_email == "owner@co.com" + assert captured_call_info.event_group == Litellm_EntityType.KEY + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_old_path_no_map(): + """Test that old single-threshold path is used when no max_budget_alert_emails in metadata""" + alert_triggered = False + captured_call_info = None + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal alert_triggered, captured_call_info + alert_triggered = True + captured_call_info = user_info + + # spend=90 is above 80% of 100 → old path should fire + valid_token = UserAPIKeyAuth( + token="test-token", + spend=90.0, + max_budget=100.0, + user_id="test-user", + key_alias="test-key", + metadata={}, + ) + + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=MockProxyLogging(), + user_obj=None, + ) + await asyncio.sleep(0.1) + + assert alert_triggered is True + assert captured_call_info is not None + assert captured_call_info.max_budget_alert_emails is None + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_old_path_below_threshold_no_alert(): + """Test that old path does NOT fire when spend is below 80% and no map is set""" + alert_triggered = False + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal alert_triggered + alert_triggered = True + + # spend=50 is below 80% of 100 → should NOT fire + valid_token = UserAPIKeyAuth( + token="test-token", + spend=50.0, + max_budget=100.0, + user_id="test-user", + key_alias="test-key", + metadata={}, + ) + + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=MockProxyLogging(), + user_obj=None, + ) + await asyncio.sleep(0.1) + + assert alert_triggered is False + + @pytest.mark.asyncio async def test_get_fuzzy_user_object_case_insensitive_email(): """Test that _get_fuzzy_user_object uses case-insensitive email lookup"""