mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
797353f13a
Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
327 lines
13 KiB
Python
327 lines
13 KiB
Python
"""Pin alerting helpers on ``ProxyLogging``.
|
|
|
|
Covers ``failed_tracking_alert``, ``budget_alerts``, ``alerting_handler``,
|
|
``failure_handler``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import Any, Dict
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
from prisma.errors import PrismaError
|
|
|
|
import litellm
|
|
from litellm._service_logger import ServiceTypes
|
|
from litellm.proxy._types import AlertType, CallInfo
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# failed_tracking_alert
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_tracking_alert_no_op_when_alerting_none(proxy_logging):
|
|
proxy_logging.alerting = None
|
|
proxy_logging.slack_alerting_instance = MagicMock(failed_tracking_alert=AsyncMock())
|
|
await proxy_logging.failed_tracking_alert(error_message="x", failing_model="m")
|
|
proxy_logging.slack_alerting_instance.failed_tracking_alert.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_tracking_alert_forwards_to_slack(proxy_logging):
|
|
proxy_logging.alerting = ["slack"]
|
|
captured: Dict[str, Any] = {}
|
|
|
|
async def fake_alert(**kwargs):
|
|
captured.update(kwargs)
|
|
|
|
proxy_logging.slack_alerting_instance = MagicMock(failed_tracking_alert=fake_alert)
|
|
await proxy_logging.failed_tracking_alert(error_message="db down", failing_model="gpt-4")
|
|
snapshot = {
|
|
"error_message": captured["error_message"],
|
|
"failing_model": captured["failing_model"],
|
|
"captured_keys": sorted(captured.keys()),
|
|
}
|
|
assert snapshot == {
|
|
"error_message": "db down",
|
|
"failing_model": "gpt-4",
|
|
"captured_keys": ["error_message", "failing_model"],
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_tracking_alert_slack_error_raises(proxy_logging):
|
|
proxy_logging.alerting = ["slack"]
|
|
proxy_logging.slack_alerting_instance = MagicMock(
|
|
failed_tracking_alert=AsyncMock(side_effect=RuntimeError("slack down"))
|
|
)
|
|
with pytest.raises(RuntimeError):
|
|
await proxy_logging.failed_tracking_alert(error_message="x", failing_model="m")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# budget_alerts
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _user_info(alert_emails=None):
|
|
return CallInfo(
|
|
spend=0.0,
|
|
max_budget=1.0,
|
|
token="tok",
|
|
user_id="u1",
|
|
team_id="t1",
|
|
team_alias=None,
|
|
user_email=None,
|
|
key_alias=None,
|
|
projected_exceeded_date=None,
|
|
projected_spend=None,
|
|
event_group="user",
|
|
event="threshold_crossed",
|
|
alert_emails=alert_emails,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_budget_alerts_no_op_when_alerting_off_and_no_emails(proxy_logging):
|
|
proxy_logging.alerting = None
|
|
proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=AsyncMock())
|
|
proxy_logging.email_logging_instance = MagicMock(budget_alerts=AsyncMock())
|
|
await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info())
|
|
proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called()
|
|
proxy_logging.email_logging_instance.budget_alerts.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_budget_alerts_slack_when_slack_alerting(proxy_logging):
|
|
proxy_logging.alerting = ["slack"]
|
|
captured: Dict[str, Any] = {}
|
|
|
|
async def fake_alert(**kwargs):
|
|
captured.update(kwargs)
|
|
|
|
proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=fake_alert)
|
|
proxy_logging.email_logging_instance = None
|
|
user_info = _user_info()
|
|
await proxy_logging.budget_alerts(type="user_budget", user_info=user_info)
|
|
snapshot = {
|
|
"type": captured["type"],
|
|
"user_info_is_callinfo": isinstance(captured["user_info"], CallInfo),
|
|
"user_id": captured["user_info"].user_id,
|
|
}
|
|
assert snapshot == {"type": "user_budget", "user_info_is_callinfo": True, "user_id": "u1"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_budget_alerts_webhook_only_forwards_to_slack_alerting_instance(proxy_logging):
|
|
proxy_logging.alerting = ["webhook"]
|
|
captured: Dict[str, Any] = {}
|
|
|
|
async def fake_alert(**kwargs):
|
|
captured.update(kwargs)
|
|
|
|
proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=fake_alert)
|
|
proxy_logging.email_logging_instance = None
|
|
await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info())
|
|
snapshot = {
|
|
"type": captured["type"],
|
|
"user_info_is_callinfo": isinstance(captured["user_info"], CallInfo),
|
|
"user_id": captured["user_info"].user_id,
|
|
}
|
|
assert snapshot == {"type": "user_budget", "user_info_is_callinfo": True, "user_id": "u1"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_budget_alerts_email_only_skips_slack_alerting_instance(proxy_logging):
|
|
proxy_logging.alerting = ["email"]
|
|
proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=AsyncMock())
|
|
proxy_logging.email_logging_instance = MagicMock(budget_alerts=AsyncMock())
|
|
await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info())
|
|
proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called()
|
|
proxy_logging.email_logging_instance.budget_alerts.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_budget_alerts_soft_budget_with_alert_emails_bypasses_global(proxy_logging):
|
|
proxy_logging.alerting = None
|
|
proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=AsyncMock())
|
|
proxy_logging.email_logging_instance = MagicMock(budget_alerts=AsyncMock())
|
|
info = _user_info(alert_emails=["a@b.c"])
|
|
await proxy_logging.budget_alerts(type="soft_budget", user_info=info)
|
|
proxy_logging.email_logging_instance.budget_alerts.assert_called_once()
|
|
proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_budget_alerts_slack_failure_raises(proxy_logging):
|
|
proxy_logging.alerting = ["slack"]
|
|
proxy_logging.slack_alerting_instance = MagicMock(
|
|
budget_alerts=AsyncMock(side_effect=ConnectionError("slack"))
|
|
)
|
|
proxy_logging.email_logging_instance = None
|
|
with pytest.raises(ConnectionError):
|
|
await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# alerting_handler
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alerting_handler_no_op_when_alerting_is_none(proxy_logging):
|
|
proxy_logging.alerting = None
|
|
proxy_logging.slack_alerting_instance = MagicMock(send_alert=AsyncMock())
|
|
await proxy_logging.alerting_handler(message="x", level="High", alert_type=AlertType.db_exceptions)
|
|
proxy_logging.slack_alerting_instance.send_alert.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alerting_handler_sends_to_slack(proxy_logging):
|
|
proxy_logging.alerting = ["slack"]
|
|
captured: Dict[str, Any] = {}
|
|
|
|
async def fake_send(**kwargs):
|
|
captured.update(kwargs)
|
|
|
|
proxy_logging.slack_alerting_instance = MagicMock(send_alert=fake_send)
|
|
await proxy_logging.alerting_handler(
|
|
message="hi", level="High", alert_type=AlertType.db_exceptions, request_data={"metadata": {}}
|
|
)
|
|
snapshot = {
|
|
"message": captured["message"],
|
|
"level": captured["level"],
|
|
"alert_type": captured["alert_type"],
|
|
"user_info": captured["user_info"],
|
|
}
|
|
assert snapshot == {
|
|
"message": "hi",
|
|
"level": "High",
|
|
"alert_type": AlertType.db_exceptions,
|
|
"user_info": None,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_alerting_handler_sentry_without_sdk_error_raises(proxy_logging, monkeypatch):
|
|
proxy_logging.alerting = ["sentry"]
|
|
monkeypatch.setattr(litellm.utils, "sentry_sdk_instance", None)
|
|
with pytest.raises(Exception, match="SENTRY_DSN"):
|
|
await proxy_logging.alerting_handler(message="x", level="Low", alert_type=AlertType.db_exceptions)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# failure_handler
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failure_handler_skips_when_db_exceptions_not_in_alert_types(proxy_logging):
|
|
proxy_logging.alert_types = ["llm_too_slow"] # type: ignore[list-item]
|
|
proxy_logging.alerting_handler = AsyncMock()
|
|
proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock())
|
|
await proxy_logging.failure_handler(original_exception=Exception("x"), duration=1.0, call_type="db_read")
|
|
proxy_logging.alerting_handler.assert_not_called()
|
|
proxy_logging.service_logging_obj.async_service_failure_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failure_handler_logs_db_error_and_calls_service_logging(proxy_logging, monkeypatch):
|
|
proxy_logging.alert_types = [AlertType.db_exceptions]
|
|
proxy_logging.alerting_handler = AsyncMock()
|
|
proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock())
|
|
monkeypatch.setattr(litellm.utils, "capture_exception", None)
|
|
await proxy_logging.failure_handler(
|
|
original_exception=HTTPException(status_code=500, detail="boom"),
|
|
duration=1.5,
|
|
call_type="db_write",
|
|
)
|
|
call_kwargs = proxy_logging.service_logging_obj.async_service_failure_hook.call_args.kwargs
|
|
snapshot = {
|
|
"service": call_kwargs["service"].value if hasattr(call_kwargs["service"], "value") else call_kwargs["service"],
|
|
"duration": call_kwargs["duration"],
|
|
"call_type": call_kwargs["call_type"],
|
|
}
|
|
assert snapshot == {
|
|
"service": "postgres",
|
|
"duration": 1.5,
|
|
"call_type": "db_write",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("call_type", ["get_data", "insert_data", "update_data", "delete_data"])
|
|
async def test_failure_handler_alerts_but_leaves_prisma_error_event_to_log_db_metrics(
|
|
proxy_logging, monkeypatch, call_type
|
|
):
|
|
proxy_logging.alert_types = [AlertType.db_exceptions]
|
|
proxy_logging.alerting_handler = AsyncMock()
|
|
proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock())
|
|
monkeypatch.setattr(litellm.utils, "capture_exception", None)
|
|
await proxy_logging.failure_handler(
|
|
original_exception=PrismaError("connection reset"), duration=1.0, call_type=call_type
|
|
)
|
|
snapshot = {
|
|
"alerting_handler_scheduled": proxy_logging.alerting_handler.called,
|
|
"service_failure_called": proxy_logging.service_logging_obj.async_service_failure_hook.called,
|
|
}
|
|
assert snapshot == {"alerting_handler_scheduled": True, "service_failure_called": False}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failure_handler_still_emits_db_event_for_wrapped_insert_error(proxy_logging, monkeypatch):
|
|
proxy_logging.alert_types = [AlertType.db_exceptions]
|
|
proxy_logging.alerting_handler = AsyncMock()
|
|
proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock())
|
|
monkeypatch.setattr(litellm.utils, "capture_exception", None)
|
|
await proxy_logging.failure_handler(
|
|
original_exception=HTTPException(status_code=400, detail={"error": "Foreign Key Constraint failed"}),
|
|
duration=1.0,
|
|
call_type="insert_data",
|
|
)
|
|
call_kwargs = proxy_logging.service_logging_obj.async_service_failure_hook.call_args.kwargs
|
|
assert (call_kwargs["service"], call_kwargs["call_type"]) == (ServiceTypes.DB, "insert_data")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failure_handler_with_capture_exception_invoked(proxy_logging, monkeypatch):
|
|
proxy_logging.alert_types = [AlertType.db_exceptions]
|
|
proxy_logging.alerting_handler = AsyncMock()
|
|
proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock())
|
|
captured: Dict[str, Any] = {}
|
|
|
|
def fake_capture(error):
|
|
captured["error"] = error
|
|
|
|
monkeypatch.setattr(litellm.utils, "capture_exception", fake_capture)
|
|
err = RuntimeError("real")
|
|
await proxy_logging.failure_handler(original_exception=err, duration=1.0, call_type="db_read")
|
|
snapshot = {
|
|
"captured_is_input": captured["error"] is err,
|
|
"service_failure_called": proxy_logging.service_logging_obj.async_service_failure_hook.called,
|
|
"alerting_handler_scheduled": proxy_logging.alerting_handler.called,
|
|
}
|
|
assert snapshot == {
|
|
"captured_is_input": True,
|
|
"service_failure_called": True,
|
|
"alerting_handler_scheduled": True,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failure_handler_propagates_service_logging_error_raises(proxy_logging, monkeypatch):
|
|
proxy_logging.alert_types = [AlertType.db_exceptions]
|
|
proxy_logging.alerting_handler = AsyncMock()
|
|
proxy_logging.service_logging_obj = MagicMock(
|
|
async_service_failure_hook=AsyncMock(side_effect=RuntimeError("svc"))
|
|
)
|
|
monkeypatch.setattr(litellm.utils, "capture_exception", None)
|
|
with pytest.raises(RuntimeError):
|
|
await proxy_logging.failure_handler(
|
|
original_exception=Exception("x"), duration=0.0, call_type="db_read"
|
|
)
|