mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
Merge 864d845990 into 60fc121ce8
This commit is contained in:
commit
80cc41079a
14 changed files with 1911 additions and 64 deletions
|
|
@ -1895,6 +1895,7 @@ USER_SPEND_ALERTS_JOB_ID: Final = "user_spend_alerts_job"
|
|||
PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job"
|
||||
SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report"
|
||||
SLACK_MODEL_DEPRECATION_LOCK_ID: Final = "slack_model_deprecation_warning"
|
||||
EMAIL_MODEL_DEPRECATION_LOCK_ID: Final = "email_model_deprecation_warning"
|
||||
SPEND_LOG_RUN_LOOPS: Final = get_env_int("SPEND_LOG_RUN_LOOPS", 500)
|
||||
SPEND_LOG_CLEANUP_BATCH_SIZE: Final = get_env_int("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)
|
||||
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = get_env_int("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import datetime
|
|||
import os
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
|
|
@ -55,6 +55,7 @@ from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
|||
from litellm.types.proxy.model_deprecation import (
|
||||
DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
|
||||
DEPRECATION_IDLE_POLL_SECONDS,
|
||||
EmailSender,
|
||||
)
|
||||
|
||||
from ..email_templates.templates import *
|
||||
|
|
@ -84,6 +85,29 @@ def _proxy_llm_router() -> Router | None:
|
|||
return llm_router
|
||||
|
||||
|
||||
def _proxy_prisma_client() -> "PrismaClient | None":
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
return prisma_client
|
||||
|
||||
|
||||
def _proxy_email_logger() -> EmailSender | None:
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
return proxy_logging_obj.email_logging_instance
|
||||
|
||||
|
||||
async def _pass_failed(pass_run: Awaitable[object], name: str) -> bool:
|
||||
try:
|
||||
await pass_run
|
||||
except Exception as e: # noqa: BLE001 # a failed pass must not kill the loop
|
||||
verbose_proxy_logger.exception("Error in model deprecation %s loop: %s", name, e)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
_CHANNELS_ADAPTER: Final = TypeAdapter(tuple[str, ...])
|
||||
|
||||
_DAILY_REPORT_TARGET: Final = "daily_report_schedule"
|
||||
|
||||
|
||||
|
|
@ -101,7 +125,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
alert_types: list[AlertType] = DEFAULT_ALERT_TYPES,
|
||||
alert_to_webhook_url: dict[AlertType, list[str] | str]
|
||||
| None = None, # if user wants to separate alerts to diff channels
|
||||
alerting_args={},
|
||||
alerting_args: Mapping[str, object] | None = None,
|
||||
default_webhook_url: str | None = None,
|
||||
alert_type_config: dict[str, dict] | None = None,
|
||||
async_http_handler: AsyncHTTPHandler | None = None,
|
||||
|
|
@ -118,7 +142,9 @@ class SlackAlerting(CustomBatchLogger):
|
|||
)
|
||||
self.alert_to_webhook_url = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url)
|
||||
self.is_running = False
|
||||
self.alerting_args = SlackAlertingArgs(**alerting_args)
|
||||
self.alerting_args = (
|
||||
SlackAlertingArgs() if alerting_args is None else SlackAlertingArgs.model_validate(alerting_args)
|
||||
)
|
||||
self.default_webhook_url = default_webhook_url
|
||||
self.flush_lock = asyncio.Lock()
|
||||
self.periodic_started = False
|
||||
|
|
@ -132,6 +158,8 @@ class SlackAlerting(CustomBatchLogger):
|
|||
self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val
|
||||
self.digest_buckets: dict[str, DigestEntry] = {}
|
||||
self.digest_lock = asyncio.Lock()
|
||||
self.deprecation_alert_backoff_until: float = 0.0
|
||||
self.deprecation_email_backoff_until: float = 0.0
|
||||
super().__init__(**kwargs, flush_lock=self.flush_lock)
|
||||
|
||||
def _ensure_periodic_flush_task(self) -> None:
|
||||
|
|
@ -158,7 +186,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
if alert_types is not None:
|
||||
self.alert_types = alert_types
|
||||
if alerting_args is not None:
|
||||
self.alerting_args = SlackAlertingArgs(**alerting_args)
|
||||
self.alerting_args = SlackAlertingArgs.model_validate(alerting_args)
|
||||
self._ensure_periodic_flush_task()
|
||||
if alert_type_config is not None:
|
||||
for key, val in alert_type_config.items():
|
||||
|
|
@ -1101,8 +1129,14 @@ Model Info:
|
|||
async def model_removed_alert(self, model_name: str):
|
||||
pass
|
||||
|
||||
def _channels(self) -> tuple[str, ...]:
|
||||
return _CHANNELS_ADAPTER.validate_python(self.alerting or ())
|
||||
|
||||
def _deprecation_alerts_enabled(self) -> bool:
|
||||
return self.alerting is not None and AlertType.model_deprecation_warnings in self.alert_types
|
||||
channels: Final = self._channels()
|
||||
return (
|
||||
"slack" in channels or MS_TEAMS_ALERTING_DESTINATION in channels
|
||||
) and AlertType.model_deprecation_warnings in self.alert_types
|
||||
|
||||
async def send_model_deprecation_alert(
|
||||
self,
|
||||
|
|
@ -1174,25 +1208,96 @@ Model Info:
|
|||
return False
|
||||
return await self.send_model_deprecation_alert(llm_router=llm_router, pod_lock_manager=pod_lock_manager)
|
||||
|
||||
def _deprecation_emails_enabled(self) -> bool:
|
||||
return (
|
||||
"email" in self._channels()
|
||||
and AlertType.model_deprecation_warnings in self.alert_types
|
||||
and bool(self.alerting_args.model_deprecation_email_thresholds)
|
||||
)
|
||||
|
||||
async def _run_deprecation_email_pass(
|
||||
self,
|
||||
llm_router: Router | None,
|
||||
pod_lock_manager: "PodLockManager | None",
|
||||
get_prisma_client: Callable[[], "PrismaClient | None"],
|
||||
get_email_logger: Callable[[], EmailSender | None],
|
||||
send_emails: "Callable[..., Awaitable[int]] | None",
|
||||
) -> int:
|
||||
"""Email team admins the milestones crossed, resolving against the DB at most once a day"""
|
||||
if llm_router is None or not self._deprecation_emails_enabled():
|
||||
return 0
|
||||
|
||||
from litellm.proxy.common_utils.model_deprecation_notifications import (
|
||||
DeprecationEmailContext,
|
||||
email_pass_done_today,
|
||||
make_email_deliverer,
|
||||
send_model_deprecation_emails,
|
||||
)
|
||||
|
||||
if await email_pass_done_today(self.internal_usage_cache):
|
||||
return 0
|
||||
prisma_client: Final = get_prisma_client()
|
||||
if prisma_client is None:
|
||||
verbose_proxy_logger.debug("model_deprecation: no database connected, skipping email pass")
|
||||
return 0
|
||||
run: Final = send_emails or send_model_deprecation_emails
|
||||
return await run(
|
||||
DeprecationEmailContext(
|
||||
llm_router=llm_router,
|
||||
prisma_client=prisma_client,
|
||||
cache=self.internal_usage_cache,
|
||||
alerting_args=self.alerting_args,
|
||||
pod_lock_manager=pod_lock_manager,
|
||||
deliver=make_email_deliverer(get_email_logger()),
|
||||
)
|
||||
)
|
||||
|
||||
async def _run_deprecation_passes(
|
||||
self,
|
||||
get_llm_router: Callable[[], Router | None],
|
||||
pod_lock_manager: "PodLockManager | None",
|
||||
get_prisma_client: Callable[[], "PrismaClient | None"],
|
||||
get_email_logger: Callable[[], EmailSender | None],
|
||||
send_emails: "Callable[..., Awaitable[int]] | None",
|
||||
now: Callable[[], float] = time.time,
|
||||
) -> None:
|
||||
"""Run the Slack pass then the email pass; a pass that raised is skipped for a day while the other keeps polling"""
|
||||
llm_router: Final = get_llm_router()
|
||||
current: Final = now()
|
||||
if current >= self.deprecation_alert_backoff_until and await _pass_failed(
|
||||
self._run_deprecation_alert_pass(llm_router, pod_lock_manager), "alert"
|
||||
):
|
||||
self.deprecation_alert_backoff_until = current + DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS
|
||||
if current >= self.deprecation_email_backoff_until and await _pass_failed(
|
||||
self._run_deprecation_email_pass(
|
||||
llm_router, pod_lock_manager, get_prisma_client, get_email_logger, send_emails
|
||||
),
|
||||
"email",
|
||||
):
|
||||
self.deprecation_email_backoff_until = current + DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS
|
||||
|
||||
async def run_scheduled_deprecation_check(
|
||||
self,
|
||||
get_llm_router: Callable[[], Router | None] = _proxy_llm_router,
|
||||
pod_lock_manager: "PodLockManager | None" = None,
|
||||
get_prisma_client: Callable[[], "PrismaClient | None"] = _proxy_prisma_client,
|
||||
get_email_logger: Callable[[], EmailSender | None] = _proxy_email_logger,
|
||||
send_emails: "Callable[..., Awaitable[int]] | None" = None,
|
||||
sleep: Callable[[float], Awaitable[None]] | None = None,
|
||||
) -> None:
|
||||
"""Poll every pass for a loaded router, the alert being on, and no alert in the last day, then alert
|
||||
"""Poll for a loaded router and run both passes every poll interval
|
||||
|
||||
A pass that could not alert (no router yet, alert type off, a sibling pod holds the daily lock, or a
|
||||
redis blip at claim time) is retried on the next poll instead of costing a day, while a pass that
|
||||
raised (a missing webhook, say) backs off a full day so a misconfiguration logs once, not every poll
|
||||
A pass that could not act (no router yet, alert type off, a sibling pod holds the daily lock, or a
|
||||
redis blip at claim time) simply runs again next poll, while a pass that raised (a missing webhook,
|
||||
say) is skipped for a full day so a misconfiguration logs once, not every poll, without stalling
|
||||
the other pass
|
||||
"""
|
||||
sleep_fn: Final = sleep or asyncio.sleep
|
||||
while True:
|
||||
try:
|
||||
await self._run_deprecation_alert_pass(get_llm_router(), pod_lock_manager)
|
||||
except Exception as e: # noqa: BLE001 # a failed alert must not kill the loop
|
||||
verbose_proxy_logger.exception("Error in model deprecation alert loop: %s", e)
|
||||
await asyncio.sleep(DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS)
|
||||
continue
|
||||
await asyncio.sleep(DEPRECATION_IDLE_POLL_SECONDS)
|
||||
await self._run_deprecation_passes(
|
||||
get_llm_router, pod_lock_manager, get_prisma_client, get_email_logger, send_emails
|
||||
)
|
||||
await sleep_fn(DEPRECATION_IDLE_POLL_SECONDS)
|
||||
|
||||
async def send_webhook_alert(self, webhook_event: WebhookEvent) -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3,10 +3,13 @@ Functions for sending Email Alerts
|
|||
"""
|
||||
|
||||
import os
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm.models.team import LiteLLM_TeamTable
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.types.integrations.slack_alerting import WebhookEvent
|
||||
|
||||
# we use this for the email header, please send a test email if you change this. verify it looks good on email
|
||||
|
|
@ -65,6 +68,18 @@ async def get_all_team_member_emails(team_id: str | None = None) -> list:
|
|||
return emails
|
||||
|
||||
|
||||
async def get_team_admin_emails(team: LiteLLM_TeamTable, prisma_client: object) -> tuple[str, ...]:
|
||||
"""Emails of the members holding the admin role: the user record's email, else the inline one, deduped"""
|
||||
admins: Final = tuple(member for member in team.members_with_roles if member.role == "admin")
|
||||
admin_ids: Final = frozenset(member.user_id for member in admins if member.user_id)
|
||||
rows: Final = await UserRepository(prisma_client).find_many_in("user_id", sorted(admin_ids))
|
||||
record_emails: Final = MappingProxyType({row.user_id: row.user_email for row in rows if row.user_email})
|
||||
inline_emails: Final = tuple(
|
||||
member.user_email for member in admins if member.user_email and member.user_id not in record_emails
|
||||
)
|
||||
return tuple(dict.fromkeys((*inline_emails, *record_emails.values())))
|
||||
|
||||
|
||||
async def send_team_budget_alert(webhook_event: WebhookEvent) -> bool:
|
||||
"""
|
||||
Send an Email Alert to All Team Members when the Team Budget is crossed
|
||||
|
|
|
|||
|
|
@ -0,0 +1,42 @@
|
|||
"""Template for the model deprecation digest sent to team admins"""
|
||||
|
||||
from typing import Final
|
||||
|
||||
MODEL_DEPRECATION_EMAIL_TEMPLATE: Final = """
|
||||
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" /> <br/><br/>
|
||||
|
||||
<p>Hi team admin,</p>
|
||||
|
||||
<p>The following models available to team <b>{team_name}</b> are being deprecated by their provider.
|
||||
Requests to a deprecated model may start failing on its deprecation date.</p>
|
||||
|
||||
<table border="1" cellpadding="6" cellspacing="0" style="border-collapse: collapse;">
|
||||
<thead>
|
||||
<tr>
|
||||
<th align="left">Model</th>
|
||||
<th align="left">Provider</th>
|
||||
<th align="left">Deprecation date</th>
|
||||
<th align="left">Days left</th>
|
||||
<th align="left">Status</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{model_rows}
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
<p>Plan a migration to a supported model. See
|
||||
<a href="https://docs.litellm.ai/docs/proxy/model_management">the LiteLLM model management docs</a> for guidance.</p>
|
||||
|
||||
<p>If you have any questions, please send an email to {email_support_contact}</p>
|
||||
|
||||
<p>Best,<br/>The LiteLLM team</p>
|
||||
"""
|
||||
|
||||
MODEL_DEPRECATION_EMAIL_ROW_TEMPLATE: Final = """ <tr>
|
||||
<td><code>{model_name}</code></td>
|
||||
<td>{provider}</td>
|
||||
<td>{deprecation_date}</td>
|
||||
<td>{days_left}</td>
|
||||
<td>{status}</td>
|
||||
</tr>"""
|
||||
426
litellm/proxy/common_utils/model_deprecation_notifications.py
Normal file
426
litellm/proxy/common_utils/model_deprecation_notifications.py
Normal file
|
|
@ -0,0 +1,426 @@
|
|||
"""Email team admins about the deprecating models their team can reach"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import date
|
||||
from types import MappingProxyType
|
||||
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 (
|
||||
LITELLM_LOGO_URL,
|
||||
LITELLM_SUPPORT_CONTACT,
|
||||
get_team_admin_emails,
|
||||
)
|
||||
from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER
|
||||
from litellm.integrations.email_templates.model_deprecation_email import (
|
||||
MODEL_DEPRECATION_EMAIL_ROW_TEMPLATE,
|
||||
MODEL_DEPRECATION_EMAIL_TEMPLATE,
|
||||
)
|
||||
from litellm.models.team import LiteLLM_TeamTable
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.auth.auth_checks import can_team_access_model
|
||||
from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.types.integrations.slack_alerting import SlackAlertingArgs, SlackAlertingCacheKeys
|
||||
from litellm.types.proxy.model_deprecation import (
|
||||
DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
|
||||
EmailSender,
|
||||
ModelDeprecationInfo,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.router import Router
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AffectedModel:
|
||||
info: ModelDeprecationInfo
|
||||
display_name: str
|
||||
milestone: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _DeploymentOwner:
|
||||
model_name: str
|
||||
team_id: str
|
||||
public_name: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamNotification:
|
||||
team_id: str
|
||||
team_alias: str | None
|
||||
recipients: tuple[str, ...]
|
||||
models: tuple[AffectedModel, ...]
|
||||
|
||||
|
||||
class DeprecationEmailCache(Protocol):
|
||||
async def async_get_cache(self, *, key: str) -> object:
|
||||
"""The value stored under key, or None"""
|
||||
|
||||
async def async_set_cache(self, *, key: str, value: float, ttl: float) -> None:
|
||||
"""Store value under key for ttl seconds"""
|
||||
|
||||
|
||||
class SmtpSend(Protocol):
|
||||
async def __call__(self, *, receiver_email: str, subject: str, html: str, raise_on_error: bool) -> object:
|
||||
"""Send one HTML email over SMTP, raising on failure when asked to"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DeprecationEmailContext:
|
||||
llm_router: Router
|
||||
prisma_client: object
|
||||
cache: DeprecationEmailCache
|
||||
alerting_args: SlackAlertingArgs
|
||||
pod_lock_manager: PodLockManager | None
|
||||
deliver: Callable[[Sequence[str], str, str], Awaitable[None]]
|
||||
|
||||
|
||||
def select_milestone(days_until: int, thresholds: Sequence[int]) -> int | None:
|
||||
"""The most urgent threshold already reached, None while the first one is still ahead"""
|
||||
return min((threshold for threshold in thresholds if days_until <= threshold), default=None)
|
||||
|
||||
|
||||
def _max_age_days(alerting_args: SlackAlertingArgs) -> int:
|
||||
return max(1, alerting_args.model_deprecation_email_ttl // (24 * 60 * 60))
|
||||
|
||||
|
||||
def _reached(
|
||||
info: ModelDeprecationInfo, thresholds: Sequence[int], max_age_days: int
|
||||
) -> tuple[ModelDeprecationInfo, int] | None:
|
||||
if info.days_until_deprecation <= -max_age_days:
|
||||
return None
|
||||
milestone: Final = select_milestone(info.days_until_deprecation, thresholds)
|
||||
return None if milestone is None else (info, milestone)
|
||||
|
||||
|
||||
class _TeamScope(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
team_id: str | None = None
|
||||
team_public_model_name: str | None = None
|
||||
|
||||
|
||||
class _OwnedDeployment(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
model_name: str
|
||||
model_info: _TeamScope = Field(default_factory=_TeamScope)
|
||||
|
||||
|
||||
def _owner_of(deployment: Mapping[str, object]) -> _DeploymentOwner | None:
|
||||
try:
|
||||
parsed: Final = _OwnedDeployment.model_validate(deployment)
|
||||
except ValidationError:
|
||||
return None
|
||||
if not parsed.model_info.team_id:
|
||||
return None
|
||||
return _DeploymentOwner(
|
||||
model_name=parsed.model_name,
|
||||
team_id=parsed.model_info.team_id,
|
||||
public_name=parsed.model_info.team_public_model_name or parsed.model_name,
|
||||
)
|
||||
|
||||
|
||||
def _deployment_owners(llm_router: Router) -> Mapping[str, _DeploymentOwner]:
|
||||
"""Team-scoped deployments keyed by their internal model name, with the name the team knows"""
|
||||
owners: Final = (_owner_of(deployment) for deployment in llm_router.get_model_list() or ())
|
||||
return MappingProxyType({owner.model_name: owner for owner in owners if owner is not None})
|
||||
|
||||
|
||||
async def _team_can_access(model_name: str, team: LiteLLM_TeamTable, llm_router: Router) -> bool:
|
||||
"""Reuse the auth check so wildcards, access groups and the empty-list rule match real requests"""
|
||||
try:
|
||||
await can_team_access_model(model=model_name, team_object=team, llm_router=llm_router)
|
||||
except ProxyException:
|
||||
return False
|
||||
except Exception as e: # noqa: BLE001 # a malformed allowlist (e.g. a wildcard with regex metacharacters) must not abort the pass
|
||||
verbose_proxy_logger.debug("model_deprecation: access check failed for team %s: %s", team.team_id, e)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def _display_name_for(
|
||||
info: ModelDeprecationInfo, team: LiteLLM_TeamTable, owner: _DeploymentOwner | None, llm_router: Router
|
||||
) -> str | None:
|
||||
if owner is not None:
|
||||
return owner.public_name if owner.team_id == team.team_id else None
|
||||
if await _team_can_access(info.model_name, team, llm_router):
|
||||
return info.model_name
|
||||
return None
|
||||
|
||||
|
||||
async def _affected_model(
|
||||
reached: tuple[ModelDeprecationInfo, int],
|
||||
team: LiteLLM_TeamTable,
|
||||
owners: Mapping[str, _DeploymentOwner],
|
||||
llm_router: Router,
|
||||
) -> AffectedModel | None:
|
||||
info, milestone = reached
|
||||
display_name: Final = await _display_name_for(info, team, owners.get(info.model_name), llm_router)
|
||||
return None if display_name is None else AffectedModel(info=info, display_name=display_name, milestone=milestone)
|
||||
|
||||
|
||||
async def _affected_for_team(
|
||||
team: LiteLLM_TeamTable,
|
||||
reached: Sequence[tuple[ModelDeprecationInfo, int]],
|
||||
owners: Mapping[str, _DeploymentOwner],
|
||||
llm_router: Router,
|
||||
) -> tuple[AffectedModel, ...]:
|
||||
candidates: Final = tuple([await _affected_model(pair, team, owners, llm_router) for pair in reached])
|
||||
return tuple(model for model in candidates if model is not None)
|
||||
|
||||
|
||||
async def resolve_affected_teams(
|
||||
infos: Sequence[ModelDeprecationInfo],
|
||||
llm_router: Router,
|
||||
teams: Sequence[LiteLLM_TeamTable],
|
||||
thresholds: Sequence[int],
|
||||
max_age_days: int,
|
||||
) -> Mapping[str, tuple[AffectedModel, ...]]:
|
||||
"""Per team id, the deprecating models it can reach that have crossed a threshold"""
|
||||
reached: Final = tuple(
|
||||
pair for pair in (_reached(info, thresholds, max_age_days) for info in infos) if pair is not None
|
||||
)
|
||||
owners: Final = _deployment_owners(llm_router)
|
||||
per_team: Final = MappingProxyType(
|
||||
{
|
||||
team.team_id: await _affected_for_team(team, reached, owners, llm_router)
|
||||
for team in teams
|
||||
if not team.blocked
|
||||
}
|
||||
)
|
||||
return MappingProxyType({team_id: models for team_id, models in per_team.items() if models})
|
||||
|
||||
|
||||
def email_sent_key(team_id: str, model_name: str, deprecation_date: date, milestone: int) -> str:
|
||||
return f"model_deprecation_email:{team_id}:{model_name}:{deprecation_date.isoformat()}:{milestone}"
|
||||
|
||||
|
||||
async def _unsent(
|
||||
team_id: str, models: Sequence[AffectedModel], cache: DeprecationEmailCache
|
||||
) -> tuple[AffectedModel, ...]:
|
||||
flags: Final = tuple(
|
||||
[
|
||||
await cache.async_get_cache(
|
||||
key=email_sent_key(team_id, model.info.model_name, model.info.deprecation_date, model.milestone)
|
||||
)
|
||||
is None
|
||||
for model in models
|
||||
]
|
||||
)
|
||||
return tuple(model for model, unsent in zip(models, flags, strict=True) if unsent)
|
||||
|
||||
|
||||
async def _notification_for(
|
||||
team: LiteLLM_TeamTable, models: Sequence[AffectedModel], cache: DeprecationEmailCache, prisma_client: object
|
||||
) -> TeamNotification | None:
|
||||
unsent: Final = await _unsent(team.team_id, models, cache)
|
||||
if not unsent:
|
||||
return None
|
||||
recipients: Final = await get_team_admin_emails(team, prisma_client)
|
||||
if not recipients:
|
||||
verbose_proxy_logger.debug("model_deprecation: team %s has no admin emails, skipping", team.team_id)
|
||||
return None
|
||||
return TeamNotification(team_id=team.team_id, team_alias=team.team_alias, recipients=recipients, models=unsent)
|
||||
|
||||
|
||||
async def build_team_notifications(
|
||||
affected: Mapping[str, Sequence[AffectedModel]],
|
||||
teams: Sequence[LiteLLM_TeamTable],
|
||||
cache: DeprecationEmailCache,
|
||||
prisma_client: object,
|
||||
) -> tuple[TeamNotification, ...]:
|
||||
"""One digest per team of the milestones not yet emailed, dropping teams with nobody to email"""
|
||||
teams_by_id: Final = MappingProxyType({team.team_id: team for team in teams})
|
||||
candidates: Final = tuple(
|
||||
[
|
||||
await _notification_for(teams_by_id[team_id], models, cache, prisma_client)
|
||||
for team_id, models in affected.items()
|
||||
]
|
||||
)
|
||||
return tuple(notification for notification in candidates if notification is not None)
|
||||
|
||||
|
||||
def _days_left_label(days_until: int) -> str:
|
||||
if days_until < 0:
|
||||
return f"deprecated {abs(days_until)}d ago"
|
||||
return "today" if days_until == 0 else f"{days_until}d"
|
||||
|
||||
|
||||
def _render_row(model: AffectedModel) -> str:
|
||||
return MODEL_DEPRECATION_EMAIL_ROW_TEMPLATE.format(
|
||||
model_name=html.escape(model.display_name),
|
||||
provider=html.escape(model.info.litellm_provider or "unknown"),
|
||||
deprecation_date=model.info.deprecation_date.isoformat(),
|
||||
days_left=_days_left_label(model.info.days_until_deprecation),
|
||||
status=model.info.status,
|
||||
)
|
||||
|
||||
|
||||
def render_model_deprecation_email(
|
||||
notification: TeamNotification, email_logo_url: str, email_support_contact: str
|
||||
) -> tuple[str, str]:
|
||||
"""(subject, html) for one team's digest, every model-sourced string escaped"""
|
||||
alias: Final = " ".join((notification.team_alias or "").split())
|
||||
team_name: Final = alias or " ".join(notification.team_id.split())
|
||||
verb: Final = (
|
||||
"deprecated" if any(model.info.days_until_deprecation < 0 for model in notification.models) else "deprecating"
|
||||
)
|
||||
subject: Final = f"[LiteLLM] {len(notification.models)} model(s) {verb} for team {team_name}"
|
||||
body: Final = MODEL_DEPRECATION_EMAIL_TEMPLATE.format(
|
||||
email_logo_url=html.escape(email_logo_url),
|
||||
team_name=html.escape(team_name),
|
||||
model_rows="\n".join(_render_row(model) for model in notification.models),
|
||||
email_support_contact=html.escape(email_support_contact),
|
||||
)
|
||||
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",
|
||||
)
|
||||
_ENTERPRISE_SMTP_EMAIL_LOGGER: Final = (
|
||||
"litellm_enterprise.enterprise_callbacks.send_emails.smtp_email",
|
||||
"SMTPEmailLogger",
|
||||
)
|
||||
|
||||
|
||||
def _class_identity(cls: type) -> tuple[str, str]:
|
||||
return (cls.__module__, cls.__qualname__)
|
||||
|
||||
|
||||
def _is_noop_email_logger(email_logger: object) -> bool:
|
||||
"""The proxy hands out the bare enterprise base logger when no email provider is configured"""
|
||||
return _class_identity(type(email_logger)) == _ENTERPRISE_BASE_EMAIL_LOGGER
|
||||
|
||||
|
||||
def _is_smtp_email_logger(email_logger: object) -> bool:
|
||||
"""The enterprise SMTP logger fires one task per recipient without awaiting it, so its failures are invisible"""
|
||||
return _ENTERPRISE_SMTP_EMAIL_LOGGER in map(_class_identity, type(email_logger).__mro__)
|
||||
|
||||
|
||||
def make_email_deliverer(
|
||||
email_logger: EmailSender | None, smtp_send: SmtpSend | None = None
|
||||
) -> Callable[[Sequence[str], str, str], Awaitable[None]]:
|
||||
"""A configured non-SMTP provider, one email per recipient; SMTP always goes through the OSS helper so failures raise"""
|
||||
if email_logger is not None and not _is_noop_email_logger(email_logger) and not _is_smtp_email_logger(email_logger):
|
||||
|
||||
async def deliver_via_logger(recipients: Sequence[str], subject: str, html_body: str) -> None:
|
||||
for recipient in recipients:
|
||||
await email_logger.send_email(
|
||||
from_email=email_logger.DEFAULT_LITELLM_EMAIL,
|
||||
to_email=(recipient,),
|
||||
subject=subject,
|
||||
html_body=html_body,
|
||||
)
|
||||
|
||||
return deliver_via_logger
|
||||
|
||||
async def deliver_over_smtp(recipients: Sequence[str], subject: str, html_body: str) -> None:
|
||||
from litellm.proxy.utils import send_email
|
||||
|
||||
send: Final = smtp_send or send_email
|
||||
for recipient in recipients:
|
||||
await send(receiver_email=recipient, subject=subject, html=html_body, raise_on_error=True)
|
||||
|
||||
return deliver_over_smtp
|
||||
|
||||
|
||||
async def _load_teams(prisma_client: object) -> tuple[LiteLLM_TeamTable, ...]:
|
||||
return tuple(await TeamRepository(prisma_client).find_many())
|
||||
|
||||
|
||||
async def _stamp_pass(cache: DeprecationEmailCache) -> None:
|
||||
await cache.async_set_cache(
|
||||
key=SlackAlertingCacheKeys.deprecation_email_pass_key.value,
|
||||
value=time.time(),
|
||||
ttl=DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
async def _claimed_email_window(pod_lock_manager: PodLockManager | None) -> bool:
|
||||
"""Without a redis backed lock there is no fleet to coordinate, so a lone pod always sends"""
|
||||
if pod_lock_manager is None:
|
||||
return True
|
||||
return (
|
||||
await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=EMAIL_MODEL_DEPRECATION_LOCK_ID,
|
||||
ttl=DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
|
||||
allow_reentrant=False,
|
||||
)
|
||||
) is not False
|
||||
|
||||
|
||||
async def _send_team_notification(notification: TeamNotification, ctx: DeprecationEmailContext) -> bool:
|
||||
subject, html_body = render_model_deprecation_email(
|
||||
notification,
|
||||
email_logo_url=os.getenv("SMTP_SENDER_LOGO", os.getenv("EMAIL_LOGO_URL", LITELLM_LOGO_URL)),
|
||||
email_support_contact=os.getenv("EMAIL_SUPPORT_CONTACT", LITELLM_SUPPORT_CONTACT),
|
||||
)
|
||||
try:
|
||||
await ctx.deliver(notification.recipients, subject, html_body)
|
||||
except Exception as e: # noqa: BLE001 # one team's mail failure must not block the rest; unstamped keys retry tomorrow
|
||||
verbose_proxy_logger.exception("model_deprecation: email to team %s failed: %s", notification.team_id, e)
|
||||
return False
|
||||
for model in notification.models:
|
||||
await ctx.cache.async_set_cache(
|
||||
key=email_sent_key(
|
||||
notification.team_id, model.info.model_name, model.info.deprecation_date, model.milestone
|
||||
),
|
||||
value=time.time(),
|
||||
ttl=ctx.alerting_args.model_deprecation_email_ttl,
|
||||
)
|
||||
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
|
||||
|
||||
An empty snapshot returns before the DB is touched. A lost lock claim (a sibling pod, or a redis blip)
|
||||
returns without a stamp, so it costs one poll rather than a day, and a failing resolution raises
|
||||
before the lock is held, so it cannot silence the other pods
|
||||
"""
|
||||
snapshot: Final = collect_model_deprecations(llm_router=ctx.llm_router)
|
||||
infos: Final = (*snapshot.deprecated, *snapshot.imminent, *snapshot.upcoming)
|
||||
if not infos:
|
||||
return 0
|
||||
teams: Final = await _load_teams(ctx.prisma_client)
|
||||
affected: Final = await resolve_affected_teams(
|
||||
infos,
|
||||
ctx.llm_router,
|
||||
teams,
|
||||
ctx.alerting_args.model_deprecation_email_thresholds,
|
||||
_max_age_days(ctx.alerting_args),
|
||||
)
|
||||
notifications: Final = await build_team_notifications(affected, teams, ctx.cache, ctx.prisma_client)
|
||||
if not notifications:
|
||||
await _stamp_pass(ctx.cache)
|
||||
return 0
|
||||
if not await _claimed_email_window(ctx.pod_lock_manager):
|
||||
return 0
|
||||
results: Final = tuple([await _send_team_notification(notification, ctx) for notification in notifications])
|
||||
await _stamp_pass(ctx.cache)
|
||||
return sum(results)
|
||||
|
|
@ -8193,6 +8193,7 @@ async def send_email(
|
|||
receiver_email: str | None = None,
|
||||
subject: str | None = None,
|
||||
html: str | None = None,
|
||||
raise_on_error: bool = False,
|
||||
):
|
||||
"""
|
||||
smtp_host,
|
||||
|
|
@ -8246,6 +8247,9 @@ async def send_email(
|
|||
)
|
||||
|
||||
except Exception as e:
|
||||
if raise_on_error:
|
||||
verbose_proxy_logger.debug("An error occurred while sending the email: %s", e)
|
||||
raise
|
||||
verbose_proxy_logger.exception("An error occurred while sending the email:" + str(e))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from datetime import datetime as dt
|
|||
from enum import Enum
|
||||
from typing import Any, Final, Literal, Optional, Union
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic import Field, field_validator
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
|
@ -128,6 +128,28 @@ class SlackAlertingArgs(LiteLLMPydanticObjectBase):
|
|||
ge=60,
|
||||
description="How often (in seconds) to check per-user spend thresholds and anomalies. Default is hourly.",
|
||||
)
|
||||
model_deprecation_email_thresholds: tuple[int, ...] = Field(
|
||||
default=(30, 7, 0),
|
||||
description=(
|
||||
"Days before a model's deprecation_date at which team admins are emailed. "
|
||||
"One email per team per threshold. An empty list disables deprecation emails."
|
||||
),
|
||||
)
|
||||
model_deprecation_email_ttl: int = Field(
|
||||
default=90 * 24 * 60 * 60,
|
||||
ge=1,
|
||||
description=(
|
||||
"Seconds each (team, model, threshold) deprecation email is remembered after sending so the same "
|
||||
"milestone is not emailed twice. Default is 90 days."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("model_deprecation_email_thresholds")
|
||||
@classmethod
|
||||
def _normalize_deprecation_email_thresholds(cls, value: tuple[int, ...]) -> tuple[int, ...]:
|
||||
if any(threshold < 0 for threshold in value):
|
||||
raise ValueError("model_deprecation_email_thresholds must be non-negative")
|
||||
return tuple(sorted(frozenset(value), reverse=True))
|
||||
|
||||
|
||||
class DeploymentMetrics(LiteLLMPydanticObjectBase):
|
||||
|
|
@ -159,6 +181,7 @@ class SlackAlertingCacheKeys(Enum):
|
|||
latency_key = "latency_daily_metrics"
|
||||
report_sent_key = "daily_metrics_report_sent"
|
||||
deprecation_alert_sent_key = "model_deprecation_alert_sent"
|
||||
deprecation_email_pass_key = "model_deprecation_email_pass_at"
|
||||
|
||||
|
||||
class AlertType(str, Enum):
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from datetime import date, datetime
|
||||
from typing import Final, Literal
|
||||
from typing import Final, Literal, Protocol
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
|
|
@ -16,6 +17,15 @@ DEPRECATION_IDLE_POLL_SECONDS: Final = 30
|
|||
DeprecationStatus = Literal["upcoming", "imminent", "deprecated"]
|
||||
|
||||
|
||||
class EmailSender(Protocol):
|
||||
"""The proxy's configured email logger, as far as deprecation digests need it"""
|
||||
|
||||
DEFAULT_LITELLM_EMAIL: str
|
||||
|
||||
async def send_email(self, from_email: str, to_email: Sequence[str], subject: str, html_body: str) -> None:
|
||||
"""Deliver one HTML email to the given recipients"""
|
||||
|
||||
|
||||
class ModelDeprecationInfo(LiteLLMBaseModel):
|
||||
model_name: str = Field(description="The public name of the model on the proxy (model_group).")
|
||||
litellm_model: str | None = Field(
|
||||
|
|
|
|||
|
|
@ -2,24 +2,25 @@
|
|||
|
||||
import asyncio
|
||||
from itertools import chain, repeat
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.constants import SLACK_MODEL_DEPRECATION_LOCK_ID
|
||||
from litellm.integrations.SlackAlerting.ms_teams import MS_TEAMS_ALERTING_DESTINATION
|
||||
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
||||
from litellm.proxy._types import AlertType
|
||||
from litellm.proxy.common_utils.model_deprecation_notifications import DeprecationEmailContext
|
||||
from litellm.types.integrations.slack_alerting import SlackAlertingCacheKeys
|
||||
from litellm.types.proxy.model_deprecation import (
|
||||
DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
|
||||
DEPRECATION_IDLE_POLL_SECONDS,
|
||||
)
|
||||
|
||||
DEAD_MODEL_COST = {
|
||||
"dead-model": {"deprecation_date": "2020-01-01", "litellm_provider": "openai"}
|
||||
}
|
||||
DEAD_MODEL_COST = {"dead-model": {"deprecation_date": "2020-01-01", "litellm_provider": "openai"}}
|
||||
DEAD_ALIAS_DEPLOYMENT = {
|
||||
"model_name": "dead-alias",
|
||||
"litellm_params": {"model": "dead-model"},
|
||||
|
|
@ -99,9 +100,7 @@ async def test_should_dispatch_high_severity_when_deprecated(monkeypatch):
|
|||
]
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
alerting, "send_alert", new_callable=AsyncMock
|
||||
) as mock_send_alert:
|
||||
with patch.object(alerting, "send_alert", new_callable=AsyncMock) as mock_send_alert:
|
||||
sent = await alerting.send_model_deprecation_alert(llm_router=router)
|
||||
|
||||
assert sent is True
|
||||
|
|
@ -175,9 +174,7 @@ async def test_should_wait_for_the_router_instead_of_sleeping_a_full_day(monkeyp
|
|||
"model_cost",
|
||||
{"dead-model": {"deprecation_date": "2020-01-01", "litellm_provider": "openai"}},
|
||||
)
|
||||
alerting = SlackAlerting(
|
||||
alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings]
|
||||
)
|
||||
alerting = SlackAlerting(alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings])
|
||||
router = _make_router(
|
||||
[
|
||||
{
|
||||
|
|
@ -204,9 +201,7 @@ async def test_should_wait_for_the_router_instead_of_sleeping_a_full_day(monkeyp
|
|||
),
|
||||
pytest.raises(asyncio.CancelledError),
|
||||
):
|
||||
await alerting.run_scheduled_deprecation_check(
|
||||
get_llm_router=lambda: next(routers)
|
||||
)
|
||||
await alerting.run_scheduled_deprecation_check(get_llm_router=lambda: next(routers))
|
||||
|
||||
assert slept == [DEPRECATION_IDLE_POLL_SECONDS] * (router_absent_passes + 1)
|
||||
mock_send_alert.assert_awaited_once()
|
||||
|
|
@ -219,18 +214,14 @@ async def test_should_wait_for_the_router_instead_of_sleeping_a_full_day(monkeyp
|
|||
ids=["lock won", "no redis lock", "another pod holds the lock"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_alert_only_from_the_pod_holding_the_daily_lock(
|
||||
monkeypatch, lock_acquired, expect_alert
|
||||
):
|
||||
async def test_should_alert_only_from_the_pod_holding_the_daily_lock(monkeypatch, lock_acquired, expect_alert):
|
||||
"""Every pod runs the loop, so a fleet must not send one identical alert per replica"""
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"model_cost",
|
||||
{"dead-model": {"deprecation_date": "2020-01-01", "litellm_provider": "openai"}},
|
||||
)
|
||||
alerting = SlackAlerting(
|
||||
alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings]
|
||||
)
|
||||
alerting = SlackAlerting(alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings])
|
||||
router = _make_router(
|
||||
[
|
||||
{
|
||||
|
|
@ -251,9 +242,7 @@ async def test_should_alert_only_from_the_pod_holding_the_daily_lock(
|
|||
),
|
||||
pytest.raises(asyncio.CancelledError),
|
||||
):
|
||||
await alerting.run_scheduled_deprecation_check(
|
||||
get_llm_router=lambda: router, pod_lock_manager=pod_lock_manager
|
||||
)
|
||||
await alerting.run_scheduled_deprecation_check(get_llm_router=lambda: router, pod_lock_manager=pod_lock_manager)
|
||||
|
||||
assert mock_send_alert.await_count == int(expect_alert)
|
||||
assert pod_lock_manager.acquire_lock.await_args.kwargs == {
|
||||
|
|
@ -267,9 +256,7 @@ async def test_should_alert_only_from_the_pod_holding_the_daily_lock(
|
|||
async def test_should_retry_on_the_next_poll_when_the_lock_claim_fails(monkeypatch):
|
||||
"""A redis blip at claim time returns False like a held lock, and must not cost every pod a day of alerts"""
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting = SlackAlerting(
|
||||
alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings]
|
||||
)
|
||||
alerting = SlackAlerting(alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings])
|
||||
router = _make_router([DEAD_ALIAS_DEPLOYMENT])
|
||||
pod_lock_manager = MagicMock()
|
||||
pod_lock_manager.acquire_lock = AsyncMock(side_effect=[False, True])
|
||||
|
|
@ -288,9 +275,7 @@ async def test_should_retry_on_the_next_poll_when_the_lock_claim_fails(monkeypat
|
|||
),
|
||||
pytest.raises(asyncio.CancelledError),
|
||||
):
|
||||
await alerting.run_scheduled_deprecation_check(
|
||||
get_llm_router=lambda: router, pod_lock_manager=pod_lock_manager
|
||||
)
|
||||
await alerting.run_scheduled_deprecation_check(get_llm_router=lambda: router, pod_lock_manager=pod_lock_manager)
|
||||
|
||||
assert slept == [DEPRECATION_IDLE_POLL_SECONDS] * 2
|
||||
assert pod_lock_manager.acquire_lock.await_count == 2
|
||||
|
|
@ -301,9 +286,7 @@ async def test_should_retry_on_the_next_poll_when_the_lock_claim_fails(monkeypat
|
|||
async def test_should_not_claim_the_lock_when_there_is_nothing_to_report(monkeypatch):
|
||||
"""An empty pass must not hold the daily lock, or a sunset added later waits out the whole window"""
|
||||
monkeypatch.setattr(litellm, "model_cost", {})
|
||||
alerting = SlackAlerting(
|
||||
alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings]
|
||||
)
|
||||
alerting = SlackAlerting(alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings])
|
||||
router = _make_router(
|
||||
[
|
||||
{
|
||||
|
|
@ -317,9 +300,7 @@ async def test_should_not_claim_the_lock_when_there_is_nothing_to_report(monkeyp
|
|||
pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
|
||||
with patch.object(alerting, "send_alert", new_callable=AsyncMock) as mock_send_alert:
|
||||
sent = await alerting.send_model_deprecation_alert(
|
||||
llm_router=router, pod_lock_manager=pod_lock_manager
|
||||
)
|
||||
sent = await alerting.send_model_deprecation_alert(llm_router=router, pod_lock_manager=pod_lock_manager)
|
||||
|
||||
assert sent is False
|
||||
pod_lock_manager.acquire_lock.assert_not_awaited()
|
||||
|
|
@ -330,9 +311,7 @@ async def test_should_not_claim_the_lock_when_there_is_nothing_to_report(monkeyp
|
|||
async def test_should_not_alert_or_claim_the_lock_within_a_day_of_a_sent_alert(monkeypatch):
|
||||
"""The shared sent stamp keeps sibling pods and restarts from re-alerting or re-asking redis for a day"""
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting = SlackAlerting(
|
||||
alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings]
|
||||
)
|
||||
alerting = SlackAlerting(alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings])
|
||||
await alerting.internal_usage_cache.async_set_cache(
|
||||
key=SlackAlertingCacheKeys.deprecation_alert_sent_key.value,
|
||||
value=1.0,
|
||||
|
|
@ -350,21 +329,17 @@ async def test_should_not_alert_or_claim_the_lock_within_a_day_of_a_sent_alert(m
|
|||
),
|
||||
pytest.raises(asyncio.CancelledError),
|
||||
):
|
||||
await alerting.run_scheduled_deprecation_check(
|
||||
get_llm_router=lambda: router, pod_lock_manager=pod_lock_manager
|
||||
)
|
||||
await alerting.run_scheduled_deprecation_check(get_llm_router=lambda: router, pod_lock_manager=pod_lock_manager)
|
||||
|
||||
pod_lock_manager.acquire_lock.assert_not_awaited()
|
||||
mock_send_alert.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_back_off_a_full_day_after_a_pass_raises(monkeypatch):
|
||||
async def test_should_skip_a_raising_pass_for_a_day_and_keep_polling(monkeypatch):
|
||||
"""A misconfigured webhook raises on every send, which must log once a day rather than every poll"""
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting = SlackAlerting(
|
||||
alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings]
|
||||
)
|
||||
alerting = SlackAlerting(alerting=["slack"], alert_types=[AlertType.model_deprecation_warnings])
|
||||
router = _make_router([DEAD_ALIAS_DEPLOYMENT])
|
||||
slept: list[float] = []
|
||||
|
||||
|
|
@ -388,5 +363,250 @@ async def test_should_back_off_a_full_day_after_a_pass_raises(monkeypatch):
|
|||
):
|
||||
await alerting.run_scheduled_deprecation_check(get_llm_router=lambda: router)
|
||||
|
||||
assert slept == [DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS] * 2
|
||||
assert mock_send_alert.await_count == 2
|
||||
assert slept == [DEPRECATION_IDLE_POLL_SECONDS] * 2
|
||||
assert mock_send_alert.await_count == 1
|
||||
|
||||
|
||||
def _email_only_alerting() -> SlackAlerting:
|
||||
return SlackAlerting(alerting=["email"], alert_types=[AlertType.model_deprecation_warnings])
|
||||
|
||||
|
||||
class _EmailRun:
|
||||
def __init__(self, result: int = 1, fail: bool = False):
|
||||
self.result = result
|
||||
self.fail = fail
|
||||
self.contexts = []
|
||||
|
||||
async def __call__(self, ctx: DeprecationEmailContext) -> int:
|
||||
self.contexts.append(ctx)
|
||||
if self.fail:
|
||||
raise ValueError("smtp exploded")
|
||||
return self.result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("alerting", "alert_types", "thresholds", "expected"),
|
||||
[
|
||||
(None, [AlertType.model_deprecation_warnings], (30,), False),
|
||||
(["slack"], [AlertType.model_deprecation_warnings], (30,), False),
|
||||
(["email"], [AlertType.llm_exceptions], (30,), False),
|
||||
(["email"], [AlertType.model_deprecation_warnings], (), False),
|
||||
(["email"], [AlertType.model_deprecation_warnings], (30,), True),
|
||||
(["slack", "email"], [AlertType.model_deprecation_warnings], (30, 7, 0), True),
|
||||
],
|
||||
)
|
||||
def test_deprecation_emails_enabled_truth_table(alerting, alert_types, thresholds, expected):
|
||||
alerting_obj: Final = SlackAlerting(
|
||||
alerting=alerting, alert_types=alert_types, alerting_args={"model_deprecation_email_thresholds": thresholds}
|
||||
)
|
||||
assert alerting_obj._deprecation_emails_enabled() is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_not_run_the_slack_pass_without_a_slack_or_teams_channel(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
pod_lock_manager: Final = MagicMock()
|
||||
pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
|
||||
sent: Final = await _email_only_alerting().send_model_deprecation_alert(
|
||||
llm_router=_make_router([DEAD_ALIAS_DEPLOYMENT]), pod_lock_manager=pod_lock_manager
|
||||
)
|
||||
|
||||
assert sent is False
|
||||
pod_lock_manager.acquire_lock.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_run_the_email_pass_with_email_only_alerting(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting: Final = _email_only_alerting()
|
||||
prisma: Final = SimpleNamespace()
|
||||
run: Final = _EmailRun(result=2)
|
||||
|
||||
await alerting._run_deprecation_passes(
|
||||
get_llm_router=lambda: _make_router([DEAD_ALIAS_DEPLOYMENT]),
|
||||
pod_lock_manager=None,
|
||||
get_prisma_client=lambda: prisma,
|
||||
get_email_logger=lambda: None,
|
||||
send_emails=run,
|
||||
)
|
||||
|
||||
assert len(run.contexts) == 1
|
||||
assert run.contexts[0].prisma_client is prisma
|
||||
assert run.contexts[0].alerting_args.model_deprecation_email_thresholds == (30, 7, 0)
|
||||
assert run.contexts[0].cache is alerting.internal_usage_cache
|
||||
assert alerting.deprecation_email_backoff_until == 0.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_still_email_when_the_slack_pass_raises_and_back_off_only_slack(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting: Final = SlackAlerting(alerting=["slack", "email"], alert_types=[AlertType.model_deprecation_warnings])
|
||||
run: Final = _EmailRun()
|
||||
|
||||
with patch.object(alerting, "send_alert", new_callable=AsyncMock, side_effect=ValueError("no webhook")):
|
||||
await alerting._run_deprecation_passes(
|
||||
get_llm_router=lambda: _make_router([DEAD_ALIAS_DEPLOYMENT]),
|
||||
pod_lock_manager=None,
|
||||
get_prisma_client=SimpleNamespace,
|
||||
get_email_logger=lambda: None,
|
||||
send_emails=run,
|
||||
)
|
||||
|
||||
assert len(run.contexts) == 1
|
||||
assert alerting.deprecation_alert_backoff_until > 0.0
|
||||
assert alerting.deprecation_email_backoff_until == 0.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_back_off_only_the_email_pass_when_it_raises(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting: Final = _email_only_alerting()
|
||||
run: Final = _EmailRun(fail=True)
|
||||
|
||||
await alerting._run_deprecation_passes(
|
||||
get_llm_router=lambda: _make_router([DEAD_ALIAS_DEPLOYMENT]),
|
||||
pod_lock_manager=None,
|
||||
get_prisma_client=SimpleNamespace,
|
||||
get_email_logger=lambda: None,
|
||||
send_emails=run,
|
||||
)
|
||||
|
||||
assert alerting.deprecation_email_backoff_until > 0.0
|
||||
assert alerting.deprecation_alert_backoff_until == 0.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_the_email_pass_within_a_day_of_the_last_one(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting: Final = _email_only_alerting()
|
||||
await alerting.internal_usage_cache.async_set_cache(
|
||||
key=SlackAlertingCacheKeys.deprecation_email_pass_key.value, value=1.0
|
||||
)
|
||||
run: Final = _EmailRun()
|
||||
|
||||
sent: Final = await alerting._run_deprecation_email_pass(
|
||||
_make_router([DEAD_ALIAS_DEPLOYMENT]), None, SimpleNamespace, lambda: None, run
|
||||
)
|
||||
|
||||
assert sent == 0
|
||||
assert run.contexts == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_the_email_pass_without_a_database(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
run: Final = _EmailRun()
|
||||
|
||||
sent: Final = await _email_only_alerting()._run_deprecation_email_pass(
|
||||
_make_router([DEAD_ALIAS_DEPLOYMENT]), None, lambda: None, lambda: None, run
|
||||
)
|
||||
|
||||
assert sent == 0
|
||||
assert run.contexts == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_the_email_pass_without_a_router():
|
||||
run: Final = _EmailRun()
|
||||
prisma_lookups: Final = []
|
||||
|
||||
def get_prisma_client():
|
||||
prisma_lookups.append(True)
|
||||
return SimpleNamespace()
|
||||
|
||||
sent: Final = await _email_only_alerting()._run_deprecation_email_pass(
|
||||
None, None, get_prisma_client, lambda: None, run
|
||||
)
|
||||
|
||||
assert sent == 0
|
||||
assert prisma_lookups == []
|
||||
assert run.contexts == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_keep_polling_every_thirty_seconds_after_an_email_failure(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting: Final = _email_only_alerting()
|
||||
slept: list[float] = []
|
||||
|
||||
async def stop_after_second_pass(seconds):
|
||||
slept.append(seconds)
|
||||
if len(slept) == 2:
|
||||
raise asyncio.CancelledError
|
||||
|
||||
run: Final = _EmailRun(fail=True)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await alerting.run_scheduled_deprecation_check(
|
||||
get_llm_router=lambda: _make_router([DEAD_ALIAS_DEPLOYMENT]),
|
||||
get_prisma_client=SimpleNamespace,
|
||||
get_email_logger=lambda: None,
|
||||
send_emails=run,
|
||||
sleep=stop_after_second_pass,
|
||||
)
|
||||
|
||||
assert slept == [DEPRECATION_IDLE_POLL_SECONDS] * 2
|
||||
assert len(run.contexts) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("alerting", "expected"),
|
||||
[
|
||||
(None, False),
|
||||
([], False),
|
||||
(["email"], False),
|
||||
(["webhook"], False),
|
||||
(["slack"], True),
|
||||
([MS_TEAMS_ALERTING_DESTINATION], True),
|
||||
(["slack", "email"], True),
|
||||
],
|
||||
)
|
||||
def test_deprecation_alerts_enabled_requires_a_slack_or_teams_channel(alerting, expected):
|
||||
alerting_obj: Final = SlackAlerting(alerting=alerting, alert_types=[AlertType.model_deprecation_warnings])
|
||||
assert alerting_obj._deprecation_alerts_enabled() is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_keep_emailing_while_a_broken_slack_pass_is_backed_off(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting: Final = SlackAlerting(alerting=["slack", "email"], alert_types=[AlertType.model_deprecation_warnings])
|
||||
run: Final = _EmailRun()
|
||||
|
||||
with patch.object(
|
||||
alerting, "send_alert", new_callable=AsyncMock, side_effect=ValueError("no webhook")
|
||||
) as send_alert:
|
||||
for _ in range(2):
|
||||
await alerting._run_deprecation_passes(
|
||||
get_llm_router=lambda: _make_router([DEAD_ALIAS_DEPLOYMENT]),
|
||||
pod_lock_manager=None,
|
||||
get_prisma_client=SimpleNamespace,
|
||||
get_email_logger=lambda: None,
|
||||
send_emails=run,
|
||||
now=lambda: 1_000.0,
|
||||
)
|
||||
|
||||
assert send_alert.await_count == 1
|
||||
assert len(run.contexts) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_retry_a_backed_off_email_pass_after_a_day(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", DEAD_MODEL_COST)
|
||||
alerting: Final = _email_only_alerting()
|
||||
run: Final = _EmailRun(fail=True)
|
||||
|
||||
for moment in (
|
||||
1_000.0,
|
||||
1_000.0 + DEPRECATION_IDLE_POLL_SECONDS,
|
||||
1_000.0 + DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
|
||||
):
|
||||
await alerting._run_deprecation_passes(
|
||||
get_llm_router=lambda: _make_router([DEAD_ALIAS_DEPLOYMENT]),
|
||||
pod_lock_manager=None,
|
||||
get_prisma_client=SimpleNamespace,
|
||||
get_email_logger=lambda: None,
|
||||
send_emails=run,
|
||||
now=lambda moment=moment: moment,
|
||||
)
|
||||
|
||||
assert len(run.contexts) == 2
|
||||
|
|
|
|||
123
tests/unit/integrations/test_email_alerting.py
Normal file
123
tests/unit/integrations/test_email_alerting.py
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
"""Tests for litellm/integrations/email_alerting.py."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.email_alerting import get_team_admin_emails
|
||||
from litellm.models.team import LiteLLM_TeamTable, Member
|
||||
|
||||
|
||||
class _FakeUserTable:
|
||||
def __init__(self, rows):
|
||||
self.rows = rows
|
||||
self.queries = []
|
||||
|
||||
async def find_many(self, where=None, take=None, skip=None, order=None):
|
||||
self.queries.append(where)
|
||||
wanted: Final = where["user_id"]["in"]
|
||||
return [row for row in self.rows if row["user_id"] in wanted]
|
||||
|
||||
|
||||
def _prisma(rows):
|
||||
return SimpleNamespace(db=SimpleNamespace(litellm_usertable=_FakeUserTable(rows)))
|
||||
|
||||
|
||||
USERS: Final = (
|
||||
{"user_id": "u-admin", "user_email": "admin@example.com"},
|
||||
{"user_id": "u-member", "user_email": "member@example.com"},
|
||||
{"user_id": "u-legacy", "user_email": "legacy@example.com"},
|
||||
{"user_id": "u-no-email", "user_email": None},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_return_only_admin_emails():
|
||||
team: Final = LiteLLM_TeamTable(
|
||||
team_id="t1",
|
||||
members_with_roles=[Member(user_id="u-admin", role="admin"), Member(user_id="u-member", role="user")],
|
||||
)
|
||||
assert await get_team_admin_emails(team, _prisma(USERS)) == ("admin@example.com",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_ignore_the_legacy_admins_list_the_proxy_does_not_authorize():
|
||||
prisma: Final = _prisma(USERS)
|
||||
team: Final = LiteLLM_TeamTable(
|
||||
team_id="t1", admins=["u-legacy"], members_with_roles=[Member(user_id="u-admin", role="admin")]
|
||||
)
|
||||
assert await get_team_admin_emails(team, prisma) == ("admin@example.com",)
|
||||
assert prisma.db.litellm_usertable.queries == [{"user_id": {"in": ["u-admin"]}}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_drop_admins_without_an_email_and_dedupe():
|
||||
team: Final = LiteLLM_TeamTable(
|
||||
team_id="t1",
|
||||
admins=["u-admin"],
|
||||
members_with_roles=[Member(user_id="u-admin", role="admin"), Member(user_id="u-no-email", role="admin")],
|
||||
)
|
||||
assert await get_team_admin_emails(team, _prisma(USERS)) == ("admin@example.com",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_not_query_when_team_has_no_admins():
|
||||
prisma: Final = _prisma(USERS)
|
||||
team: Final = LiteLLM_TeamTable(team_id="t1", members_with_roles=[Member(user_id="u-member", role="user")])
|
||||
assert await get_team_admin_emails(team, prisma) == ()
|
||||
assert prisma.db.litellm_usertable.queries == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_include_email_only_admins_and_dedupe_ids_and_emails():
|
||||
prisma: Final = _prisma(
|
||||
(
|
||||
{"user_id": "u-admin", "user_email": "admin@example.com"},
|
||||
{"user_id": "u-legacy", "user_email": "admin@example.com"},
|
||||
)
|
||||
)
|
||||
team: Final = LiteLLM_TeamTable(
|
||||
team_id="t1",
|
||||
members_with_roles=[
|
||||
Member(user_id="u-admin", role="admin"),
|
||||
Member(user_id="u-legacy", role="admin"),
|
||||
Member(user_email="mail-only@example.com", role="admin"),
|
||||
],
|
||||
)
|
||||
assert await get_team_admin_emails(team, prisma) == ("mail-only@example.com", "admin@example.com")
|
||||
assert prisma.db.litellm_usertable.queries == [{"user_id": {"in": ["u-admin", "u-legacy"]}}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_email_an_email_only_admin_without_querying_the_database():
|
||||
prisma: Final = _prisma(USERS)
|
||||
team: Final = LiteLLM_TeamTable(
|
||||
team_id="t1",
|
||||
members_with_roles=[
|
||||
Member(user_email="mail-only@example.com", role="admin"),
|
||||
Member(user_id="u-member", role="user"),
|
||||
],
|
||||
)
|
||||
assert await get_team_admin_emails(team, prisma) == ("mail-only@example.com",)
|
||||
assert prisma.db.litellm_usertable.queries == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_fall_back_to_the_inline_email_when_the_user_record_has_none():
|
||||
team: Final = LiteLLM_TeamTable(
|
||||
team_id="t1",
|
||||
members_with_roles=[Member(user_id="u-no-email", user_email="inline@example.com", role="admin")],
|
||||
)
|
||||
assert await get_team_admin_emails(team, _prisma(USERS)) == ("inline@example.com",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_prefer_the_user_record_email_over_a_differing_inline_one():
|
||||
team: Final = LiteLLM_TeamTable(
|
||||
team_id="t1",
|
||||
members_with_roles=[Member(user_id="u-admin", user_email="stale@example.com", role="admin")],
|
||||
)
|
||||
assert await get_team_admin_emails(team, _prisma(USERS)) == ("admin@example.com",)
|
||||
|
|
@ -0,0 +1,818 @@
|
|||
"""Tests for the model deprecation email notifications module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
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
|
||||
from litellm.proxy.common_utils.model_deprecation_notifications import (
|
||||
AffectedModel,
|
||||
DeprecationEmailContext,
|
||||
TeamNotification,
|
||||
build_team_notifications,
|
||||
email_pass_done_today,
|
||||
email_sent_key,
|
||||
make_email_deliverer,
|
||||
render_model_deprecation_email,
|
||||
resolve_affected_teams,
|
||||
select_milestone,
|
||||
send_model_deprecation_emails,
|
||||
)
|
||||
from litellm.types.integrations.slack_alerting import SlackAlertingArgs, SlackAlertingCacheKeys
|
||||
from litellm.types.proxy.model_deprecation import DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS, ModelDeprecationInfo
|
||||
|
||||
|
||||
class TestSelectMilestone:
|
||||
@pytest.mark.parametrize(
|
||||
("days_until", "expected"),
|
||||
[(45, None), (30, 30), (25, 30), (7, 7), (3, 7), (0, 0), (-12, 0)],
|
||||
)
|
||||
def test_should_pick_the_most_urgent_threshold_reached(self, days_until, expected):
|
||||
assert select_milestone(days_until, (30, 7, 0)) == expected
|
||||
|
||||
def test_should_return_none_when_no_thresholds_configured(self):
|
||||
assert select_milestone(-5, ()) is None
|
||||
|
||||
def test_should_not_depend_on_threshold_order(self):
|
||||
assert select_milestone(5, (0, 7, 30)) == 7
|
||||
|
||||
|
||||
def _info(
|
||||
model_name: str,
|
||||
days_until: int,
|
||||
provider: str | None = "openai",
|
||||
deprecation_date: date = date(2026, 10, 1),
|
||||
) -> ModelDeprecationInfo:
|
||||
return ModelDeprecationInfo(
|
||||
model_name=model_name,
|
||||
litellm_model=model_name,
|
||||
deprecation_date=deprecation_date,
|
||||
days_until_deprecation=days_until,
|
||||
status="deprecated" if days_until < 0 else ("imminent" if days_until <= 30 else "upcoming"),
|
||||
litellm_provider=provider,
|
||||
)
|
||||
|
||||
|
||||
def _router(deployments: Sequence[Mapping[str, object]]) -> MagicMock:
|
||||
router: Final = MagicMock()
|
||||
router.get_model_list.return_value = deployments
|
||||
router.get_model_access_groups.return_value = {}
|
||||
router.model_group_alias = {}
|
||||
return router
|
||||
|
||||
|
||||
def _team(team_id: str, models: Sequence[str] = (), blocked: bool = False) -> LiteLLM_TeamTable:
|
||||
return LiteLLM_TeamTable(team_id=team_id, models=list(models), blocked=blocked)
|
||||
|
||||
|
||||
THRESHOLDS: Final = (30, 7, 0)
|
||||
|
||||
|
||||
class TestResolveAffectedTeams:
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_include_team_listing_the_model_explicitly(self):
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("gpt-old", 5),), _router([]), (_team("t1", ("gpt-old",)), _team("t2", ("other",))), THRESHOLDS, 90
|
||||
)
|
||||
assert set(affected) == {"t1"}
|
||||
assert affected["t1"] == (AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_include_unrestricted_team(self):
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("gpt-old", 5),), _router([]), (_team("t1"),), THRESHOLDS, 90
|
||||
)
|
||||
assert set(affected) == {"t1"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_match_wildcard_allowlists(self):
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("openai/gpt-old", 5),),
|
||||
_router([]),
|
||||
(_team("t1", ("openai/*",)), _team("t2", ("azure/*",))),
|
||||
THRESHOLDS,
|
||||
90,
|
||||
)
|
||||
assert set(affected) == {"t1"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_blocked_teams(self):
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("gpt-old", 5),), _router([]), (_team("t1", blocked=True),), THRESHOLDS, 90
|
||||
)
|
||||
assert dict(affected) == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_models_that_have_not_reached_a_threshold(self):
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("gpt-old", 45),), _router([]), (_team("t1"),), THRESHOLDS, 90
|
||||
)
|
||||
assert dict(affected) == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_attribute_team_scoped_deployment_to_its_owner_only(self):
|
||||
deployments: Final = [
|
||||
{
|
||||
"model_name": "model_name_t2_abc",
|
||||
"litellm_params": {"model": "openai/gpt-old"},
|
||||
"model_info": {"id": "1", "team_id": "t2", "team_public_model_name": "gpt-old"},
|
||||
}
|
||||
]
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("model_name_t2_abc", -3),), _router(deployments), (_team("t1"), _team("t2")), THRESHOLDS, 90
|
||||
)
|
||||
assert set(affected) == {"t2"}
|
||||
assert affected["t2"][0].display_name == "gpt-old"
|
||||
assert affected["t2"][0].milestone == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_collect_every_reachable_model_per_team(self):
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("gpt-old", 5), _info("gpt-older", -1)),
|
||||
_router([]),
|
||||
(_team("t1", ("gpt-old", "gpt-older")),),
|
||||
THRESHOLDS,
|
||||
90,
|
||||
)
|
||||
assert [m.info.model_name for m in affected["t1"]] == ["gpt-old", "gpt-older"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_treat_an_ordinary_deployment_as_shared(self):
|
||||
deployments: Final = [
|
||||
{"model_name": "gpt-old", "litellm_params": {"model": "openai/gpt-old"}, "model_info": {"id": "1"}}
|
||||
]
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("gpt-old", 5),), _router(deployments), (_team("t1", ("gpt-old",)),), THRESHOLDS, 90
|
||||
)
|
||||
assert set(affected) == {"t1"}
|
||||
assert affected["t1"][0].display_name == "gpt-old"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_fall_back_to_the_internal_name_when_a_scoped_deployment_has_no_public_name(self):
|
||||
deployments: Final = [
|
||||
{
|
||||
"model_name": "model_name_t2_abc",
|
||||
"litellm_params": {"model": "openai/gpt-old"},
|
||||
"model_info": {"id": "1", "team_id": "t2"},
|
||||
}
|
||||
]
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("model_name_t2_abc", -3),), _router(deployments), (_team("t2"),), THRESHOLDS, 90
|
||||
)
|
||||
assert affected["t2"][0].display_name == "model_name_t2_abc"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_models_deprecated_longer_ago_than_the_marker_lifetime(self):
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("ancient", -200), _info("gpt-old", -10)), _router([]), (_team("t1"),), THRESHOLDS, 90
|
||||
)
|
||||
assert [model.info.model_name for model in affected["t1"]] == ["gpt-old"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_exclude_a_model_exactly_at_the_marker_lifetime_boundary(self):
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("edge-out", -90), _info("edge-in", -89)), _router([]), (_team("t1"),), THRESHOLDS, 90
|
||||
)
|
||||
assert [model.info.model_name for model in affected["t1"]] == ["edge-in"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_deny_a_team_whose_allowlist_is_malformed_and_keep_going(self):
|
||||
affected: Final = await resolve_affected_teams(
|
||||
(_info("openai/gpt-old", 5),),
|
||||
_router([]),
|
||||
(_team("t1", ("openai/[*",)), _team("t2", ("openai/gpt-old",))),
|
||||
THRESHOLDS,
|
||||
90,
|
||||
)
|
||||
assert set(affected) == {"t2"}
|
||||
|
||||
|
||||
class TestRenderModelDeprecationEmail:
|
||||
def _notification(self, models: tuple[AffectedModel, ...]) -> TeamNotification:
|
||||
return TeamNotification(team_id="t1", team_alias="Data Team", recipients=("a@example.com",), models=models)
|
||||
|
||||
def test_should_render_one_row_per_model_with_escaped_names(self):
|
||||
hostile: Final = _info("<img src=x onerror=alert(1)>", 5, provider="op&en")
|
||||
subject, html = render_model_deprecation_email(
|
||||
self._notification((AffectedModel(info=hostile, display_name=hostile.model_name, milestone=7),)),
|
||||
email_logo_url="https://logo",
|
||||
email_support_contact="help@example.com",
|
||||
)
|
||||
assert "<img src=x" not in html
|
||||
assert "<img src=x onerror=alert(1)>" in html
|
||||
assert "op&en" in html
|
||||
assert "2026-10-01" in html
|
||||
assert "<td>5d</td>" in html
|
||||
assert "https://logo" in html
|
||||
assert "help@example.com" in html
|
||||
assert "Data Team" in html
|
||||
assert subject == "[LiteLLM] 1 model(s) deprecating for team Data Team"
|
||||
|
||||
def test_should_say_deprecated_in_subject_once_a_model_reaches_day_zero(self):
|
||||
models: Final = (
|
||||
AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),
|
||||
AffectedModel(info=_info("gpt-older", -4), display_name="gpt-older", milestone=0),
|
||||
)
|
||||
subject, html = render_model_deprecation_email(
|
||||
self._notification(models), email_logo_url="https://logo", email_support_contact="help@example.com"
|
||||
)
|
||||
assert subject == "[LiteLLM] 2 model(s) deprecated for team Data Team"
|
||||
assert "deprecated 4d ago" in html
|
||||
assert "<td>deprecated</td>" in html
|
||||
assert html.index("gpt-old") < html.index("gpt-older")
|
||||
|
||||
def test_should_fall_back_to_team_id_when_alias_missing(self):
|
||||
notification: Final = TeamNotification(
|
||||
team_id="t1",
|
||||
team_alias=None,
|
||||
recipients=("a@example.com",),
|
||||
models=(AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),),
|
||||
)
|
||||
subject, _ = render_model_deprecation_email(
|
||||
notification, email_logo_url="https://logo", email_support_contact="help@example.com"
|
||||
)
|
||||
assert subject.endswith("for team t1")
|
||||
|
||||
def test_should_escape_the_team_alias(self):
|
||||
notification: Final = TeamNotification(
|
||||
team_id="t1",
|
||||
team_alias="<b>Ops</b>",
|
||||
recipients=("a@example.com",),
|
||||
models=(AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),),
|
||||
)
|
||||
_, html = render_model_deprecation_email(
|
||||
notification, email_logo_url="https://logo", email_support_contact="help@example.com"
|
||||
)
|
||||
assert "<b>Ops</b>" not in html
|
||||
assert "<b>Ops</b>" in html
|
||||
|
||||
def test_should_render_unknown_when_the_provider_is_missing(self):
|
||||
models: Final = (AffectedModel(info=_info("gpt-old", 5, provider=None), display_name="gpt-old", milestone=7),)
|
||||
_, html = render_model_deprecation_email(
|
||||
self._notification(models), email_logo_url="https://logo", email_support_contact="help@example.com"
|
||||
)
|
||||
assert "<td>unknown</td>" in html
|
||||
|
||||
def test_should_describe_day_zero_as_today_and_still_deprecating(self):
|
||||
models: Final = (AffectedModel(info=_info("gpt-old", 0), display_name="gpt-old", milestone=0),)
|
||||
subject, html = render_model_deprecation_email(
|
||||
self._notification(models), email_logo_url="https://logo", email_support_contact="help@example.com"
|
||||
)
|
||||
assert subject == "[LiteLLM] 1 model(s) deprecating for team Data Team"
|
||||
assert "<td>today</td>" in html
|
||||
assert "<td>imminent</td>" in html
|
||||
|
||||
def test_should_append_the_shared_footer_and_escape_the_support_contact(self):
|
||||
models: Final = (AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),)
|
||||
_, html = render_model_deprecation_email(
|
||||
self._notification(models), email_logo_url="https://logo", email_support_contact="<help>@example.com"
|
||||
)
|
||||
assert "github.com/BerriAI/litellm" in html
|
||||
assert "<help>@example.com" in html
|
||||
assert "<help>@example.com" not in html
|
||||
|
||||
def test_should_escape_the_logo_url(self):
|
||||
models: Final = (AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),)
|
||||
_, html = render_model_deprecation_email(
|
||||
self._notification(models), email_logo_url='https://logo" onerror="x', email_support_contact="h@x.io"
|
||||
)
|
||||
assert '" onerror="' not in html
|
||||
assert "https://logo" onerror="x" in html
|
||||
|
||||
def test_should_collapse_newlines_in_the_subject_team_name(self):
|
||||
notification: Final = TeamNotification(
|
||||
team_id="t1",
|
||||
team_alias="Data\r\n Team",
|
||||
recipients=("a@example.com",),
|
||||
models=(AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),),
|
||||
)
|
||||
subject, _ = render_model_deprecation_email(
|
||||
notification, email_logo_url="https://logo", email_support_contact="h@x.io"
|
||||
)
|
||||
assert subject == "[LiteLLM] 1 model(s) deprecating for team Data Team"
|
||||
|
||||
def test_should_collapse_newlines_in_the_team_id_fallback(self):
|
||||
notification: Final = TeamNotification(
|
||||
team_id="t\r\n1",
|
||||
team_alias=None,
|
||||
recipients=("a@example.com",),
|
||||
models=(AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),),
|
||||
)
|
||||
subject, _ = render_model_deprecation_email(
|
||||
notification, email_logo_url="https://logo", email_support_contact="h@x.io"
|
||||
)
|
||||
assert subject.endswith("for team t 1")
|
||||
|
||||
|
||||
def _prisma_with_users(rows: Sequence[Mapping[str, object]]) -> SimpleNamespace:
|
||||
async def find_many(where=None, take=None, skip=None, order=None):
|
||||
wanted: Final = where["user_id"]["in"]
|
||||
return [row for row in rows if row["user_id"] in wanted]
|
||||
|
||||
return SimpleNamespace(db=SimpleNamespace(litellm_usertable=SimpleNamespace(find_many=find_many)))
|
||||
|
||||
|
||||
ADMIN_USERS: Final = ({"user_id": "u1", "user_email": "admin@example.com"},)
|
||||
|
||||
|
||||
def _admin_team(team_id: str, alias: str | None = None) -> LiteLLM_TeamTable:
|
||||
return LiteLLM_TeamTable(team_id=team_id, team_alias=alias, members_with_roles=[{"user_id": "u1", "role": "admin"}])
|
||||
|
||||
|
||||
class TestBuildTeamNotifications:
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_build_one_digest_per_team_with_admin_recipients(self):
|
||||
models: Final = (AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),)
|
||||
notifications: Final = await build_team_notifications(
|
||||
{"t1": models}, (_admin_team("t1", "Data"),), DualCache(), _prisma_with_users(ADMIN_USERS)
|
||||
)
|
||||
assert notifications == (
|
||||
TeamNotification(team_id="t1", team_alias="Data", recipients=("admin@example.com",), models=models),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_milestones_already_sent(self):
|
||||
cache: Final = DualCache()
|
||||
await cache.async_set_cache(key=email_sent_key("t1", "gpt-old", date(2026, 10, 1), 7), value=1.0)
|
||||
sent: Final = AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7)
|
||||
fresh: Final = AffectedModel(info=_info("gpt-older", -1), display_name="gpt-older", milestone=0)
|
||||
notifications: Final = await build_team_notifications(
|
||||
{"t1": (sent, fresh)}, (_admin_team("t1"),), cache, _prisma_with_users(ADMIN_USERS)
|
||||
)
|
||||
assert notifications[0].models == (fresh,)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_drop_team_when_everything_was_already_sent(self):
|
||||
cache: Final = DualCache()
|
||||
await cache.async_set_cache(key=email_sent_key("t1", "gpt-old", date(2026, 10, 1), 7), value=1.0)
|
||||
models: Final = (AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),)
|
||||
notifications: Final = await build_team_notifications(
|
||||
{"t1": models}, (_admin_team("t1"),), cache, _prisma_with_users(ADMIN_USERS)
|
||||
)
|
||||
assert notifications == ()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_drop_team_without_admin_emails(self):
|
||||
models: Final = (AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7),)
|
||||
no_admins: Final = LiteLLM_TeamTable(team_id="t1", members_with_roles=[{"user_id": "u1", "role": "user"}])
|
||||
notifications: Final = await build_team_notifications(
|
||||
{"t1": models}, (no_admins,), DualCache(), _prisma_with_users(ADMIN_USERS)
|
||||
)
|
||||
assert notifications == ()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_send_a_second_deployment_sharing_a_model_name_with_a_later_date(self):
|
||||
cache: Final = DualCache()
|
||||
await cache.async_set_cache(key=email_sent_key("t1", "gpt-old", date(2026, 10, 1), 7), value=1.0)
|
||||
sent: Final = AffectedModel(info=_info("gpt-old", 5), display_name="gpt-old", milestone=7)
|
||||
other_date: Final = AffectedModel(
|
||||
info=_info("gpt-old", 5, deprecation_date=date(2027, 3, 1)), display_name="gpt-old", milestone=7
|
||||
)
|
||||
notifications: Final = await build_team_notifications(
|
||||
{"t1": (sent, other_date)}, (_admin_team("t1"),), cache, _prisma_with_users(ADMIN_USERS)
|
||||
)
|
||||
assert notifications[0].models == (other_date,)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_key_the_sent_marker_by_team_model_date_and_milestone(self):
|
||||
assert (
|
||||
email_sent_key("t1", "gpt-old", date(2026, 10, 1), 7) == "model_deprecation_email:t1:gpt-old:2026-10-01:7"
|
||||
)
|
||||
assert email_sent_key("t1", "gpt-old", date(2026, 10, 1), 7) != email_sent_key(
|
||||
"t2", "gpt-old", date(2026, 10, 1), 7
|
||||
)
|
||||
assert email_sent_key("t1", "gpt-old", date(2026, 10, 1), 7) != email_sent_key(
|
||||
"t1", "gpt-old", date(2026, 10, 1), 0
|
||||
)
|
||||
assert email_sent_key("t1", "gpt-old", date(2026, 10, 1), 7) != email_sent_key(
|
||||
"t1", "gpt-old", date(2027, 3, 1), 7
|
||||
)
|
||||
|
||||
|
||||
TODAY: Final = datetime.now(timezone.utc).date()
|
||||
DEAD_DATE: Final = TODAY - timedelta(days=10)
|
||||
LONG_DEAD_DATE: Final = TODAY - timedelta(days=200)
|
||||
UPCOMING_DATE: Final = TODAY + timedelta(days=20)
|
||||
FAR_DATE: Final = TODAY + timedelta(days=40)
|
||||
|
||||
DEAD_DEPLOYMENT: Final = {
|
||||
"model_name": "dead-alias",
|
||||
"litellm_params": {"model": "openai/dead-model"},
|
||||
"model_info": {"id": "1", "deprecation_date": DEAD_DATE.isoformat(), "litellm_provider": "openai"},
|
||||
}
|
||||
LONG_DEAD_DEPLOYMENT: Final = {
|
||||
"model_name": "ancient-alias",
|
||||
"litellm_params": {"model": "openai/ancient-model"},
|
||||
"model_info": {"id": "4", "deprecation_date": LONG_DEAD_DATE.isoformat(), "litellm_provider": "openai"},
|
||||
}
|
||||
UPCOMING_DEPLOYMENT: Final = {
|
||||
"model_name": "sunset-alias",
|
||||
"litellm_params": {"model": "openai/sunset-model"},
|
||||
"model_info": {"id": "2", "deprecation_date": UPCOMING_DATE.isoformat(), "litellm_provider": "openai"},
|
||||
}
|
||||
FRESH_DEPLOYMENT: Final = {
|
||||
"model_name": "fresh",
|
||||
"litellm_params": {"model": "openai/fresh-model-with-no-cost-map-entry"},
|
||||
"model_info": {"id": "x"},
|
||||
}
|
||||
TODAY_DEPLOYMENT: Final = {
|
||||
"model_name": "today-alias",
|
||||
"litellm_params": {"model": "openai/today-model"},
|
||||
"model_info": {"id": "5", "deprecation_date": TODAY.isoformat(), "litellm_provider": "openai"},
|
||||
}
|
||||
FAR_DEPLOYMENT: Final = {
|
||||
"model_name": "far-alias",
|
||||
"litellm_params": {"model": "openai/far-model"},
|
||||
"model_info": {"id": "3", "deprecation_date": FAR_DATE.isoformat(), "litellm_provider": "openai"},
|
||||
}
|
||||
|
||||
|
||||
class _Deliverer:
|
||||
def __init__(self, fail_for: Sequence[str] = ()):
|
||||
self.fail_for = fail_for
|
||||
self.sent = []
|
||||
|
||||
async def __call__(self, recipients: Sequence[str], subject: str, html_body: str) -> None:
|
||||
if any(recipient in self.fail_for for recipient in recipients):
|
||||
raise ConnectionError("smtp down")
|
||||
self.sent.append((tuple(recipients), subject))
|
||||
|
||||
|
||||
class _Lock:
|
||||
def __init__(self, result: bool | None):
|
||||
self.result = result
|
||||
self.calls = []
|
||||
|
||||
async def acquire_lock(self, cronjob_id: str, ttl: int | None = None, allow_reentrant: bool = True) -> bool | None:
|
||||
self.calls.append({"cronjob_id": cronjob_id, "ttl": ttl, "allow_reentrant": allow_reentrant})
|
||||
return self.result
|
||||
|
||||
|
||||
class _RecordingCache:
|
||||
def __init__(self):
|
||||
self.values = {}
|
||||
self.ttls = {}
|
||||
|
||||
async def async_get_cache(self, *, key: str) -> object:
|
||||
return self.values.get(key)
|
||||
|
||||
async def async_set_cache(self, *, key: str, value: float, ttl: float) -> None:
|
||||
self.values[key] = value
|
||||
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,
|
||||
"team_alias": alias,
|
||||
"models": [],
|
||||
"members_with_roles": [{"user_id": f"admin-{team_id}", "role": "admin"}],
|
||||
}
|
||||
|
||||
|
||||
def _prisma(team_rows: Sequence[Mapping[str, object]], user_rows: Sequence[Mapping[str, object]]) -> SimpleNamespace:
|
||||
async def find_many_teams(where=None, take=None, skip=None, order=None):
|
||||
return list(team_rows)
|
||||
|
||||
async def find_many_users(where=None, take=None, skip=None, order=None):
|
||||
wanted: Final = where["user_id"]["in"]
|
||||
return [row for row in user_rows if row["user_id"] in wanted]
|
||||
|
||||
return SimpleNamespace(
|
||||
db=SimpleNamespace(
|
||||
litellm_teamtable=SimpleNamespace(find_many=find_many_teams),
|
||||
litellm_usertable=SimpleNamespace(find_many=find_many_users),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
TWO_TEAMS: Final = (_team_row("t1", "Alpha"), _team_row("t2", "Beta"))
|
||||
TWO_ADMINS: Final = (
|
||||
{"user_id": "admin-t1", "user_email": "a@x.io"},
|
||||
{"user_id": "admin-t2", "user_email": "b@x.io"},
|
||||
)
|
||||
|
||||
|
||||
def _context(router, prisma, deliver, cache=None, pod_lock_manager=None, alerting_args=None) -> DeprecationEmailContext:
|
||||
return DeprecationEmailContext(
|
||||
llm_router=router,
|
||||
prisma_client=prisma,
|
||||
cache=cache or DualCache(),
|
||||
alerting_args=alerting_args or SlackAlertingArgs(),
|
||||
pod_lock_manager=pod_lock_manager,
|
||||
deliver=deliver,
|
||||
)
|
||||
|
||||
|
||||
class TestSendModelDeprecationEmails:
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_email_each_affected_team_once_and_stamp_the_milestones(self):
|
||||
cache: Final = DualCache()
|
||||
deliverer: Final = _Deliverer()
|
||||
ctx: Final = _context(_router([DEAD_DEPLOYMENT]), _prisma(TWO_TEAMS, TWO_ADMINS), deliverer, cache)
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == 2
|
||||
assert sorted(recipients for recipients, _ in deliverer.sent) == [("a@x.io",), ("b@x.io",)]
|
||||
assert all("deprecated" in subject for _, subject in deliverer.sent)
|
||||
assert await cache.async_get_cache(key=email_sent_key("t1", "dead-alias", DEAD_DATE, 0)) is not None
|
||||
assert await cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value) is not None
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == 0
|
||||
assert len(deliverer.sent) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_touch_nothing_when_no_model_deprecates(self):
|
||||
cache: Final = DualCache()
|
||||
deliverer: Final = _Deliverer()
|
||||
lock: Final = _Lock(True)
|
||||
|
||||
async def explode(where=None, take=None, skip=None, order=None):
|
||||
raise AssertionError("teams must not be loaded when nothing deprecates")
|
||||
|
||||
prisma: Final = _prisma((), ())
|
||||
prisma.db.litellm_teamtable.find_many = explode
|
||||
|
||||
assert (
|
||||
await send_model_deprecation_emails(_context(_router([FRESH_DEPLOYMENT]), prisma, deliverer, cache, lock))
|
||||
== 0
|
||||
)
|
||||
assert deliverer.sent == []
|
||||
assert lock.calls == []
|
||||
assert await cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_stamp_the_pass_but_send_nothing_when_every_milestone_was_already_sent(self):
|
||||
cache: Final = DualCache()
|
||||
await cache.async_set_cache(key=email_sent_key("t1", "dead-alias", DEAD_DATE, 0), value=1.0)
|
||||
lock: Final = _Lock(True)
|
||||
deliverer: Final = _Deliverer()
|
||||
ctx: Final = _context(_router([DEAD_DEPLOYMENT]), _prisma(TWO_TEAMS[:1], TWO_ADMINS), deliverer, cache, lock)
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == 0
|
||||
assert deliverer.sent == []
|
||||
assert lock.calls == []
|
||||
assert await cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value) is not None
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("lock_result", "expected_sent"),
|
||||
[(True, 1), (None, 1), (False, 0)],
|
||||
ids=["lock won", "no redis lock", "another pod holds the lock"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_send_only_from_the_pod_holding_the_daily_lock(self, lock_result, expected_sent):
|
||||
cache: Final = DualCache()
|
||||
lock: Final = _Lock(lock_result)
|
||||
deliverer: Final = _Deliverer()
|
||||
ctx: Final = _context(_router([DEAD_DEPLOYMENT]), _prisma(TWO_TEAMS[:1], TWO_ADMINS), deliverer, cache, lock)
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == expected_sent
|
||||
assert len(deliverer.sent) == expected_sent
|
||||
assert lock.calls == [
|
||||
{
|
||||
"cronjob_id": EMAIL_MODEL_DEPRECATION_LOCK_ID,
|
||||
"ttl": DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
|
||||
"allow_reentrant": False,
|
||||
}
|
||||
]
|
||||
stamped: Final = await cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value)
|
||||
assert (stamped is not None) is (expected_sent == 1)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_not_send_or_stamp_when_the_lock_is_lost(self):
|
||||
cache: Final = DualCache()
|
||||
deliverer: Final = _Deliverer()
|
||||
prisma: Final = _prisma(TWO_TEAMS[:1], TWO_ADMINS)
|
||||
|
||||
assert (
|
||||
await send_model_deprecation_emails(
|
||||
_context(_router([DEAD_DEPLOYMENT]), prisma, deliverer, cache, _Lock(False))
|
||||
)
|
||||
== 0
|
||||
)
|
||||
assert deliverer.sent == []
|
||||
assert await cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_raise_before_claiming_the_lock_when_loading_teams_fails(self):
|
||||
cache: Final = DualCache()
|
||||
lock: Final = _Lock(True)
|
||||
|
||||
async def explode(where=None, take=None, skip=None, order=None):
|
||||
raise RuntimeError("db down")
|
||||
|
||||
prisma: Final = _prisma((), ())
|
||||
prisma.db.litellm_teamtable.find_many = explode
|
||||
|
||||
with pytest.raises(RuntimeError, match="db down"):
|
||||
await send_model_deprecation_emails(_context(_router([DEAD_DEPLOYMENT]), prisma, _Deliverer(), cache, lock))
|
||||
|
||||
assert lock.calls == []
|
||||
assert await cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_email_the_thirty_day_milestone_for_an_upcoming_model(self):
|
||||
cache: Final = DualCache()
|
||||
deliverer: Final = _Deliverer()
|
||||
ctx: Final = _context(_router([UPCOMING_DEPLOYMENT]), _prisma(TWO_TEAMS[:1], TWO_ADMINS), deliverer, cache)
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == 1
|
||||
assert deliverer.sent[0][1] == "[LiteLLM] 1 model(s) deprecating for team Alpha"
|
||||
assert await cache.async_get_cache(key=email_sent_key("t1", "sunset-alias", UPCOMING_DATE, 30)) is not None
|
||||
assert await cache.async_get_cache(key=email_sent_key("t1", "sunset-alias", UPCOMING_DATE, 7)) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_email_an_upcoming_model_when_a_threshold_reaches_past_the_slack_window(self):
|
||||
cache: Final = DualCache()
|
||||
deliverer: Final = _Deliverer()
|
||||
ctx: Final = _context(
|
||||
_router([FAR_DEPLOYMENT]),
|
||||
_prisma(TWO_TEAMS[:1], TWO_ADMINS),
|
||||
deliverer,
|
||||
cache,
|
||||
alerting_args=SlackAlertingArgs(model_deprecation_email_thresholds=(45, 7, 0)),
|
||||
)
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == 1
|
||||
assert deliverer.sent[0][1] == "[LiteLLM] 1 model(s) deprecating for team Alpha"
|
||||
assert await cache.async_get_cache(key=email_sent_key("t1", "far-alias", FAR_DATE, 45)) is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_stamp_sent_keys_with_the_configured_ttl_and_the_pass_for_a_day(self):
|
||||
cache: Final = _RecordingCache()
|
||||
deliverer: Final = _Deliverer()
|
||||
ctx: Final = _context(_router([DEAD_DEPLOYMENT]), _prisma(TWO_TEAMS[:1], TWO_ADMINS), deliverer, cache)
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == 1
|
||||
assert (
|
||||
cache.ttls[email_sent_key("t1", "dead-alias", DEAD_DATE, 0)]
|
||||
== SlackAlertingArgs().model_deprecation_email_ttl
|
||||
)
|
||||
assert (
|
||||
cache.ttls[SlackAlertingCacheKeys.deprecation_email_pass_key.value]
|
||||
== DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_not_email_a_model_deprecated_longer_ago_than_the_marker_lifetime(self):
|
||||
cache: Final = DualCache()
|
||||
deliverer: Final = _Deliverer()
|
||||
lock: Final = _Lock(True)
|
||||
ctx: Final = _context(
|
||||
_router([LONG_DEAD_DEPLOYMENT]), _prisma(TWO_TEAMS[:1], TWO_ADMINS), deliverer, cache, lock
|
||||
)
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == 0
|
||||
assert deliverer.sent == []
|
||||
assert lock.calls == []
|
||||
assert await cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value) is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_still_email_the_day_zero_milestone_when_the_marker_lifetime_is_under_a_day(self):
|
||||
deliverer: Final = _Deliverer()
|
||||
ctx: Final = _context(
|
||||
_router([TODAY_DEPLOYMENT, DEAD_DEPLOYMENT]),
|
||||
_prisma(TWO_TEAMS[:1], TWO_ADMINS),
|
||||
deliverer,
|
||||
alerting_args=SlackAlertingArgs(model_deprecation_email_ttl=3600),
|
||||
)
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == 1
|
||||
assert len(deliverer.sent) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_keep_going_and_leave_keys_unstamped_when_one_team_fails(self):
|
||||
cache: Final = DualCache()
|
||||
deliverer: Final = _Deliverer(fail_for=("a@x.io",))
|
||||
ctx: Final = _context(_router([DEAD_DEPLOYMENT]), _prisma(TWO_TEAMS, TWO_ADMINS), deliverer, cache)
|
||||
|
||||
assert await send_model_deprecation_emails(ctx) == 1
|
||||
assert [recipients for recipients, _ in deliverer.sent] == [("b@x.io",)]
|
||||
assert await cache.async_get_cache(key=email_sent_key("t1", "dead-alias", DEAD_DATE, 0)) is None
|
||||
assert await cache.async_get_cache(key=email_sent_key("t2", "dead-alias", DEAD_DATE, 0)) is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_report_whether_the_daily_pass_already_ran():
|
||||
cache: Final = DualCache()
|
||||
assert await email_pass_done_today(cache) is False
|
||||
await cache.async_set_cache(key=SlackAlertingCacheKeys.deprecation_email_pass_key.value, value=1.0)
|
||||
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):
|
||||
calls: Final = []
|
||||
|
||||
async def send_email(from_email, to_email, subject, html_body):
|
||||
calls.append((from_email, tuple(to_email), subject, html_body))
|
||||
|
||||
logger: Final = SimpleNamespace(DEFAULT_LITELLM_EMAIL="noreply@litellm.ai", send_email=send_email)
|
||||
deliver: Final = make_email_deliverer(logger)
|
||||
|
||||
await deliver(("a@x.io", "b@x.io"), "subj", "<p>hi</p>")
|
||||
assert calls == [
|
||||
("noreply@litellm.ai", ("a@x.io",), "subj", "<p>hi</p>"),
|
||||
("noreply@litellm.ai", ("b@x.io",), "subj", "<p>hi</p>"),
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_fall_back_to_smtp_per_recipient_without_a_logger(self):
|
||||
calls: Final = []
|
||||
|
||||
async def smtp_send(*, receiver_email, subject, html, raise_on_error):
|
||||
assert raise_on_error is True
|
||||
calls.append((receiver_email, subject, html))
|
||||
|
||||
deliver: Final = make_email_deliverer(None, smtp_send=smtp_send)
|
||||
|
||||
await deliver(("a@x.io", "b@x.io"), "subj", "<p>hi</p>")
|
||||
assert calls == [("a@x.io", "subj", "<p>hi</p>"), ("b@x.io", "subj", "<p>hi</p>")]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_route_the_enterprise_smtp_logger_through_the_oss_helper(self):
|
||||
smtp_email: Final = pytest.importorskip("litellm_enterprise.enterprise_callbacks.send_emails.smtp_email")
|
||||
calls: Final = []
|
||||
|
||||
async def smtp_send(*, receiver_email, subject, html, raise_on_error):
|
||||
assert raise_on_error is True
|
||||
calls.append(receiver_email)
|
||||
|
||||
deliver: Final = make_email_deliverer(smtp_email.SMTPEmailLogger(), smtp_send=smtp_send)
|
||||
|
||||
await deliver(("a@x.io", "b@x.io"), "subj", "<p>hi</p>")
|
||||
assert calls == ["a@x.io", "b@x.io"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_raise_when_smtp_is_not_configured(self, monkeypatch):
|
||||
monkeypatch.delenv("SMTP_SENDER_EMAIL", raising=False)
|
||||
deliver: Final = make_email_deliverer(None)
|
||||
|
||||
with pytest.raises(ValueError, match="SMTP_SENDER_EMAIL"):
|
||||
await deliver(("a@x.io",), "subj", "<p>hi</p>")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_skip_the_no_op_enterprise_base_logger(self):
|
||||
base_email: Final = pytest.importorskip("litellm_enterprise.enterprise_callbacks.send_emails.base_email")
|
||||
calls: Final = []
|
||||
|
||||
async def smtp_send(*, receiver_email, subject, html, raise_on_error):
|
||||
assert raise_on_error is True
|
||||
calls.append(receiver_email)
|
||||
|
||||
deliver: Final = make_email_deliverer(base_email.BaseEmailLogger(), smtp_send=smtp_send)
|
||||
|
||||
await deliver(("a@x.io",), "subj", "<p>hi</p>")
|
||||
assert calls == ["a@x.io"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_use_a_configured_enterprise_logger_subclass(self):
|
||||
base_email: Final = pytest.importorskip("litellm_enterprise.enterprise_callbacks.send_emails.base_email")
|
||||
calls: Final = []
|
||||
|
||||
class ConfiguredLogger(base_email.BaseEmailLogger):
|
||||
async def send_email(self, from_email, to_email, subject, html_body):
|
||||
calls.append((from_email, tuple(to_email)))
|
||||
|
||||
deliver: Final = make_email_deliverer(ConfiguredLogger())
|
||||
|
||||
await deliver(("a@x.io", "b@x.io"), "subj", "<p>hi</p>")
|
||||
assert calls == [
|
||||
(base_email.BaseEmailLogger.DEFAULT_LITELLM_EMAIL, ("a@x.io",)),
|
||||
(base_email.BaseEmailLogger.DEFAULT_LITELLM_EMAIL, ("b@x.io",)),
|
||||
]
|
||||
|
|
@ -22,6 +22,7 @@ from litellm.proxy.utils import (
|
|||
ProxyLogging,
|
||||
_failure_fields_to_lift,
|
||||
handle_exception_on_proxy,
|
||||
send_email,
|
||||
)
|
||||
from litellm.router_utils.fallback_event_handlers import carry_superseded_attempt_usage
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
|
@ -4041,6 +4042,21 @@ class TestPrismaClientTokenAuthBehindThePool:
|
|||
assert client.db.reader.iam_token_db_auth is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_email_raises_transport_errors_only_when_asked(monkeypatch):
|
||||
monkeypatch.setenv("SMTP_HOST", "127.0.0.1")
|
||||
monkeypatch.setenv("SMTP_PORT", "9")
|
||||
monkeypatch.setenv("SMTP_SENDER_EMAIL", "proxy@example.com")
|
||||
monkeypatch.setenv("SMTP_TIMEOUT", "5")
|
||||
monkeypatch.delenv("SMTP_USERNAME", raising=False)
|
||||
monkeypatch.delenv("SMTP_PASSWORD", raising=False)
|
||||
|
||||
with pytest.raises(ConnectionRefusedError):
|
||||
await send_email(receiver_email="a@x.io", subject="s", html="<p>x</p>", raise_on_error=True)
|
||||
|
||||
assert await send_email(receiver_email="a@x.io", subject="s", html="<p>x</p>") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"])
|
||||
def test_mcp_conversion_preserves_request_policy_and_isolates_guardrail_data(bucket):
|
||||
from copy import deepcopy
|
||||
|
|
|
|||
0
tests/unit/types/integrations/__init__.py
Normal file
0
tests/unit/types/integrations/__init__.py
Normal file
44
tests/unit/types/integrations/test_slack_alerting.py
Normal file
44
tests/unit/types/integrations/test_slack_alerting.py
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
"""Tests for the model deprecation email settings on SlackAlertingArgs."""
|
||||
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.constants import EMAIL_MODEL_DEPRECATION_LOCK_ID
|
||||
from litellm.types.integrations.slack_alerting import (
|
||||
SlackAlertingArgs,
|
||||
SlackAlertingCacheKeys,
|
||||
)
|
||||
|
||||
|
||||
def test_should_default_thresholds_to_30_7_0():
|
||||
assert SlackAlertingArgs().model_deprecation_email_thresholds == (30, 7, 0)
|
||||
|
||||
|
||||
def test_should_default_ttl_to_90_days():
|
||||
assert SlackAlertingArgs().model_deprecation_email_ttl == 90 * 24 * 60 * 60
|
||||
|
||||
|
||||
def test_should_dedupe_and_sort_thresholds_descending():
|
||||
args: Final = SlackAlertingArgs(model_deprecation_email_thresholds=[0, 7, 7, 30])
|
||||
assert args.model_deprecation_email_thresholds == (30, 7, 0)
|
||||
|
||||
|
||||
def test_should_allow_empty_thresholds_to_disable_emails():
|
||||
assert SlackAlertingArgs(model_deprecation_email_thresholds=[]).model_deprecation_email_thresholds == ()
|
||||
|
||||
|
||||
def test_should_reject_negative_thresholds():
|
||||
with pytest.raises(ValidationError, match="non-negative"):
|
||||
SlackAlertingArgs(model_deprecation_email_thresholds=[30, -1])
|
||||
|
||||
|
||||
def test_should_reject_non_positive_ttl():
|
||||
with pytest.raises(ValidationError, match="greater than or equal to 1"):
|
||||
SlackAlertingArgs(model_deprecation_email_ttl=0)
|
||||
|
||||
|
||||
def test_should_expose_pass_cache_key_and_lock_id():
|
||||
assert SlackAlertingCacheKeys.deprecation_email_pass_key.value == "model_deprecation_email_pass_at"
|
||||
assert EMAIL_MODEL_DEPRECATION_LOCK_ID == "email_model_deprecation_warning"
|
||||
Loading…
Add table
Reference in a new issue