From c3fb28654d5895fa1f15083fcf4be6040d385b40 Mon Sep 17 00:00:00 2001 From: Joseph Barker <156112794+seph-barker@users.noreply.github.com> Date: Wed, 1 Jul 2026 20:57:33 -0700 Subject: [PATCH] fix: report the blocked LLM response's real token usage (#31217) When a guardrail blocks a post-call response, the synthetic violation response reported hard-coded zero usage, discarding the token usage the upstream call had already consumed. Fix the root cause rather than re-counting tokens: - Add an optional `original_response` field to ModifyResponseException. - The unified guardrail's post-call success hook attaches the blocked LLM response to the exception. - The /v1/messages and OpenAI-format (/v1/chat/completions, /v1/completions) block handlers report `original_response.usage` directly. Pre-call blocks never invoked the LLM, so usage is zero. Mock-based tests cover the helper (returns original usage / zero), the success hook attaching original_response, and the endpoint reporting it end-to-end. Co-authored-by: Claude Opus 4.8 (1M context) --- litellm/exceptions.py | 6 ++ .../proxy/anthropic_endpoints/endpoints.py | 34 +++++++- .../unified_guardrail/unified_guardrail.py | 27 ++++-- litellm/proxy/proxy_server.py | 33 +++++--- .../anthropic_endpoints/test_endpoints.py | 76 ++++++++++++++--- .../proxy/test_blocked_response_usage.py | 84 +++++++++++++++++++ 6 files changed, 229 insertions(+), 31 deletions(-) create mode 100644 tests/test_litellm/proxy/test_blocked_response_usage.py diff --git a/litellm/exceptions.py b/litellm/exceptions.py index d97ba347b07..adf7b3ef05a 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1165,12 +1165,18 @@ class ModifyResponseException(Exception): request_data: Dict[str, Any], guardrail_name: Optional[str] = None, detection_info: Optional[Dict[str, Any]] = None, + original_response: Optional[Any] = None, ): self.message = message self.model = model self.request_data = request_data self.guardrail_name = guardrail_name self.detection_info = detection_info or {} + # The LLM response that was blocked (post-call). Carries the real token + # usage the upstream call consumed, so the synthetic block response can + # report it instead of discarding it. None for pre-call blocks (the LLM + # was never invoked). + self.original_response = original_response super().__init__(message) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 71acc1f3106..1008880d6db 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -19,6 +19,7 @@ from litellm.proxy.common_request_processing import ( create_response, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.utils import TokenCountResponse router = APIRouter() @@ -58,6 +59,33 @@ def _strip_total_tokens_from_anthropic_response(response: Any) -> None: usage.pop("total_tokens", None) +def _blocked_response_usage(original_response: Optional[Any]) -> AnthropicUsage: + """ + Token usage for a synthetic guardrail-blocked response. + + A post-call block replaces the LLM's response with the violation message, + but the upstream call already consumed tokens -- report that real usage + (carried on ``ModifyResponseException.original_response``) rather than + discarding it. Pre-call blocks never invoked the LLM (no original_response), + so usage is zero. + """ + usage_obj: Any = None + if isinstance(original_response, dict): + usage_obj = original_response.get("usage") + elif original_response is not None: + usage_obj = getattr(original_response, "usage", None) + + def _tokens(key: str) -> int: + if isinstance(usage_obj, dict): + return int(usage_obj.get(key, 0) or 0) + return int(getattr(usage_obj, key, 0) or 0) + + return AnthropicUsage( + input_tokens=_tokens("input_tokens"), + output_tokens=_tokens("output_tokens"), + ) + + @router.post( "/v1/messages", tags=["[beta] Anthropic `/v1/messages`"], @@ -134,6 +162,10 @@ async def anthropic_response( from litellm.types.utils import AnthropicMessagesResponse + # Report the blocked LLM response's real token usage (carried on the + # exception) instead of discarding it; zero for pre-call blocks. + _usage = _blocked_response_usage(e.original_response) + _anthropic_response = AnthropicMessagesResponse( id=f"msg_{str(uuid.uuid4())}", type="message", @@ -141,7 +173,7 @@ async def anthropic_response( content=[{"type": "text", "text": e.message}], model=e.model, stop_reason="end_turn", - usage={"input_tokens": 0, "output_tokens": 0}, + usage=_usage, ) if data.get("stream", None) is not None and data["stream"] is True: diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 79c29670d93..83b787aaaf4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -197,6 +197,10 @@ class UnifiedLLMGuardrails(CustomLogger): ) from litellm.types.guardrails import GuardrailEventHooks + # Local import avoids a module-level cyclic import with + # litellm.integrations.custom_guardrail. + from litellm.integrations.custom_guardrail import ModifyResponseException + guardrail_to_apply: CustomGuardrail = data.pop("guardrail_to_apply", None) if guardrail_to_apply is None: @@ -238,13 +242,22 @@ class UnifiedLLMGuardrails(CustomLogger): endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() - response = await endpoint_translation.process_output_response( - response=response, # type: ignore - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=data.get("litellm_logging_obj"), - user_api_key_dict=user_api_key_dict, - request_data=data, - ) + try: + response = await endpoint_translation.process_output_response( + response=response, # type: ignore + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, + request_data=data, + ) + except ModifyResponseException as e: + # The guardrail blocked the response. Attach the original LLM + # response so the endpoint handler can report its real token usage + # instead of discarding it (the block replaces the content, but the + # upstream call already consumed those tokens). + if e.original_response is None: + e.original_response = response + raise # Add guardrail to applied guardrails header add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_to_apply.guardrail_name) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a29acf3b06f..bb925799bb8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8361,6 +8361,22 @@ async def model_info( ) +def _blocked_response_usage(original_response: Optional[Any]) -> "litellm.Usage": + """ + Token usage for a synthetic guardrail-blocked response. + + A post-call block replaces the LLM's response with the violation message, + but the upstream call already consumed tokens -- report that real usage + (carried on ``ModifyResponseException.original_response``) rather than + discarding it. Pre-call blocks never invoked the LLM (no original_response), + so usage is zero. + """ + usage = getattr(original_response, "usage", None) if original_response is not None else None + if isinstance(usage, litellm.Usage): + return usage + return litellm.Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + + @router.post( "/v1/chat/completions", dependencies=[Depends(user_api_key_auth)], @@ -8467,6 +8483,9 @@ async def chat_completion( _chat_response.model = e.model # type: ignore _chat_response.choices[0].message.content = e.message # type: ignore _chat_response.choices[0].finish_reason = "content_filter" # type: ignore + # Report the blocked LLM response's real usage (set before the stream + # branch so both paths carry it); zero for pre-call blocks. + _chat_response.usage = _blocked_response_usage(e.original_response) # type: ignore if data.get("stream", None) is not None and data["stream"] is True: _iterator = litellm.utils.ModelResponseIterator(model_response=_chat_response, convert_to_delta=True) @@ -8488,8 +8507,6 @@ async def chat_completion( media_type="text/event-stream", status_code=200, # Return 200 for passthrough mode ) - _usage = litellm.Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) - _chat_response.usage = _usage # type: ignore return _chat_response except RejectedRequestError as e: _data = e.request_data @@ -8618,11 +8635,7 @@ async def completion( # Set text attribute dynamically for text completion format setattr(_text_response.choices[0], "text", e.message) _text_response.model = e.model # type: ignore[assignment] - _usage = litellm.Usage( - prompt_tokens=0, - completion_tokens=0, - total_tokens=0, - ) + _usage = _blocked_response_usage(e.original_response) # Set usage attribute dynamically (ModelResponse accepts usage in __init__ but it's not in type definition) setattr(_text_response, "usage", _usage) _iterator = litellm.utils.ModelResponseIterator(model_response=_text_response, convert_to_delta=True) @@ -8647,11 +8660,7 @@ async def completion( _response = litellm.TextCompletionResponse() _response.choices[0].text = e.message _response.model = e.model # type: ignore - _usage = litellm.Usage( - prompt_tokens=0, - completion_tokens=0, - total_tokens=0, - ) + _usage = _blocked_response_usage(e.original_response) _response.usage = _usage # type: ignore return _response except RejectedRequestError as e: diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index a4da4587b7f..69d90a8b59b 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -58,15 +58,71 @@ class TestAnthropicEndpoints(unittest.TestCase): self.assertEqual(result, expected_result) # Assert safe_dumps was called for dictionary objects - mock_safe_dumps.assert_any_call( - {"type": "message_start", "message": {"id": "msg_123"}} + mock_safe_dumps.assert_any_call({"type": "message_start", "message": {"id": "msg_123"}}) + mock_safe_dumps.assert_any_call({"type": "content_block_delta", "delta": {"text": "more data"}}) + assert mock_safe_dumps.call_count == 2 # Called twice, once for each dict object + + +class TestBlockedResponseUsage: + """Blocked responses report the blocked LLM response's real usage.""" + + def test_uses_original_response_usage(self): + from litellm.proxy.anthropic_endpoints.endpoints import _blocked_response_usage + + # original_response is the AnthropicMessagesResponse the LLM produced + # before the guardrail blocked it; its usage is real. + original = {"usage": {"input_tokens": 31, "output_tokens": 9}} + assert _blocked_response_usage(original) == { + "input_tokens": 31, + "output_tokens": 9, + } + + def test_zero_usage_when_no_original_response(self): + from litellm.proxy.anthropic_endpoints.endpoints import _blocked_response_usage + + # Pre-call blocks never invoked the LLM -> nothing consumed. + assert _blocked_response_usage(None) == { + "input_tokens": 0, + "output_tokens": 0, + } + + @pytest.mark.asyncio + async def test_blocked_endpoint_response_carries_original_usage(self): + """The /v1/messages block handler reports the blocked response's real + usage, carried on ModifyResponseException.original_response.""" + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.integrations.custom_guardrail import ModifyResponseException + + exc = ModifyResponseException( + message="blocked by guardrail", + model="claude-3-5-sonnet-20240620", + request_data={"messages": [{"role": "user", "content": "hi"}]}, + guardrail_name="rubrik", + original_response={"usage": {"input_tokens": 12, "output_tokens": 5}}, ) - mock_safe_dumps.assert_any_call( - {"type": "content_block_delta", "delta": {"text": "more data"}} - ) - assert ( - mock_safe_dumps.call_count == 2 - ) # Called twice, once for each dict object + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})), + patch.object( + ep.ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + new=AsyncMock(side_effect=exc), + ), + patch.object(proxy_server, "proxy_logging_obj") as mock_logging, + ): + mock_logging.post_call_failure_hook = AsyncMock() + response = await ep.anthropic_response( + fastapi_response=MagicMock(), + request=MagicMock(), + user_api_key_dict=MagicMock(), + ) + + assert response["content"][0]["text"] == "blocked by guardrail" + assert response["usage"] == {"input_tokens": 12, "output_tokens": 5} + mock_logging.post_call_failure_hook.assert_awaited_once() class TestEventLoggingBatchEndpoint: @@ -159,9 +215,7 @@ class TestStripTotalTokens(unittest.TestCase): # SimpleNamespace mimics the .usage attribute access pattern; the # helper's contract: if .usage is dict-shaped, strip total_tokens. - response = SimpleNamespace( - usage={"input_tokens": 100, "output_tokens": 50, "total_tokens": 150} - ) + response = SimpleNamespace(usage={"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}) _strip_total_tokens_from_anthropic_response(response) assert "total_tokens" not in response.usage assert response.usage == {"input_tokens": 100, "output_tokens": 50} diff --git a/tests/test_litellm/proxy/test_blocked_response_usage.py b/tests/test_litellm/proxy/test_blocked_response_usage.py new file mode 100644 index 00000000000..d486431ca3e --- /dev/null +++ b/tests/test_litellm/proxy/test_blocked_response_usage.py @@ -0,0 +1,84 @@ +""" +Token usage on synthetic guardrail-blocked responses for the OpenAI-format +proxy endpoints (/v1/chat/completions and /v1/completions). + +A post-call block replaces the LLM response with the violation message, but the +upstream call already consumed tokens. `_blocked_response_usage` reports that +real usage (carried on `ModifyResponseException.original_response`) rather than +zero; a pre-call block never invoked the LLM, so usage is zero. +""" + +import pytest + +import litellm +from litellm.proxy.proxy_server import _blocked_response_usage + + +def test_uses_original_response_usage(): + resp = litellm.ModelResponse() + resp.usage = litellm.Usage(prompt_tokens=42, completion_tokens=7, total_tokens=49) + + usage = _blocked_response_usage(resp) + + assert usage.prompt_tokens == 42 + assert usage.completion_tokens == 7 + assert usage.total_tokens == 49 + + +def test_zero_usage_when_no_original_response(): + usage = _blocked_response_usage(None) + + assert usage.prompt_tokens == 0 + assert usage.completion_tokens == 0 + assert usage.total_tokens == 0 + + +@pytest.mark.asyncio +async def test_success_hook_attaches_original_response_on_block(): + """The unified guardrail's post-call success hook must attach the blocked + LLM response to ModifyResponseException so its real usage isn't discarded.""" + from unittest.mock import AsyncMock, MagicMock, patch + + import litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail as ug + from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import CallTypes + + response = litellm.ModelResponse() + response.usage = litellm.Usage(prompt_tokens=15, completion_tokens=3, total_tokens=18) + + guardrail = MagicMock() + guardrail.should_run_guardrail.return_value = True + guardrail.guardrail_name = "rubrik" + + # The translation layer raises a block without pre-setting original_response. + translation = MagicMock() + translation.process_output_response = AsyncMock( + side_effect=ModifyResponseException( + message="blocked", + model="gpt-4o", + request_data={}, + guardrail_name="rubrik", + ) + ) + + unified = ug.UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions") + data = {"guardrail_to_apply": guardrail, "model": "gpt-4o"} + + # Inject our translation for the inferred call type (the module global is + # cached across tests, so patch it directly rather than the loader). + with patch.object( + ug, + "endpoint_guardrail_translation_mappings", + { + CallTypes.acompletion: lambda: translation, + CallTypes.completion: lambda: translation, + }, + ): + with pytest.raises(ModifyResponseException) as excinfo: + await unified.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + + assert excinfo.value.original_response is response