diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 586ccc77985..8208a269b84 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -956,11 +956,24 @@ LITELLM_EXCEPTION_TYPES = [ class BudgetExceededError(Exception): def __init__( - self, current_cost: float, max_budget: float, message: Optional[str] = None + self, + current_cost: float, + max_budget: float, + message: Optional[str] = None, + llm_provider: Optional[str] = None, ): self.current_cost = current_cost self.max_budget = max_budget self.status_code = 429 + self.llm_provider = llm_provider or "" + # Surface unified rate-limit fields without joining the RateLimitError + # hierarchy so existing `except BudgetExceededError:` handlers keep + # working; custom callbacks reading StandardLoggingPayload pick these + # up via the same `category` / `rate_limit_type` attributes the rest + # of the unified rate-limit error path uses. Stored as plain strings + # to match the normalization RateLimitError.__init__ performs. + self.category: str = RateLimitErrorCategory.LITELLM_RATE_LIMIT.value + self.rate_limit_type: str = RateLimitType.BUDGET.value message = ( message or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f919bc7705d..4ce4014d077 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2685,15 +2685,16 @@ class PrometheusLogger(CustomLogger): ) -> Tuple[Optional[str], Optional[str]]: """ Pull the unified ``category`` / ``rate_limit_type`` fields off any - :class:`litellm.RateLimitError` (vendor-side or litellm-internal) so - Prometheus can split 429s by source + dimension without the consumer - parsing free-text error messages. + exception that declares them (``litellm.RateLimitError`` and bare- + Exception subclasses like ``BudgetExceededError`` that set these + attributes directly) so Prometheus can split 429s by source + + dimension without the consumer parsing free-text error messages. - Returns ``(None, None)`` for non-rate-limit exceptions. Both values are - sanitized to plain ``str`` so str-enum subclasses don't leak their - repr into the label. + Returns ``(None, None)`` for exceptions that don't declare these + fields. Both classes normalize their values to plain ``str`` at + construction, so this helper only needs to coerce defensively. """ - if exception is None or not isinstance(exception, litellm.RateLimitError): + if exception is None: return None, None def _coerce(value: Any) -> Optional[str]: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3bcb5aa5fc8..e1585398831 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5159,22 +5159,18 @@ class StandardLoggingPayloadSetup: # Get additional error details error_message = str(original_exception) - # For rate-limit errors (litellm.RateLimitError + the proxy-side - # ProxyRateLimitError subclass), surface the unified `category` and - # `rate_limit_type` fields so callbacks can distinguish vendor vs. - # litellm rate limits AND split by the dimension that was exceeded - # (requests / tokens / concurrent_requests / budget / max_iterations) - # without reaching for the raw exception object. - is_rate_limit_error = isinstance(original_exception, litellm.RateLimitError) - rate_limit_category: Optional[str] = ( - getattr(original_exception, "category", None) - if is_rate_limit_error - else None + # Surface the unified `category` and `rate_limit_type` fields off any + # exception that opts in by setting them. Duck-typed rather than + # isinstance-gated on RateLimitError so bare-Exception subclasses like + # `litellm.BudgetExceededError` can participate without joining the + # RateLimitError hierarchy (which would break `except BudgetExceededError`). + # Both RateLimitError and BudgetExceededError normalize their values to + # plain strings at construction. + rate_limit_category: Optional[str] = getattr( + original_exception, "category", None ) - rate_limit_type: Optional[str] = ( - getattr(original_exception, "rate_limit_type", None) - if is_rate_limit_error - else None + rate_limit_type: Optional[str] = getattr( + original_exception, "rate_limit_type", None ) return StandardLoggingPayloadErrorInformation( diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py index 25030a11ff8..76ab363ed6f 100644 --- a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py @@ -94,6 +94,20 @@ def test_should_return_none_for_none_exception(): assert PrometheusLogger._extract_rate_limit_labels(None) == (None, None) +def test_should_extract_budget_dimension_for_budget_exceeded_error(): + # Virtual-key / team / org / end-user budget caps raise + # `litellm.BudgetExceededError` (a bare Exception subclass), which sets + # the same `.category` / `.rate_limit_type` attributes as the unified + # RateLimitError path so Prometheus can split budget 429s from other + # 429s without the customer parsing free-text error messages. + import litellm + + err = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels(err) + assert category == "litellm_rate_limit" + assert rate_limit_type == "budget" + + @pytest.mark.parametrize( "category_enum,rate_limit_enum,expected_category,expected_type", [ diff --git a/tests/test_litellm/test_rate_limit_error_unification.py b/tests/test_litellm/test_rate_limit_error_unification.py index 30c08690e59..2913492ed63 100644 --- a/tests/test_litellm/test_rate_limit_error_unification.py +++ b/tests/test_litellm/test_rate_limit_error_unification.py @@ -1419,3 +1419,91 @@ class TestProxyHooksWireTypeCorrectly: e = exc_info.value assert e.rate_limit_type == "requests" assert e.category == RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT + + +class TestBudgetExceededErrorSurfacesUnifiedFields: + """ + The hot path for virtual-key / team / org / end-user max_budget caps + raises :class:`litellm.BudgetExceededError`, which historically had no + relationship to :class:`RateLimitError` and therefore left the unified + `error_rate_limit_category` / `error_rate_limit_type` fields empty. + Test 2 of the QA pass surfaced this gap; this class pins the fix. + + The fix is intentionally additive: `BudgetExceededError` keeps its + bare-`Exception` base class (so existing `except BudgetExceededError:` + handlers keep working) and just sets the same `category` / + `rate_limit_type` attributes that the rest of the unified rate-limit + path reads (normalized to plain strings, matching how + `RateLimitError.__init__` stores its own values). Duck-typed dispatch + in `get_error_information` picks them up automatically. + """ + + def test_should_carry_litellm_rate_limit_category(self): + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + # Stored as the plain string value (matches RateLimitError behavior), + # but equality with the enum still works because the enum subclasses + # str. + assert e.category == "litellm_rate_limit" + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + + def test_should_carry_budget_rate_limit_type(self): + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert e.rate_limit_type == "budget" + assert e.rate_limit_type == RateLimitType.BUDGET + + def test_should_default_llm_provider_to_empty_string(self): + # `llm_provider` is read off the exception in `get_error_information` + # — it must always be a string so the StandardLoggingPayload field + # stays serializable. Default to "" when no caller passes one. + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert e.llm_provider == "" + + def test_should_accept_llm_provider_kwarg(self): + # Callers that have the resolved provider in scope (e.g. the + # auth-checks budget enforcement paths) can thread it through. + e = litellm.BudgetExceededError( + current_cost=0.5, max_budget=0.1, llm_provider="anthropic" + ) + assert e.llm_provider == "anthropic" + + def test_should_keep_existing_status_code_and_message(self): + # Backward-compat guard: existing callers depend on `status_code=429` + # and the canonical message format. + e = litellm.BudgetExceededError(current_cost=0.000109, max_budget=0.0001) + assert e.status_code == 429 + assert "Current cost: 0.000109" in e.message + assert "Max budget: 0.0001" in e.message + + def test_should_still_be_catchable_as_exception_not_rate_limit_error(self): + # Critical: we deliberately did NOT make BudgetExceededError a + # RateLimitError subclass. Existing `except BudgetExceededError:` + # handlers must keep catching it, and `except RateLimitError:` + # handlers must NOT start catching it (which would surprise callers + # who rely on the two being distinct). + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert isinstance(e, Exception) + assert isinstance(e, litellm.BudgetExceededError) + assert not isinstance(e, RateLimitError) + + def test_should_propagate_category_to_standard_logging_payload(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_category"] == "litellm_rate_limit" + assert info["error_rate_limit_type"] == "budget" + assert info["error_code"] == "429" + assert info["error_class"] == "BudgetExceededError" + + def test_should_propagate_llm_provider_to_standard_logging_payload(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = litellm.BudgetExceededError( + current_cost=0.5, max_budget=0.1, llm_provider="bedrock" + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["llm_provider"] == "bedrock"