scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(guardrails): write TrendAI redactions through structured messages

Browse source
This commit is contained in:
Oliver Fei 2026-10-27 13:53:16 -04:00
parent 2753325618
commit 409b522762
4 changed files with 185 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py
View file

@ -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

29
litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py
View file

@ -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)))

20
litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py
View file

@ -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,

181
tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py
View file

@ -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"),
[