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 a29f0a1b43a..93276065cce 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -39,6 +39,7 @@ from litellm.integrations.email_templates.templates import ( from litellm.integrations.email_templates.user_invitation_email import ( USER_INVITATION_EMAIL_TEMPLATE, ) +from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_threshold from litellm.proxy._types import ( CallInfo, InvitationNew, @@ -454,6 +455,7 @@ class BaseEmailLogger(CustomLogger): "projected_limit_exceeded", ], user_info: CallInfo, + budget_alert_thresholds: tuple[int, ...] | None = None, ): """ Send a budget alert via email @@ -555,8 +557,17 @@ class BaseEmailLogger(CustomLogger): ) return - alert_threshold = ( - user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + configured_threshold: Final = ( + get_budget_alert_threshold(user_info.spend, user_info.max_budget, budget_alert_thresholds) + if budget_alert_thresholds is not None + else None + ) + if budget_alert_thresholds is not None and configured_threshold is None: + return + alert_threshold = user_info.max_budget * ( + configured_threshold / 100 + if configured_threshold is not None + else EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE ) # Only alert if we've crossed the threshold but haven't exceeded max_budget yet @@ -566,7 +577,12 @@ class BaseEmailLogger(CustomLogger): ): # Generate cache key based on event type and identifier _id = user_info.token or user_info.user_id or "default_id" - _cache_key = f"email_budget_alerts:max_budget_alert:{_id}" + cache_id: Final = f"{_id}:{configured_threshold}" if configured_threshold is not None else _id + _cache_key = f"email_budget_alerts:max_budget_alert:{cache_id}" + if configured_threshold is not None and await _cache.async_get_cache( + key=f"email_budget_alerts:max_budget_alert:{_id}" + ) is not None: + return send_count = await _cache.async_increment_cache( key=_cache_key, @@ -575,8 +591,10 @@ class BaseEmailLogger(CustomLogger): ) if send_count is None or send_count <= 1: # Calculate percentage - percentage = int( - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100 + percentage = ( + configured_threshold + if configured_threshold is not None + else int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100) ) # Create WebhookEvent for max budget alert @@ -601,7 +619,10 @@ class BaseEmailLogger(CustomLogger): ) try: - await self.send_max_budget_alert_email(webhook_event) + if configured_threshold is None: + await self.send_max_budget_alert_email(webhook_event) + else: + await self.send_max_budget_alert_email(webhook_event, threshold_pct=configured_threshold) except Exception as e: verbose_proxy_logger.error( f"Error sending max budget alert email: {e}", diff --git a/litellm/integrations/SlackAlerting/budget_alert_types.py b/litellm/integrations/SlackAlerting/budget_alert_types.py index 4fe833acecc..32645eadd68 100644 --- a/litellm/integrations/SlackAlerting/budget_alert_types.py +++ b/litellm/integrations/SlackAlerting/budget_alert_types.py @@ -1,9 +1,17 @@ +import math from abc import ABC, abstractmethod +from collections.abc import Sequence from typing import Final, Literal from litellm.proxy._types import CallInfo, Litellm_EntityType +def get_budget_alert_threshold(spend: float, max_budget: float | None, thresholds: Sequence[int]) -> int | None: + if max_budget is None or not math.isfinite(max_budget) or max_budget <= 0 or not math.isfinite(spend): + return None + return max((threshold for threshold in thresholds if spend >= max_budget * (threshold / 100)), default=None) + + class BaseBudgetAlertType(ABC): """Base class for different budget alert types""" diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 4c2b722ef90..a67aac59269 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -7,6 +7,7 @@ import random import time from collections.abc import Callable from datetime import timedelta +from fnmatch import fnmatchcase from typing import TYPE_CHECKING, Any, Final, Literal from openai import APIError @@ -25,7 +26,7 @@ from litellm.constants import ( SLACK_MODEL_DEPRECATION_LOCK_ID, ) from litellm.integrations.custom_batch_logger import CustomBatchLogger -from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type +from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_threshold, get_budget_alert_type from litellm.integrations.SlackAlerting.hanging_request_check import ( AlertingHangingRequestCheck, ) @@ -118,7 +119,7 @@ class SlackAlerting(CustomBatchLogger): ) self.alert_to_webhook_url = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url) self.is_running = False - self.alerting_args = SlackAlertingArgs(**alerting_args) + self.alerting_args = SlackAlertingArgs.model_validate(alerting_args) self.default_webhook_url = default_webhook_url self.flush_lock = asyncio.Lock() self.periodic_started = False @@ -158,7 +159,7 @@ class SlackAlerting(CustomBatchLogger): if alert_types is not None: self.alert_types = alert_types if alerting_args is not None: - self.alerting_args = SlackAlertingArgs(**alerting_args) + self.alerting_args = SlackAlertingArgs.model_validate(alerting_args) self._ensure_periodic_flush_task() if alert_type_config is not None: for key, val in alert_type_config.items(): @@ -539,6 +540,7 @@ class SlackAlerting(CustomBatchLogger): "project_budget", ], user_info: CallInfo, + send_threshold_email: bool = True, ): """ Send a budget alert on slack or webhook @@ -589,9 +591,26 @@ class SlackAlerting(CustomBatchLogger): # send alert if event is not None and user_info.event_group is not None: - _cache_key: Final = f"budget_alerts:{event}:{_id}" + threshold: Final = ( + get_budget_alert_threshold( + user_info.spend, user_info.max_budget, self.alerting_args.budget_alert_thresholds + ) + if event == "threshold_crossed" and self.alerting_args.budget_alert_thresholds is not None + else None + ) + cache_id: Final = f"{_id}:{threshold}" if threshold is not None else _id + _cache_key: Final = f"budget_alerts:{event}:{cache_id}" + if threshold is not None and await _cache.async_get_cache(key=f"budget_alerts:{event}:{_id}") == "SENT": + return result: Final = await _cache.async_get_cache(key=_cache_key) - if result is None: + slack_cache_key: Final = f"budget_alerts:slack:{event}:{cache_id}" + slack_due: Final[bool] = ( + "slack" in self.alerting + and self._slack_budget_alert_allowed(user_info) + and result != "SENT" + and await _cache.async_get_cache(key=slack_cache_key) is None + ) + if result is None or slack_due: webhook_event = WebhookEvent( event=event, event_message=event_message, @@ -612,18 +631,28 @@ class SlackAlerting(CustomBatchLogger): alert_emails=user_info.alert_emails, max_budget_alert_emails=user_info.max_budget_alert_emails, ) - await self.send_alert( + slack_accepted: Final = await self.send_alert( message=event_message + "\n\n" + user_info_str, level="High", alert_type=AlertType.budget_alerts, user_info=webhook_event, alerting_metadata={}, + budget_alert_destination=("slack" if result is not None else "all") if slack_due else "non_slack", + budget_alert_email=send_threshold_email or event != "threshold_crossed", ) - await _cache.async_set_cache( - key=_cache_key, - value="SENT", - ttl=self.alerting_args.budget_alert_ttl, - ) + if slack_accepted: + await _cache.async_set_cache( + key=slack_cache_key, + value="SENT", + ttl=self.alerting_args.budget_alert_ttl, + ) + if result is None: + # Legacy SENT includes Slack; new markers use the independent Slack window. + await _cache.async_set_cache( + key=_cache_key, + value="SENT_WITH_SLACK_DEDUP", + ttl=self.alerting_args.budget_alert_ttl, + ) return return @@ -665,6 +694,14 @@ class SlackAlerting(CustomBatchLogger): if user_info.spend >= user_info.max_budget: event = "budget_crossed" event_message += f"Budget Crossed\n Total Budget:`{user_info.max_budget}`" + elif self.alerting_args.budget_alert_thresholds is not None: + if event in ("soft_budget_crossed", "projected_limit_exceeded"): + return event, event_message + threshold: Final = get_budget_alert_threshold( + user_info.spend, user_info.max_budget, self.alerting_args.budget_alert_thresholds + ) + if threshold is not None: + return "threshold_crossed", event_message + f"{threshold}% of budget consumed" elif percent_left <= SLACK_ALERTING_THRESHOLD_5_PERCENT: event = "threshold_crossed" event_message += "5% or less of budget remaining" @@ -1431,17 +1468,29 @@ Model Info: return False + def _slack_budget_alert_allowed(self, user_info: CallInfo | WebhookEvent | None) -> bool: + patterns: Final = self.alerting_args.slack_budget_alert_key_aliases + return patterns is None or ( + user_info is not None + and user_info.event_group == Litellm_EntityType.KEY + and user_info.key_alias is not None + and user_info.key_alias != "" + and any(fnmatchcase(user_info.key_alias, pattern) for pattern in patterns) + ) + async def send_alert( self, message: str, level: Literal["Low", "Medium", "High"], alert_type: AlertType, - alerting_metadata: dict, + alerting_metadata: dict[str, object], user_info: WebhookEvent | None = None, request_model: str | None = None, api_base: str | None = None, - **kwargs, - ): + budget_alert_destination: Literal["all", "slack", "non_slack"] = "all", + budget_alert_email: bool = True, + **kwargs: object, + ) -> bool: """ Alerting based on thresholds: - https://github.com/BerriAI/litellm/issues/1298 @@ -1459,25 +1508,36 @@ Model Info: api_base: Optional[str] - api base for digest grouping """ if self.alerting is None: - return + return False # Start periodic flush if not already started if self.alerting is not None and len(self.alerting) > 0: self._ensure_periodic_flush_task() - if "webhook" in self.alerting and alert_type == "budget_alerts" and user_info is not None: + if ( + budget_alert_destination != "slack" + and "webhook" in self.alerting + and alert_type == "budget_alerts" + and user_info is not None + ): await self.send_webhook_alert(webhook_event=user_info) - if "email" in self.alerting and alert_type == "budget_alerts" and user_info is not None: + if ( + budget_alert_destination != "slack" + and budget_alert_email + and "email" in self.alerting + and alert_type == "budget_alerts" + and user_info is not None + ): # only send budget alerts over Email await self.send_email_alert_using_smtp(webhook_event=user_info, alert_type=alert_type) - send_to_slack: Final = "slack" in self.alerting - send_to_ms_teams: Final = MS_TEAMS_ALERTING_DESTINATION in self.alerting + send_to_slack: Final = "slack" in self.alerting and budget_alert_destination != "non_slack" + send_to_ms_teams: Final = MS_TEAMS_ALERTING_DESTINATION in self.alerting and budget_alert_destination != "slack" if not send_to_slack and not send_to_ms_teams: - return + return False if alert_type not in self.alert_types: - return + return False from datetime import datetime @@ -1504,10 +1564,12 @@ Model Info: if send_to_ms_teams: self._enqueue_ms_teams_alert(formatted_message=formatted_message, alert_type=alert_type) - if not send_to_slack: + if not send_to_slack or ( + alert_type == AlertType.budget_alerts and not self._slack_budget_alert_allowed(user_info) + ): if len(self.log_queue) >= self.batch_size: await self.flush_queue() - return + return False # Check if digest mode is enabled for this alert type alert_type_name_str: Final = getattr(alert_type, "value", str(alert_type)) @@ -1522,6 +1584,8 @@ Model Info: _digest_webhook = os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERTING_WEBHOOK_URL") if _digest_webhook is None: raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment") + if _digest_webhook == []: + return False digest_key: Final = f"{alert_type_name_str}:{request_model or ''}:{api_base or ''}" @@ -1542,7 +1606,7 @@ Model Info: last_time=now, webhook_url=_digest_webhook, ) - return # Suppress immediate alert; will be emitted by _flush_digest_buckets + return True # check if we find the slack webhook url in self.alert_to_webhook_url if self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url: @@ -1554,6 +1618,8 @@ Model Info: if slack_webhook_url is None: raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment") + if slack_webhook_url == []: + return False payload: Final = {"text": formatted_message} headers: Final = {"Content-type": "application/json"} @@ -1579,6 +1645,7 @@ Model Info: if len(self.log_queue) >= self.batch_size: await self.flush_queue() + return True def _enqueue_ms_teams_alert(self, formatted_message: str, alert_type: AlertType) -> None: ms_teams_webhook_url: Final = get_ms_teams_webhook_url() diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dfdf7dc4e66..6a7e9b3688b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5720,6 +5720,7 @@ async def _virtual_key_max_budget_alert_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, user_obj: LiteLLM_UserTable | None = None, + budget_alert_thresholds: tuple[int, ...] | None = None, ): """ Triggers a budget alert if the token has reached EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE @@ -5765,8 +5766,13 @@ async def _virtual_key_max_budget_alert_check( ) ) else: - # Old path: existing single 80% threshold — completely unchanged - alert_threshold: Final = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + if budget_alert_thresholds == (): + return + alert_threshold: Final = valid_token.max_budget * ( + min(budget_alert_thresholds) / 100 + if budget_alert_thresholds is not None + else 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( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index cb39801a8e0..33fdef1a226 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2461,10 +2461,16 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e # Check 4. Max Budget Alert Check (runs before budget enforcement # so multi-threshold 100% alerts fire on the request that crosses # max_budget, before BudgetExceededError is raised below) + configured_thresholds: Final = ( + proxy_logging_obj.slack_alerting_instance.alerting_args.budget_alert_thresholds + ) await _virtual_key_max_budget_alert_check( valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, user_obj=user_obj, + budget_alert_thresholds=( + tuple(configured_thresholds) if configured_thresholds is not None else None + ), ) # Check 5. Token Spend is under budget diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1c738aefb97..ba28527a677 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -157,6 +157,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, CallInfo, + Litellm_EntityType, LiteLLM_VerificationTokenView, Member, UserAPIKeyAuth, @@ -1319,7 +1320,7 @@ class ProxyLogging: if alert_to_webhook_url is not None: self.alert_to_webhook_url = alert_to_webhook_url updated_slack_alerting = True - if alert_type_config is not None: + if alert_type_config is not None or alerting_args is not None: updated_slack_alerting = True if updated_slack_alerting is True: @@ -3082,14 +3083,22 @@ class ProxyLogging: # do nothing if alerting is not switched on (unless it's a soft_budget alert with team-specific emails) return + configured_thresholds: Final = self.slack_alerting_instance.alerting_args.budget_alert_thresholds if self.alerting is not None and ( "slack" in self.alerting or "ms_teams" in self.alerting or "webhook" in self.alerting ): if self.slack_alerting_instance is not None: - await self.slack_alerting_instance.budget_alerts( - type=type, - user_info=user_info, - ) + if configured_thresholds is None: + await self.slack_alerting_instance.budget_alerts(type=type, user_info=user_info) + else: + await self.slack_alerting_instance.budget_alerts( + type=type, + user_info=user_info, + send_threshold_email=( + self.email_logging_instance is None + or user_info.event_group not in (Litellm_EntityType.KEY, Litellm_EntityType.TEAM_MEMBER) + ), + ) # Call email_logging_instance if: # 1. "email" is in alerting config, OR @@ -3097,10 +3106,14 @@ class ProxyLogging: should_send_email = (self.alerting is not None and "email" in self.alerting) or is_soft_budget_with_alert_emails if should_send_email and self.email_logging_instance is not None: - await self.email_logging_instance.budget_alerts( - type=type, - user_info=user_info, - ) + if configured_thresholds is None: + await self.email_logging_instance.budget_alerts(type=type, user_info=user_info) + else: + await self.email_logging_instance.budget_alerts( + type=type, + user_info=user_info, + budget_alert_thresholds=tuple(configured_thresholds), + ) async def alerting_handler( self, diff --git a/litellm/types/integrations/slack_alerting.py b/litellm/types/integrations/slack_alerting.py index 770746196e4..41a56d76f81 100644 --- a/litellm/types/integrations/slack_alerting.py +++ b/litellm/types/integrations/slack_alerting.py @@ -1,11 +1,11 @@ import os import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime as dt from enum import Enum -from typing import Any, Final, Literal, Optional, Union +from typing import Annotated, Final -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.types.utils import LiteLLMPydanticObjectBase @@ -68,6 +68,26 @@ class SlackAlertingArgs(LiteLLMPydanticObjectBase): default=SlackAlertingArgsEnum.budget_alert_ttl.value, description="Cache ttl for budgets alerts. Prevents spamming same alert, each time budget is crossed. Value is in seconds.", ) # 24 hours + slack_budget_alert_key_aliases: ( # mutable-ok: public alerting config accepts and serializes a list + Annotated[list[Annotated[str, Field(strict=True, min_length=1)]], Field(strict=True)] | None + ) = Field( + default=None, + description="Case-sensitive key alias glob patterns for Slack budget alerts. Null allows all budget alerts; an empty list disables them.", + ) + budget_alert_thresholds: ( # mutable-ok: public configuration accepts and serializes a list + Annotated[list[Annotated[int, Field(strict=True, ge=1, le=99)]], Field(strict=True)] | None + ) = Field( + default=None, + description="Consumed budget percentages for enabled alert destinations. Null preserves legacy thresholds; an empty list disables percentage warnings.", + ) + + @field_validator("budget_alert_thresholds") + @classmethod + def validate_budget_alert_thresholds(cls, thresholds: Sequence[int] | None) -> Sequence[int] | None: + if thresholds is not None and len(thresholds) != len(set(thresholds)): + raise ValueError("Budget alert thresholds must be unique") + return thresholds + outage_alert_ttl: int = Field( default=SlackAlertingArgsEnum.outage_alert_ttl.value, description="Cache ttl for model outage alerts. Sets time-window for errors. Default is 1 minute. Value is in seconds.", diff --git a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py index 52e44ca5448..73682ec5991 100644 --- a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -2,25 +2,117 @@ import asyncio import json import os import unittest.mock as mock -from unittest.mock import patch +from typing import Final +from unittest.mock import AsyncMock, patch +import httpx import pytest from fastapi.testclient import TestClient - -from litellm.caching.caching import DualCache from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( BaseEmailLogger, ) - +from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import ( + SendGridEmailLogger, +) from litellm_enterprise.types.enterprise_callbacks.send_emails import ( EmailEvent, SendKeyCreatedEmailEvent, SendKeyRotatedEmailEvent, ) +from litellm.caching.caching import DualCache +from litellm.constants import EMAIL_BUDGET_ALERT_TTL from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER from litellm.proxy._types import CallInfo, Litellm_EntityType, WebhookEvent -from litellm.constants import EMAIL_BUDGET_ALERT_TTL + + +@pytest.mark.asyncio +async def test_common_budget_thresholds_use_real_email_rendering_and_per_level_dedup( + monkeypatch, +) -> None: + monkeypatch.setenv("SENDGRID_API_KEY", "synthetic-not-a-secret") + transport: Final = AsyncMock() + transport.post.return_value = httpx.Response(202) + logger: Final = SendGridEmailLogger() + logger.async_httpx_client = transport + info: Final = CallInfo( + spend=70, + max_budget=100, + token="synthetic", + user_email="owner@example.test", + event_group=Litellm_EntityType.KEY, + ) + for spend, count in ( + (69, 0), + (70, 1), + (70, 1), + (85, 2), + (96, 3), + (96, 3), + (100, 3), + ): + await logger.budget_alerts( + type="max_budget_alert", + user_info=info.model_copy(update={"spend": spend}), + budget_alert_thresholds=(95, 70, 85), + ) + assert transport.post.await_count == count + for call, pct in zip(transport.post.await_args_list, (70, 85, 95)): + payload: Final = call.kwargs["json"] + assert payload["personalizations"][0]["to"] == [{"email": "owner@example.test"}] + assert f"{pct}%" in payload["personalizations"][0]["subject"] + assert f"{pct}%" in payload["content"][0]["value"] + await logger.budget_alerts( + type="max_budget_alert", user_info=info, budget_alert_thresholds=() + ) + assert transport.post.await_count == 3 + await logger.budget_alerts( + type="max_budget_alert", + user_info=info.model_copy( + update={"max_budget_alert_emails": {"50": ["finance@example.test"]}} + ), + budget_alert_thresholds=(), + ) + assert transport.post.await_count == 4 + assert { + email["email"] + for email in transport.post.await_args.kwargs["json"]["personalizations"][0][ + "to" + ] + } == { + "owner@example.test", + "finance@example.test", + } + assert "50%" in transport.post.await_args.kwargs["json"]["content"][0]["value"] + + +@pytest.mark.asyncio +async def test_common_email_threshold_failed_send_releases_claim(monkeypatch) -> None: + monkeypatch.setenv("SENDGRID_API_KEY", "synthetic-not-a-secret") + transport: Final = AsyncMock() + transport.post.side_effect = ( + httpx.ConnectError("synthetic failure"), + httpx.Response(202), + ) + logger: Final = SendGridEmailLogger() + logger.async_httpx_client = transport + info: Final = CallInfo( + spend=70, + max_budget=100, + token="synthetic", + user_email="owner@example.test", + event_group=Litellm_EntityType.KEY, + ) + await logger.budget_alerts( + type="max_budget_alert", user_info=info, budget_alert_thresholds=(70,) + ) + await logger.budget_alerts( + type="max_budget_alert", user_info=info, budget_alert_thresholds=(70,) + ) + await logger.budget_alerts( + type="max_budget_alert", user_info=info, budget_alert_thresholds=(70,) + ) + assert transport.post.await_count == 2 @pytest.fixture(autouse=True) diff --git a/tests/unit/integrations/SlackAlerting/test_slack_alerting.py b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py index 47e55c8476e..1d77e08a427 100644 --- a/tests/unit/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/unit/integrations/SlackAlerting/test_slack_alerting.py @@ -3,12 +3,13 @@ import datetime import json import time import unittest -from typing import Final, List, Optional, Tuple +from collections.abc import Mapping +from typing import Final, List, Literal, Optional, Tuple from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch import httpx import pytest -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -17,8 +18,135 @@ from litellm.caching.caching import DualCache from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import CallInfo, Litellm_EntityType -from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType, SlackAlertingCacheKeys +from litellm.proxy._types import CallInfo, Litellm_EntityType, WebhookEvent +from litellm.types.integrations.slack_alerting import ( + AlertQueueItem, + AlertType, + SlackAlertingArgs, + SlackAlertingCacheKeys, +) + + +@pytest.mark.parametrize("thresholds", ([0], [100], [70, 70], [True], [70.0], ["70"], "70", (70,), {})) +def test_budget_alert_thresholds_reject_invalid_configuration(thresholds: object) -> None: + with pytest.raises(ValidationError): + SlackAlertingArgs.model_validate({"budget_alert_thresholds": thresholds}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("digest", (False, True)) +async def test_configured_budget_thresholds_notify_each_level_and_destination(digest: bool, monkeypatch) -> None: + http_handler: Final = _webhook_accepting_posts() + monkeypatch.setenv("WEBHOOK_URL", "https://example.test/budget") + monkeypatch.setenv("MS_TEAMS_WEBHOOK_URL", "https://example.test/teams") + alerts: Final = SlackAlerting( + alerting=["slack", "webhook", "ms_teams"], + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"budget_alert_thresholds": [95, 70, 85]}, + alert_type_config={"budget_alerts": {"digest": digest, "digest_interval": 0}}, + async_http_handler=http_handler, + ) + for spend, expected_posts in ((69, 0), (70, 3), (70, 3), (85, 6), (96, 9), (96, 9), (100, 12)): + await alerts.budget_alerts( + type="token_budget", + user_info=CallInfo( + spend=spend, + max_budget=100, + token="synthetic", + key_alias="example-api", + event_group=Litellm_EntityType.KEY, + ), + ) + await alerts._flush_digest_buckets() + await alerts.flush_queue() + assert http_handler.post.await_count == expected_posts + webhook_events: Final = tuple( + json.loads(call.kwargs["data"]) + for call in http_handler.post.await_args_list + if call.kwargs["url"] == "https://example.test/budget" + ) + assert tuple(event["event"] for event in webhook_events) == ( + "threshold_crossed", + "threshold_crossed", + "threshold_crossed", + "budget_crossed", + ) + assert tuple(event["event_message"] for event in webhook_events[:3]) == tuple( + f"Key Budget: {pct}% of budget consumed" for pct in (70, 85, 95) + ) + + +@pytest.mark.asyncio +async def test_budget_threshold_jump_empty_reload_and_soft_budget() -> None: + http_handler: Final = _webhook_accepting_posts() + alerts: Final = SlackAlerting( + alerting=["slack"], + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"budget_alert_thresholds": [70, 85, 95]}, + async_http_handler=http_handler, + ) + info: Final = CallInfo(spend=96, max_budget=100, token="synthetic", event_group=Litellm_EntityType.KEY) + await alerts.budget_alerts(type="token_budget", user_info=info) + await alerts.flush_queue() + assert len(_posted_slack_bodies(http_handler)) == 1 + assert "95% of budget consumed" in _posted_slack_bodies(http_handler)[0]["text"] + alerts.update_values(alerting_args={"budget_alert_thresholds": []}) + await alerts.budget_alerts(type="token_budget", user_info=info.model_copy(update={"spend": 99})) + await alerts.flush_queue() + assert http_handler.post.await_count == 1 + await alerts.budget_alerts(type="soft_budget", user_info=info.model_copy(update={"spend": 60, "soft_budget": 50})) + await alerts.budget_alerts(type="token_budget", user_info=info.model_copy(update={"spend": 100})) + await alerts.flush_queue() + assert http_handler.post.await_count == 3 + alerts.update_values(alerting_args={"budget_alert_thresholds": [50]}) + await alerts.budget_alerts(type="token_budget", user_info=info.model_copy(update={"spend": 50})) + await alerts.flush_queue() + assert "50% of budget consumed" in _posted_slack_bodies(http_handler)[-1]["text"] + + +@pytest.mark.asyncio +async def test_configured_thresholds_preserve_soft_and_projected_events() -> None: + http_handler: Final = _webhook_accepting_posts() + alerts: Final = SlackAlerting( + alerting=["slack"], + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"budget_alert_thresholds": [70, 95]}, + async_http_handler=http_handler, + ) + info: Final = CallInfo( + spend=70, max_budget=100, soft_budget=50, token="synthetic", event_group=Litellm_EntityType.KEY + ) + await alerts.budget_alerts(type="soft_budget", user_info=info, send_threshold_email=False) + await alerts.budget_alerts(type="token_budget", user_info=info, send_threshold_email=False) + await alerts.budget_alerts( + type="projected_limit_exceeded", + user_info=info.model_copy(update={"soft_budget": 90}), + send_threshold_email=False, + ) + await alerts.flush_queue() + assert http_handler.post.await_count == 2 + assert "Soft Budget Crossed" in _posted_slack_bodies(http_handler)[0]["text"] + assert "Projected Limit Exceeded" in _posted_slack_bodies(http_handler)[1]["text"] + + +@pytest.mark.asyncio +async def test_configured_threshold_respects_pre_upgrade_sent_marker() -> None: + cache: Final = DualCache() + await cache.async_set_cache("budget_alerts:threshold_crossed:synthetic", "SENT", ttl=86400) + http_handler: Final = _webhook_accepting_posts() + alerts: Final = SlackAlerting( + alerting=["slack"], + internal_usage_cache=cache, + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"budget_alert_thresholds": [70]}, + async_http_handler=http_handler, + ) + await alerts.budget_alerts( + type="token_budget", + user_info=CallInfo(spend=70, max_budget=100, token="synthetic", event_group=Litellm_EntityType.KEY), + ) + await alerts.flush_queue() + http_handler.post.assert_not_awaited() class TestSlackAlerting(unittest.TestCase): @@ -556,6 +684,455 @@ async def test_async_send_batch_collapses_only_identical_alerts() -> None: ) +def _budget_webhook_event( + key_alias: str | None, event_group: Litellm_EntityType = Litellm_EntityType.KEY +) -> WebhookEvent: + return WebhookEvent( + spend=85.0, + max_budget=100.0, + token="hashed_key", + key_alias=key_alias, + event="threshold_crossed", + event_message="15% or less of budget remaining", + event_group=event_group, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "alerting_args, key_alias, event_group, delivered", + [ + ({}, None, None, True), + ({}, None, Litellm_EntityType.KEY, True), + ({"slack_budget_alert_key_aliases": None}, None, None, True), + ({"slack_budget_alert_key_aliases": None}, None, Litellm_EntityType.TEAM, True), + ({"slack_budget_alert_key_aliases": []}, "github-example-api", Litellm_EntityType.KEY, False), + ( + {"slack_budget_alert_key_aliases": ["github-example-api"]}, + "github-example-api", + Litellm_EntityType.KEY, + True, + ), + ( + {"slack_budget_alert_key_aliases": ["github-example-api"]}, + "github-example-api-extra", + Litellm_EntityType.KEY, + False, + ), + ({"slack_budget_alert_key_aliases": ["github-example-*"]}, "github-example-api", Litellm_EntityType.KEY, True), + ({"slack_budget_alert_key_aliases": ["github-example-*"]}, "other-api", Litellm_EntityType.KEY, False), + ({"slack_budget_alert_key_aliases": ["github-example-*"]}, "GitHub-example-api", Litellm_EntityType.KEY, False), + ({"slack_budget_alert_key_aliases": ["github-example-*"]}, None, Litellm_EntityType.KEY, False), + ({"slack_budget_alert_key_aliases": ["*"]}, "", Litellm_EntityType.KEY, False), + ({"slack_budget_alert_key_aliases": ["*"]}, None, None, False), + ( + {"slack_budget_alert_key_aliases": ["other-*", "github-example-?"]}, + "github-example-a", + Litellm_EntityType.KEY, + True, + ), + ( + {"slack_budget_alert_key_aliases": ["github-example-?", "other-*"]}, + "github-example-a", + Litellm_EntityType.KEY, + True, + ), + ( + {"slack_budget_alert_key_aliases": ["github-example-?", "other-*"]}, + "github-example-ab", + Litellm_EntityType.KEY, + False, + ), + ({"slack_budget_alert_key_aliases": ["github-example-[ab]"]}, "github-example-b", Litellm_EntityType.KEY, True), + ], +) +async def test_slack_budget_key_alias_filter_delivery( + alerting_args: Mapping[str, object], + key_alias: str | None, + event_group: Litellm_EntityType | None, + delivered: bool, +) -> None: + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args=dict(alerting_args), + async_http_handler=http_handler, + ) + event: Final = _budget_webhook_event(key_alias, event_group) if event_group is not None else None + await slack_alerting.send_alert( + message=THRESHOLD_ALERT, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={"key_alias": "github-example-api"}, + user_info=event, + ) + + assert len(slack_alerting.log_queue) == int(delivered) + await slack_alerting.flush_queue() + assert http_handler.post.await_count == int(delivered) + if delivered: + assert THRESHOLD_ALERT in _posted_slack_bodies(http_handler)[0]["text"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_group", tuple(group for group in Litellm_EntityType if group != Litellm_EntityType.KEY)) +async def test_slack_budget_key_alias_filter_excludes_non_key_entities(event_group: Litellm_EntityType) -> None: + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"slack_budget_alert_key_aliases": ["github-example-*"]}, + async_http_handler=http_handler, + ) + await slack_alerting.send_alert( + message=THRESHOLD_ALERT, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + user_info=_budget_webhook_event("github-example-api", event_group), + ) + + assert slack_alerting.log_queue == [] + await slack_alerting.flush_queue() + http_handler.post.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "alert_type", + ( + AlertType.llm_exceptions, + AlertType.failed_tracking_spend, + AlertType.user_spend_thresholds, + AlertType.user_spend_anomalies, + ), +) +async def test_slack_budget_key_alias_filter_leaves_other_alert_types_unchanged(alert_type: AlertType) -> None: + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + alert_types=[alert_type], + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"slack_budget_alert_key_aliases": []}, + async_http_handler=http_handler, + ) + await slack_alerting.send_alert( + message="other alert", + level="High", + alert_type=alert_type, + alerting_metadata={}, + ) + await slack_alerting.flush_queue() + + http_handler.post.assert_awaited_once() + assert "other alert" in _posted_slack_bodies(http_handler)[0]["text"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("batch_size", (1, 100)) +async def test_slack_budget_key_alias_filter_preserves_webhook_and_teams_delivery( + monkeypatch: pytest.MonkeyPatch, batch_size: int +) -> None: + monkeypatch.setenv("WEBHOOK_URL", "https://webhook.example/budget") + monkeypatch.setenv("MS_TEAMS_WEBHOOK_URL", "https://teams.example/budget") + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting( + alerting=["slack", "webhook", "ms_teams"], + alerting_args={"slack_budget_alert_key_aliases": []}, + async_http_handler=http_handler, + batch_size=batch_size, + ) + event: Final = _budget_webhook_event("github-example-api") + await slack_alerting.send_alert( + message=THRESHOLD_ALERT, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + user_info=event, + ) + assert tuple(call.kwargs["url"] for call in http_handler.post.call_args_list) == ( + ("https://webhook.example/budget", "https://teams.example/budget") + if batch_size == 1 + else ("https://webhook.example/budget",) + ) + await slack_alerting.flush_queue() + + assert tuple(call.kwargs["url"] for call in http_handler.post.call_args_list) == ( + "https://webhook.example/budget", + "https://teams.example/budget", + ) + assert WebhookEvent.model_validate_json(http_handler.post.call_args_list[0].kwargs["data"]) == event + teams_body: Final = json.loads(http_handler.post.call_args_list[1].kwargs["data"]) + assert THRESHOLD_ALERT in teams_body["attachments"][0]["content"]["body"][0]["text"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key_alias, delivered", (("github-example-api", True), ("other-api", False))) +async def test_slack_budget_key_alias_filter_retains_native_thresholds_and_dedup( + key_alias: str, delivered: bool +) -> None: + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"slack_budget_alert_key_aliases": ["github-example-*"]}, + async_http_handler=http_handler, + ) + below_threshold: Final = CallInfo( + spend=50.0, max_budget=100.0, token="hashed_key", key_alias=key_alias, event_group=Litellm_EntityType.KEY + ) + await slack_alerting.budget_alerts(type="token_budget", user_info=below_threshold) + assert slack_alerting.log_queue == [] + assert ( + await slack_alerting.internal_usage_cache.async_get_cache("budget_alerts:threshold_crossed:hashed_key") is None + ) + + at_threshold: Final = below_threshold.model_copy(update={"spend": 85.0}) + await slack_alerting.budget_alerts(type="token_budget", user_info=at_threshold) + assert len(slack_alerting.log_queue) == int(delivered) + assert ( + await slack_alerting.internal_usage_cache.async_get_cache("budget_alerts:threshold_crossed:hashed_key") + == "SENT_WITH_SLACK_DEDUP" + ) + assert ( + await slack_alerting.internal_usage_cache.async_get_cache("budget_alerts:slack:threshold_crossed:hashed_key") + ) == ("SENT" if delivered else None) + await slack_alerting.budget_alerts(type="token_budget", user_info=at_threshold) + assert len(slack_alerting.log_queue) == int(delivered) + await slack_alerting.flush_queue() + assert http_handler.post.await_count == int(delivered) + if delivered: + assert "15% or less of budget remaining" in _posted_slack_bodies(http_handler)[0]["text"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key_alias, delivered", (("github-example-api", True), ("other-api", False))) +async def test_slack_budget_key_alias_filter_applies_before_digest(key_alias: str, delivered: bool) -> None: + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"slack_budget_alert_key_aliases": ["github-example-*"]}, + alert_type_config={"budget_alerts": {"digest": True, "digest_interval": 0}}, + async_http_handler=http_handler, + ) + await slack_alerting.send_alert( + message=THRESHOLD_ALERT, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + user_info=_budget_webhook_event(key_alias), + ) + assert slack_alerting.log_queue == [] + assert len(slack_alerting.digest_buckets) == int(delivered) + await slack_alerting._flush_digest_buckets() + await slack_alerting.flush_queue() + assert http_handler.post.await_count == int(delivered) + if delivered: + assert THRESHOLD_ALERT in _posted_slack_bodies(http_handler)[0]["text"] + + +@pytest.mark.parametrize("patterns", (None, [], ["github-example-*", "exact-alias", "?"])) +def test_slack_budget_key_alias_patterns_validate_and_round_trip(patterns: object) -> None: + args: Final = SlackAlertingArgs.model_validate({"slack_budget_alert_key_aliases": patterns}) + assert args.slack_budget_alert_key_aliases == patterns + assert args.model_dump()["slack_budget_alert_key_aliases"] == patterns + + +@pytest.mark.parametrize( + "patterns", ("github-example-*", 123, {}, ("alias",), [""], [None], [123], [True], [["alias"]]) +) +def test_slack_budget_key_alias_patterns_reject_invalid_config(patterns: object) -> None: + with pytest.raises(ValidationError, match="slack_budget_alert_key_aliases"): + SlackAlertingArgs.model_validate({"slack_budget_alert_key_aliases": patterns}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("patterns, delivered", (([], False), (["github-example-*"], True))) +async def test_slack_budget_key_alias_filter_reload_validates_and_changes_delivery( + patterns: list[str], delivered: bool +) -> None: + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + default_webhook_url=SLACK_WEBHOOK_URL, + async_http_handler=http_handler, + ) + slack_alerting.update_values(alerting_args={"slack_budget_alert_key_aliases": patterns}) + with pytest.raises(ValidationError, match="slack_budget_alert_key_aliases"): + slack_alerting.update_values(alerting_args={"slack_budget_alert_key_aliases": "github-example-*"}) + await slack_alerting.send_alert( + message=THRESHOLD_ALERT, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + user_info=_budget_webhook_event("github-example-api"), + ) + await slack_alerting.flush_queue() + assert http_handler.post.await_count == int(delivered) + if delivered: + assert THRESHOLD_ALERT in _posted_slack_bodies(http_handler)[0]["text"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("digest", (False, True)) +@pytest.mark.parametrize("patterns", (["github-example-*"], None)) +@pytest.mark.parametrize( + "budget_type, spend, max_budget, soft_budget, event", + ( + ("token_budget", 85.0, 100.0, None, "threshold_crossed"), + ("max_budget_alert", 100.0, 100.0, None, "budget_crossed"), + ("soft_budget", 50.0, None, 40.0, "soft_budget_crossed"), + ("projected_limit_exceeded", 50.0, 100.0, None, "projected_limit_exceeded"), + ), +) +async def test_budget_filter_reload_sends_slack_without_repeating_other_destinations( + monkeypatch: pytest.MonkeyPatch, + digest: bool, + patterns: list[str] | None, + budget_type: Literal["token_budget", "max_budget_alert", "soft_budget", "projected_limit_exceeded"], + spend: float, + max_budget: float | None, + soft_budget: float | None, + event: str, +) -> None: + monkeypatch.setenv("WEBHOOK_URL", "https://webhook.example/budget") + monkeypatch.setenv("MS_TEAMS_WEBHOOK_URL", "https://teams.example/budget") + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting( + alerting=["slack", "webhook", "ms_teams"], + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"slack_budget_alert_key_aliases": []}, + alert_type_config={"budget_alerts": {"digest": digest, "digest_interval": 0}}, + async_http_handler=http_handler, + ) + info: Final = CallInfo( + spend=spend, + max_budget=max_budget, + soft_budget=soft_budget, + token="hashed_key", + key_alias="github-example-api", + event_group=Litellm_EntityType.KEY, + ) + await slack_alerting.budget_alerts(type=budget_type, user_info=info) + await slack_alerting.flush_queue() + assert tuple(c.kwargs["url"] for c in http_handler.post.call_args_list) == ( + "https://webhook.example/budget", + "https://teams.example/budget", + ) + slack_alerting.update_values(alerting_args={"slack_budget_alert_key_aliases": patterns}) + await slack_alerting.budget_alerts(type=budget_type, user_info=info) + await slack_alerting.budget_alerts(type=budget_type, user_info=info) + await slack_alerting._flush_digest_buckets() + await slack_alerting.flush_queue() + assert tuple(c.kwargs["url"] for c in http_handler.post.call_args_list) == ( + "https://webhook.example/budget", + "https://teams.example/budget", + SLACK_WEBHOOK_URL, + ) + assert ( + "github-example-api" + in _SLACK_WEBHOOK_BODY.validate_json(http_handler.post.call_args_list[-1].kwargs["data"])["text"] + ) + assert ( + await slack_alerting.internal_usage_cache.async_get_cache(f"budget_alerts:slack:{event}:hashed_key") == "SENT" + ) + + +@pytest.mark.asyncio +async def test_budget_slack_and_other_destination_windows_expire_independently( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("WEBHOOK_URL", "https://webhook.example/budget") + http_handler: Final = _webhook_accepting_posts() + cache: Final = DualCache() + slack_alerting: Final = SlackAlerting( + alerting=["slack", "webhook"], + internal_usage_cache=cache, + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"slack_budget_alert_key_aliases": []}, + async_http_handler=http_handler, + ) + info: Final = CallInfo( + spend=85.0, + max_budget=100.0, + token="hashed_key", + key_alias="github-example-api", + event_group=Litellm_EntityType.KEY, + ) + await slack_alerting.budget_alerts(type="token_budget", user_info=info) + slack_alerting.update_values(alerting_args={"slack_budget_alert_key_aliases": None}) + await slack_alerting.budget_alerts(type="token_budget", user_info=info) + await slack_alerting.flush_queue() + cache.delete_cache("budget_alerts:threshold_crossed:hashed_key") + await slack_alerting.budget_alerts(type="token_budget", user_info=info) + await slack_alerting.flush_queue() + cache.delete_cache("budget_alerts:slack:threshold_crossed:hashed_key") + await slack_alerting.budget_alerts(type="token_budget", user_info=info) + await slack_alerting.flush_queue() + assert tuple(c.kwargs["url"] for c in http_handler.post.call_args_list) == ( + "https://webhook.example/budget", + SLACK_WEBHOOK_URL, + "https://webhook.example/budget", + SLACK_WEBHOOK_URL, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("digest", (False, True)) +async def test_budget_empty_channel_mapping_does_not_consume_slack_dedup(digest: bool) -> None: + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + alert_to_webhook_url={AlertType.budget_alerts: []}, + alert_type_config={"budget_alerts": {"digest": digest, "digest_interval": 0}}, + async_http_handler=http_handler, + ) + info: Final = CallInfo( + spend=85.0, + max_budget=100.0, + token="hashed_key", + key_alias="github-example-api", + event_group=Litellm_EntityType.KEY, + ) + await slack_alerting.budget_alerts(type="token_budget", user_info=info) + assert slack_alerting.log_queue == [] + assert slack_alerting.digest_buckets == {} + slack_alerting.update_values(alert_to_webhook_url={AlertType.budget_alerts: SLACK_WEBHOOK_URL}) + await slack_alerting.budget_alerts(type="token_budget", user_info=info) + await slack_alerting._flush_digest_buckets() + await slack_alerting.flush_queue() + assert "github-example-api" in _posted_slack_bodies(http_handler)[0]["text"] + http_handler.post.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_legacy_budget_sent_marker_does_not_repeat_slack() -> None: + http_handler: Final = _webhook_accepting_posts() + cache: Final = DualCache() + await cache.async_set_cache("budget_alerts:threshold_crossed:hashed_key", "SENT", ttl=86400) + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + internal_usage_cache=cache, + default_webhook_url=SLACK_WEBHOOK_URL, + alerting_args={"slack_budget_alert_key_aliases": ["github-example-*"]}, + async_http_handler=http_handler, + ) + await slack_alerting.budget_alerts( + type="token_budget", + user_info=CallInfo( + spend=85.0, + max_budget=100.0, + token="hashed_key", + key_alias="github-example-api", + event_group=Litellm_EntityType.KEY, + ), + ) + await slack_alerting.flush_queue() + http_handler.post.assert_not_awaited() + + def _periodic_flush_tasks() -> list[asyncio.Task[object]]: return [ t diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 059cff0c385..2231dfcb8e7 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -3792,6 +3792,43 @@ async def test_virtual_key_soft_budget_check_scenarios(spend, soft_budget, expec ) +@pytest.mark.asyncio +@pytest.mark.parametrize("thresholds,spend,expected", (([70, 95], 69, 0), ([70, 95], 70, 1), ([], 85, 0))) +async def test_common_email_threshold_auth_gate(thresholds, spend: int, expected: int, monkeypatch) -> None: + from unittest.mock import AsyncMock + import httpx + from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import SendGridEmailLogger + + monkeypatch.setenv("SENDGRID_API_KEY", "synthetic-not-a-secret") + transport: Final = AsyncMock() + transport.post.return_value = httpx.Response(202) + from litellm.proxy.utils import ProxyLogging + + proxy_logging: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + proxy_logging.alerting = ["email"] + proxy_logging.slack_alerting_instance = SlackAlerting( + alerting=["email"], alerting_args={"budget_alert_thresholds": thresholds} + ) + email_logger: Final = SendGridEmailLogger() + email_logger.async_httpx_client = transport + proxy_logging.email_logging_instance = email_logger + await _virtual_key_max_budget_alert_check( + valid_token=UserAPIKeyAuth(token="synthetic", spend=spend, max_budget=100), + proxy_logging_obj=proxy_logging, + user_obj=LiteLLM_UserTable(user_id="owner", user_email="owner@example.test"), + budget_alert_thresholds=tuple(thresholds), + ) + pending: Final = tuple( + task + for task in asyncio.all_tasks() + if task is not asyncio.current_task() and task.get_coro().__qualname__ == "ProxyLogging.budget_alerts" + ) + if pending: + await asyncio.gather(*pending) + assert transport.post.await_count == expected + + @pytest.mark.asyncio async def test_virtual_key_max_budget_alert_check_with_user_obj(): """Test _virtual_key_max_budget_alert_check includes user_email when user_obj is provided""" @@ -10123,7 +10160,9 @@ async def test_authoritative_access_group_reads_writer_despite_stale_allow_cache from litellm.proxy._types import LiteLLM_AccessGroupTable from litellm.proxy.auth.auth_checks import get_access_object - stale: Final = LiteLLM_AccessGroupTable(access_group_id="group", access_group_name="Policy", access_model_names=["old"]) + stale: Final = LiteLLM_AccessGroupTable( + access_group_id="group", access_group_name="Policy", access_model_names=["old"] + ) current: Final = stale.model_copy(update={"access_model_names": ["new"] if allowed else []}) client: Final = MagicMock() client.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(return_value=current) @@ -10275,9 +10314,12 @@ async def test_authoritative_key_cannot_keep_grants_when_permission_is_unavailab from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache database: Final = MagicMock() - database.get_data = AsyncMock(return_value=UserAPIKeyAuth( - object_permission_id="grant", object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="grant", agents=["allowed"]) - )) + database.get_data = AsyncMock( + return_value=UserAPIKeyAuth( + object_permission_id="grant", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="grant", agents=["allowed"]), + ) + ) database.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock( return_value=None, side_effect=None if missing else RuntimeError("writer unavailable") ) @@ -10300,7 +10342,9 @@ async def test_authoritative_group_grants_propagate_policy_outages( database: Final = MagicMock() database.db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) - database.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + database.writer_db.litellm_accessgrouptable.find_unique = AsyncMock( + side_effect=RuntimeError("database unavailable") + ) monkeypatch.setattr(proxy_server, "prisma_client", database) monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) if strict: diff --git a/tests/unit/proxy/utils/proxy_logging/test_alerting.py b/tests/unit/proxy/utils/proxy_logging/test_alerting.py index 43ee1094330..7aafa931d3c 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_alerting.py +++ b/tests/unit/proxy/utils/proxy_logging/test_alerting.py @@ -7,16 +7,65 @@ Covers ``failed_tracking_alert``, ``budget_alerts``, ``alerting_handler``, from __future__ import annotations from datetime import datetime -from typing import Any, Dict +from typing import Any, Dict, Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from fastapi import HTTPException +from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import SendGridEmailLogger from prisma.errors import PrismaError import litellm from litellm._service_logger import ServiceTypes -from litellm.proxy._types import AlertType, CallInfo +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.proxy._types import AlertType, CallInfo, Litellm_EntityType + + +@pytest.mark.asyncio +@pytest.mark.parametrize("destinations", (["email"], ["email", "slack", "webhook", "ms_teams"])) +async def test_common_thresholds_propagate_to_email_without_duplicate_smtp( + proxy_logging, monkeypatch, destinations +) -> None: + monkeypatch.setenv("SENDGRID_API_KEY", "synthetic-not-a-secret") + monkeypatch.setenv("WEBHOOK_URL", "https://example.test/budget") + monkeypatch.setenv("MS_TEAMS_WEBHOOK_URL", "https://example.test/teams") + transport: Final = AsyncMock() + transport.post.return_value = httpx.Response(200) + proxy_logging.alerting = destinations + proxy_logging.slack_alerting_instance = SlackAlerting( + alerting=destinations, + default_webhook_url="https://example.test/slack", + alerting_args={"budget_alert_thresholds": [70, 95]}, + async_http_handler=transport, + ) + email_logger: Final = SendGridEmailLogger() + email_logger.async_httpx_client = transport + proxy_logging.email_logging_instance = email_logger + info: Final = CallInfo( + spend=70, max_budget=100, token="synthetic", user_email="owner@example.test", event_group=Litellm_EntityType.KEY + ) + for spend in (70, 70, 96, 96): + event: Final = info.model_copy(update={"spend": spend}) + await proxy_logging.budget_alerts(type="token_budget", user_info=event) + await proxy_logging.budget_alerts(type="max_budget_alert", user_info=event) + await proxy_logging.slack_alerting_instance.flush_queue() + email_posts: Final = tuple(call for call in transport.post.await_args_list if "json" in call.kwargs) + assert len(email_posts) == 2 + assert "70%" in email_posts[0].kwargs["json"]["personalizations"][0]["subject"] + assert "95%" in email_posts[1].kwargs["json"]["personalizations"][0]["subject"] + assert transport.post.await_count == (2 if destinations == ["email"] else 8) + + +@pytest.mark.asyncio +async def test_alerting_args_only_reload_validates_and_preserves_last_valid(proxy_logging) -> None: + proxy_logging.update_values(alerting_args={"budget_alert_thresholds": [70, 95]}) + assert proxy_logging.slack_alerting_instance.alerting_args.budget_alert_thresholds == [70, 95] + from pydantic import ValidationError + + with pytest.raises(ValidationError): + proxy_logging.update_values(alerting_args={"budget_alert_thresholds": [True]}) + assert proxy_logging.slack_alerting_instance.alerting_args.budget_alert_thresholds == [70, 95] # --------------------------------------------------------------------------- @@ -160,9 +209,7 @@ async def test_budget_alerts_soft_budget_with_alert_emails_bypasses_global(proxy @pytest.mark.asyncio async def test_budget_alerts_slack_failure_raises(proxy_logging): proxy_logging.alerting = ["slack"] - proxy_logging.slack_alerting_instance = MagicMock( - budget_alerts=AsyncMock(side_effect=ConnectionError("slack")) - ) + proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=AsyncMock(side_effect=ConnectionError("slack"))) proxy_logging.email_logging_instance = None with pytest.raises(ConnectionError): await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info()) @@ -317,11 +364,7 @@ async def test_failure_handler_with_capture_exception_invoked(proxy_logging, mon async def test_failure_handler_propagates_service_logging_error_raises(proxy_logging, monkeypatch): proxy_logging.alert_types = [AlertType.db_exceptions] proxy_logging.alerting_handler = AsyncMock() - proxy_logging.service_logging_obj = MagicMock( - async_service_failure_hook=AsyncMock(side_effect=RuntimeError("svc")) - ) + proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock(side_effect=RuntimeError("svc"))) monkeypatch.setattr(litellm.utils, "capture_exception", None) with pytest.raises(RuntimeError): - await proxy_logging.failure_handler( - original_exception=Exception("x"), duration=0.0, call_type="db_read" - ) + await proxy_logging.failure_handler(original_exception=Exception("x"), duration=0.0, call_type="db_read") diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a72aec07766..13a1bda8cbe 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -2627,7 +2627,9 @@ class TestProxyLoggingBudgetAlerts: proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging.alerting = ["slack"] - proxy_logging.slack_alerting_instance = AsyncMock() + proxy_logging.slack_alerting_instance = AsyncMock( + alerting_args=proxy_logging.slack_alerting_instance.alerting_args + ) user_info = MagicMock() @@ -2675,7 +2677,9 @@ class TestProxyLoggingBudgetAlerts: proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging.alerting = ["slack", "email"] - proxy_logging.slack_alerting_instance = AsyncMock() + proxy_logging.slack_alerting_instance = AsyncMock( + alerting_args=proxy_logging.slack_alerting_instance.alerting_args + ) proxy_logging.email_logging_instance = AsyncMock() user_info = MagicMock() @@ -2708,7 +2712,9 @@ class TestProxyLoggingBudgetAlerts: proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging.alerting = ["slack", "email"] - proxy_logging.slack_alerting_instance = AsyncMock() + proxy_logging.slack_alerting_instance = AsyncMock( + alerting_args=proxy_logging.slack_alerting_instance.alerting_args + ) proxy_logging.email_logging_instance = AsyncMock() user_info = MagicMock() @@ -2736,7 +2742,9 @@ class TestProxyLoggingBudgetAlerts: proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging.alerting = None # Global alerting is disabled - proxy_logging.slack_alerting_instance = AsyncMock() + proxy_logging.slack_alerting_instance = AsyncMock( + alerting_args=proxy_logging.slack_alerting_instance.alerting_args + ) proxy_logging.email_logging_instance = AsyncMock() # Create CallInfo with alert_emails set (simulating team metadata extraction)