mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
Dedupe team soft budget alerts by team_id instead of token
_team_soft_budget_check sends type="soft_budget" alerts with event_group=TEAM, but SoftBudgetAlert.get_id always returned the request token. The alert cache key was therefore scoped per virtual key, so every active key in a team over its soft budget fired its own alert within budget_alert_ttl. Branch on event_group so team-level alerts dedupe by team_id, matching TeamBudgetAlert, while key and project level alerts keep per-token dedupe. Fixes #27398.
This commit is contained in:
parent
e15b37a18e
commit
9a5ce6d395
2 changed files with 28 additions and 1 deletions
|
|
@ -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"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue