scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(guardrails): fail closed on TrustGuard transforms that empty or miscount messages

Browse source 
When TrustGuard transformed a response into messages whose content was
empty or null, the hook dropped those messages while rebuilding texts and
then fell back to the original inputs, so a fully redacted completion
reached the client unredacted. Returning an empty list would not help
either: the chat translation handler skips the write-back when the
returned texts are empty

Texts are now rebuilt one per returned message, with "" for empty or
non-string content, and the fallback is gone. That keeps the positional
alignment the handlers rely on when they map texts back onto choices, so
a redaction that empties only the first of two choices no longer shifts
the second choice's text onto the first. When the handler sent no texts
at all (a tool-call-only completion) the hook keeps texts as sent instead
of inventing one for the placeholder message, which the Anthropic and
Responses handlers would index out of range

A transform whose message count differs from what the hook sent is now
rejected with the same 400 the tool-call count mismatch already raises.
Fewer messages used to leave trailing choices unredacted and more used to
crash the handler write-back with an IndexError

Regression tests cover the empty and null cases in both directions, the
two-choice alignment through the OpenAI chat translation handler, the
tool-call-only reply, and both count mismatches
This commit is contained in:
albertbausili 2026-09-13 12:14:30 +02:00
parent 9d06efa6a9
commit e9b056045a
2 changed files with 205 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

61
litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py
View file

