From 409b522762927a7cec277a13f30513382a030fd2 Mon Sep 17 00:00:00 2001 From: Oliver Fei Date: Tue, 27 Oct 2026 13:53:16 -0400 Subject: [PATCH] fix(guardrails): write TrendAI redactions through structured messages --- .../guardrail_hooks/trendai/_models.py | 3 +- .../guardrail_hooks/trendai/_text.py | 29 +-- .../guardrail_hooks/trendai/trendai.py | 20 +- .../guardrail_hooks/trendai/test_trendai.py | 181 ++++++++++++++++-- 4 files changed, 185 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py index 05d30d8cb20..8172113ae7a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py @@ -127,4 +127,5 @@ class TrendAIWindowScan: @dataclass(frozen=True, slots=True) class TrendAIRequestPrompt: prompt: str - text_range: tuple[int, int] | None + texts: tuple[str, ...] + message_index: int | None diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py index 62d9bd65b33..936fb103264 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py @@ -2,10 +2,9 @@ # This file has been modified for integration into LiteLLM. from collections.abc import Iterator, Sequence -from itertools import chain from typing import Final -from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_text_slot_count +from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts from litellm.types.llms.openai import AllMessageValues from ._models import TrendAIRequestPrompt, TrendAITextWindow @@ -65,10 +64,7 @@ def locate_request_prompt( texts: Sequence[str], structured_messages: Sequence[AllMessageValues] | None, ) -> TrendAIRequestPrompt | None: - """Locate the last user turn by message order, never by matching its text. - - A different flat-text layout can still be scanned, but cannot be safely rewritten. - """ + """Locate the last user turn and its text slots, or use flat texts when no rows are supplied.""" messages: Final = tuple(structured_messages or ()) user_index: Final = next( (index for index in reversed(range(len(messages))) if messages[index]["role"] == "user"), None @@ -79,32 +75,23 @@ def locate_request_prompt( prompt: Final = "".join(parts).strip() if not prompt: return None - start: Final = sum(map(message_text_slot_count, messages[:user_index])) if user_index is not None else 0 - normalized_texts: Final = ( - tuple(chain.from_iterable(map(message_slot_texts, messages))) if messages else tuple(texts) - ) - text_range: Final = (start, start + len(parts)) if normalized_texts == tuple(texts) else None - return TrendAIRequestPrompt(prompt=prompt, text_range=text_range) + return TrendAIRequestPrompt(prompt=prompt, texts=tuple(parts), message_index=user_index) -def redact_request_texts(texts: Sequence[str], prompt: TrendAIRequestPrompt, redacted: str) -> tuple[str, ...] | None: +def redact_request_texts(texts: Sequence[str], redacted: str) -> tuple[str, ...]: """Write a redacted prompt back over the ``texts`` entries the prompt was assembled from. A single-part prompt is replaced outright. A multi-part prompt whose redaction kept its length is split back into the original parts; otherwise the first part carries the whole redaction and the rest are blanked, so no original part can leak. """ - if prompt.text_range is None: - return None - start, end = prompt.text_range - parts: Final = tuple(texts[start:end]) + parts: Final = tuple(texts) if len(parts) == 1: - return (*texts[:start], redacted, *texts[end:]) + return (redacted,) joined: Final = "".join(parts) leading: Final = len(joined) - len(joined.lstrip()) aligned: Final = f"{joined[:leading]}{redacted}{joined[len(joined.rstrip()) :]}" if len(aligned) != len(joined): - return (*texts[:start], redacted, *("" for _ in parts[1:]), *texts[end:]) + return (redacted, *("" for _ in parts[1:])) offsets: Final = tuple(sum(len(part) for part in parts[:index]) for index in range(len(parts) + 1)) - split: Final = tuple(aligned[offsets[index] : offsets[index + 1]] for index in range(len(parts))) - return (*texts[:start], *split, *texts[end:]) + return tuple(aligned[offsets[index] : offsets[index + 1]] for index in range(len(parts))) diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py index 1cb8569bdbc..8dc82f37536 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # legacy decorator has an untyped signature ) +from litellm.llms.base_llm.guardrail_translation.utils import message_with_slot_texts from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # legacy client factory has an untyped params map httpxSpecialProvider, @@ -161,7 +162,8 @@ class TrendAIGuardrail(CustomGuardrail): texts: tuple[str, ...], request_data: dict[str, object], # mutable-ok: CustomGuardrail hook contract is a plain dict ) -> GenericGuardrailAPIInputs: - prompt: Final = locate_request_prompt(texts, inputs.get("structured_messages")) + messages: Final = tuple(inputs.get("structured_messages") or ()) + prompt: Final = locate_request_prompt(texts, messages) if prompt is None: verbose_proxy_logger.debug("Trend AI Guard: no user prompt to scan in request inputs") return inputs @@ -171,14 +173,20 @@ class TrendAIGuardrail(CustomGuardrail): redacted: Final = self._enforce(result) if redacted is None: return inputs - redacted_texts: Final = redact_request_texts(texts, prompt, redacted) - if redacted_texts is None: + redacted_texts: Final = redact_request_texts(prompt.texts, redacted) + if prompt.message_index is None: + return {**inputs, "texts": list(redacted_texts)} + rewritten_message: Final = message_with_slot_texts(messages[prompt.message_index], redacted_texts) + if rewritten_message is None: self._raise_redaction_failure(request_data, started_at, event_type=None) - redacted_inputs: Final[GenericGuardrailAPIInputs] = { + return { **inputs, - "texts": list(redacted_texts), + "structured_messages": [ + *messages[: prompt.message_index], + rewritten_message, + *messages[prompt.message_index + 1 :], + ], } - return redacted_inputs async def _apply_response_guardrail( self, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py b/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py index ff5a8533d32..c886a0aff43 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py @@ -9,8 +9,11 @@ import pytest import litellm from litellm.exceptions import GuardrailRaisedException +from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.bedrock.passthrough.guardrail_translation.handler import BedrockPassthroughGuardrailHandler from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler +from litellm.llms.openai.completion.guardrail_translation.handler import OpenAITextCompletionHandler +from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler from litellm.proxy.guardrails.guardrail_hooks.trendai import TrendAIGuardrail, initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.trendai._models import ( TrendAIAllow, @@ -340,7 +343,12 @@ async def test_request_scans_only_the_last_user_turn_and_writes_redaction_back() ) assert scanned == ["mail a@b.com now"] - assert result["texts"] == ["system prompt", "first turn a@b.com", "reply", "mail [EMAIL] now"] + assert result["structured_messages"] == [ + {"role": "system", "content": "system prompt"}, + {"role": "user", "content": "first turn a@b.com"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "mail [EMAIL] now"}, + ] (record,) = _guardrail_records(request_data) assert record["guardrail_status"] == "success" assert record["masked_entity_count"] == {"EMAIL": 1} @@ -404,40 +412,79 @@ async def test_request_redaction_targets_the_last_user_even_when_later_messages_ @pytest.mark.asyncio -async def test_converse_redaction_with_unaligned_text_slots_fails_closed() -> None: +@pytest.mark.parametrize("action", ["converse", "converse-stream"]) +@pytest.mark.parametrize("tool_result", ["found", "a@b.com"]) +@pytest.mark.parametrize("tool_input", [{}, {"query": "original"}]) +async def test_converse_redaction_preserves_native_field_positions( + action: str, tool_result: str, tool_input: Mapping[str, str] +) -> None: respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"}) + history: Final = [ + {"role": "user", "content": [{"text": "lookup"}]}, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "call-1", "name": "lookup", "input": tool_input}}], + }, + ] body: Final = { "messages": [ - {"role": "user", "content": [{"text": "lookup"}]}, - { - "role": "assistant", - "content": [{"toolUse": {"toolUseId": "call-1", "name": "lookup", "input": {}}}], - }, + *history, { "role": "user", "content": [ {"text": "mail a@b.com"}, - {"toolResult": {"toolUseId": "call-1", "content": [{"text": "found"}]}}, + {"toolResult": {"toolUseId": "call-1", "content": [{"text": tool_result}]}}, ], }, ], - "additionalModelRequestFields": {"system": "extra"}, + "additionalModelRequestFields": {"system": "a@b.com"}, } - original_body: Final = deepcopy(body) + original_history: Final = deepcopy(history) request_data: Final[dict[str, object]] = { - "endpoint": "/model/test-model/converse", + "endpoint": f"/model/test-model/{action}", "data": body, "metadata": {}, } async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: - with pytest.raises(GuardrailRaisedException, match="could not apply the requested redaction"): - await BedrockPassthroughGuardrailHandler().process_input_messages( - request_data, _guardrail(async_handler=client, fallback_on_error="allow") - ) + await BedrockPassthroughGuardrailHandler().process_input_messages( + request_data, _guardrail(async_handler=client) + ) - assert scanned == ["mail a@b.comextra"] - assert request_data["data"] == original_body - assert _guardrail_records(request_data)[-1]["guardrail_status"] == "guardrail_failed_to_respond" + assert scanned == ["mail a@b.coma@b.com"] + assert request_data["data"] == { + "messages": [ + *original_history, + { + "role": "user", + "content": [ + {"text": "mail [EMAIL]"}, + {"toolResult": {"toolUseId": "call-1", "content": [{"text": tool_result}]}}, + ], + }, + ], + "additionalModelRequestFields": {"system": "[EMAIL]"}, + } + + +@pytest.mark.asyncio +async def test_converse_redaction_rewrites_guard_content_without_flat_text_slots() -> None: + respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"}) + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + out: Final = await BedrockPassthroughGuardrailHandler().process_input_messages( + data={ + "endpoint": "/model/test-model/converse", + "data": { + "messages": [{"role": "user", "content": [{"guardContent": {"text": {"text": "mail a@b.com"}}}]}] + }, + "metadata": {}, + }, + guardrail_to_apply=_guardrail(async_handler=client), + ) + + assert scanned == ["mail a@b.com"] + assert out["data"] == { + "messages": [{"role": "user", "content": [{"guardContent": {"text": {"text": "mail [EMAIL]"}}}]}] + } @pytest.mark.asyncio @@ -463,7 +510,16 @@ async def test_request_redaction_is_split_back_across_multipart_user_content() - ) assert scanned == ["card 4111 ok"] - assert result["texts"] == ["card ", "#### ok"] + assert result["structured_messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": "card "}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}, + {"type": "text", "text": "#### ok"}, + ], + } + ] @pytest.mark.asyncio @@ -480,7 +536,9 @@ async def test_request_scans_text_slots_from_normalized_message_parts() -> None: ) assert scanned == ["mail a@b.com"] - assert result["texts"] == ["mail [EMAIL]"] + assert result["structured_messages"] == [ + {"role": "user", "content": [{"type": "input_text", "text": "mail [EMAIL]"}]} + ] @pytest.mark.asyncio @@ -657,6 +715,89 @@ async def test_request_redaction_flows_through_the_chat_completions_handler() -> assert out["messages"][1]["content"][0]["text"] == "mail [EMAIL]" +@pytest.mark.asyncio +async def test_responses_redaction_keeps_instructions_and_tool_items() -> None: + respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"}) + history: Final = [ + {"role": "user", "content": "mail a@b.com"}, + {"type": "function_call", "call_id": "call-1", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call-1", "output": "a@b.com"}, + ] + original_history: Final = deepcopy(history) + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + out: Final = await OpenAIResponsesHandler().process_input_messages( + data={ + "instructions": "be terse", + "input": [ + *history, + {"role": "user", "content": [{"type": "input_text", "text": "mail a@b.com"}]}, + ], + "metadata": {}, + }, + guardrail_to_apply=_guardrail(async_handler=client), + ) + + assert scanned == ["mail a@b.com"] + assert out["instructions"] == "be terse" + assert out["input"] == [ + *original_history, + {"role": "user", "content": [{"type": "input_text", "text": "mail [EMAIL]"}]}, + ] + + +@pytest.mark.asyncio +async def test_anthropic_redaction_preserves_system_and_image_content() -> None: + respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"}) + image: Final = {"type": "image", "source": {"type": "url", "url": "https://example.com/image.png"}} + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + out: Final = await AnthropicMessagesHandler().process_input_messages( + data={ + "model": "claude-sonnet-5-5", + "system": "be terse", + "messages": [ + {"role": "user", "content": "first a@b.com"}, + {"role": "assistant", "content": "reply"}, + { + "role": "user", + "content": [{"type": "text", "text": "mail "}, image, {"type": "text", "text": "a@b.com"}], + }, + ], + "metadata": {}, + }, + guardrail_to_apply=_guardrail(async_handler=client), + ) + + assert scanned == ["mail a@b.com"] + assert out["system"] == "be terse" + assert out["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "first a@b.com"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "reply"}]}, + { + "role": "user", + "content": [{"type": "text", "text": "mail "}, image, {"type": "text", "text": "[EMAIL]"}], + }, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("prompt", "expected"), + [("mail a@b.com", "mail [EMAIL]"), (["mail ", "a@b.com"], ["mail ", "[EMAIL]"])], +) +async def test_text_completion_redaction_still_rewrites_the_prompt( + prompt: str | list[str], expected: str | list[str] +) -> None: + respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"}) + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + out: Final = await OpenAITextCompletionHandler().process_input_messages( + data={"prompt": prompt, "metadata": {}}, + guardrail_to_apply=_guardrail(async_handler=client), + ) + + assert scanned == ["mail a@b.com"] + assert out["prompt"] == expected + + @pytest.mark.parametrize( ("content", "chunk_size_bytes", "overlap", "expected"), [