mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-27 03:34:10 +00:00
(testing) add router unit testing for send_llm_exception_alert
, router_cooldown_event_callback
, cooldown utils (#6258)
* add router unit testing for send_llm_exception_alert * test router_cooldown_event_callback * test test_router_cooldown_event_callback_no_prometheus * test_router_cooldown_event_callback_no_deployment * test_router_cooldown_event_callback_no_deployment * add testing for test_should_run_cooldown_logic * test_increment_deployment_successes_for_current_minute_does_not_write_to_redis * test test_should_cooldown_deployment_allowed_fails_set_on_router * use helper for _is_allowed_fails_set_on_router * add complete testing for cooldown utils * move router unit tests * move router handle error * fix test_send_llm_exception_alert_no_logger
This commit is contained in:
parent
dee6de0105
commit
3ab2b86062
5 changed files with 564 additions and 26 deletions
|
@ -12,8 +12,10 @@ if TYPE_CHECKING:
|
|||
from litellm.router import Router as _Router
|
||||
|
||||
LitellmRouter = _Router
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
else:
|
||||
LitellmRouter = Any
|
||||
PrometheusLogger = Any
|
||||
|
||||
|
||||
async def router_cooldown_event_callback(
|
||||
|
@ -56,34 +58,38 @@ async def router_cooldown_event_callback(
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
# Trigger cooldown on Prometheus
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
get_custom_logger_compatible_class,
|
||||
)
|
||||
|
||||
# get the prometheus logger from in memory loggers
|
||||
prometheusLogger: Optional[CustomLogger] = get_custom_logger_compatible_class(
|
||||
logging_integration="prometheus",
|
||||
prometheusLogger: Optional[PrometheusLogger] = (
|
||||
_get_prometheus_logger_from_callbacks()
|
||||
)
|
||||
|
||||
if prometheusLogger is not None:
|
||||
prometheusLogger.set_deployment_complete_outage(
|
||||
litellm_model_name=_model_name,
|
||||
model_id=model_id,
|
||||
api_base=_api_base,
|
||||
api_provider=llm_provider,
|
||||
)
|
||||
|
||||
if isinstance(prometheusLogger, PrometheusLogger):
|
||||
prometheusLogger.set_deployment_complete_outage(
|
||||
litellm_model_name=_model_name,
|
||||
model_id=model_id,
|
||||
api_base=_api_base,
|
||||
api_provider=llm_provider,
|
||||
)
|
||||
|
||||
prometheusLogger.increment_deployment_cooled_down(
|
||||
litellm_model_name=_model_name,
|
||||
model_id=model_id,
|
||||
api_base=_api_base,
|
||||
api_provider=llm_provider,
|
||||
exception_status=str(exception_status),
|
||||
)
|
||||
prometheusLogger.increment_deployment_cooled_down(
|
||||
litellm_model_name=_model_name,
|
||||
model_id=model_id,
|
||||
api_base=_api_base,
|
||||
api_provider=llm_provider,
|
||||
exception_status=str(exception_status),
|
||||
)
|
||||
|
||||
return
|
||||
|
||||
|
||||
def _get_prometheus_logger_from_callbacks() -> Optional[PrometheusLogger]:
|
||||
"""
|
||||
Checks if prometheus is a initalized callback, if yes returns it
|
||||
"""
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
for _callback in litellm.callbacks:
|
||||
if isinstance(_callback, PrometheusLogger):
|
||||
return _callback
|
||||
|
||||
return None
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue