fix(guardrails): moderate tool and system text in its own call so a long payload cannot hide a flagged message

enterprise/enterprise_hooks/google_text_moderation.py

@@ -12,7 +12,7 @@ import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.proxy._types import UserAPIKeyAuth
-from litellm.proxy.guardrails._content_utils import iter_message_text
+from litellm.proxy.guardrails._content_utils import split_message_text
 from litellm.types.utils import CallTypesLiteral
 
 
@@ -95,36 +95,42 @@ class ENTERPRISE_GoogleTextModeration(CustomLogger):
         - Calls Google's Text Moderation API
         - Rejects request if it fails safety check
         """
-        # Covers multimodal list content + Responses-API input.
-        text = "".join(iter_message_text(data))
-        if text:
-            document = self.language_document(content=text, type_=self.document_type)
-
-            request = self.moderate_text_request(
-                document=document,
-            )
-
-            # Make the request
-            response = self.client.moderate_text(request=request)
-            for category in response.moderation_categories:
-                category_name = category.name
-                category_name = category_name.lower()
-                category_name = category_name.replace("&", "and")
-                category_name = category_name.replace(",", "")
-                category_name = category_name.replace(
-                    " ", "_"
-                )  # e.g. go from 'Firearms & Weapons' to 'firearms_and_weapons'
-                if category.confidence > getattr(
-                    self, f"{category_name}_confidence_threshold"
-                ):
-                    raise HTTPException(
-                        status_code=400,
-                        detail={
-                            "error": f"Violated content safety policy. Category={category}"
-                        },
-                    )
+        # Covers multimodal list content + Responses-API input. System, instructions
+        # and tool text go in a second document after the message text.
+        message_text, added_text = split_message_text(data)
+        texts = [text for text in ("".join(message_text), "".join(added_text)) if text]
+        if texts:
+            for text in texts:
+                self._moderate_text(text)
             # Handle the response
             return data
+
+    def _moderate_text(self, text: str) -> None:
+        document = self.language_document(content=text, type_=self.document_type)
+
+        request = self.moderate_text_request(
+            document=document,
+        )
+
+        # Make the request
+        response = self.client.moderate_text(request=request)
+        for category in response.moderation_categories:
+            category_name = category.name
+            category_name = category_name.lower()
+            category_name = category_name.replace("&", "and")
+            category_name = category_name.replace(",", "")
+            category_name = category_name.replace(
+                " ", "_"
+            )  # e.g. go from 'Firearms & Weapons' to 'firearms_and_weapons'
+            if category.confidence > getattr(
+                self, f"{category_name}_confidence_threshold"
+            ):
+                raise HTTPException(
+                    status_code=400,
+                    detail={
+                        "error": f"Violated content safety policy. Category={category}"
+                    },
+                )
 _ENTERPRISE_GoogleTextModeration = ENTERPRISE_GoogleTextModeration
 
 

enterprise/enterprise_hooks/openai_moderation.py

@@ -20,7 +20,7 @@ from litellm._logging import verbose_proxy_logger
 from litellm.constants import DEFAULT_OPENAI_MODERATIONS_MODEL
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.proxy._types import UserAPIKeyAuth
-from litellm.proxy.guardrails._content_utils import iter_message_text
+from litellm.proxy.guardrails._content_utils import split_message_text
 from litellm.types.utils import CallTypesLiteral
 
 
@@ -37,22 +37,24 @@ class ENTERPRISE_OpenAI_Moderation(CustomLogger):
         user_api_key_dict: UserAPIKeyAuth,
         call_type: CallTypesLiteral,
     ):
-        # Covers multimodal list content + Responses-API input.
-        text = "".join(iter_message_text(data))
+        # Covers multimodal list content + Responses-API input. System, instructions
+        # and tool text go in a second call after the message text.
+        message_text, added_text = split_message_text(data)
+        texts = ("".join(message_text), "".join(added_text)) if added_text else ("".join(message_text),)
 
         from litellm.proxy.proxy_server import llm_router
 
         if llm_router is None:
             return
 
-        moderation_response = await llm_router.amoderation(
-            model=self.model_name, input=text
-        )
-
-        verbose_proxy_logger.debug("Moderation response: %s", moderation_response)
-        if moderation_response and moderation_response.results[0].flagged is True:
-            raise HTTPException(
-                status_code=403, detail={"error": "Violated content safety policy"}
+        for text in texts:
+            moderation_response = await llm_router.amoderation(
+                model=self.model_name, input=text
             )
-        pass
+
+            verbose_proxy_logger.debug("Moderation response: %s", moderation_response)
+            if moderation_response and moderation_response.results[0].flagged is True:
+                raise HTTPException(
+                    status_code=403, detail={"error": "Violated content safety policy"}
+                )
 _ENTERPRISE_OpenAI_Moderation = ENTERPRISE_OpenAI_Moderation

litellm/proxy/guardrails/_content_utils.py

@@ -315,8 +315,18 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]:
     system, instructions and tool text. Walks every role (user, assistant,
     system, …) — guardrails inspect the entire conversation, not just user turns.
     """
