diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ddf31cb1d8a..afbf75ef39e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -82,7 +82,11 @@ from litellm._logging import _redact_string, verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes from litellm.caching.caching import DualCache, RedisCache from litellm.caching.dual_cache import LimitedSizeOrderedDict -from litellm.exceptions import RejectedRequestError, SensitiveDataRouteException +from litellm.exceptions import ( + GuardrailRaisedException, + RejectedRequestError, + SensitiveDataRouteException, +) from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -2628,6 +2632,7 @@ class ProxyLogging: - Authentication Errors from user_api_key_auth - HTTP HTTPException (rate limit errors) - ProxyException (guardrail blocks, budget / rate-limit errors) + - GuardrailRaisedException (guardrail blocks / guardrail failures) """ ######################################################### @@ -2642,9 +2647,9 @@ class ProxyLogging: if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)): return False - return isinstance(original_exception, (HTTPException, ProxyException)) or ( - error_type == ProxyErrorTypes.auth_error - ) + return isinstance( + original_exception, (HTTPException, ProxyException, GuardrailRaisedException) + ) or (error_type == ProxyErrorTypes.auth_error) async def _handle_logging_proxy_only_error( self, @@ -3181,7 +3186,7 @@ class ProxyLogging: except (GeneratorExit, asyncio.CancelledError): raise except Exception: - ProxyLogging._fire_deferred_stream_logging(request_data) + ProxyLogging._discard_deferred_stream_logging_for_failure(request_data) raise ProxyLogging._fire_deferred_stream_logging(request_data) return @@ -3241,7 +3246,7 @@ class ProxyLogging: except (GeneratorExit, asyncio.CancelledError): raise except Exception: - ProxyLogging._fire_deferred_stream_logging(request_data) + ProxyLogging._discard_deferred_stream_logging_for_failure(request_data) raise # Fire deferred logging AFTER all guardrail end-of-stream blocks @@ -3272,6 +3277,44 @@ class ProxyLogging: logging_obj._deferred_stream_complete_args = None asyncio.create_task(_deferred_cb(*_args)) + @staticmethod + def _discard_deferred_stream_logging_for_failure(request_data: dict) -> None: + """Discard the deferred stream-complete dispatch when the stream ends in + a failure (e.g. an end-of-stream guardrail block raising out of the + callback chain). The deferred dispatch is the success logging path — + firing it here would record the blocked request as a success callback + and a ``status=success`` spend row before the outer generator's + ``post_call_failure_hook`` writes the failure row. The CSW shape parks + ``(assembled ModelResponse, cache_hit)``; record its partial usage so + the failure row bills what the stream consumed instead of zero. The + native /v1/messages and responses shapes park ``(coroutine,)`` and + still need the flush (no success row is produced without it), so they + keep the existing fire behaviour. + """ + logging_obj: Final = request_data.get("litellm_logging_obj") + if logging_obj is None: + return + _deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr( + logging_obj, "_on_deferred_stream_complete", None + ) + _args: Final[tuple[object, ...] | None] = getattr( + logging_obj, "_deferred_stream_complete_args", None + ) + if _deferred_cb is None or _args is None: + return + assembled: Final = _args[0] + if not isinstance(assembled, ModelResponse): + ProxyLogging._fire_deferred_stream_logging(request_data) + return + logging_obj._on_deferred_stream_complete = None + logging_obj._deferred_stream_complete_args = None + usage: Final[Usage | None] = getattr(assembled, "usage", None) + if isinstance(usage, Usage): + logging_obj.record_partial_usage_for_failure( + usage, + logging_obj._response_cost_calculator(result=assembled) or 0.0, + ) + async def _arelease_max_parallel_requests_on_disconnect( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py index a2a57931d26..793132bf6ed 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -11,6 +11,7 @@ import pytest from fastapi import HTTPException import litellm +from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import AlertType, ProxyErrorTypes from litellm.proxy.utils import ProxyLogging @@ -49,12 +50,17 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging): error_type=ProxyErrorTypes.auth_error, route="/chat/completions", ), + "guardrail_raised_on_llm_route": proxy_logging._is_proxy_only_llm_api_error( + original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"), + route="/chat/completions", + ), } assert snapshot == { "no_route": False, "non_llm_route": False, "http_on_llm_route": True, "auth_short_circuit": True, + "guardrail_raised_on_llm_route": True, } @@ -267,3 +273,49 @@ async def test_handle_logging_proxy_only_path_propagates_async_failure_raises( route="/chat/completions", original_exception=Exception("x"), ) + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """A ``GuardrailRaisedException`` on an LLM route must reach the logging + object's ``async_failure_handler`` so custom loggers see a ``failure`` + status - without this, guardrail blocks produce only + ``post_call_failure_hook`` and no failure logging event.""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + recorded: dict[str, Any] = {} + + class _StatusRecorder(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + recorded["status"] = (kwargs.get("standard_logging_object") or {}).get("status") + + monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()]) + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test_guardrail_block_failure_cb", + function_id="test_guardrail_block_failure_cb", + ) + request_data = { + "litellm_logging_obj": logging_obj, + "litellm_call_id": "test_guardrail_block_failure_cb", + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {}, + } + proxy_logging.alert_types = [] + await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"), + user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"), + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + assert recorded["status"] == "failure" diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index ec5b994f147..68e37da8e3a 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -479,6 +479,89 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_ assert logging_obj._deferred_stream_complete_args is None +@pytest.mark.asyncio +async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + On /chat/completions streams the CSW shape parks + ``(assembled ModelResponse, cache_hit)`` as the deferred args. A guardrail + that raises ``GuardrailRaisedException`` at end of stream must NOT dispatch + that deferred success logging - the request is logged via the failure path + instead, with the consumed usage carried over so the failure row bills + correctly. + """ + from litellm.exceptions import GuardrailRaisedException + from litellm.types.utils import Usage + + events: List[Any] = [] + request_data: Dict[str, Any] = {"metadata": {}} + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test_chat_stream_guardrail_block", + function_id="test_chat_stream_guardrail_block", + ) + logging_obj.optional_params = {} + logging_obj.litellm_params = {} + logging_obj.standard_built_in_tools_params = None + + async def _dispatch_deferred_logging(*args): + events.append("success_dispatched") + + logging_obj._on_deferred_stream_complete = _dispatch_deferred_logging + request_data["litellm_logging_obj"] = logging_obj + + assembled = litellm.ModelResponse( + model="gpt-4o-mini", + choices=[{"index": 0, "message": {"role": "assistant", "content": "BANANA"}}], + usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8), + ) + + async def _upstream(): + yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "BAN"}}]} + yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "ANA"}}]} + logging_obj._deferred_stream_complete_args = (assembled, False) + + class _BlockingGuardrail(CustomLogger): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for chunk in response: + yield chunk + raise GuardrailRaisedException( + guardrail_name="g", message="blocked", blocked_content=True + ) + + monkeypatch.setattr(litellm, "callbacks", [_BlockingGuardrail()]) + + with pytest.raises(GuardrailRaisedException): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=_upstream(), + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + await asyncio.sleep(0) + await asyncio.sleep(0) + + snapshot = { + "events": events, + "callback_cleared": logging_obj._on_deferred_stream_complete is None, + "args_cleared": logging_obj._deferred_stream_complete_args is None, + "combined_usage_total_tokens": logging_obj.model_call_details["combined_usage_object"].total_tokens, + "response_cost_positive": logging_obj.model_call_details["response_cost"] > 0, + } + assert snapshot == { + "events": [], + "callback_cleared": True, + "args_cleared": True, + "combined_usage_total_tokens": 8, + "response_cost_positive": True, + } + + # --------------------------------------------------------------------------- # _fire_deferred_stream_logging # ---------------------------------------------------------------------------