From b4e1b6e52681f9da60e57bf355ce95e6ca6640a2 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 23 Sep 2026 02:58:45 +0000 Subject: [PATCH] fix(alerting): label team member budget crossings as team member budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/SlackAlerting/Readme.md | 4 ++-- .../SlackAlerting/budget_alert_types.py | 20 ++++++++++--------- .../SlackAlerting/slack_alerting.py | 2 +- .../SlackAlerting/test_budget_alert_types.py | 9 +++++++++ .../SlackAlerting/test_slack_alerting.py | 6 +++--- 5 files changed, 26 insertions(+), 15 deletions(-) diff --git a/litellm/integrations/SlackAlerting/Readme.md b/litellm/integrations/SlackAlerting/Readme.md index 1719941d0da..aef8074f410 100644 --- a/litellm/integrations/SlackAlerting/Readme.md +++ b/litellm/integrations/SlackAlerting/Readme.md @@ -15,7 +15,7 @@ The `budget_alert_types.py` module provides a flexible framework for handling di - `BaseBudgetAlertType`: An abstract base class with abstract methods that all alert types must implement: - `get_event_group()`: Returns the Litellm_EntityType for the alert - - `get_event_message()`: Returns the message prefix for the alert + - `get_event_message(user_info)`: Returns the message prefix for the alert - `get_id(user_info)`: Returns the ID to use for caching/tracking the alert Concrete implementations include: @@ -36,7 +36,7 @@ budget_alert_class = get_budget_alert_type("user_budget") # Use the handler methods event_group = budget_alert_class.get_event_group() # Returns Litellm_EntityType.USER -event_message = budget_alert_class.get_event_message() # Returns "User Budget: " +event_message = budget_alert_class.get_event_message(user_info) # Returns "User Budget: " cache_id = budget_alert_class.get_id(user_info) # Returns user_id ``` diff --git a/litellm/integrations/SlackAlerting/budget_alert_types.py b/litellm/integrations/SlackAlerting/budget_alert_types.py index 4fe833acecc..a32964d17e4 100644 --- a/litellm/integrations/SlackAlerting/budget_alert_types.py +++ b/litellm/integrations/SlackAlerting/budget_alert_types.py @@ -8,7 +8,7 @@ class BaseBudgetAlertType(ABC): """Base class for different budget alert types""" @abstractmethod - def get_event_message(self) -> str: + def get_event_message(self, user_info: CallInfo) -> str: """Return the event message for this alert type""" @abstractmethod @@ -17,7 +17,7 @@ class BaseBudgetAlertType(ABC): class ProxyBudgetAlert(BaseBudgetAlertType): - def get_event_message(self) -> str: + def get_event_message(self, user_info: CallInfo) -> str: return "Proxy Budget: " def get_id(self, user_info: CallInfo) -> str: @@ -25,7 +25,7 @@ class ProxyBudgetAlert(BaseBudgetAlertType): class SoftBudgetAlert(BaseBudgetAlertType): - def get_event_message(self) -> str: + def get_event_message(self, user_info: CallInfo) -> str: return "Soft Budget Crossed: " def get_id(self, user_info: CallInfo) -> str: @@ -35,7 +35,7 @@ class SoftBudgetAlert(BaseBudgetAlertType): class UserBudgetAlert(BaseBudgetAlertType): - def get_event_message(self) -> str: + def get_event_message(self, user_info: CallInfo) -> str: return "User Budget: " def get_id(self, user_info: CallInfo) -> str: @@ -43,7 +43,7 @@ class UserBudgetAlert(BaseBudgetAlertType): class TeamBudgetAlert(BaseBudgetAlertType): - def get_event_message(self) -> str: + def get_event_message(self, user_info: CallInfo) -> str: return "Team Budget: " def get_id(self, user_info: CallInfo) -> str: @@ -51,7 +51,7 @@ class TeamBudgetAlert(BaseBudgetAlertType): class OrganizationBudgetAlert(BaseBudgetAlertType): - def get_event_message(self) -> str: + def get_event_message(self, user_info: CallInfo) -> str: return "Organization Budget: " def get_id(self, user_info: CallInfo) -> str: @@ -59,7 +59,9 @@ class OrganizationBudgetAlert(BaseBudgetAlertType): class TokenBudgetAlert(BaseBudgetAlertType): - def get_event_message(self) -> str: + def get_event_message(self, user_info: CallInfo) -> str: + if user_info.event_group == Litellm_EntityType.TEAM_MEMBER: + return "Team Member Budget: " return "Key Budget: " def get_id(self, user_info: CallInfo) -> str: @@ -69,7 +71,7 @@ class TokenBudgetAlert(BaseBudgetAlertType): class ProjectedLimitExceededAlert(BaseBudgetAlertType): - def get_event_message(self) -> str: + def get_event_message(self, user_info: CallInfo) -> str: return "Key Budget: Projected Limit Exceeded" def get_id(self, user_info: CallInfo) -> str: @@ -77,7 +79,7 @@ class ProjectedLimitExceededAlert(BaseBudgetAlertType): class ProjectBudgetAlert(BaseBudgetAlertType): - def get_event_message(self) -> str: + def get_event_message(self, user_info: CallInfo) -> str: return "Project Budget: " def get_id(self, user_info: CallInfo) -> str: diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 8d0d044ff93..46c19771510 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -551,7 +551,7 @@ class SlackAlerting(CustomBatchLogger): budget_alert_class: Final = get_budget_alert_type(type) _id: Final = budget_alert_class.get_id(user_info) user_info_str: Final = self._get_user_info_str(user_info) - event_message = budget_alert_class.get_event_message() + event_message = budget_alert_class.get_event_message(user_info) # Set default event unless we're in projected_limit_exceeded event: ( diff --git a/tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py b/tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py index 99dcdd9d4c3..22e0f81a474 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py @@ -94,3 +94,12 @@ class TestTokenBudgetAlert: ) assert alert.get_id(user_info) == "hashed_key" + assert alert.get_event_message(user_info) == "Key Budget: " + + def test_get_event_message_labels_team_member_alerts(self): + alert = TokenBudgetAlert() + user_info = CallInfo( + spend=8.0, max_budget=10.0, user_id="member_1", team_id="team_a", event_group=Litellm_EntityType.TEAM_MEMBER + ) + + assert alert.get_event_message(user_info) == "Team Member Budget: " diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 2d5eb78950c..0ad1b8746e1 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -67,7 +67,6 @@ class TestSlackAlerting(unittest.TestCase): def test_get_event_and_event_message_max_budget(self): event = None - event_message = get_budget_alert_type("user_budget").get_event_message() # Test case 1: When spend exceeds max_budget user_info = CallInfo( @@ -76,32 +75,33 @@ class TestSlackAlerting(unittest.TestCase): soft_budget=None, event_group=Litellm_EntityType.KEY, ) + event_message = get_budget_alert_type("user_budget").get_event_message(user_info) event, event_message = self.slack_alerting._get_event_and_event_message( user_info=user_info, event=event, event_message=event_message ) self.assertEqual(event, "budget_crossed") self.assertTrue("Budget Crossed" in event_message) - event_message = get_budget_alert_type("user_budget").get_event_message() user_info = CallInfo( max_budget=100.0, spend=95.0, soft_budget=None, event_group=Litellm_EntityType.KEY, ) + event_message = get_budget_alert_type("user_budget").get_event_message(user_info) event, event_message = self.slack_alerting._get_event_and_event_message( user_info=user_info, event=event, event_message=event_message ) self.assertEqual(event, "threshold_crossed") self.assertEqual(event_message, "User Budget: 5% or less of budget remaining") - event_message = get_budget_alert_type("user_budget").get_event_message() user_info = CallInfo( max_budget=100.0, spend=85.0, soft_budget=None, event_group=Litellm_EntityType.KEY, ) + event_message = get_budget_alert_type("user_budget").get_event_message(user_info) event, event_message = self.slack_alerting._get_event_and_event_message( user_info=user_info, event=event, event_message=event_message )