diff --git a/litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py b/litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py index f8469186c49..61a8ec8940d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py +++ b/litellm/proxy/guardrails/guardrail_hooks/neuraltrust/neuraltrust.py @@ -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) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_neuraltrust.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_neuraltrust.py index 20993b1f3b5..7f325790059 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_neuraltrust.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_neuraltrust.py @@ -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()