diff --git a/litellm/__init__.py b/litellm/__init__.py index eebd2dad91e..743d4a363c3 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -399,6 +399,7 @@ budget_duration: Optional[str] = ( ) default_soft_budget: float = DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0 budget_exceeded_throttle_percentage: Optional[float] = None +budget_exceeded_error_message: Optional[str] = None forward_traceparent_to_llm_provider: bool = False diff --git a/litellm/constants.py b/litellm/constants.py index 23e92d26a59..cc8867454f2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1640,6 +1640,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ "cost_margin_config", "block_requests_for_models_without_pricing", "budget_exceeded_throttle_percentage", + "budget_exceeded_error_message", # Every field editable from the Admin UI (proxy_server._GENERAL_SETTINGS_UI_LITELLM_FIELDS) # must be listed here so a DB write from one worker overrides the live litellm attribute on # the others when config reloads; otherwise peer workers stay on their startup value. diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 286f7528896..e2dee778062 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -985,6 +985,12 @@ class BudgetExceededError(Exception): self.message = message super().__init__(message) + @property + def client_facing_message(self) -> str: + import litellm + + return litellm.budget_exceeded_error_message or self.message + ## DEPRECATED ## class InvalidRequestError(openai.BadRequestError): diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 7d85f3c4908..e784a5782a5 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1137,7 +1137,7 @@ class MCPRequestHandler: except (HTTPException, ProxyException): raise except litellm.BudgetExceededError as e: - raise HTTPException(status_code=getattr(e, "status_code", 429), detail=str(e)) from None + raise HTTPException(status_code=getattr(e, "status_code", 429), detail=e.client_facing_message) from None except Exception as e: # noqa: BLE001 # untyped gate failure: retryable 503 for a DB outage, else fail closed 401 MCPRequestHandler._raise_503_if_db_unavailable(e) raise HTTPException(status_code=401, detail="Invalid or expired credential") from None diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index a42187b3a44..f20887021c9 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -169,7 +169,7 @@ class UserAPIKeyAuthExceptionHandler: if isinstance(e, litellm.BudgetExceededError): raise ProxyException( - message=e.message, + message=e.client_facing_message, type=ProxyErrorTypes.budget_exceeded, param=None, code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS), diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index fb633870d21..cbc7b0b686c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -45,10 +45,11 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth + from litellm.litellm_core_utils.streaming_handler import ( backfill_missing_cache_usage_fields, ) -from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model from litellm.proxy.auth.auth_utils import check_response_size_is_safe from litellm.proxy.common_utils.callback_utils import ( @@ -3291,6 +3292,15 @@ class ProxyBaseLLMRequestProcessing: code=status.HTTP_400_BAD_REQUEST, headers=safe_headers, ) + if isinstance(e, litellm.BudgetExceededError): + raise ProxyException( + message=e.client_facing_message, + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=e.status_code, + headers=safe_headers, + ) + # Extract status_code from the exception if it carries one. # Provider exceptions (NotFoundError, BadRequestError, GeminiError, # VertexAIError, etc.) all have a status_code attribute reflecting diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 1322f50d4af..d6a8f5f6f3d 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -237,7 +237,7 @@ async def _authorize_models_this_test_can_call( ) except BudgetExceededError as e: raise ProxyException( - message=e.message, + message=e.client_facing_message, type=ProxyErrorTypes.budget_exceeded, param=None, code=status.HTTP_400_BAD_REQUEST, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 697c9b018ec..6a54c0a991c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -5669,15 +5669,17 @@ class TestMCPDcrBridgeDelegateAdmission: ) return exc_info.value - async def test_over_budget_admission_surfaces_429_not_401(self): + async def test_over_budget_admission_surfaces_429_not_401(self, monkeypatch): """A validly-authenticated but over-budget identity surfaces the standard pipeline's 429, not a misleading 401. Flattening budget to 401 told the caller their credential was invalid, which on a DCR client reads as broken auth and triggers a re-authorize that cannot fix a budget problem. Regression for the status-flattening finding on the live-policy gate.""" import litellm + monkeypatch.setattr(litellm, "budget_exceeded_error_message", "Allowance reached") mapped = await self._enforce_with_gate_error(litellm.BudgetExceededError(current_cost=10.0, max_budget=1.0)) assert mapped.status_code == 429 + assert mapped.detail == "Allowance reached" async def test_db_outage_during_policy_surfaces_503_not_401(self): """A transient database outage during the live-policy gate surfaces a retryable 503, not a 401 diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 90b3b29d919..c872123422d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -25,6 +25,7 @@ from prisma.errors import ( ) +import litellm from litellm._logging import verbose_proxy_logger from litellm.exceptions import BudgetExceededError from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth @@ -353,6 +354,41 @@ async def test_handle_authentication_error_budget_exceeded(): assert int(exc_info.value.code) == status.HTTP_429_TOO_MANY_REQUESTS +@pytest.mark.asyncio +async def test_handle_authentication_error_budget_exceeded_custom_message(monkeypatch): + from litellm.exceptions import BudgetExceededError + + monkeypatch.setattr( + litellm, + "budget_exceeded_error_message", + "Your AI usage allowance has been reached. Please contact the AI team.", + ) + + budget_error = BudgetExceededError( + message="Budget has been exceeded! Current cost: 10.51, Max budget: 10.00", + current_cost=10.51, + max_budget=10.00, + ) + + with pytest.raises(ProxyException) as exc_info: + await UserAPIKeyAuthExceptionHandler()._handle_authentication_error( + budget_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-test", + ) + + assert ( + exc_info.value.message + == "Your AI usage allowance has been reached. Please contact the AI team." + ) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert int(exc_info.value.code) == status.HTTP_429_TOO_MANY_REQUESTS + assert "10.51" in budget_error.message + + @pytest.mark.asyncio async def test_route_passed_to_post_call_failure_hook(): """ diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index b595c44d2ce..e5cf02660c3 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3314,6 +3314,28 @@ class TestHandleLLMApiExceptionRetryAfter: assert proxy_exc.headers["retry-after"] == "43" assert proxy_exc.headers["x-custom"] == "1" + async def test_handle_llm_api_exception_budget_exceeded_uses_custom_message( + self, monkeypatch + ): + monkeypatch.setattr( + litellm, "budget_exceeded_error_message", "Allowance reached, contact the AI team." + ) + exc = litellm.BudgetExceededError(current_cost=10.51, max_budget=10.0) + + proxy_exc = await self._invoke(exc) + + assert proxy_exc.message == "Allowance reached, contact the AI team." + assert proxy_exc.type == "budget_exceeded" + assert proxy_exc.code == "429" + + async def test_handle_llm_api_exception_budget_exceeded_defaults_to_detail(self): + exc = litellm.BudgetExceededError(current_cost=10.51, max_budget=10.0) + + proxy_exc = await self._invoke(exc) + + assert "10.51" in proxy_exc.message + assert proxy_exc.type == "budget_exceeded" + class TestHandleLLMApiExceptionFramingHeaders: """HTTP-framing headers on the provider exception must be stripped before the