diff --git a/litellm/proxy/common_utils/model_deprecation_notifications.py b/litellm/proxy/common_utils/model_deprecation_notifications.py index d54183fb452..f912df11138 100644 --- a/litellm/proxy/common_utils/model_deprecation_notifications.py +++ b/litellm/proxy/common_utils/model_deprecation_notifications.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import BaseModel, ConfigDict, Field, ValidationError +from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.constants import EMAIL_MODEL_DEPRECATION_LOCK_ID from litellm.integrations.email_alerting import ( @@ -288,6 +289,9 @@ def render_model_deprecation_email( return subject, body + EMAIL_FOOTER +_DEPRECATION_EMAIL_TARGET: Final = "model_deprecation_emails" + + _ENTERPRISE_BASE_EMAIL_LOGGER: Final = ( "litellm_enterprise.enterprise_callbacks.send_emails.base_email", "BaseEmailLogger", @@ -351,6 +355,7 @@ async def _stamp_pass(cache: DeprecationEmailCache) -> None: ) +@with_service_target(_DEPRECATION_EMAIL_TARGET) async def email_pass_done_today(cache: DeprecationEmailCache) -> bool: return (await cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value)) is not None @@ -390,6 +395,7 @@ async def _send_team_notification(notification: TeamNotification, ctx: Deprecati return True +@with_service_target(_DEPRECATION_EMAIL_TARGET) async def send_model_deprecation_emails(ctx: DeprecationEmailContext) -> int: """One pass: resolve affected teams, claim the day only when there is something to send, deliver, then stamp diff --git a/tests/unit/proxy/common_utils/test_model_deprecation_notifications.py b/tests/unit/proxy/common_utils/test_model_deprecation_notifications.py index 01d6854999f..e3bd45f28a0 100644 --- a/tests/unit/proxy/common_utils/test_model_deprecation_notifications.py +++ b/tests/unit/proxy/common_utils/test_model_deprecation_notifications.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock import pytest +from litellm._internal_context import current_service_target from litellm.caching.caching import DualCache from litellm.constants import EMAIL_MODEL_DEPRECATION_LOCK_ID from litellm.models.team import LiteLLM_TeamTable @@ -469,6 +470,20 @@ class _RecordingCache: self.ttls[key] = ttl +class _TargetRecordingCache(_RecordingCache): + def __init__(self): + super().__init__() + self.targets = [] + + async def async_get_cache(self, *, key: str) -> object: + self.targets.append(current_service_target()) + return await super().async_get_cache(key=key) + + async def async_set_cache(self, *, key: str, value: float, ttl: float) -> None: + self.targets.append(current_service_target()) + await super().async_set_cache(key=key, value=value, ttl=ttl) + + def _team_row(team_id: str, alias: str = "Team") -> Mapping[str, object]: return { "team_id": team_id, @@ -707,6 +722,18 @@ async def test_should_report_whether_the_daily_pass_already_ran(): assert await email_pass_done_today(cache) is True +@pytest.mark.asyncio +async def test_should_tag_every_cache_read_and_write_with_the_deprecation_email_family(): + cache: Final = _TargetRecordingCache() + ctx: Final = _context(_router([DEAD_DEPLOYMENT]), _prisma(TWO_TEAMS[:1], TWO_ADMINS), _Deliverer(), cache) + + assert await email_pass_done_today(cache) is False + assert await send_model_deprecation_emails(ctx) == 1 + assert len(cache.targets) >= 3 + assert set(cache.targets) == {"model_deprecation_emails"} + assert current_service_target() is None + + class TestMakeEmailDeliverer: @pytest.mark.asyncio async def test_should_use_the_configured_email_logger_once_per_recipient(self):