scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-30 06:43:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 249

fix(guardrails): correct TrendAI initialization and request redaction

Browse source
This commit is contained in:
Oliver Fei 2026-10-20 17:17:34 -04:00
parent bf2ae95cbb
commit 15a5371e06
5 changed files with 174 additions and 41 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py
View file

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

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

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

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

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

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

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

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

@ -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": "####"})