From 15a5371e0699dcc2b83fef1ffc99e8f4c4026437 Mon Sep 17 00:00:00 2001 From: Oliver Fei Date: Tue, 20 Oct 2026 17:17:34 -0400 Subject: [PATCH] fix(guardrails): correct TrendAI initialization and request redaction --- .../guardrail_hooks/trendai/__init__.py | 2 +- .../guardrail_hooks/trendai/_models.py | 3 +- .../guardrail_hooks/trendai/_text.py | 58 +++---- .../guardrail_hooks/trendai/trendai.py | 8 +- .../guardrail_hooks/trendai/test_trendai.py | 144 +++++++++++++++++- 5 files changed, 174 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py index 14faca620ae..7976f4fc734 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py @@ -21,7 +21,7 @@ def _normalize_event_hook( def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> TrendAIGuardrail: - settings: Final = TrendAISettings.model_validate(litellm_params.model_dump(mode="python")) + settings: Final = TrendAISettings.model_validate(litellm_params.model_dump(mode="python", exclude_none=True)) guardrail_name: Final = guardrail["guardrail_name"] callback: Final = TrendAIGuardrail( api_key=litellm_params.api_key, diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py index 9ba9ed2e10a..05d30d8cb20 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py @@ -127,5 +127,4 @@ class TrendAIWindowScan: @dataclass(frozen=True, slots=True) class TrendAIRequestPrompt: prompt: str - start: int - end: int + text_range: tuple[int, int] | None diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py index 19e655b7302..62d9bd65b33 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py @@ -2,9 +2,10 @@ # 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 +from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_text_slot_count from litellm.types.llms.openai import AllMessageValues from ._models import TrendAIRequestPrompt, TrendAITextWindow @@ -60,61 +61,50 @@ def apply_window_redaction(content: str, window: TrendAITextWindow, redacted: st return f"{content[: window.start]}{merged}{content[end:]}" -def _last_user_text_parts(structured_messages: Sequence[AllMessageValues]) -> tuple[str, ...] | None: - last_user_message: Final = next( - (message for message in reversed(structured_messages) if message["role"] == "user"), None - ) - return message_slot_texts(last_user_message) if last_user_message is not None else None - - -def _last_occurrence(texts: Sequence[str], parts: Sequence[str]) -> int | None: - return next( - ( - start - for start in range(len(texts) - len(parts), -1, -1) - if tuple(texts[start : start + len(parts)]) == tuple(parts) - ), - None, - ) - - def locate_request_prompt( texts: Sequence[str], structured_messages: Sequence[AllMessageValues] | None, ) -> TrendAIRequestPrompt | None: - """Pick the text Trend AI Guard scans for a request and where it lives in ``texts``. + """Locate the last user turn by message order, never by matching its text. - The scanned prompt is the last user turn, matching the upstream plugin. Its text parts are - located as a run inside ``texts`` so a redacted prompt can be written back to exactly those - entries. Without structured messages the whole ``texts`` list is the prompt. + A different flat-text layout can still be scanned, but cannot be safely rewritten. """ - parts: Final = _last_user_text_parts(structured_messages) if structured_messages else tuple(texts) - if parts is None or not parts: + messages: Final = tuple(structured_messages or ()) + user_index: Final = next( + (index for index in reversed(range(len(messages))) if messages[index]["role"] == "user"), None + ) + if messages and user_index is None: return None + parts: Final = message_slot_texts(messages[user_index]) if user_index is not None else tuple(texts) prompt: Final = "".join(parts).strip() if not prompt: return None - start: Final = _last_occurrence(texts, parts) - if start is None: - return None - return TrendAIRequestPrompt(prompt=prompt, start=start, end=start + len(parts)) + 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) -def redact_request_texts(texts: Sequence[str], prompt: TrendAIRequestPrompt, redacted: str) -> tuple[str, ...]: +def redact_request_texts(texts: Sequence[str], prompt: TrendAIRequestPrompt, redacted: str) -> tuple[str, ...] | None: """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. """ - parts: Final = tuple(texts[prompt.start : prompt.end]) + if prompt.text_range is None: + return None + start, end = prompt.text_range + parts: Final = tuple(texts[start:end]) if len(parts) == 1: - return (*texts[: prompt.start], redacted, *texts[prompt.end :]) + return (*texts[:start], redacted, *texts[end:]) 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[: prompt.start], redacted, *("" for _ in parts[1:]), *texts[prompt.end :]) + return (*texts[:start], redacted, *("" for _ in parts[1:]), *texts[end:]) 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[: prompt.start], *split, *texts[prompt.end :]) + return (*texts[:start], *split, *texts[end:]) diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py index d2d70cabaf6..1cb8569bdbc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py @@ -135,6 +135,9 @@ class TrendAIGuardrail(CustomGuardrail): GuardrailEventHooks.logging_only, ] + def supports_scan_only_tool_results(self) -> bool: + return False + @log_guardrail_information async def apply_guardrail( self, @@ -168,9 +171,12 @@ 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: + self._raise_redaction_failure(request_data, started_at, event_type=None) redacted_inputs: Final[GenericGuardrailAPIInputs] = { **inputs, - "texts": list(redact_request_texts(texts, prompt, redacted)), + "texts": list(redacted_texts), } return redacted_inputs 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 1a17db0a8bc..ff5a8533d32 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py @@ -1,14 +1,17 @@ import json from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy from functools import reduce -from typing import Literal +from typing import Final, Literal import httpx import pytest +import litellm from litellm.exceptions import GuardrailRaisedException +from litellm.llms.bedrock.passthrough.guardrail_translation.handler import BedrockPassthroughGuardrailHandler from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler -from litellm.proxy.guardrails.guardrail_hooks.trendai import TrendAIGuardrail +from litellm.proxy.guardrails.guardrail_hooks.trendai import TrendAIGuardrail, initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.trendai._models import ( TrendAIAllow, TrendAIBlock, @@ -20,7 +23,8 @@ from litellm.proxy.guardrails.guardrail_hooks.trendai.trendai import ( OPENAI_CHAT_COMPLETION_RESPONSE_V1, PLUGIN_VERSION, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.utils import ( CallTypes, Choices, @@ -223,6 +227,47 @@ def test_guardrail_is_discovered_by_global_registries() -> None: assert "trendai" in guardrail_initializer_registry +@pytest.mark.parametrize( + ("timeout_config", "expected_timeout"), + [({}, 5.0), ({"timeout": None}, 5.0), ({"timeout": 2.5}, 2.5)], + ids=["omitted", "null", "configured"], +) +def test_initializer_defaults_null_timeout_and_preserves_configured_timeout( + timeout_config: Mapping[str, float | None], expected_timeout: float +) -> None: + params: Final = LitellmParams.model_validate( + { + "guardrail": "trendai", + "mode": "pre_call", + "api_key": "test-key", + "api_base": "https://guard.example.com", + **timeout_config, + } + ) + callback: Final = initialize_guardrail(params, {"guardrail_name": "trendai-timeout"}) + try: + assert callback.timeout == expected_timeout + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(callback) + + +def test_registry_rejects_tool_result_only_scanning() -> None: + with pytest.raises(ValueError, match="never scans tool results"): + InMemoryGuardrailHandler().initialize_guardrail( + { + "guardrail_name": "trendai-tool-results", + "litellm_params": { + "guardrail": "trendai", + "mode": "pre_call", + "api_key": "test-key", + "api_base": "https://guard.example.com", + "timeout": 5.0, + "scan_only_tool_results": True, + }, + } + ) + + Responder = Callable[[httpx.Request], httpx.Response] @@ -302,6 +347,99 @@ async def test_request_scans_only_the_last_user_turn_and_writes_redaction_back() assert record["guardrail_response"] == {"action": "allow", "redacted": True} +@pytest.mark.asyncio +@pytest.mark.parametrize("repeated_role", ["assistant", "tool"]) +@pytest.mark.parametrize("multipart", [False, True], ids=["string", "multipart"]) +async def test_request_redaction_targets_the_last_user_even_when_later_messages_repeat_it( + repeated_role: Literal["assistant", "tool"], multipart: bool +) -> None: + respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"}) + user_content: Final = ( + [ + {"type": "text", "text": "mail "}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}, + {"type": "text", "text": "a@b.com"}, + ] + if multipart + else "mail a@b.com" + ) + repeated_message: Final = { + "role": repeated_role, + "content": ( + [{"type": "text", "text": "mail "}, {"type": "text", "text": "a@b.com"}] if multipart else "mail a@b.com" + ), + **({"tool_call_id": "call-1"} if repeated_role == "tool" else {}), + } + expected_repeated_message: Final = deepcopy(repeated_message) + expected_content: Final = ( + [ + {"type": "text", "text": "mail "}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}, + {"type": "text", "text": "[EMAIL]"}, + ] + if multipart + else "mail [EMAIL]" + ) + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + out: Final = await OpenAIChatCompletionsHandler().process_input_messages( + data={ + "messages": [ + {"role": "user", "content": "mail a@b.com"}, + {"role": "assistant", "content": "previous reply"}, + {"role": "user", "content": user_content}, + repeated_message, + ], + "metadata": {}, + }, + guardrail_to_apply=_guardrail(async_handler=client), + ) + + assert scanned == ["mail a@b.com"] + assert out["messages"] == [ + {"role": "user", "content": "mail a@b.com"}, + {"role": "assistant", "content": "previous reply"}, + {"role": "user", "content": expected_content}, + expected_repeated_message, + ] + + +@pytest.mark.asyncio +async def test_converse_redaction_with_unaligned_text_slots_fails_closed() -> None: + respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"}) + body: Final = { + "messages": [ + {"role": "user", "content": [{"text": "lookup"}]}, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "call-1", "name": "lookup", "input": {}}}], + }, + { + "role": "user", + "content": [ + {"text": "mail a@b.com"}, + {"toolResult": {"toolUseId": "call-1", "content": [{"text": "found"}]}}, + ], + }, + ], + "additionalModelRequestFields": {"system": "extra"}, + } + original_body: Final = deepcopy(body) + request_data: Final[dict[str, object]] = { + "endpoint": "/model/test-model/converse", + "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") + ) + + assert scanned == ["mail a@b.comextra"] + assert request_data["data"] == original_body + assert _guardrail_records(request_data)[-1]["guardrail_status"] == "guardrail_failed_to_respond" + + @pytest.mark.asyncio async def test_request_redaction_is_split_back_across_multipart_user_content() -> None: respond, scanned = _engine(redact={"4111": "####"})