From 7c19ec02d56d69826dbe1abc278f92f64d019dfc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 23:56:43 -0700 Subject: [PATCH] fix(guardrails): apply the safeguard_results and stop_details rules to policy pipeline stream rewrites --- .../chat/guardrail_translation/handler.py | 62 ++++++++++++ .../anthropic/pass_through/messages/utils.py | 67 ++++++++++++- litellm/proxy/guardrails/anthropic_sse.py | 63 +++--------- .../test_anthropic_guardrail_handler.py | 95 ++++++++++++++++++- .../guardrail_hooks/test_tool_permission.py | 32 ++++--- .../proxy/guardrails/test_anthropic_sse.py | 17 +++- 6 files changed, 267 insertions(+), 69 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 6b6974becac..7654b56ad9a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -30,6 +30,7 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, is_provider_native_tool_dict, ) +from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details from litellm.llms.anthropic.pass_through.stream_assembly import ( build_complete_streaming_response, build_usage_only_response_from_chunks, @@ -1456,6 +1457,9 @@ class AnthropicMessagesHandler(BaseTranslation): return _SSEFieldRewrite("delta", "text", next(replacements)) cls._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta) + cls._redact_ended_stream_message_delta( + responses_so_far, kept_tool_use_ids=frozenset(cls._streamed_tool_use_ids(responses_so_far).values()) + ) @classmethod def _write_ended_stream_tool_call_rewrites( @@ -1523,6 +1527,64 @@ class AnthropicMessagesHandler(BaseTranslation): return None cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use) + streamed_tool_use_ids: Final = cls._streamed_tool_use_ids(responses_so_far) + cls._redact_ended_stream_message_delta( + responses_so_far, + kept_tool_use_ids=frozenset( + tool_use_id for index, tool_use_id in streamed_tool_use_ids.items() if index not in rewrites_by_block + ), + ) + + @classmethod + def _streamed_tool_use_ids(cls, responses_so_far: Sequence[object]) -> Mapping[int, str]: + """The id of each buffered ``tool_use`` block, by its content block index.""" + events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far) + tool_use_blocks: Final = ( + (event.get("index"), block) + for event in events + if event.get("type") == "content_block_start" + and (block := as_str_mapping(event.get("content_block"))) is not None + and block.get("type") == "tool_use" + ) + return MappingProxyType( + { + index: tool_use_id + for index, block in tool_use_blocks + if isinstance(index, int) and isinstance(tool_use_id := block.get("id"), str) + } + ) + + @classmethod + def _redact_ended_stream_message_delta( + cls, + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place + *, + kept_tool_use_ids: frozenset[str], + ) -> None: + """Hold the buffered ``message_delta`` to the rules a re-emitted rewrite follows: a verdict survives + only for a tool call delivered unchanged, and no free-text ``explanation`` describes the text the + guardrail rewrote.""" + + def rewrite_stop_details(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None + if delta is None or "stop_details" not in delta: + return None + return _SSEFieldRewrite( + "delta", "stop_details", redacted_stop_details(delta["stop_details"], keep_explanations=False) + ) + + def rewrite_safeguard_results(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None + if delta is None or "safeguard_results" not in delta: + return None + return _SSEFieldRewrite( + "delta", + "safeguard_results", + filtered_safeguard_results(delta["safeguard_results"], kept_tool_use_ids, keep_explanations=False), + ) + + cls._rewrite_ended_stream_events(responses_so_far, rewrite_stop_details) + cls._rewrite_ended_stream_events(responses_so_far, rewrite_safeguard_results) @staticmethod def _rewrite_ended_stream_events( diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index 67d3affd7c0..6f8fe4f49c0 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/litellm/llms/anthropic/pass_through/messages/utils.py @@ -2,7 +2,7 @@ from collections.abc import Iterable, Mapping, Sequence from functools import lru_cache from typing import TYPE_CHECKING, Final, cast, get_type_hints -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter, ValidationError from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.litellm_core_utils.core_helpers import as_str_mapping @@ -73,6 +73,71 @@ def anthropic_message_delta( } +_SAFEGUARD_RESULT_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def redacted_stop_details(raw: object, keep_explanations: bool) -> object: + """``stop_details`` as a rewritten stream may deliver it: its free-text ``explanation`` can quote + the text a guardrail rewrote, so it is nulled unless the rewrite changed no text.""" + stop_details: Final = as_str_mapping(raw) + if keep_explanations or stop_details is None or "explanation" not in stop_details: + return raw + return {**stop_details, "explanation": None} + + +def _without_explanation(raw: object) -> object: + mapping: Final = as_str_mapping(raw) + if mapping is None: + return raw + return {key: value for key, value in mapping.items() if key != "explanation"} + + +def _status_without_explanations(raw: object) -> object: + status: Final = as_str_mapping(_without_explanation(raw)) + if status is None: + return raw + tool_uses: Final = as_str_mapping(status.get("tool_uses")) + if tool_uses is None: + return status + return { + **status, + "tool_uses": {tool_use_id: _without_explanation(verdict) for tool_use_id, verdict in tool_uses.items()}, + } + + +def _result_without_explanations(raw: object) -> object: + result: Final = as_str_mapping(_without_explanation(raw)) + if result is None: + return raw + if "status" not in result: + return result + return {**result, "status": _status_without_explanations(result["status"])} + + +def _kept_tool_use_verdicts(raw: object, kept_tool_use_ids: frozenset[str]) -> object: + result: Final = as_str_mapping(raw) + status: Final = as_str_mapping(result.get("status")) if result is not None else None + tool_uses: Final = as_str_mapping(status.get("tool_uses")) if status is not None else None + if result is None or status is None or tool_uses is None: + return raw + kept_tool_uses: Final = { + tool_use_id: verdict for tool_use_id, verdict in tool_uses.items() if tool_use_id in kept_tool_use_ids + } + return {**result, "status": {**status, "tool_uses": kept_tool_uses}} + + +def filtered_safeguard_results(raw: object, kept_tool_use_ids: frozenset[str], keep_explanations: bool) -> object: + """``safeguard_results`` as a rewritten stream may deliver them: a ``tool_uses`` verdict survives only + for a tool call the client gets with the input the classifier judged, and the free-text ``explanation`` of a + result, its status and each verdict is dropped unless the rewrite changed no text.""" + try: + results: Final = _SAFEGUARD_RESULT_ITEMS.validate_python(raw) + except ValidationError: + return raw + kept: Final = tuple(_kept_tool_use_verdicts(result, kept_tool_use_ids) for result in results) + return kept if keep_explanations else tuple(_result_without_explanations(result) for result in kept) + + def _mapping_field(container: object, key: str) -> object | None: """One key of a raw provider payload, or None when the payload is not a mapping.""" if not isinstance(container, Mapping): diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index 8e363ff92f0..2baeabdce7b 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -12,9 +12,8 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final -from pydantic import TypeAdapter, ValidationError - from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details from litellm.llms.anthropic.pass_through.stream_assembly import anthropic_message_delta_fields from litellm.llms.anthropic.pass_through.stream_assembly import ( assemble_anthropic_sse_stream as assemble_anthropic_sse_stream, # noqa: PLC0414 # public re-export @@ -55,52 +54,6 @@ from litellm.llms.anthropic.pass_through.stream_assembly import ( # noqa: E402 is_sse_error_stream as is_sse_error_stream, # noqa: PLC0414 # public re-export ) -_ITEMS: Final = TypeAdapter(tuple[object, ...]) - - -def _stop_details(raw: object, keep_explanations: bool) -> object: - stop_details: Final = as_str_mapping(raw) - if keep_explanations or stop_details is None or "explanation" not in stop_details: - return raw - return {**stop_details, "explanation": None} - - -def _verdict(raw: object, keep_explanations: bool) -> object: - verdict: Final = as_str_mapping(raw) - if keep_explanations or verdict is None: - return raw - return {key: value for key, value in verdict.items() if key != "explanation"} - - -def _status(raw: object, unchanged_tool_use_ids: frozenset[str], keep_explanations: bool) -> object: - status: Final = as_str_mapping(raw) - tool_uses: Final = as_str_mapping(status.get("tool_uses")) if status is not None else None - if status is None or tool_uses is None: - return raw - return { - **status, - "tool_uses": { - tool_use_id: _verdict(verdict, keep_explanations) - for tool_use_id, verdict in tool_uses.items() - if tool_use_id in unchanged_tool_use_ids - }, - } - - -def _safeguard_result(raw: object, unchanged_tool_use_ids: frozenset[str], keep_explanations: bool) -> object: - result: Final = as_str_mapping(raw) - if result is None or "status" not in result: - return raw - return {**result, "status": _status(result["status"], unchanged_tool_use_ids, keep_explanations)} - - -def _safeguard_results(raw: object, unchanged_tool_use_ids: frozenset[str], keep_explanations: bool) -> object: - try: - results: Final = _ITEMS.validate_python(raw) - except ValidationError: - return raw - return tuple(_safeguard_result(result, unchanged_tool_use_ids, keep_explanations) for result in results) - def _tool_use_inputs(content: Sequence[object]) -> Mapping[str, object]: blocks: Final = (block for raw in content if (block := as_str_mapping(raw)) is not None) @@ -130,9 +83,17 @@ def _rewritten_message_delta_fields( safeguard_results: Final = source_fields.get("safeguard_results") return MappingProxyType( { - **({"stop_details": _stop_details(stop_details, keep_explanations)} if stop_details is not None else {}), **( - {"safeguard_results": _safeguard_results(safeguard_results, unchanged_tool_use_ids, keep_explanations)} + {"stop_details": redacted_stop_details(stop_details, keep_explanations)} + if stop_details is not None + else {} + ), + **( + { + "safeguard_results": filtered_safeguard_results( + safeguard_results, unchanged_tool_use_ids, keep_explanations + ) + } if safeguard_results is not None else {} ), @@ -141,7 +102,7 @@ def _rewritten_message_delta_fields( def anthropic_sse_chunks_from_response( - assembled: ModelResponse, source_chunks: Sequence[object], *, keep_explanations: bool = False + assembled: ModelResponse, source_chunks: Sequence[object] = (), *, keep_explanations: bool = False ) -> tuple[bytes, ...]: """Re-emit a rewritten response as the SSE events of the ``source_chunks`` stream it replaces. diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 892736b5043..0aca60351bf 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -6,7 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic """ import json -from typing import Any, Literal, Optional +from typing import Any, Literal, Optional, Final from unittest.mock import MagicMock, patch import pytest @@ -18,7 +18,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import Streami from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs class MockPassThroughGuardrail(CustomGuardrail): @@ -383,6 +383,97 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert '"name": "lookup_fruit"' not in raw assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"} + @staticmethod + def _message_delta(chunks: list[bytes]) -> dict[str, object]: + lines: Final = b"".join(chunks).decode().split("\n") + events: Final = (json.loads(line[len("data:") :]) for line in lines if line.startswith("data:")) + return next(event["delta"] for event in events if event.get("type") == "message_delta") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_text_rewrite_drops_the_upstream_explanations(self): + delta_fields: Final = { + "stop_details": {"type": "refusal", "category": "cyber", "explanation": "the reply says world"}, + "safeguard_results": [ + {"type": "prompt_injection", "status": {"type": "available", "explanation": "quotes world"}} + ], + } + source: Final[list[bytes]] = self._ended_sse_chunks() + chunks: Final = [ + chunk.replace( + b'"stop_sequence": null}', b'"stop_sequence": null, ' + json.dumps(delta_fields)[1:-1].encode() + b"}" + ) + for chunk in source + ] + + await AnthropicMessagesHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert self._delta_texts(chunks) == ["hello [MASKED]", ""] + assert self._message_delta(chunks) == { + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": {"type": "refusal", "category": "cyber", "explanation": None}, + "safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}], + } + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_use_rewrite_drops_the_verdict_of_the_rewritten_call_only(self): + class MaskPersimmon(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + for tool_call in inputs.get("tool_calls", []): + if isinstance(tool_call, ChatCompletionMessageToolCall) and "persimmon" in tool_call.function.arguments: + tool_call.function.arguments = '{"fruit": "[MASKED]"}' + return inputs + + first_call: Final[list[bytes]] = self._ended_tool_use_sse_chunks() + second_call: Final = [ + chunk.replace(b'"index": 0', b'"index": 1') + .replace(b"toolu_1", b"toolu_2") + .replace(b"persim", b"pear") + .replace(b'mon\\"}', b'\\"}') + for chunk in first_call[1:6] + ] + verdicts: Final = { + "toolu_1": {"type": "evaluated", "outcome": "flagged", "explanation": "looks up persimmon"}, + "toolu_2": {"type": "evaluated", "outcome": "not_flagged", "explanation": "looks up pear"}, + } + safeguard_results: Final = [ + {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": verdicts}} + ] + message_delta: Final = first_call[6].replace( + b'"stop_sequence": null}', + b'"stop_sequence": null, "safeguard_results": ' + json.dumps(safeguard_results).encode() + b"}", + ) + chunks: Final = [*first_call[:6], *second_call, message_delta, first_call[7]] + + await AnthropicMessagesHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MaskPersimmon(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert self._partial_jsons(chunks)[:3] == ['{"fruit": "[MASKED]"}', "", ""] + assert self._message_delta(chunks)["safeguard_results"] == [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {"toolu_2": {"type": "evaluated", "outcome": "not_flagged"}}, + }, + } + ] + @pytest.mark.asyncio async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index f439bd40059..2ee4c8c72c0 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -5,6 +5,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) import json import logging import re +from typing import Final from unittest.mock import patch import pytest @@ -1383,8 +1384,16 @@ class TestToolPermissionGuardrailAnthropicMessages: async def test_rewrite_mode_keeps_the_allowed_tool_use_verdict_with_its_explanation(self): """Claude Code auto mode reads the classifier verdict from message_delta. The rewrite changes no text the explanation could repeat, so the allowed call keeps its verdict whole.""" - bash_verdict = {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"} - safeguard_results = [ + default_on_rewriting: Final = ToolPermissionGuardrail( + guardrail_name="anthropic-rewrite-default-on", + rules=self.rules, + default_action="deny", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + bash_verdict: Final = {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"} + safeguard_results: Final = [ { "type": "dangerous_tool_use", "status": { @@ -1393,28 +1402,23 @@ class TestToolPermissionGuardrailAnthropicMessages: }, } ] - bash_chunks = self._sse_chunks("Bash", "tu_ok") - read_chunks = self._sse_chunks("Read", "tu_bad") - read_tool_use = [ - chunk.replace(b'"index": 1', b'"index": 2') for chunk in read_chunks[4:7] - ] - message_delta = { + bash_chunks: Final = self._sse_chunks("Bash", "tu_ok") + read_chunks: Final = self._sse_chunks("Read", "tu_bad") + read_tool_use: Final = [chunk.replace(b'"index": 1', b'"index": 2') for chunk in read_chunks[4:7]] + message_delta: Final = { "type": "message_delta", "delta": {"stop_reason": "tool_use", "safeguard_results": safeguard_results}, "usage": {"output_tokens": 5}, } - chunks = [ + chunks: Final = [ *bash_chunks[:7], *read_tool_use, f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode(), bash_chunks[-1], ] - with patch.object(self.rewriting, "should_run_guardrail", return_value=True): - out = await self._drain(self.rewriting, chunks) - - body = b"".join(out).decode() - delta = next( + body: Final = b"".join(await self._drain(default_on_rewriting, chunks)).decode() + delta: Final = next( json.loads(line[6:])["delta"] for line in body.splitlines() if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_delta" diff --git a/tests/unit/proxy/guardrails/test_anthropic_sse.py b/tests/unit/proxy/guardrails/test_anthropic_sse.py index 000ecaadcb9..41f058a764d 100644 --- a/tests/unit/proxy/guardrails/test_anthropic_sse.py +++ b/tests/unit/proxy/guardrails/test_anthropic_sse.py @@ -169,8 +169,23 @@ SSN: Final = "123-45-6789" (("toolu_bash", "Bash"),), {"safeguard_results": _safeguard_results({"toolu_bash": {"type": "evaluated", "outcome": "flagged"}})}, ), + ( + { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": [ + { + "type": "prompt_injection", + "status": {"type": "available", "explanation": f"quotes SSN {SSN}"}, + "explanation": f"saw SSN {SSN}", + } + ], + }, + (), + {"safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}]}, + ), ], - ids=["stop_details_explanation", "verdict_explanation"], + ids=["stop_details_explanation", "verdict_explanation", "unknown_shape_explanation"], ) def test_masked_stream_drops_the_unmasked_upstream_explanations( upstream_delta: dict[str, object], tool_uses: tuple[tuple[str, str], ...], expected_fields: dict[str, object]