@ -46,9 +46,9 @@ class _TrustGuardUnreachable(Exception):
"""Transport or availability failure; eligible for unreachable_fallback."""
def _message_text(message: Mapping[str, object]) -> str | None:
def _message_text(message: Mapping[str, object]) -> str:
content: Final = message.get("content")
return content if isinstance(content, str) and content else None
return content if isinstance(content, str) else ""
def _copy_message(value: object) -> Mapping[str, object] | None:
@ -63,7 +63,7 @@ def _copy_messages(messages: Sequence[object]) -> tuple[Mapping[str, object], ..
def _texts_from_messages(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
return tuple(text for message in messages if (text := _message_text(message)) is not None)
return tuple(_message_text(message) for message in messages)
def _tool_calls_in_message(message: Mapping[str, object]) -> tuple[object, ...] | None:
@ -118,13 +118,24 @@ def _assistant_messages(texts: Sequence[str], tool_calls: object) -> tuple[Mappi
return tuple(_assistant_message(text, tool_calls if index == last else None) for index, text in enumerate(texts))
def _sent_messages(
inputs: GenericGuardrailAPIInputs,
input_type: Literal["request", "response"],
) -> Sequence[Mapping[str, object]]:
if input_type == "response":
return _assistant_messages(tuple(inputs.get("texts") or ()), inputs.get("tool_calls"))
structured: Final = inputs.get("structured_messages")
if structured:
return structured
return tuple({"role": "user", "content": text} for text in (inputs.get("texts") or ())) # mutable-ok: outbound JSON
def _inputs_with_messages(
inputs: GenericGuardrailAPIInputs,
messages: Sequence[Mapping[str, object]],
*,
replace_tool_calls: bool,
) -> GenericGuardrailAPIInputs:
texts: Final = _texts_from_messages(messages)
extracted: Final = _tool_calls_from_messages(messages) if replace_tool_calls else None
original_tool_calls: Final = inputs.get("tool_calls")
if extracted is not None and original_tool_calls is not None and len(extracted) != len(original_tool_calls):
@ -132,11 +143,15 @@ def _inputs_with_messages(
merged: Final[GenericGuardrailAPIInputs] = { # mutable-ok: GenericGuardrailAPIInputs is a TypedDict
**inputs,
"structured_messages": list(messages), # mutable-ok: GenericGuardrailAPIInputs.structured_messages is a list
"texts": list(texts) if texts else inputs.get("texts"), # mutable-ok: GenericGuardrailAPIInputs.texts is a list
}
rebuilt: Final[GenericGuardrailAPIInputs] = (
{**merged, "texts": list(_texts_from_messages(messages))} # mutable-ok: TypedDict field is a list
if inputs.get("texts")
else merged
)
if extracted is None:
return merged
return {**merged, "tool_calls": list(extracted)} # mutable-ok: GenericGuardrailAPIInputs.tool_calls is a list
return rebuilt
return {**rebuilt, "tool_calls": list(extracted)} # mutable-ok: GenericGuardrailAPIInputs.tool_calls is a list
class NeuralTrustGuardrail(CustomGuardrail):
@ -217,7 +232,11 @@ class NeuralTrustGuardrail(CustomGuardrail):
},
)
if status == STATUS_TRANSFORM:
return self._apply_transform(inputs, result.get("transformed_payload"))
return self._apply_transform(
inputs,
result.get("transformed_payload"),
sent_count=len(_sent_messages(inputs, input_type)),
)
if status == STATUS_REPORT:
verbose_proxy_logger.info("TrustGuard report-only findings trace_id=%s", result.get("trace_id"))
return inputs
@ -247,23 +266,11 @@ class NeuralTrustGuardrail(CustomGuardrail):
inputs: GenericGuardrailAPIInputs,
input_type: Literal["request", "response"],
) -> Mapping[str, object]:
if input_type == "request":
structured: Final = inputs.get("structured_messages")
messages: Final = (
structured
if structured
else tuple(
{"role": "user", "content": text} # mutable-ok: outbound JSON
for text in (inputs.get("texts") or ())
)
)
tools: Final = inputs.get("tools")
if tools:
return {"messages": messages, "tools": tools} # mutable-ok: outbound JSON
return {"messages": messages} # mutable-ok: outbound JSON
output_messages: Final = _assistant_messages(tuple(inputs.get("texts") or ()), inputs.get("tool_calls"))
return {"messages": output_messages} # mutable-ok: outbound JSON
messages: Final = _sent_messages(inputs, input_type)
tools: Final = inputs.get("tools") if input_type == "request" else None
if tools:
return {"messages": messages, "tools": tools} # mutable-ok: outbound JSON
return {"messages": messages} # mutable-ok: outbound JSON
async def _call_evaluate(self, body: dict[str, object]) -> dict[str, object]: # mutable-ok: TrustGuard JSON
url: Final = f"{self.api_base}{EVALUATE_PATH}"
@ -335,6 +342,8 @@ class NeuralTrustGuardrail(CustomGuardrail):
def _apply_transform(
inputs: GenericGuardrailAPIInputs,
transformed: object,
*,
sent_count: int,
) -> GenericGuardrailAPIInputs:
if not isinstance(transformed, Mapping):
raise HTTPException(status_code=400, detail=TRANSFORM_MISSING)
@ -342,7 +351,7 @@ class NeuralTrustGuardrail(CustomGuardrail):
raw_messages: Final = transformed.get("messages")
if isinstance(raw_messages, list) and raw_messages:
rewritten_messages: Final = _copy_messages(raw_messages)
if rewritten_messages is None:
if rewritten_messages is None or len(rewritten_messages) != sent_count:
raise HTTPException(status_code=400, detail=TRANSFORM_MISSING)
return _inputs_with_messages(inputs, rewritten_messages, replace_tool_calls=True)

171
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_neuraltrust.py
View file

@ -9,10 +9,11 @@ from httpx import Request, Response
from litellm.exceptions import Timeout
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
from litellm.proxy.guardrails.guardrail_hooks.neuraltrust.neuraltrust import (
NeuralTrustGuardrail,
)
from litellm.types.utils import GenericGuardrailAPIInputs
from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message, ModelResponse
def _response(payload: object, status_code: int = 200) -> Response:
@ -289,6 +290,174 @@ class TestNeuralTrustGuardrail:
assert result["tool_calls"] is not original_tool_calls
assert result["structured_messages"][0]["tool_calls"] == rewritten_tool_calls
@pytest.mark.asyncio
@pytest.mark.parametrize("emptied", ["", None])
async def test_transform_emptied_output_blanks_text_instead_of_restoring_original(
self, emptied: str | None
) -> None:
guardrail = _guardrail(event_hook="post_call")
mock_post = AsyncMock(
return_value=_response(
{
"status": "transform",
"transformed_payload": {"messages": [{"role": "assistant", "content": emptied}]},
}
)
)
with patch.object(guardrail.async_handler, "post", mock_post):
result = await guardrail.apply_guardrail(
inputs={"texts": ["my ssn is 123-45-6789"]},
request_data={},
input_type="response",
logging_obj=_logging(),
)
assert result["texts"] == [""]
@pytest.mark.asyncio
async def test_transform_emptied_output_keeps_choice_alignment(self) -> None:
guardrail = _guardrail(event_hook="post_call")
rewritten = [
{"role": "assistant", "content": ""},
{"role": "assistant", "content": "card ending [REDACTED]"},
]
mock_post = AsyncMock(
return_value=_response({"status": "transform", "transformed_payload": {"messages": rewritten}})
)
with patch.object(guardrail.async_handler, "post", mock_post):
result = await guardrail.apply_guardrail(
inputs={"texts": ["ssn 123-45-6789", "card ending 4242"]},
request_data={},
input_type="response",
logging_obj=_logging(),
)
assert result["texts"] == ["", "card ending [REDACTED]"]
@pytest.mark.asyncio
async def test_transform_emptied_output_reaches_client_blank_and_aligned(self) -> None:
guardrail = _guardrail(event_hook="post_call")
rewritten = [
{"role": "assistant", "content": ""},
{"role": "assistant", "content": "card ending [REDACTED]"},
]
mock_post = AsyncMock(
return_value=_response({"status": "transform", "transformed_payload": {"messages": rewritten}})
)
response = ModelResponse(
id="chatcmpl-1",
created=1,
model="gpt-4o-mini",
object="chat.completion",
choices=[
Choices(finish_reason="stop", index=0, message=Message(content="ssn 123-45-6789", role="assistant")),
Choices(finish_reason="stop", index=1, message=Message(content="card ending 4242", role="assistant")),
],
)
with patch.object(guardrail.async_handler, "post", mock_post):
processed = await OpenAIChatCompletionsHandler().process_output_response(response, guardrail)
assert processed.choices[0].message.content == ""
assert processed.choices[1].message.content == "card ending [REDACTED]"
@pytest.mark.asyncio
@pytest.mark.parametrize("sent_texts", [{}, {"texts": []}])
async def test_transform_tool_call_only_output_adds_no_text(self, sent_texts: GenericGuardrailAPIInputs) -> None:
guardrail = _guardrail(event_hook="post_call")
original_tool_calls = [
{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": '{"ssn":"123-45-6789"}'}}
]
rewritten_tool_calls = [
{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": '{"ssn":"[REDACTED]"}'}}
]
mock_post = AsyncMock(
return_value=_response(
{
"status": "transform",
"transformed_payload": {
"messages": [{"role": "assistant", "content": None, "tool_calls": rewritten_tool_calls}]
},
}
)
)
with patch.object(guardrail.async_handler, "post", mock_post):
result = await guardrail.apply_guardrail(
inputs={**sent_texts, "tool_calls": original_tool_calls},
request_data={},
input_type="response",
logging_obj=_logging(),
)
assert not result.get("texts")
assert result["tool_calls"] == rewritten_tool_calls
@pytest.mark.asyncio
@pytest.mark.parametrize("emptied", ["", None])
async def test_transform_emptied_input_blanks_text_and_message(self, emptied: str | None) -> None:
guardrail = _guardrail()
mock_post = AsyncMock(
return_value=_response(
{
"status": "transform",
"transformed_payload": {"messages": [{"role": "user", "content": emptied}]},
}
)
)
with patch.object(guardrail.async_handler, "post", mock_post):
result = await guardrail.apply_guardrail(
inputs={
"texts": ["my ssn is 123-45-6789"],
"structured_messages": [{"role": "user", "content": "my ssn is 123-45-6789"}],
},
request_data={},
input_type="request",
logging_obj=_logging(),
)
assert result["texts"] == [""]
assert result["structured_messages"] == [{"role": "user", "content": emptied}]
@pytest.mark.asyncio
@pytest.mark.parametrize("returned", [1, 3])
async def test_transform_output_message_count_mismatch_fail_closed(self, returned: int) -> None:
guardrail = _guardrail(event_hook="post_call")
rewritten = [{"role": "assistant", "content": "[REDACTED]"} for _ in range(returned)]
mock_post = AsyncMock(
return_value=_response({"status": "transform", "transformed_payload": {"messages": rewritten}})
)
with patch.object(guardrail.async_handler, "post", mock_post):
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs={"texts": ["ssn 111-11-1111", "ssn 222-22-2222"]},
request_data={},
input_type="response",
logging_obj=_logging(),
)
assert exc_info.value.status_code == 400
assert "transform missing payload" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_transform_input_message_count_mismatch_fail_closed(self) -> None:
guardrail = _guardrail()
mock_post = AsyncMock(
return_value=_response(
{
"status": "transform",
"transformed_payload": {"messages": [{"role": "user", "content": "ssn is [REDACTED]"}]},
}
)
)
with patch.object(guardrail.async_handler, "post", mock_post):
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs={
"texts": ["you are a helpful assistant", "ssn is 123-45-6789"],
"structured_messages": [
{"role": "system", "content": "you are a helpful assistant"},
{"role": "user", "content": "ssn is 123-45-6789"},
],
},
request_data={},
input_type="request",
logging_obj=_logging(),
)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_transform_messages_keeps_tool_calls_when_omitted(self) -> None:
guardrail = _guardrail()