scripts/retry_after_callback.py

"""LiteLLM proxy callback: inject Retry-After on 429 responses.

When all deployments are exhausted, LiteLLM returns 429 but may not
include a Retry-After header. This callback catches rate-limit failures
and ensures the header is present, set to ``cooldown_time`` (300 s).

Registered via ``litellm_settings.callbacks`` in the generated config.
"""

from __future__ import annotations

from typing import Optional

from fastapi import HTTPException

COOLDOWN_SECONDS = 300


class RetryAfterCallback:
    """CustomLogger hook that ensures 429 responses carry Retry-After."""

    async def async_post_call_success_hook(self, *args, **kwargs) -> None:
        return None

    async def async_post_call_failure_hook(
        self,
        request_data: dict,
        original_exception: Exception,
        user_api_key_dict: object,
        traceback_str: Optional[str] = None,
    ) -> Optional[HTTPException]:
        status = _extract_status(original_exception)
        if status != 429:
            return None

        existing_headers: dict = _extract_headers(original_exception)
        if existing_headers.get("retry-after"):
            return None

        return HTTPException(
            status_code=429,
            detail={"error": {
                "message": str(original_exception),
                "type": "rate_limit_error",
                "code": 429,
            }},
            headers={"retry-after": str(COOLDOWN_SECONDS)},
        )


def _extract_status(exc: Exception) -> int | None:
    for attr in ("status_code", "code"):
        val = getattr(exc, attr, None)  # nosemgrep: python-no-getattr -- Probing exception attributes from upstream library where type is unknown at compile time
        if isinstance(val, int):
            return val
    return None


def _extract_headers(exc: Exception) -> dict:
    headers = getattr(exc, "headers", None)  # nosemgrep: python-no-getattr -- Probing exception headers from upstream library
    if isinstance(headers, dict):
        return headers
    return {}


proxy_handler_instance = RetryAfterCallback()