diff --git a/litellm/integrations/SlackAlerting/budget_alert_types.py b/litellm/integrations/SlackAlerting/budget_alert_types.py index ea80b258540..2a19ec0b7fa 100644 --- a/litellm/integrations/SlackAlerting/budget_alert_types.py +++ b/litellm/integrations/SlackAlerting/budget_alert_types.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Literal -from litellm.proxy._types import CallInfo +from litellm.proxy._types import CallInfo, Litellm_EntityType class BaseBudgetAlertType(ABC): @@ -31,6 +31,8 @@ class SoftBudgetAlert(BaseBudgetAlertType): return "Soft Budget Crossed: " def get_id(self, user_info: CallInfo) -> str: + if user_info.event_group == Litellm_EntityType.TEAM: + return user_info.team_id or "default_id" return user_info.token or "default_id" 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 efb8c1c4b28..52b7cc983a7 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_budget_alert_types.py @@ -28,6 +28,31 @@ class TestSoftBudgetAlert: result = alert.get_id(user_info) assert result == "default_id" + def test_get_id_returns_team_id_for_team_event_group(self): + """Team soft budget alerts dedupe by team, not by the calling key's token""" + alert = SoftBudgetAlert() + user_info = CallInfo( + spend=120.0, + token="test_token_123", + team_id="team_456", + event_group=Litellm_EntityType.TEAM, + ) + + result = alert.get_id(user_info) + assert result == "team_456" + + def test_get_id_returns_default_id_for_team_event_group_without_team_id(self): + alert = SoftBudgetAlert() + user_info = CallInfo( + spend=120.0, + token="test_token_123", + team_id=None, + event_group=Litellm_EntityType.TEAM, + ) + + result = alert.get_id(user_info) + assert result == "default_id" + def test_get_id_with_empty_token(self): """Test that get_id returns 'default_id' when token is empty string""" alert = SoftBudgetAlert()