diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1dc8a77db21..d7b5b57285e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1439,21 +1439,40 @@ class OpenAIResponsesHandler(BaseTranslation): rewrites_by_position=MappingProxyType(dict(zip(placeable_positions, chain((rewritten_text,), repeat(""))))), ) - async def deliver_text_rewrite( + def message_part_texts(self, response: "ResponsesAPIResponse") -> Mapping[tuple[int, int], str]: + """Text of every message content part, keyed by ``(output_index, content_index)``.""" + texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists + positions: Final[list[tuple[int, int]]] = [] # mutable-ok: the same extractor's position sink + for output_idx, output_item in enumerate(response.output or ()): + self._extract_output_text_and_images( + output_item=output_item, + output_idx=output_idx, + texts_to_check=texts, + images_to_check=[], + task_mappings=positions, + ) + return MappingProxyType( + {position: text for position, text in zip(positions, texts) if position[1] != _COMPACTION_SUMMARY_POSITION} + ) + + async def spread_text_rewrite( self, response: "ResponsesAPIResponse", - position: tuple[int, int], rewritten_text: str, stream_events: Sequence[object] = (), ) -> None: - """Write one rewritten text into the response at ``position`` and into any buffered events - that spell out the same text, so a client reading deltas sees what the body says.""" + """Replace all message text in ``response`` with one rewrite, for a guardrail that answers + one string for everything it scanned: the whole rewrite goes on the first part and every + later part is blanked, in the body and in any buffered events spelling out the same text.""" + rewrites_by_position: Final = MappingProxyType( + dict(zip(self.message_part_texts(response), chain((rewritten_text,), repeat("")))) + ) await self._apply_guardrail_responses_to_output( - response=response, responses=[rewritten_text], task_mappings=[position] - ) - self._sync_stream_events_with_rewrites( - stream_events=stream_events, rewrites_by_position=MappingProxyType({position: rewritten_text}) + response=response, + responses=list(rewrites_by_position.values()), + task_mappings=list(rewrites_by_position), ) + self._sync_stream_events_with_rewrites(stream_events=stream_events, rewrites_by_position=rewrites_by_position) @staticmethod def _write_event_field(event: object, field: str, value: object) -> None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 38c11dab063..503089fe790 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,7 +1,7 @@ import time from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence from enum import Enum, auto -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Protocol, runtime_checkable import httpx from fastapi import HTTPException @@ -10,6 +10,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel import json @@ -174,12 +175,6 @@ def _redact_scanned_content(payload: object, depth: int = 0) -> object: return payload -class _ResponsesBodyScan(NamedTuple): - texts: tuple[str, ...] - positions: tuple[tuple[int, int], ...] - tool_calls: tuple[ChatCompletionToolCallChunk, ...] - - class ModelArmorGuardrail(CustomGuardrail, VertexBase): """ Google Cloud Model Armor Guardrail integration for LiteLLM. @@ -548,7 +543,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return False, None if isinstance(assembled_response, ModelResponse) or ( isinstance(assembled_response, ResponsesAPIResponse) - and self._sole_text_position(assembled_response, content) is not None + and self._responses_message_text(assembled_response) == content ): return False, sanitized_content verbose_proxy_logger.warning( @@ -919,7 +914,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return # Extract content from response - content: Final = self._extract_scannable_content(response) + content: Final = ( + self._responses_message_text(response) + if isinstance(response, ResponsesAPIResponse) + else self._extract_content_from_response(response) + ) if not content: verbose_proxy_logger.debug("Model Armor: No text content to process in response, skipping guardrail") return @@ -973,13 +972,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if self.mask_response_content: sanitized_content: Final = self._get_sanitized_content(armor_response) if sanitized_content and sanitized_content != content: - if isinstance(response, ResponsesAPIResponse) and not await self._apply_sanitized_responses_content( - response, content, sanitized_content - ): - raise HTTPException( - status_code=400, - detail=self._build_block_error_detail("Response blocked by Model Armor", armor_response), - ) + if isinstance(response, ResponsesAPIResponse): + await self._responses_handler().spread_text_rewrite(response, sanitized_content) if isinstance(response, litellm.ModelResponse): self._apply_sanitized_content(response, sanitized_content) @@ -1039,26 +1033,17 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): ) @staticmethod - def _scan_responses_api_response(response: ResponsesAPIResponse) -> _ResponsesBodyScan: - """Texts of a Responses API response, where each one sits, and its tool calls.""" + def _responses_handler() -> "OpenAIResponsesHandler": from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) - texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists - positions: Final[list[tuple[int, int]]] = [] # mutable-ok: the same extractor's position sink - tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] # mutable-ok: the same extractor's tool-call sink - handler: Final = OpenAIResponsesHandler() - for output_idx, output_item in enumerate(response.output or ()): - handler._extract_output_text_and_images( # pyright: ignore[reportPrivateUsage] # the shared Responses output extractor; forking it would duplicate per-item parsing - output_item=output_item, - output_idx=output_idx, - texts_to_check=texts, - images_to_check=[], - task_mappings=positions, - tool_calls_to_check=tool_calls, - ) - return _ResponsesBodyScan(texts=tuple(texts), positions=tuple(positions), tool_calls=tuple(tool_calls)) + return OpenAIResponsesHandler() + + @staticmethod + def _responses_message_text(response: ResponsesAPIResponse) -> str: + """Message text of a Responses API response, every part joined as Model Armor scans it.""" + return "".join(ModelArmorGuardrail._responses_handler().message_part_texts(response).values()) @staticmethod def _responses_api_response_text(response: ResponsesAPIResponse) -> str: @@ -1067,47 +1052,25 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): Tool calls are folded in because ``get_content_from_model_response`` folds them into what the chat surface scans, and a Responses turn can carry its whole payload in them. """ - scan: Final = ModelArmorGuardrail._scan_responses_api_response(response) - return "".join((*scan.texts, *(json.dumps(tool_call) for tool_call in scan.tool_calls))) + texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists + tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] # mutable-ok: the same extractor's tool-call sink + handler: Final = ModelArmorGuardrail._responses_handler() + for output_idx, output_item in enumerate(response.output or ()): + handler._extract_output_text_and_images( # pyright: ignore[reportPrivateUsage] # the shared Responses output extractor; forking it would duplicate per-item parsing + output_item=output_item, + output_idx=output_idx, + texts_to_check=texts, + images_to_check=[], + task_mappings=[], + tool_calls_to_check=tool_calls, + ) + return "".join((*texts, *(json.dumps(tool_call) for tool_call in tool_calls))) - @staticmethod - def _sole_text_position(response: ResponsesAPIResponse, content: str) -> tuple[int, int] | None: - """Where one sanitized string can land: the part ``content`` came from, when that is one part. - - A scan that also joined a second part, tool-call arguments or reasoning deltas comes back as - one string with no boundaries, so it has no single place to go. - """ - scan: Final = ModelArmorGuardrail._scan_responses_api_response(response) - return scan.positions[0] if scan.texts == (content,) else None - - async def _apply_sanitized_responses_content( - self, - response: ResponsesAPIResponse, - content: str, - sanitized_content: str, - stream_events: Sequence[object] = (), - ) -> bool: - """Write the sanitized text into the response and its buffered events, if it has one place to go.""" - from litellm.llms.openai.responses.guardrail_translation.handler import ( - OpenAIResponsesHandler, - ) - - position: Final = self._sole_text_position(response, content) - if position is None: - return False - await OpenAIResponsesHandler().deliver_text_rewrite( - response=response, - position=position, - rewritten_text=sanitized_content, - stream_events=stream_events, - ) - return True - - def _extract_scannable_content(self, response: object) -> str: - """Text to scan from a response, for every endpoint shape this hook serves.""" - if isinstance(response, ResponsesAPIResponse): - return self._responses_api_response_text(response) - return self._extract_content_from_response(response) + def _extract_streaming_content(self, assembled_response: object) -> str: + """Text to scan from an assembled stream, for every endpoint shape this hook serves.""" + if isinstance(assembled_response, ResponsesAPIResponse): + return self._responses_api_response_text(assembled_response) + return self._extract_content_from_response(assembled_response) @staticmethod def _responses_delta_field(chunk: object) -> tuple[str, ...]: @@ -1145,7 +1108,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): delta events that its terminal body never repeats, so every delta field the body does not already carry is scanned after it. """ - content: Final = self._extract_scannable_content(assembled_response) + content: Final = self._extract_streaming_content(assembled_response) if surface is not _StreamSurface.RESPONSES: return content unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content) @@ -1316,8 +1279,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): ) if sanitized_content is not None and isinstance(assembled_response, ResponsesAPIResponse): - await self._apply_sanitized_responses_content( - assembled_response, content, sanitized_content, stream_events=all_chunks + await self._responses_handler().spread_text_rewrite( + assembled_response, sanitized_content, stream_events=all_chunks ) if sanitized_content is not None and isinstance(assembled_response, ModelResponse): diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index e3089f69413..abcb48c8c02 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -4966,6 +4966,34 @@ class TestCompactionOutputSummaryReachesOutputGuardrails: assert self._summary(result.output[0]) == "The user lives at 1 [MASKED]." + @pytest.mark.asyncio + async def test_spreading_one_rewrite_over_message_text_leaves_the_summary_alone(self) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=[ + self._compaction_item(), + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "first", "annotations": []}, + {"type": "refusal", "refusal": "second"}, + ], + }, + ], + ) + + await OpenAIResponsesHandler().spread_text_rewrite(response, "rewritten") + + assert self._summary(response.output[0]) == self.SUMMARY + assert [response.output[1].content[0].text, response.output[1].content[1].refusal] == ["rewritten", ""] + @pytest.mark.asyncio async def test_ended_stream_summary_is_masked_in_the_done_item_and_the_envelope(self) -> None: events: Final = [ diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index caa6ee34372..94e990c1005 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3892,9 +3892,17 @@ async def _anthropic_sse_stream(): _SECOND_TEXT_PART: Final = {"type": "output_text", "text": " and more"} +_TOOL_CALL_ITEM: Final = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "save_card", + "arguments": '{"card": "4111-1111-1111-1111"}', + "status": "completed", +} -def _responses_api_events(*extra_parts: dict): +def _responses_api_events(extra_parts: tuple[dict, ...] = (), extra_items: tuple[dict, ...] = ()): from litellm.types.llms.openai import ( OutputTextDeltaEvent, ResponseCompletedEvent, @@ -3914,7 +3922,8 @@ def _responses_api_events(*extra_parts: dict): "status": "completed", "role": "assistant", "content": [{"type": "output_text", "text": "my card is 4111-1111-1111-1111"}, *extra_parts], - } + }, + *extra_items, ], parallel_tool_calls=False, tool_choice="auto", @@ -4050,9 +4059,9 @@ async def test_streaming_masking_re_emits_anthropic_sse_with_sanitized_text(): @pytest.mark.asyncio -async def test_streaming_masking_blocks_a_responses_stream_with_no_single_part_to_rewrite(): - """One sanitized string cannot be split back over several parts, so releasing the stream would - ship the content the guardrail just rewrote. It is blocked instead.""" +async def test_streaming_masking_blocks_a_responses_stream_whose_scan_includes_a_tool_call(): + """Tool-call arguments are scanned on a stream but one sanitized string has no way back into + them, so releasing the stream would ship the content the guardrail just rewrote.""" from litellm.types.llms.openai import ErrorEvent guardrail = _surface_guardrail(mask_response_content=True) @@ -4060,7 +4069,7 @@ async def test_streaming_masking_blocks_a_responses_stream_with_no_single_part_t with patch.object( guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_SANITIZED) ): - delivered = await _drain_surface_hook(guardrail, _responses_api_events(_SECOND_TEXT_PART)) + delivered = await _drain_surface_hook(guardrail, _responses_api_events(extra_items=(_TOOL_CALL_ITEM,))) assert len(delivered) == 1 assert isinstance(delivered[0], ErrorEvent) @@ -4538,9 +4547,9 @@ async def test_streaming_deidentify_match_without_a_rewrite_fails_closed(chunks, @pytest.mark.asyncio async def test_streaming_status_records_a_surface_that_cannot_carry_the_rewrite_as_blocked(): - """A multi-part Responses body has no single part to rewrite, so a de-identify match ends as a - refusal. The status stamped on metadata feeds the spend log and has to say so rather than - reporting the success the block check alone would have implied.""" + """A Responses stream that scanned a tool call has nowhere to put the rewrite, so a de-identify + match ends as a refusal. The status stamped on metadata feeds the spend log and has to say so + rather than reporting the success the block check alone would have implied.""" guardrail = _surface_guardrail(mask_response_content=True) request_data = { "model": "gpt-4o-mini", @@ -4550,7 +4559,7 @@ async def test_streaming_status_records_a_surface_that_cannot_carry_the_rewrite_ with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)): delivered = await _drain_surface_hook( - guardrail, _responses_api_events(_SECOND_TEXT_PART), request_data=request_data + guardrail, _responses_api_events(extra_items=(_TOOL_CALL_ITEM,)), request_data=request_data ) body = _delivered_bytes(delivered) @@ -5477,7 +5486,7 @@ async def test_post_call_masks_plain_text_in_a_responses_body(): @pytest.mark.asyncio -async def test_post_call_blocks_a_responses_body_whose_parts_one_redaction_cannot_split_across(): +async def test_post_call_masks_every_part_of_a_responses_body_through_the_first(): guardrail = _model_armor_post_call_guardrail(mask_response_content=True) post = AsyncMock(return_value=_armor_http_response(_armor_redaction("card [REDACTED] and more"))) response = _responses_body( @@ -5485,11 +5494,38 @@ async def test_post_call_blocks_a_responses_body_whose_parts_one_redaction_canno {"type": "output_text", "text": " and more", "annotations": []}, ) - with patch.object(guardrail.async_handler, "post", post), pytest.raises(HTTPException): + with patch.object(guardrail.async_handler, "post", post): await guardrail.async_post_call_success_hook( data=dict(_POST_CALL_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=response ) + assert [part.text for part in response.output[0].content] == ["card [REDACTED] and more", ""] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("mask_response_content", "flagged"), + [(True, _armor_redaction("saving [REDACTED]")), (False, _armor_block())], +) +async def test_post_call_answers_a_tool_call_beside_responses_text_unchanged(mask_response_content, flagged): + from litellm.types.llms.openai import ResponsesAPIResponse + + async def flag_the_card(**kwargs): + verdict = flagged if "4111" in str(kwargs["json"]) else {"filterMatchState": "NO_MATCH_FOUND"} + return _armor_http_response(verdict) + + guardrail = _model_armor_post_call_guardrail(mask_response_content=mask_response_content) + body = _responses_body({"type": "output_text", "text": "saving it", "annotations": []}).model_dump() + response = ResponsesAPIResponse.model_validate({**body, "output": [*body["output"], _TOOL_CALL_ITEM]}) + + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=flag_the_card)): + await guardrail.async_post_call_success_hook( + data=dict(_POST_CALL_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert response.output[0].content[0].text == "saving it" + assert response.output[1].arguments == _TOOL_CALL_ITEM["arguments"] + @pytest.mark.asyncio @pytest.mark.parametrize( @@ -5540,3 +5576,25 @@ async def test_streaming_masks_every_event_that_spells_out_a_responses_part(part assert "4111" not in str(released) assert "".join(event.delta for event in released[:2]) == "I cannot help card [REDACTED]" assert getattr(released[2], field) == "I cannot help card [REDACTED]" + + +@pytest.mark.asyncio +async def test_streaming_masks_a_responses_body_with_several_parts(): + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + guardrail = _surface_guardrail(mask_response_content=True) + first_delta, terminal = _responses_api_events(extra_parts=(_SECOND_TEXT_PART,)) + second_delta = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=1, + delta=_SECOND_TEXT_PART["text"], + ) + + with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)): + delivered = await _drain_surface_hook(guardrail, (first_delta, second_delta, terminal)) + + assert b"4111-1111-1111-1111" not in _delivered_bytes(delivered) + assert [delivered[0].delta, delivered[1].delta] == ["my card is [REDACTED]", ""] + assert [part["text"] for part in delivered[-1].response.output[0]["content"]] == ["my card is [REDACTED]", ""]