diff --git a/enterprise/dist/litellm_enterprise-0.1.26-py3-none-any.whl b/enterprise/dist/litellm_enterprise-0.1.26-py3-none-any.whl new file mode 100644 index 00000000000..e4cfac65530 Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.26-py3-none-any.whl differ diff --git a/enterprise/dist/litellm_enterprise-0.1.26.tar.gz b/enterprise/dist/litellm_enterprise-0.1.26.tar.gz new file mode 100644 index 00000000000..c8e0081ff11 Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.26.tar.gz differ diff --git a/enterprise/dist/litellm_enterprise-0.1.27-py3-none-any.whl b/enterprise/dist/litellm_enterprise-0.1.27-py3-none-any.whl new file mode 100644 index 00000000000..0274d62e16e Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.27-py3-none-any.whl differ diff --git a/enterprise/dist/litellm_enterprise-0.1.27.tar.gz b/enterprise/dist/litellm_enterprise-0.1.27.tar.gz new file mode 100644 index 00000000000..d802b5a89d5 Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.27.tar.gz differ 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 1fe82c2c188..61e0745bab1 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -5,7 +5,7 @@ Base class for sending emails to user after creating keys or invite links import json import os -from typing import List, Optional +from typing import List, Literal, Optional from litellm_enterprise.types.enterprise_callbacks.send_emails import ( EmailEvent, @@ -15,6 +15,7 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import ( ) from litellm._logging import verbose_proxy_logger +from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER from litellm.integrations.email_templates.key_created_email import ( @@ -26,9 +27,17 @@ from litellm.integrations.email_templates.key_rotated_email import ( from litellm.integrations.email_templates.user_invitation_email import ( USER_INVITATION_EMAIL_TEMPLATE, ) -from litellm.proxy._types import InvitationNew, UserAPIKeyAuth, WebhookEvent +from litellm.integrations.email_templates.templates import ( + MAX_BUDGET_ALERT_EMAIL_TEMPLATE, + SOFT_BUDGET_ALERT_EMAIL_TEMPLATE, +) +from litellm.proxy._types import CallInfo, InvitationNew, UserAPIKeyAuth, WebhookEvent from litellm.secret_managers.main import get_secret_bool from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL +from litellm.constants import ( + EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, + EMAIL_BUDGET_ALERT_TTL, +) class BaseEmailLogger(CustomLogger): @@ -40,6 +49,21 @@ class BaseEmailLogger(CustomLogger): EmailEvent.virtual_key_rotated: "LiteLLM: {event_message}", } + def __init__( + self, + internal_usage_cache: Optional[DualCache] = None, + **kwargs, + ): + """ + Initialize BaseEmailLogger + + Args: + internal_usage_cache: DualCache instance for preventing duplicate alerts + **kwargs: Additional arguments passed to CustomLogger + """ + super().__init__(**kwargs) + self.internal_usage_cache = internal_usage_cache or DualCache() + async def send_user_invitation_email(self, event: WebhookEvent): """ Send email to user after inviting them to the team @@ -154,6 +178,218 @@ class BaseEmailLogger(CustomLogger): ) pass + async def send_soft_budget_alert_email(self, event: WebhookEvent): + """ + Send email to user when soft budget is crossed + """ + email_params = await self._get_email_params( + email_event=EmailEvent.soft_budget_crossed, # Reuse existing event type for subject template + user_id=event.user_id, + user_email=event.user_email, + event_message=event.event_message, + ) + + verbose_proxy_logger.debug( + f"send_soft_budget_alert_email_event: {json.dumps(event.model_dump(exclude_none=True), indent=4, default=str)}" + ) + + # Format budget values + soft_budget_str = f"${event.soft_budget}" if event.soft_budget is not None else "N/A" + spend_str = f"${event.spend}" if event.spend is not None else "$0.00" + max_budget_info = "" + if event.max_budget is not None: + max_budget_info = f"Maximum Budget: ${event.max_budget}
" + + email_html_content = SOFT_BUDGET_ALERT_EMAIL_TEMPLATE.format( + email_logo_url=email_params.logo_url, + recipient_email=email_params.recipient_email, + soft_budget=soft_budget_str, + spend=spend_str, + max_budget_info=max_budget_info, + 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 + + async def send_max_budget_alert_email(self, event: WebhookEvent): + """ + 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, + ) + + verbose_proxy_logger.debug( + f"send_max_budget_alert_email_event: {json.dumps(event.model_dump(exclude_none=True), indent=4, default=str)}" + ) + + # Format budget values + spend_str = f"${event.spend}" if event.spend is not None else "$0.00" + max_budget_str = f"${event.max_budget}" if event.max_budget is not None else "N/A" + + # Calculate percentage and alert threshold + percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100) + alert_threshold_str = f"${event.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE:.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 + + async def budget_alerts( + self, + type: Literal[ + "token_budget", + "soft_budget", + "max_budget_alert", + "user_budget", + "team_budget", + "organization_budget", + "proxy_budget", + "projected_limit_exceeded", + ], + user_info: CallInfo, + ): + """ + Send a budget alert via email + + Args: + type: The type of budget alert to send + user_info: The user info to send the alert for + """ + ## PREVENTITIVE ALERTING ## + # - Alert once within 24hr period + # - Cache this information + # - Don't re-alert, if alert already sent + _cache: DualCache = self.internal_usage_cache + + # percent of max_budget left to spend + if user_info.max_budget is None and user_info.soft_budget is None: + return + + # For soft_budget alerts, check if we've already sent an alert + if type == "soft_budget": + if user_info.soft_budget is not None and user_info.spend >= user_info.soft_budget: + # 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:soft_budget_crossed:{_id}" + + # Check if we've already sent this alert + result = await _cache.async_get_cache(key=_cache_key) + if result is None: + # Create WebhookEvent for soft budget alert + event_message = f"Soft Budget Crossed - Total Soft Budget: ${user_info.soft_budget}" + webhook_event = WebhookEvent( + event="soft_budget_crossed", + 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_soft_budget_alert_email(webhook_event) + + # Cache the alert to prevent duplicate sends + 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 soft budget alert email: {e}", + exc_info=True, + ) + return + + # 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: + alert_threshold = user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + + # Only alert if we've crossed the threshold but haven't exceeded max_budget yet + if user_info.spend >= alert_threshold and user_info.spend < user_info.max_budget: + # 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}" + + # Check if we've already sent this alert + result = await _cache.async_get_cache(key=_cache_key) + if result is None: + # Calculate percentage + percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100) + + # Create WebhookEvent for max budget alert + event_message = f"Max Budget Alert - {percentage}% 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) + + # Cache the alert to prevent duplicate sends + 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 max budget alert email: {e}", + exc_info=True, + ) + return + async def _get_email_params( self, email_event: EmailEvent, diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/resend_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/resend_email.py index 8119e4a7ef5..7593e66aa47 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/resend_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/resend_email.py @@ -19,7 +19,8 @@ RESEND_API_ENDPOINT = "https://api.resend.com/emails" class ResendEmailLogger(BaseEmailLogger): - def __init__(self): + def __init__(self, internal_usage_cache=None, **kwargs): + super().__init__(internal_usage_cache=internal_usage_cache, **kwargs) self.async_httpx_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py index dfde9ce329a..8fc2d66d531 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py @@ -27,7 +27,8 @@ class SendGridEmailLogger(BaseEmailLogger): - SENDGRID_API_KEY """ - def __init__(self): + def __init__(self, internal_usage_cache=None, **kwargs): + super().__init__(internal_usage_cache=internal_usage_cache, **kwargs) self.async_httpx_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/smtp_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/smtp_email.py index 4ede8ee59fe..8efdaf231b7 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/smtp_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/smtp_email.py @@ -21,7 +21,8 @@ class SMTPEmailLogger(BaseEmailLogger): - SMTP_SENDER_EMAIL """ - def __init__(self): + def __init__(self, internal_usage_cache=None, **kwargs): + super().__init__(internal_usage_cache=internal_usage_cache, **kwargs) verbose_logger.debug("SMTP Email Logger initialized....") async def send_email( diff --git a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py index 736aaff1f75..380b0a6facb 100644 --- a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py +++ b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py @@ -36,6 +36,8 @@ class EmailEvent(str, enum.Enum): virtual_key_created = "Virtual Key Created" new_user_invitation = "New User Invitation" virtual_key_rotated = "Virtual Key Rotated" + soft_budget_crossed = "Soft Budget Crossed" + max_budget_alert = "Max Budget Alert" class EmailEventSettings(BaseModel): event: EmailEvent @@ -51,6 +53,8 @@ class DefaultEmailSettings(BaseModel): EmailEvent.virtual_key_created: True, # On by default EmailEvent.new_user_invitation: True, # On by default EmailEvent.virtual_key_rotated: True, # On by default + EmailEvent.soft_budget_crossed: True, # On by default + EmailEvent.max_budget_alert: True, # On by default } ) def to_dict(self) -> Dict[str, bool]: diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 2bcd8d33adc..1f3da432574 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-enterprise" -version = "0.1.25" +version = "0.1.27" description = "Package for LiteLLM Enterprise features" authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.1.25" +version = "0.1.27" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-enterprise==", diff --git a/litellm/constants.py b/litellm/constants.py index 38d3e8a1753..23766874841 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -313,6 +313,8 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE = os.getenv( "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield" ) +EMAIL_BUDGET_ALERT_TTL = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) # 24 hours in seconds +EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float(os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)) # 80% of max budget ############### LLM Provider Constants ############### ### ANTHROPIC CONSTANTS ### ANTHROPIC_SKILLS_API_BETA_VERSION = "skills-2025-10-02" diff --git a/litellm/integrations/email_templates/templates.py b/litellm/integrations/email_templates/templates.py index 7029e8ce12a..5de23db0f24 100644 --- a/litellm/integrations/email_templates/templates.py +++ b/litellm/integrations/email_templates/templates.py @@ -60,3 +60,51 @@ USER_INVITED_EMAIL_TEMPLATE = """ Best,
The LiteLLM team
""" + +SOFT_BUDGET_ALERT_EMAIL_TEMPLATE = """ + LiteLLM Logo + +

Hi {recipient_email},
+ + Your LiteLLM API key has crossed its soft budget limit of {soft_budget}.

+ + Current Spend: {spend}
+ Soft Budget: {soft_budget}
+ {max_budget_info} + +

+ ⚠️ Note: Your API requests will continue to work, but you should monitor your usage closely. + If you reach your maximum budget, requests will be rejected. +

+ + You can view your usage and manage your budget in the LiteLLM Dashboard.

+ + If you have any questions, please send an email to {email_support_contact}

+ + Best,
+ The LiteLLM team
+""" + +MAX_BUDGET_ALERT_EMAIL_TEMPLATE = """ + LiteLLM Logo + +

Hi {recipient_email},
+ + Your LiteLLM API key has reached {percentage}% of its maximum budget.

+ + Current Spend: {spend}
+ Maximum Budget: {max_budget}
+ Alert Threshold: {alert_threshold} ({percentage}%)
+ +

+ ⚠️ Warning: You are approaching your maximum budget limit. + Once you reach your maximum budget of {max_budget}, all API requests will be rejected. +

+ + You can view your usage and manage your budget in the LiteLLM Dashboard.

+ + If you have any questions, please send an email to {email_support_contact}

+ + Best,
+ The LiteLLM team
+""" \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8f9929597fe..b4588db6252 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2484,6 +2484,7 @@ class CallInfo(LiteLLMPydanticObjectBase): class WebhookEvent(CallInfo): event: Literal[ "budget_crossed", + "max_budget_alert", "soft_budget_crossed", "threshold_crossed", "projected_limit_exceeded", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 309bd577606..6a0223b35ba 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -24,6 +24,7 @@ from litellm.constants import ( DEFAULT_IN_MEMORY_TTL, DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, DEFAULT_MAX_RECURSE_DEPTH, + EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, ) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.proxy._types import ( @@ -1911,6 +1912,7 @@ async def _virtual_key_max_budget_check( 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, organization_id=valid_token.org_id, @@ -1939,6 +1941,7 @@ async def _virtual_key_max_budget_check( async def _virtual_key_soft_budget_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, + user_obj: Optional[LiteLLM_UserTable] = None, ): """ Triggers a budget alert if the token is over it's soft budget. @@ -1961,10 +1964,11 @@ async def _virtual_key_soft_budget_check( team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=valid_token.org_id, - user_email=None, + user_email=user_obj.user_email if user_obj else None, key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, ) + asyncio.create_task( proxy_logging_obj.budget_alerts( type="soft_budget", @@ -1973,6 +1977,56 @@ async def _virtual_key_soft_budget_check( ) +async def _virtual_key_max_budget_alert_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: Optional[LiteLLM_UserTable] = None, +): + """ + Triggers a budget alert if the token has reached EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + (default 80%) of its max budget. + This is a warning alert before the token actually exceeds the max budget. + + """ + + if ( + valid_token.max_budget is not None + 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 + + # 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, + ) + 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=user_obj.user_email if user_obj else None, + 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 _team_max_budget_check( team_object: Optional[LiteLLM_TeamTable], valid_token: Optional[UserAPIKeyAuth], diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d0c284e921c..495d4db304c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -29,6 +29,7 @@ from litellm.proxy.auth.auth_checks import ( _get_user_role, _is_user_proxy_admin, _virtual_key_max_budget_check, + _virtual_key_max_budget_alert_check, _virtual_key_soft_budget_check, can_key_call_model, common_checks, @@ -1062,10 +1063,18 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 user_obj=user_obj, ) - # Check 5. Soft Budget Check + # Check 5. Max Budget Alert Check + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=user_obj, + ) + + # Check 6. Soft Budget Check await _virtual_key_soft_budget_check( valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, + user_obj=user_obj, ) # Check 5. Token Model Spend is under Model budget diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index af74853d82d..a773e934ef1 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -3,7 +3,8 @@ model_list: litellm_params: model: gemini/* - - litellm_settings: - callbacks: ["langfuse"] + callbacks: ["dynamic_rate_limiter_v3"] + priority_reservation: + "prod": 0.9 # 90% reserved for production + "dev": 0.1 # 10% reserved for development diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 19f3b2a87d7..3927380c8ae 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2795,19 +2795,25 @@ class ProxyConfig: verbose_proxy_logger.debug(f"_alerting_callbacks: {general_settings}") if _alerting_callbacks is None: return + + # Ensure proxy_logging_obj.alerting is set for all alerting types + _alerting_value = general_settings.get("alerting", None) + verbose_proxy_logger.debug(f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}") + proxy_logging_obj.update_values( + alerting=_alerting_value, + alerting_threshold=general_settings.get("alerting_threshold", 600), + alert_types=general_settings.get("alert_types", None), + alert_to_webhook_url=general_settings.get( + "alert_to_webhook_url", None + ), + alerting_args=general_settings.get("alerting_args", None), + redis_cache=redis_usage_cache, + ) + for _alert in _alerting_callbacks: if _alert == "slack": - # [OLD] v0 implementation - proxy_logging_obj.update_values( - alerting=general_settings.get("alerting", None), - alerting_threshold=general_settings.get("alerting_threshold", 600), - alert_types=general_settings.get("alert_types", None), - alert_to_webhook_url=general_settings.get( - "alert_to_webhook_url", None - ), - alerting_args=general_settings.get("alerting_args", None), - redis_cache=redis_usage_cache, - ) + # [OLD] v0 implementation - already handled by update_values above + pass else: # [NEW] v1 implementation - init as a custom logger if _alert in litellm._known_custom_logger_compatible_callbacks: @@ -3274,6 +3280,7 @@ class ProxyConfig: proxy_logging_obj: ProxyLogging """ _general_settings = config_data.get("general_settings", {}) + if _general_settings is not None and "alerting" in _general_settings: if ( general_settings is not None @@ -3282,29 +3289,36 @@ class ProxyConfig: and _general_settings.get("alerting", None) is not None and isinstance(_general_settings["alerting"], list) ): - verbose_proxy_logger.debug( - "Overriding Default 'alerting' values with db 'alerting' values." - ) - general_settings["alerting"] = _general_settings[ - "alerting" - ] # override yaml values with db - proxy_logging_obj.alerting = general_settings["alerting"] - proxy_logging_obj.slack_alerting_instance.alerting = general_settings[ - "alerting" + # Merge DB and YAML/config alerting values instead of overriding + _yaml_alerting = set(general_settings["alerting"]) + _db_alerting = set(_general_settings["alerting"]) + _merged_alerting = list(_yaml_alerting.union(_db_alerting)) + # Preserve order: YAML values first, then DB values + _merged_alerting = list(general_settings["alerting"]) + [ + item for item in _general_settings["alerting"] + if item not in general_settings["alerting"] ] + verbose_proxy_logger.debug( + f"Merging alerting values: YAML={general_settings['alerting']}, DB={_general_settings['alerting']}, Merged={_merged_alerting}" + ) + general_settings["alerting"] = _merged_alerting + # Use update_values to properly set alerting for both slack and email + proxy_logging_obj.update_values( + alerting=general_settings["alerting"], + ) elif general_settings is None: general_settings = {} general_settings["alerting"] = _general_settings["alerting"] - proxy_logging_obj.alerting = general_settings["alerting"] - proxy_logging_obj.slack_alerting_instance.alerting = general_settings[ - "alerting" - ] + # Use update_values to properly set alerting for both slack and email + proxy_logging_obj.update_values( + alerting=general_settings["alerting"], + ) elif isinstance(general_settings, dict): general_settings["alerting"] = _general_settings["alerting"] - proxy_logging_obj.alerting = general_settings["alerting"] - proxy_logging_obj.slack_alerting_instance.alerting = general_settings[ - "alerting" - ] + # Use update_values to properly set alerting for both slack and email + proxy_logging_obj.update_values( + alerting=general_settings["alerting"], + ) if _general_settings is not None and "alert_types" in _general_settings: general_settings["alert_types"] = _general_settings["alert_types"] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 275baa88da8..6f42c485f59 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -35,6 +35,17 @@ from litellm.proxy._types import ( from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes, CallTypesLiteral +try: + 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.enterprise_callbacks.send_emails.smtp_email import SMTPEmailLogger + from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger +except ImportError: + BaseEmailLogger = None # type: ignore + SendGridEmailLogger = None # type: ignore + SMTPEmailLogger = None # type: ignore + ResendEmailLogger = None # type: ignore + try: import backoff except ImportError: @@ -128,6 +139,33 @@ def print_verbose(print_statement): print(f"LiteLLM Proxy: {print_statement}") # noqa +def _get_email_logger_class(): + """ + Determine which email logger class to use based on environment variables. + Priority: SendGrid > Resend > SMTP > BaseEmailLogger (fallback) + + Returns: + The email logger class to use, or None if BaseEmailLogger is not available + """ + if BaseEmailLogger is None: + return None + + # Check for SendGrid API key + if SendGridEmailLogger is not None and os.getenv("SENDGRID_API_KEY"): + return SendGridEmailLogger + + # Check for Resend API key + if ResendEmailLogger is not None and os.getenv("RESEND_API_KEY"): + return ResendEmailLogger + + # Check for SMTP configuration + if SMTPEmailLogger is not None and os.getenv("SMTP_HOST"): + return SMTPEmailLogger + + # Fallback to BaseEmailLogger (though it won't actually send emails) + return BaseEmailLogger + + class InternalUsageCache: def __init__(self, dual_cache: DualCache): self.dual_cache: DualCache = dual_cache @@ -266,6 +304,14 @@ class ProxyLogging: alerting=self.alerting, internal_usage_cache=self.internal_usage_cache.dual_cache, ) + self.email_logging_instance: Optional[Any] = None + if BaseEmailLogger is not None: + email_logger_class = _get_email_logger_class() + if email_logger_class is not None: + # All email logger classes now accept internal_usage_cache + self.email_logging_instance = email_logger_class( + internal_usage_cache=self.internal_usage_cache.dual_cache, + ) self.premium_user = premium_user self.service_logging_obj = ServiceLogging() self.db_spend_update_writer = DBSpendUpdateWriter() @@ -1149,6 +1195,7 @@ class ProxyLogging: "token_budget", "user_budget", "soft_budget", + "max_budget_alert", "team_budget", "organization_budget", "proxy_budget", @@ -1159,10 +1206,18 @@ class ProxyLogging: if self.alerting is None: # do nothing if alerting is not switched on return - await self.slack_alerting_instance.budget_alerts( - type=type, - user_info=user_info, - ) + + if "slack" in self.alerting: + await self.slack_alerting_instance.budget_alerts( + type=type, + user_info=user_info, + ) + + if "email" in self.alerting and self.email_logging_instance is not None: + await self.email_logging_instance.budget_alerts( + type=type, + user_info=user_info, + ) async def alerting_handler( self, diff --git a/pyproject.toml b/pyproject.toml index 3ac63bc214a..6020e9b6eb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3. mcp = {version = "^1.21.2", optional = true, python = ">=3.10"} litellm-proxy-extras = {version = "0.4.14", optional = true} rich = {version = "13.7.1", optional = true} -litellm-enterprise = {version = "0.1.25", optional = true} +litellm-enterprise = {version = "0.1.27", optional = true} diskcache = {version = "^5.6.1", optional = true} polars = {version = "^1.31.0", optional = true, python = ">=3.10"} semantic-router = {version = ">=0.1.12", optional = true, python = ">=3.9,<3.14"} diff --git a/requirements.txt b/requirements.txt index 69eaaff0835..c36f94f0752 100644 --- a/requirements.txt +++ b/requirements.txt @@ -67,4 +67,4 @@ soundfile==0.12.1 # for audio file processing ######################## # LITELLM ENTERPRISE DEPENDENCIES ######################## -litellm-enterprise==0.1.25 +litellm-enterprise==0.1.27 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 6b140d489cf..744195dfb6f 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 @@ -19,7 +19,8 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import ( ) from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER -from litellm.proxy._types import Litellm_EntityType, WebhookEvent +from litellm.proxy._types import CallInfo, Litellm_EntityType, WebhookEvent +from litellm.constants import EMAIL_BUDGET_ALERT_TTL @pytest.fixture(autouse=True) @@ -605,4 +606,276 @@ async def test_get_email_params_default_templates(monkeypatch): ) assert key_params.subject == "LiteLLM: API Key Created" - assert key_params.signature == EMAIL_FOOTER \ No newline at end of file + assert key_params.signature == EMAIL_FOOTER + + +@pytest.mark.asyncio +async def test_send_soft_budget_alert_email( + base_email_logger, mock_send_email, mock_lookup_user_email +): + """Test that send_soft_budget_alert_email sends an email with the correct parameters and content""" + event = WebhookEvent( + user_id="test_user", + user_email="test@example.com", + event_group=Litellm_EntityType.USER, + event="soft_budget_crossed", + event_message="Soft Budget Crossed - Total Soft Budget: $100.0", + spend=105.0, + max_budget=200.0, + soft_budget=100.0, + ) + + with mock.patch.dict( + os.environ, + { + "EMAIL_LOGO_URL": "https://litellm-listing.s3.amazonaws.com/litellm_logo.png", + "EMAIL_SUPPORT_CONTACT": "support@berri.ai", + "PROXY_BASE_URL": "http://test.com", + }, + ): + await base_email_logger.send_soft_budget_alert_email(event) + + mock_send_email.assert_called_once() + call_args = mock_send_email.call_args[1] + assert call_args["from_email"] == BaseEmailLogger.DEFAULT_LITELLM_EMAIL + assert call_args["to_email"] == ["test@example.com"] + assert call_args["subject"] == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0" + assert "$100.0" in call_args["html_body"] # soft_budget + assert "$105.0" in call_args["html_body"] # spend + assert "$200.0" in call_args["html_body"] # max_budget + + +@pytest.mark.asyncio +async def test_send_soft_budget_alert_email_no_max_budget( + base_email_logger, mock_send_email, mock_lookup_user_email +): + """Test that send_soft_budget_alert_email handles missing max_budget correctly""" + event = WebhookEvent( + user_id="test_user", + user_email="test@example.com", + event_group=Litellm_EntityType.USER, + event="soft_budget_crossed", + event_message="Soft Budget Crossed - Total Soft Budget: $100.0", + spend=105.0, + max_budget=None, + soft_budget=100.0, + ) + + with mock.patch.dict( + os.environ, + { + "PROXY_BASE_URL": "http://test.com", + }, + ): + await base_email_logger.send_soft_budget_alert_email(event) + + mock_send_email.assert_called_once() + call_args = mock_send_email.call_args[1] + assert "$100.0" in call_args["html_body"] # soft_budget + assert "$105.0" in call_args["html_body"] # spend + assert "Maximum Budget" not in call_args["html_body"] # max_budget should not be shown + + +@pytest.mark.asyncio +async def test_budget_alerts_soft_budget_crossed( + base_email_logger, mock_send_email +): + """Test that budget_alerts sends email when soft budget is crossed""" + user_info = CallInfo( + user_id="test_user", + user_email="test@example.com", + spend=105.0, + max_budget=200.0, + soft_budget=100.0, + event_group=Litellm_EntityType.USER, + ) + + # Mock the cache to return None (no previous alert sent) + 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="soft_budget", user_info=user_info) + + # Verify email was sent + mock_send_email.assert_called_once() + call_args = mock_send_email.call_args[1] + assert call_args["to_email"] == ["test@example.com"] + + # Verify cache was set to prevent duplicate alerts + mock_cache.async_set_cache.assert_called_once() + cache_call_args = mock_cache.async_set_cache.call_args[1] + assert cache_call_args["key"] == "email_budget_alerts:soft_budget_crossed:test_user" + assert cache_call_args["value"] == "SENT" + assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL + + +@pytest.mark.asyncio +async def test_budget_alerts_soft_budget_not_crossed( + base_email_logger, mock_send_email +): + """Test that budget_alerts does not send email when soft budget is not crossed""" + user_info = CallInfo( + user_id="test_user", + user_email="test@example.com", + spend=50.0, + max_budget=200.0, + soft_budget=100.0, + event_group=Litellm_EntityType.USER, + ) + + mock_cache = mock.AsyncMock() + base_email_logger.internal_usage_cache = mock_cache + + await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info) + + # Verify email was NOT sent + mock_send_email.assert_not_called() + + +@pytest.mark.asyncio +async def test_budget_alerts_soft_budget_duplicate_prevention( + base_email_logger, mock_send_email +): + """Test that budget_alerts does not send duplicate alerts within TTL period""" + user_info = CallInfo( + user_id="test_user", + user_email="test@example.com", + spend=105.0, + max_budget=200.0, + soft_budget=100.0, + event_group=Litellm_EntityType.USER, + ) + + # Mock the cache to return "SENT" (previous alert already sent) + mock_cache = mock.AsyncMock() + mock_cache.async_get_cache = mock.AsyncMock(return_value="SENT") + base_email_logger.internal_usage_cache = mock_cache + + await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info) + + # Verify email was NOT sent (duplicate prevention) + mock_send_email.assert_not_called() + + +@pytest.mark.asyncio +async def test_budget_alerts_no_budgets( + base_email_logger, mock_send_email +): + """Test that budget_alerts returns early when no budgets are set""" + user_info = CallInfo( + user_id="test_user", + user_email="test@example.com", + spend=50.0, + max_budget=None, + soft_budget=None, + event_group=Litellm_EntityType.USER, + ) + + await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info) + + # Verify email was NOT sent + mock_send_email.assert_not_called() + + +@pytest.mark.asyncio +async def test_budget_alerts_uses_token_for_cache_key( + base_email_logger, mock_send_email +): + """Test that budget_alerts uses token for cache key when available""" + user_info = CallInfo( + user_id="test_user", + user_email="test@example.com", + token="hashed_token_123", + spend=105.0, + max_budget=200.0, + soft_budget=100.0, + event_group=Litellm_EntityType.KEY, + ) + + # Mock the cache to return None (no previous alert sent) + 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="soft_budget", user_info=user_info) + + # Verify cache key uses token instead of user_id + mock_cache.async_set_cache.assert_called_once() + cache_call_args = mock_cache.async_set_cache.call_args[1] + assert cache_call_args["key"] == "email_budget_alerts:soft_budget_crossed:hashed_token_123" + + +@pytest.mark.asyncio +async def test_get_email_params_soft_budget_crossed( + base_email_logger, mock_lookup_user_email +): + """Test that _get_email_params handles soft_budget_crossed event correctly""" + with mock.patch.dict( + os.environ, + { + "PROXY_BASE_URL": "http://test.com", + }, + ): + result = await base_email_logger._get_email_params( + email_event=EmailEvent.soft_budget_crossed, + user_email="test@example.com", + event_message="Soft Budget Crossed - Total Soft Budget: $100.0", + ) + + # Should use default subject template for soft_budget_crossed + assert result.subject == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0" + assert result.recipient_email == "test@example.com" + assert result.base_url == "http://test.com" + + +@pytest.mark.asyncio +async def test_budget_alerts_max_budget_alert_crossed( + base_email_logger, mock_send_email +): + """Test that budget_alerts sends email when max budget alert threshold is crossed""" + 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"] + assert "Max Budget Alert" in call_args["subject"] + + mock_cache.async_set_cache.assert_called_once() + 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 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 3d4b68ce441..807559207e6 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -14,9 +14,11 @@ import pytest import litellm from litellm.proxy._types import ( + CallInfo, LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable, + Litellm_EntityType, LitellmUserRoles, ProxyErrorTypes, ProxyException, @@ -27,6 +29,8 @@ from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _can_object_call_vector_stores, _get_team_db_check, + _virtual_key_max_budget_alert_check, + _virtual_key_soft_budget_check, get_user_object, vector_store_access_check, ) @@ -988,3 +992,288 @@ async def test_reject_clientside_metadata_tags_non_llm_route(): ) assert result is True + + +@pytest.mark.asyncio +async def test_virtual_key_soft_budget_check_with_user_obj(): + """Test _virtual_key_soft_budget_check includes user_email when user_obj is provided""" + 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 + assert type == "soft_budget" + assert isinstance(user_info, CallInfo) + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=100.0, + soft_budget=50.0, + user_id="test-user", + team_id="test-team", + team_alias="test-team-alias", + org_id="test-org", + key_alias="test-key", + max_budget=200.0, + ) + + user_obj = LiteLLM_UserTable( + user_id="test-user", + user_email="test@example.com", + max_budget=None, + ) + + proxy_logging_obj = MockProxyLogging() + + await _virtual_key_soft_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + 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.user_email == "test@example.com" + assert captured_call_info.token == "test-token" + assert captured_call_info.spend == 100.0 + assert captured_call_info.soft_budget == 50.0 + assert captured_call_info.max_budget == 200.0 + assert captured_call_info.user_id == "test-user" + assert captured_call_info.team_id == "test-team" + assert captured_call_info.team_alias == "test-team-alias" + assert captured_call_info.organization_id == "test-org" + assert captured_call_info.key_alias == "test-key" + assert captured_call_info.event_group == Litellm_EntityType.KEY + + +@pytest.mark.asyncio +async def test_virtual_key_soft_budget_check_without_user_obj(): + """Test _virtual_key_soft_budget_check sets user_email to None when user_obj is not provided""" + 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 + assert type == "soft_budget" + assert isinstance(user_info, CallInfo) + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=100.0, + soft_budget=50.0, + user_id="test-user", + team_id="test-team", + key_alias="test-key", + ) + + proxy_logging_obj = MockProxyLogging() + + await _virtual_key_soft_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=None, + ) + + await asyncio.sleep(0.1) + + assert alert_triggered is True + assert captured_call_info is not None + assert captured_call_info.user_email is None + + +@pytest.mark.parametrize( + "spend, soft_budget, expect_alert", + [ + (100.0, 50.0, True), # Over soft budget + (50.0, 50.0, True), # At soft budget + (25.0, 50.0, False), # Under soft budget + (100.0, None, False), # No soft budget set + ], +) +@pytest.mark.asyncio +async def test_virtual_key_soft_budget_check_scenarios( + spend, soft_budget, expect_alert +): + """Test _virtual_key_soft_budget_check with various spend and soft_budget scenarios""" + alert_triggered = False + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal alert_triggered + alert_triggered = True + assert type == "soft_budget" + assert isinstance(user_info, CallInfo) + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=spend, + soft_budget=soft_budget, + user_id="test-user", + key_alias="test-key", + ) + + proxy_logging_obj = MockProxyLogging() + + await _virtual_key_soft_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=None, + ) + + await asyncio.sleep(0.1) + + assert ( + alert_triggered == expect_alert + ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + + +@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""" + 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 + assert type == "max_budget_alert" + assert isinstance(user_info, CallInfo) + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=90.0, + max_budget=100.0, + user_id="test-user", + team_id="test-team", + team_alias="test-team-alias", + org_id="test-org", + key_alias="test-key", + soft_budget=50.0, + ) + + user_obj = LiteLLM_UserTable( + user_id="test-user", + user_email="test@example.com", + max_budget=None, + ) + + proxy_logging_obj = MockProxyLogging() + + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + 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.user_email == "test@example.com" + assert captured_call_info.token == "test-token" + assert captured_call_info.spend == 90.0 + assert captured_call_info.max_budget == 100.0 + assert captured_call_info.soft_budget == 50.0 + assert captured_call_info.user_id == "test-user" + assert captured_call_info.team_id == "test-team" + assert captured_call_info.team_alias == "test-team-alias" + assert captured_call_info.organization_id == "test-org" + assert captured_call_info.key_alias == "test-key" + assert captured_call_info.event_group == Litellm_EntityType.KEY + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_without_user_obj(): + """Test _virtual_key_max_budget_alert_check sets user_email to None when user_obj is not provided""" + 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 + assert type == "max_budget_alert" + assert isinstance(user_info, CallInfo) + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=90.0, + max_budget=100.0, + user_id="test-user", + team_id="test-team", + key_alias="test-key", + ) + + proxy_logging_obj = MockProxyLogging() + + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=None, + ) + + await asyncio.sleep(0.1) + + assert alert_triggered is True + assert captured_call_info is not None + assert captured_call_info.user_email is None + + +@pytest.mark.parametrize( + "spend, max_budget, expect_alert", + [ + (80.0, 100.0, True), # At 80% threshold (alert threshold) + (90.0, 100.0, True), # Above threshold, below max_budget + (79.0, 100.0, False), # Below threshold + (100.0, 100.0, False), # At max_budget (not below, so no alert) + (110.0, 100.0, False), # Above max_budget (already exceeded) + (100.0, None, False), # No max_budget set + (0.0, 100.0, False), # Spend is 0 + ], +) +@pytest.mark.asyncio +async def test_virtual_key_max_budget_alert_check_scenarios( + spend, max_budget, expect_alert +): + """Test _virtual_key_max_budget_alert_check with various spend and max_budget scenarios""" + alert_triggered = False + + class MockProxyLogging: + async def budget_alerts(self, type, user_info): + nonlocal alert_triggered + alert_triggered = True + assert type == "max_budget_alert" + assert isinstance(user_info, CallInfo) + + valid_token = UserAPIKeyAuth( + token="test-token", + spend=spend, + max_budget=max_budget, + user_id="test-user", + key_alias="test-key", + ) + + proxy_logging_obj = MockProxyLogging() + + await _virtual_key_max_budget_alert_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + user_obj=None, + ) + + await asyncio.sleep(0.1) + + assert ( + alert_triggered == expect_alert + ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 46b828e7940..7dba7c99916 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,7 +1,7 @@ import json import os import sys -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from jsonschema import validate @@ -2604,6 +2604,132 @@ class TestIsCachedMessage: assert is_cached_message(message) is False +@pytest.mark.asyncio +class TestProxyLoggingBudgetAlerts: + """Test budget_alerts method in ProxyLogging class.""" + + async def test_budget_alerts_when_alerting_is_none(self): + """Test that budget_alerts returns early when alerting is None.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alerting = None + proxy_logging.slack_alerting_instance = AsyncMock() + proxy_logging.email_logging_instance = AsyncMock() + + user_info = MagicMock() + + # Should return without calling any alerting instances + await proxy_logging.budget_alerts(type="user_budget", user_info=user_info) + + # Verify no calls were made + proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called() + proxy_logging.email_logging_instance.budget_alerts.assert_not_called() + + async def test_budget_alerts_with_slack_only(self): + """Test that budget_alerts calls slack_alerting_instance when slack is in alerting.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alerting = ["slack"] + proxy_logging.slack_alerting_instance = AsyncMock() + + user_info = MagicMock() + + await proxy_logging.budget_alerts(type="token_budget", user_info=user_info) + + proxy_logging.slack_alerting_instance.budget_alerts.assert_called_once_with( + type="token_budget", user_info=user_info + ) + + async def test_budget_alerts_with_email_only(self): + """Test that budget_alerts calls email_logging_instance when email is in alerting.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alerting = ["email"] + proxy_logging.email_logging_instance = AsyncMock() + + user_info = MagicMock() + + await proxy_logging.budget_alerts(type="team_budget", user_info=user_info) + + proxy_logging.email_logging_instance.budget_alerts.assert_called_once_with( + type="team_budget", user_info=user_info + ) + + async def test_budget_alerts_with_email_when_instance_is_none(self): + """Test that budget_alerts does not call email_logging_instance when it is None.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alerting = ["email"] + proxy_logging.email_logging_instance = None + + user_info = MagicMock() + + # Should not raise an error + await proxy_logging.budget_alerts(type="organization_budget", user_info=user_info) + + async def test_budget_alerts_with_both_slack_and_email(self): + """Test that budget_alerts calls both slack and email instances when both are in alerting.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alerting = ["slack", "email"] + proxy_logging.slack_alerting_instance = AsyncMock() + proxy_logging.email_logging_instance = AsyncMock() + + user_info = MagicMock() + + await proxy_logging.budget_alerts(type="proxy_budget", user_info=user_info) + + proxy_logging.slack_alerting_instance.budget_alerts.assert_called_once_with( + type="proxy_budget", user_info=user_info + ) + proxy_logging.email_logging_instance.budget_alerts.assert_called_once_with( + type="proxy_budget", user_info=user_info + ) + + @pytest.mark.parametrize( + "alert_type", + [ + "token_budget", + "user_budget", + "soft_budget", + "team_budget", + "organization_budget", + "proxy_budget", + "projected_limit_exceeded", + ], + ) + async def test_budget_alerts_with_all_alert_types(self, alert_type): + """Test that budget_alerts works with all supported alert types.""" + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alerting = ["slack", "email"] + proxy_logging.slack_alerting_instance = AsyncMock() + proxy_logging.email_logging_instance = AsyncMock() + + user_info = MagicMock() + + await proxy_logging.budget_alerts(type=alert_type, user_info=user_info) + + proxy_logging.slack_alerting_instance.budget_alerts.assert_called_once_with( + type=alert_type, user_info=user_info + ) + proxy_logging.email_logging_instance.budget_alerts.assert_called_once_with( + type=alert_type, user_info=user_info + ) + + def test_azure_ai_claude_provider_config(): """Test that Azure AI Claude models return AzureAnthropicConfig for proper tool transformation.""" from litellm import AzureAnthropicConfig, AzureAIStudioConfig