diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 2814ed4eafd..8a45dfb84d5 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -104,6 +104,9 @@ class AnthropicMessagesHandler(BaseTranslation): from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( FakeAnthropicMessagesStreamIterator, ) + from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_response_usage, + ) from litellm.types.utils import AnthropicMessagesResponse block_response = AnthropicMessagesResponse( @@ -113,7 +116,7 @@ class AnthropicMessagesHandler(BaseTranslation): content=[{"type": "text", "text": exc.message}], model=exc.model, stop_reason="end_turn", - usage={"input_tokens": 0, "output_tokens": 0}, + usage=blocked_response_usage(getattr(exc, "original_response", None)), ) return list(FakeAnthropicMessagesStreamIterator(response=block_response)) @@ -122,9 +125,16 @@ class AnthropicMessagesHandler(BaseTranslation): append the block message as a new text block, then end the message -- without a second message_start.""" + from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_response_usage, + ) + def _sse(event_type: str, payload: dict) -> bytes: return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() + output_tokens = blocked_response_usage(getattr(exc, "original_response", None))[ + "output_tokens" + ] open_index, max_index = self._content_block_state(responses_so_far) new_index = (max_index + 1) if max_index is not None else 0 chunks: list[bytes] = [] @@ -153,7 +163,7 @@ class AnthropicMessagesHandler(BaseTranslation): { "type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, - "usage": {"output_tokens": 0}, + "usage": {"output_tokens": output_tokens}, }, ), _sse("message_stop", {"type": "message_stop"}), diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 97ece6b5eab..aade7180617 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,10 +1,38 @@ from __future__ import annotations -from typing import Any, List +from typing import Any, List, Optional +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues +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"), + ) + + def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool: per = getattr(guardrail_to_apply, "skip_system_message_in_guardrail", None) if per is not None: diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1008880d6db..92279a9e685 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -12,6 +12,9 @@ from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.anthropic.experimental_pass_through.context_management import ( AnthropicContextManagementError, ) +from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_response_usage as _blocked_response_usage, +) from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ( @@ -19,7 +22,6 @@ 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() @@ -59,33 +61,6 @@ 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`"], 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 c325c8ba186..531d9a0c412 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -319,8 +319,9 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply: CustomGuardrail = request_data.pop("guardrail_to_apply", None) - # Get streaming configuration, with precedence: guardrail attribute -> - # guardrail_config dict -> this callback's optional_params. + # Get streaming configuration. Resolution order (later wins): default + # < guardrail attribute < guardrail_config dict < this callback's + # optional_params. def _streaming_flag(name: str, default: Any) -> Any: value = default if guardrail_to_apply is not None: @@ -336,11 +337,25 @@ class UnifiedLLMGuardrails(CustomLogger): # Withhold every chunk until end-of-stream moderation passes, then # release the original chunks (clean) or only the block message # (blocked) -- moderating the whole response *before* any content - # reaches the client. Intended for allow/block guardrails: on release - # the original chunks are replayed as-is, so content-rewriting - # guardrails (e.g. PII masking) are not applied to a buffered stream. + # reaches the client. Only safe for allow/block guardrails: on + # release the original chunks are replayed as-is, so a + # content-rewriting guardrail (e.g. PII masking) would leak + # unredacted content. Guarded below via mask_response_content. buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", False) + if ( + buffer_until_moderated + and guardrail_to_apply is not None + and getattr(guardrail_to_apply, "mask_response_content", False) + ): + verbose_proxy_logger.warning( + "UnifiedLLMGuardrails: streaming_buffer_until_moderated is disabled for %s " + "because mask_response_content=True -- buffered replay would release " + "unredacted original chunks instead of the moderated output.", + guardrail_to_apply.guardrail_name, + ) + buffer_until_moderated = False + # Buffering can only moderate the assembled response, so it always # defers to end-of-stream. if buffer_until_moderated: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index 857389c6af2..4a2d8efc12f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py @@ -157,3 +157,19 @@ async def test_buffered_clean_releases_all_content(): raw.rstrip().endswith('event: message_stop\ndata: {"type": "message_stop"}'.rstrip()) or "message_stop" in raw ) assert BLOCK_MESSAGE not in raw + + +@pytest.mark.asyncio +async def test_buffered_mode_disabled_for_content_rewriting_guardrail(): + """Buffered replay yields the withheld *original* chunks verbatim, which + is unsafe for a guardrail that rewrites response text (e.g. PII masking): + the client would get the unredacted original instead of the moderated + output. mask_response_content=True must force buffering off so the + request falls back to the (correctly moderated) non-buffered path.""" + guardrail = _PassingGuardrail( + guardrail_name="masker", event_hook="post_call", mask_response_content=True + ) + raw = await _run(guardrail) + assert guardrail.streaming_buffer_until_moderated is True # request asked for buffering + assert ORIGINAL_MARKER in raw + assert BLOCK_MESSAGE not in raw