diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 6c681c322f1..8fdd8b23cb2 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -143,6 +143,30 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: return visited +def apply_redacted_messages_back( + data: Dict[str, Any], redacted_messages: List[Dict[str, Any]] +) -> None: + """Write redacted messages back to whichever field(s) the caller used. + + Mask/anonymize paths take a synthesised messages list (from + :func:`build_inspection_messages`), get a redacted version back from a + third-party guardrail, and need to rewrite the request body. Writing + only to ``data["messages"]`` leaves the Responses-API ``data["input"]`` + field untouched, so the unredacted text still reaches the LLM. + + This helper updates both fields when both are present. + """ + if "messages" in data: + data["messages"] = redacted_messages + if isinstance(data.get("input"), str): + text_parts: List[str] = [] + for msg in redacted_messages: + if not isinstance(msg, dict): + continue + text_parts.extend(_iter_text_parts_in_content(msg.get("content"))) + data["input"] = "\n".join(text_parts) + + def has_non_string_content(data: Dict[str, Any]) -> bool: """Return True if any inspected content is not a plain string. diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 749bd4acf2b..03a73f9fc9d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -23,6 +23,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import ( + apply_redacted_messages_back, build_inspection_messages, has_non_string_content, ) @@ -156,13 +157,17 @@ class AimGuardrail(CustomGuardrail): "or rely on block-mode policies." ), ) - data["messages"] = [ + redacted_messages = [ { "role": message["role"], "content": message["content"], } for message in redacted_chat["all_redacted_messages"] ] + # Write back to ``messages`` AND ``input``. The Responses-API + # backend reads ``input``; writing only to ``messages`` would let + # unredacted text reach the LLM for ``/v1/responses`` calls. + apply_redacted_messages_back(data, redacted_messages) return data async def call_aim_guardrail_on_output( diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 4d6e6a302b5..c6af4d3c428 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -14,6 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import ( + apply_redacted_messages_back, build_inspection_messages, has_non_string_content, ) @@ -251,11 +252,15 @@ class LakeraAIGuardrail(CustomGuardrail): self._is_only_pii_violation(lakera_guardrail_response) and not is_multimodal_input ): - data["messages"] = self._mask_pii_in_messages( + redacted_messages = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, masked_entity_count=masked_entity_count, ) + # Write back to ``messages`` AND ``input``. The Responses-API + # backend reads ``input``; writing only to ``messages`` + # would let unredacted PII reach the LLM for /v1/responses. + apply_redacted_messages_back(data, list(redacted_messages)) # type: ignore[arg-type] verbose_proxy_logger.debug( "Lakera AI: Masked PII in messages instead of blocking request" ) @@ -325,11 +330,15 @@ class LakeraAIGuardrail(CustomGuardrail): self._is_only_pii_violation(lakera_guardrail_response) and not is_multimodal_input ): - data["messages"] = self._mask_pii_in_messages( + redacted_messages = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, masked_entity_count=masked_entity_count, ) + # Write back to ``messages`` AND ``input``. The Responses-API + # backend reads ``input``; writing only to ``messages`` + # would let unredacted PII reach the LLM for /v1/responses. + apply_redacted_messages_back(data, list(redacted_messages)) # type: ignore[arg-type] verbose_proxy_logger.debug( "Lakera AI: Masked PII in messages instead of blocking request" ) diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 23683abd1cc..099fca78a62 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -1,6 +1,7 @@ """Tests for the shared guardrail content extraction helpers.""" from litellm.proxy.guardrails._content_utils import ( + apply_redacted_messages_back, build_inspection_messages, has_non_string_content, iter_message_text, @@ -263,3 +264,40 @@ def test_has_non_string_content_empty_data(): assert has_non_string_content({}) is False assert has_non_string_content({"messages": []}) is False assert has_non_string_content({"input": ""}) is False + + +# ── apply_redacted_messages_back ────────────────────────────────────────────── + + +def test_apply_redacted_messages_back_chat_completion(): + data = {"messages": [{"role": "user", "content": "secret"}]} + apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) + assert data["messages"] == [{"role": "user", "content": "[REDACTED]"}] + assert "input" not in data + + +def test_apply_redacted_messages_back_responses_api_string_input(): + """A Responses-API request reads ``data["input"]``; writing only to + ``messages`` would let unredacted text reach the LLM.""" + data = {"input": "secret payload"} + apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) + assert data["input"] == "[REDACTED]" + + +def test_apply_redacted_messages_back_both_fields(): + """Defensive: when both fields are present, both are updated.""" + data = { + "messages": [{"role": "user", "content": "old"}], + "input": "old", + } + apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) + assert data["messages"] == [{"role": "user", "content": "[REDACTED]"}] + assert data["input"] == "[REDACTED]" + + +def test_apply_redacted_messages_back_skips_input_when_not_string(): + """List ``input`` (multimodal Responses-API) is left alone — the + multimodal-degrades-to-block guard runs upstream.""" + data = {"input": [{"type": "text", "text": "leak"}]} + apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) + assert data["input"] == [{"type": "text", "text": "leak"}] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py index df321b80749..04bd57fa32a 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -161,6 +161,81 @@ async def test_lakera_v2_inspects_responses_api_input(user_api_key, monkeypatch) assert seen_messages == [[{"role": "user", "content": "responses-api content"}]] +@pytest.mark.asyncio +async def test_lakera_v2_responses_api_input_redacted_writeback( + user_api_key, monkeypatch +): + """Greptile P1: when input arrives via Responses-API ``data["input"]`` + (string) and Lakera flags PII, the redacted content must be written + back to ``data["input"]`` — the Responses-API backend reads from + ``input``, so writing only to ``messages`` would let unredacted PII + reach the LLM.""" + monkeypatch.setenv("LAKERA_API_KEY", "lk-test") + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( + LakeraAIGuardrail, + ) + + guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="block") + + async def fake_call_v2_guard(messages, request_data, event_type): + return ({"flagged": True, "payload": []}, {"EMAIL": 1}) + + def fake_mask(messages, lakera_response, masked_entity_count): + return [{"role": "user", "content": "[REDACTED EMAIL]"}] + + with ( + patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard), + patch.object(guard, "_is_only_pii_violation", return_value=True), + patch.object(guard, "_mask_pii_in_messages", side_effect=fake_mask), + ): + data = {"input": "user@example.com leaked"} + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data=data, + call_type="responses", + ) + + assert data["input"] == "[REDACTED EMAIL]" + + +@pytest.mark.asyncio +async def test_aim_responses_api_input_anonymize_writeback(user_api_key, monkeypatch): + """Greptile P1: Aim's anonymize action must redact ``data["input"]`` + for Responses-API requests, not just ``data["messages"]``.""" + monkeypatch.setenv("AIM_API_KEY", "hs-aim-key") + from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail + + guard = AimGuardrail() + + aim_response_body = { + "required_action": {"action_type": "anonymize_action"}, + "redacted_chat": { + "all_redacted_messages": [ + {"role": "user", "content": "[REDACTED] anonymised"} + ] + }, + } + + async def capture(url, headers, json): + return Response( + status_code=200, + json=aim_response_body, + request=Request("POST", "https://api.aim.security/fw/v1/analyze"), + ) + + with patch.object(guard.async_handler, "post", side_effect=capture): + data = {"input": "user@example.com leaked"} + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data=data, + call_type="responses", + ) + + assert data["input"] == "[REDACTED] anonymised" + + @pytest.mark.asyncio async def test_lakera_v2_multimodal_pii_degrades_to_block(user_api_key, monkeypatch): """Mask-in-place uses Lakera offsets and cannot preserve image/audio