-    message_text: Final = tuple(_iter_message_parts_text(data))
+    message_text, added_text = split_message_text(data)
     yield from message_text
+    yield from added_text
+
+
+def split_message_text(data: Mapping[str, object]) -> tuple[tuple[str, ...], tuple[str, ...]]:
+    """Return the message text parts, then the system, instructions and tool text, in :func:`iter_message_text` order."""
+    message_text: Final = tuple(_iter_message_parts_text(data))
+    return message_text, tuple(_iter_added_text(data, message_text))
+
+
+def _iter_added_text(data: Mapping[str, object], message_text: tuple[str, ...]) -> Iterator[str]:
     pending: Final = Counter(message_text)
     for text in _iter_request_text(data):
         if pending[text]:

tests/integration/security/test_tool_text_guardrail_consumers.py

@@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter
 
 KEYWORD: Final = "zebracornflux"
 FLAGGED: Final = "moderation-flagged-fragment"
+MODERATION_LIMIT: Final = 24_000
 AZURE_FLAGGED: Final = "azure-flagged-fragment"
 AZURE_LIMIT: Final = 10_000
 AZURE_ANALYZE: Final = "/contentsafety/text:analyze"
@@ -38,7 +39,7 @@ def _vendors(request: Request) -> Reply:
         return Reply()
     body: Final = json.loads(request.body)
     if request.target.endswith("/moderations"):
