From e4f2ea12bc5311f6c3ce18f3b45ddf92eaa36a42 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:04:27 -0700 Subject: [PATCH] fix(responses_api): map bridged chat usage on guardrail-blocked replies Move the blocked-usage mapping for /v1/responses next to blocked_response_usage in guardrail_translation utils, map bridged chat prompt/completion tokens to Responses API input/output tokens, and let raise_passthrough_exception attach the blocked response so post-call guardrail blocks report real usage --- litellm/integrations/custom_guardrail.py | 6 ++ .../base_llm/guardrail_translation/utils.py | 54 +++++++++--- .../proxy/response_api_endpoints/endpoints.py | 12 +-- .../response_api_endpoints/test_endpoints.py | 86 +++++++++++++++++++ .../proxy/test_blocked_response_usage.py | 46 +++++++++- 5 files changed, 180 insertions(+), 24 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2e91e082bd4..f721e01e2c8 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -198,6 +198,7 @@ class CustomGuardrail(CustomLogger): violation_message: str, request_data: dict[str, Any], detection_info: dict[str, Any] | None = None, + original_response: object = None, ) -> None: """ Raise a passthrough exception for guardrail violations. @@ -213,6 +214,10 @@ class CustomGuardrail(CustomLogger): violation_message: The formatted violation message to return to the user request_data: The original request data dictionary detection_info: Optional dictionary with detection metadata (scores, rules, etc.) + original_response: The blocked LLM response when raising from a post-call + hook. It carries the real token usage the upstream call consumed, so + the synthetic block response reports it instead of zeros. Leave None + for pre-call/during-call blocks (the LLM was never invoked). Raises: ModifyResponseException: Always raises this exception to short-circuit @@ -235,6 +240,7 @@ class CustomGuardrail(CustomLogger): request_data=request_data, guardrail_name=self.guardrail_name, detection_info=detection_info, + original_response=original_response, ) def raise_sensitive_data_route_exception( diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index f1ddf21cd3c..1546adbb0bd 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -5,7 +5,7 @@ from collections.abc import Callable, Iterator, Sequence from typing import Any, Final, TypeVar from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage def _anthropic_stream_chunk_events(item: Any) -> list[dict]: @@ -65,6 +65,20 @@ def _usage_from_anthropic_stream_chunks(original_response: list[Any]) -> Anthrop return AnthropicUsage(input_tokens=input_tokens, output_tokens=output_tokens) +def _blocked_usage_obj(original_response: object) -> object: + if isinstance(original_response, dict): + return original_response.get("usage") + if original_response is not None and not isinstance(original_response, list): + return getattr(original_response, "usage", None) + return None + + +def _usage_tokens(usage_obj: object, key: str, fallback_key: str) -> int: + if isinstance(usage_obj, dict): + return int(usage_obj.get(key, usage_obj.get(fallback_key, 0)) or 0) + return int(getattr(usage_obj, key, getattr(usage_obj, fallback_key, 0)) or 0) + + def blocked_response_usage(original_response: Any | None) -> AnthropicUsage: """ Token usage for a synthetic guardrail-blocked response. @@ -75,24 +89,38 @@ def blocked_response_usage(original_response: Any | None) -> AnthropicUsage: discarding it. Pre-call blocks never invoked the LLM (no original_response), so usage is zero. """ - usage_obj: Any = None if isinstance(original_response, list): stream_usage: Final = _usage_from_anthropic_stream_chunks(original_response) if stream_usage is not None: return stream_usage - elif 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, fallback_key: str) -> int: - if isinstance(usage_obj, dict): - return int(usage_obj.get(key, usage_obj.get(fallback_key, 0)) or 0) - return int(getattr(usage_obj, key, getattr(usage_obj, fallback_key, 0)) or 0) + usage_obj: Final = _blocked_usage_obj(original_response) return AnthropicUsage( - input_tokens=_tokens("input_tokens", "prompt_tokens"), - output_tokens=_tokens("output_tokens", "completion_tokens"), + input_tokens=_usage_tokens(usage_obj, "input_tokens", "prompt_tokens"), + output_tokens=_usage_tokens(usage_obj, "output_tokens", "completion_tokens"), + ) + + +def blocked_responses_api_usage(original_response: object) -> ResponseAPIUsage: + """ + Token usage for a synthetic guardrail-blocked /v1/responses reply. + + Same contract as ``blocked_response_usage`` in Responses API shape: a + native ``ResponsesAPIResponse`` usage passes through unchanged, a bridged + chat ``ModelResponse`` usage maps prompt/completion tokens to input/output + tokens, and a pre-call block (no original_response) reports zeros. + """ + usage_obj: Final = _blocked_usage_obj(original_response) + if isinstance(usage_obj, ResponseAPIUsage): + return usage_obj + + input_tokens: Final = _usage_tokens(usage_obj, "input_tokens", "prompt_tokens") + output_tokens: Final = _usage_tokens(usage_obj, "output_tokens", "completion_tokens") + total_tokens: Final = _usage_tokens(usage_obj, "total_tokens", "total_tokens") + return ResponseAPIUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens or input_tokens + output_tokens, ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index d85dda71d81..5e56e822484 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -12,6 +12,9 @@ from starlette.websockets import WebSocket, WebSocketDisconnect from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_responses_api_usage as _blocked_responses_api_usage, +) from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import ( UserAPIKeyAuth, @@ -23,7 +26,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_set_request_parsed_body, ) -from litellm.types.llms.openai import REASONING_EFFORT, ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.llms.openai import REASONING_EFFORT, ResponsesAPIResponse from litellm.types.responses.main import DeleteResponseResult if TYPE_CHECKING: @@ -164,13 +167,6 @@ async def _resolve_cursor_model_variant_before_auth(request: Request) -> None: _safe_set_request_parsed_body(request=request, parsed_body=resolved) -def _blocked_responses_api_usage(original_response: Any) -> ResponseAPIUsage: - usage: Final = getattr(original_response, "usage", None) if original_response is not None else None - if isinstance(usage, ResponseAPIUsage): - return usage - return ResponseAPIUsage(input_tokens=0, output_tokens=0, total_tokens=0) - - @router.post( "/v1/responses", dependencies=[Depends(user_api_key_auth)], diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 079454d963f..9177944df2d 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -1748,3 +1748,89 @@ class TestCursorGateRecognizesRoutingGroups: resolved = _resolve_cursor_model_variant(body, router) assert resolved["model"] == "grouped-thinking-high" assert "reasoning_effort" not in resolved + + +class TestGuardrailBlockedResponsesUsage: + """Regression tests for https://github.com/BerriAI/litellm/issues/36880. + + The ModifyResponseException handler in responses_api hardcoded the synthetic + blocked reply's usage to zeros, discarding the real token counts the blocked + upstream call consumed. The blocked reply must carry the usage from + e.original_response, exactly like /v1/chat/completions already does.""" + + def _post_blocked_responses(self, original_response): + from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + exc = ModifyResponseException( + message="Content flagged by policy, response withheld", + model="gpt-4o-mini", + request_data={"model": "gpt-4o-mini", "input": "hi"}, + guardrail_name="zero-usage-regression", + original_response=original_response, + ) + mock_proxy_logging = MagicMock() + mock_proxy_logging.post_call_failure_hook = AsyncMock() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="sk-test", request_route="/v1/responses" + ) + try: + with ( + patch( + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new=AsyncMock(side_effect=exc), + ), + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), + ): + client = TestClient(app) + return client.post( + "/v1/responses", + json={"model": "gpt-4o-mini", "input": "Write a haiku about token accounting"}, + headers={"Authorization": "Bearer sk-1234"}, + ) + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + def test_post_call_block_reports_real_upstream_usage(self): + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + original = ResponsesAPIResponse( + id="resp_upstream", + created_at=1, + model="gpt-4o-mini", + object="response", + output=[], + status="completed", + usage=ResponseAPIUsage(input_tokens=14, output_tokens=20, total_tokens=34), + ) + + response = self._post_blocked_responses(original) + + assert response.status_code == 200, response.text + body = response.json() + assert body["output"][0]["content"][0]["text"] == "Content flagged by policy, response withheld" + assert body["usage"]["input_tokens"] == 14 + assert body["usage"]["output_tokens"] == 20 + assert body["usage"]["total_tokens"] == 34 + + def test_post_call_block_maps_bridged_chat_usage(self): + original = litellm.ModelResponse() + original.usage = litellm.Usage(prompt_tokens=14, completion_tokens=18, total_tokens=32) + + response = self._post_blocked_responses(original) + + assert response.status_code == 200, response.text + usage = response.json()["usage"] + assert usage["input_tokens"] == 14 + assert usage["output_tokens"] == 18 + assert usage["total_tokens"] == 32 + + def test_pre_call_block_reports_zero_usage(self): + response = self._post_blocked_responses(None) + + assert response.status_code == 200, response.text + usage = response.json()["usage"] + assert usage["input_tokens"] == 0 + assert usage["output_tokens"] == 0 + assert usage["total_tokens"] == 0 diff --git a/tests/test_litellm/proxy/test_blocked_response_usage.py b/tests/test_litellm/proxy/test_blocked_response_usage.py index b02f6541f23..37aea8fe3aa 100644 --- a/tests/test_litellm/proxy/test_blocked_response_usage.py +++ b/tests/test_litellm/proxy/test_blocked_response_usage.py @@ -1,10 +1,11 @@ """ Token usage on synthetic guardrail-blocked responses for the OpenAI-format -proxy endpoints (/v1/chat/completions and /v1/completions). +proxy endpoints (/v1/chat/completions, /v1/completions, and /v1/responses). 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 +upstream call already consumed tokens. `_blocked_response_usage` (and its +Responses API counterpart `_blocked_responses_api_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. """ @@ -124,3 +125,42 @@ def test_responses_api_blocked_reply_zero_usage_when_no_original_response(): assert usage.input_tokens == 0 assert usage.output_tokens == 0 assert usage.total_tokens == 0 + + +def test_responses_api_blocked_reply_maps_bridged_chat_usage(): + """A chat model bridged through /v1/responses blocks with a ModelResponse whose + Usage fields must map prompt_tokens -> input_tokens and completion_tokens -> output_tokens.""" + from litellm.proxy.response_api_endpoints.endpoints import ( + _blocked_responses_api_usage, + ) + + resp = litellm.ModelResponse() + resp.usage = litellm.Usage(prompt_tokens=14, completion_tokens=18, total_tokens=32) + + usage = _blocked_responses_api_usage(resp) + + assert usage.input_tokens == 14 + assert usage.output_tokens == 18 + assert usage.total_tokens == 32 + + +def test_raise_passthrough_exception_attaches_original_response(): + """Post-call guardrails raising through the blessed helper must be able to + attach the blocked response so its real usage reaches the synthetic reply.""" + from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, + ) + + resp = litellm.ModelResponse() + resp.usage = litellm.Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7) + guardrail = CustomGuardrail(guardrail_name="passthrough-usage") + + with pytest.raises(ModifyResponseException) as excinfo: + guardrail.raise_passthrough_exception( + violation_message="blocked", + request_data={"model": "gpt-4o"}, + original_response=resp, + ) + + assert excinfo.value.original_response is resp