scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-28 01:32:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 43

fix(alerting): label team member budget crossings as team member budget

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-23 02:58:45 +00:00
parent 30e1ce2819
commit b4e1b6e526
5 changed files with 26 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/integrations/SlackAlerting/Readme.md
View file

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

20
litellm/integrations/SlackAlerting/budget_alert_types.py
View file

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

2
litellm/integrations/SlackAlerting/slack_alerting.py
View file

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

9
tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py
View file

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

6
tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py
View file

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