-        flagged: Final = FLAGGED in body["input"]
+        flagged: Final = FLAGGED in body["input"] and len(body["input"]) <= MODERATION_LIMIT
         return Reply(
             body=json.dumps(
                 {
@@ -209,25 +210,28 @@ def test_deep_tool_result_chain_cannot_hide_a_later_keyword(rig: Rig) -> None:
 
 
 @pytest.mark.parametrize(
-    ("path", "body", "fragments"),
+    ("path", "body", "message", "added"),
     (
         pytest.param(
             "/v1/messages",
             _messages("tool result"),
-            ("read notes", "summarize", "system note", "notes.txt\ntool input", "tool result"),
+            ("read notes", "summarize"),
+            ("system note", "notes.txt\ntool input", "tool result"),
             id="messages",
         ),
         pytest.param(
             "/v1/chat/completions",
             _chat(json.dumps({"path": "n.txt", "content": "argument text"})),
-            ("write notes", "written", "say hi", "n.txt\nargument text"),
+            ("write notes", "written", "say hi"),
+            ("n.txt\nargument text",),
             id="chat",
         ),
     ),
 )
 def test_read_only_vendors_receive_tool_and_system_text(
-    rig: Rig, path: str, body: dict[str, JsonValue], fragments: tuple[str, ...]
+    rig: Rig, path: str, body: dict[str, JsonValue], message: tuple[str, ...], added: tuple[str, ...]
 ) -> None:
+    fragments: Final = (*message, *added)
     rig.wire.drain()
     response: Final = rig.proxy.request("POST", path, body)
     assert response.status_code == 200, response.text
@@ -235,12 +239,28 @@ def test_read_only_vendors_receive_tool_and_system_text(
     moderated: Final = tuple(json.loads(r.body)["input"] for r in sent if r.target.endswith("/moderations"))
     detected: Final = tuple(json.loads(r.body)["contents"] for r in sent if "/api/v1/text/contents" in r.target)
     orchestrated: Final = tuple(json.loads(r.body)["content"] for r in sent if r.target.endswith("/detection/content"))
-    assert moderated == ("".join(fragments),)
+    assert moderated == ("".join(message), "".join(added))
     assert detected == (list(fragments),)
     assert orchestrated == fragments
     assert len(_provider_calls(sent)) == 1
 
 
+def test_moderation_flag_in_message_text_blocks_next_to_long_system_text(rig: Rig) -> None:
+    rig.wire.drain()
+    system: Final = ("Follow the repository conventions and write tests. " * 600)[:30_000]
+    body: Final = {
+        "model": "claude",
+        "max_tokens": 32,
+        "cache": {"no-cache": True},
+        "system": system,
+        "messages": [{"role": "user", "content": f"say {FLAGGED}"}],
+    }
+    response: Final = rig.proxy.request("POST", "/v1/messages", body)
+    assert response.status_code == 403, response.text
+    moderated: Final = tuple(json.loads(r.body)["input"] for r in rig.wire.drain() if r.target.endswith("/moderations"))
+    assert moderated == (f"say {FLAGGED}",)
+
+
 def test_moderation_flag_in_tool_result_blocks(rig: Rig) -> None:
     rig.wire.drain()
     response: Final = rig.proxy.request("POST", "/v1/messages", _messages(FLAGGED))

tests/unit/proxy/guardrails/test_guardrail_coverage.py

@@ -18,7 +18,7 @@ from collections.abc import Mapping
 from functools import reduce
 from types import SimpleNamespace
 from typing import Any, Dict, Final, NamedTuple
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, call, patch
 
 import pytest
 from fastapi import HTTPException
@@ -851,6 +851,62 @@ async def test_openai_moderation_inspects_multimodal_content(monkeypatch, user_a
     assert seen_inputs == ["alpha beta"]
 
 
+_MODERATION_FLAG: Final = "I am going to kill you tonight."
+_MODERATION_LIMIT: Final = 24_000
+_LONG_BENIGN: Final = ("Follow the repository conventions and write tests. " * 600)[:30_000]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("data", "added"),
+    [
+        ({"instructions": _LONG_BENIGN, "input": _MODERATION_FLAG}, _LONG_BENIGN),
+        ({"system": _LONG_BENIGN, "messages": [{"role": "user", "content": _MODERATION_FLAG}]}, _LONG_BENIGN),
+        (
+            {
+                "messages": [
+                    {
+                        "role": "assistant",
+                        "tool_calls": [
+                            {
+                                "id": "call_1",
+                                "type": "function",
+                                "function": {"name": "write_file", "arguments": json.dumps({"content": _LONG_BENIGN})},
+                            }
+                        ],
+                    },
+                    {"role": "user", "content": _MODERATION_FLAG},
+                ]
+            },
+            _LONG_BENIGN,
+        ),
+    ],
+    ids=("responses-instructions", "messages-system", "chat-tool-arguments"),
+)
+async def test_openai_moderation_sends_tool_and_system_text_in_its_own_call(
+    monkeypatch: pytest.MonkeyPatch, user_api_key: UserAPIKeyAuth, data: dict[str, JsonValue], added: str
+) -> None:
+    from enterprise.enterprise_hooks.openai_moderation import ENTERPRISE_OpenAI_Moderation
+
+    seen: Final[list[str]] = []
+
+    async def amoderation(model: str, input: str) -> SimpleNamespace:
+        seen.append(input)
+        flagged: Final = _MODERATION_FLAG in input and len(input) <= _MODERATION_LIMIT
+        return SimpleNamespace(results=[SimpleNamespace(flagged=flagged)])
+
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", SimpleNamespace(amoderation=amoderation))
+    with pytest.raises(HTTPException) as rejection:
+        await ENTERPRISE_OpenAI_Moderation().async_moderation_hook(data, user_api_key, "acompletion")
+    assert rejection.value.status_code == 403
+    assert seen == [_MODERATION_FLAG]
+
+    seen.clear()
+    clean: Final = json.loads(json.dumps(data).replace(_MODERATION_FLAG, "say hi"))
+    assert await ENTERPRISE_OpenAI_Moderation().async_moderation_hook(clean, user_api_key, "acompletion") is None
+    assert seen == ["say hi", added]
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     ("configured_after_init", "expected_model"),
@@ -1004,11 +1060,11 @@ class _GoogleRequest(NamedTuple):
 @pytest.mark.asyncio
 @pytest.mark.parametrize("blocked", (False, True), ids=("allow", "block"))
 @pytest.mark.parametrize(
-    ("data", "document_text"),
+    ("data", "documents"),
     (
         pytest.param(
             {"messages": [{"role": "user", "content": "ordinary blocked fragment"}]},
-            "ordinary blocked fragment",
+            ("ordinary blocked fragment",),
             id="ordinary-text-control",
         ),
         pytest.param(
@@ -1016,7 +1072,7 @@ class _GoogleRequest(NamedTuple):
                 "system": [{"type": "text", "text": "system blocked fragment"}],
                 "messages": [{"role": "user", "content": "ordinary control"}],
             },
-            "ordinary controlsystem blocked fragment",
+            ("ordinary control", "system blocked fragment"),
             id="native-system",
         ),
         pytest.param(
@@ -1039,7 +1095,7 @@ class _GoogleRequest(NamedTuple):
                     },
                 ]
             },
-            "ordinary controlinput first \ninput blocked fragment",
+            ("ordinary control", "input first \ninput blocked fragment"),
             id="native-tool-input",
         ),
         pytest.param(
@@ -1062,7 +1118,7 @@ class _GoogleRequest(NamedTuple):
                     },
                 ]
             },
-            "ordinary controlresult first result blocked fragment",
+            ("ordinary control", "result first result blocked fragment"),
             id="native-tool-result",
         ),
         pytest.param(
@@ -1089,7 +1145,7 @@ class _GoogleRequest(NamedTuple):
                     },
                 ]
             },
-            "ordinary controlargument first \nargument blocked fragment",
+            ("ordinary control", "argument first \nargument blocked fragment"),
             id="chat-tool-arguments",
         ),
     ),
@@ -1097,7 +1153,7 @@ class _GoogleRequest(NamedTuple):
 async def test_google_text_moderation_inspects_tool_and_system_payloads(
     user_api_key: UserAPIKeyAuth,
     data: Mapping[str, JsonValue],
-    document_text: str,
+    documents: tuple[str, ...],
     blocked: bool,
 ) -> None:
     sdk: Final = MagicMock()
@@ -1121,7 +1177,8 @@ async def test_google_text_moderation_inspects_tool_and_system_payloads(
         else:
             assert await guard.async_moderation_hook(request_data, user_api_key, "acompletion") is request_data
     sdk.LanguageServiceClient.assert_called_once_with()
-    sdk.types.Document.assert_called_once_with(content=document_text, type_=sdk.types.Document.Type.PLAIN_TEXT)
-    client.moderate_text.assert_called_once_with(
-        request=_GoogleRequest(_GoogleDocument(document_text, sdk.types.Document.Type.PLAIN_TEXT))
-    )
+    plain_text: Final = sdk.types.Document.Type.PLAIN_TEXT
+    assert sdk.types.Document.call_args_list == [call(content=text, type_=plain_text) for text in documents]
+    assert client.moderate_text.call_args_list == [
+        call(request=_GoogleRequest(_GoogleDocument(text, plain_text))) for text in documents
+    ]