diff --git a/enterprise/enterprise_hooks/google_text_moderation.py b/enterprise/enterprise_hooks/google_text_moderation.py index b3ef5a3d496..1841e507aa2 100644 --- a/enterprise/enterprise_hooks/google_text_moderation.py +++ b/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 diff --git a/enterprise/enterprise_hooks/openai_moderation.py b/enterprise/enterprise_hooks/openai_moderation.py index d8536ad1ac3..fdf4d5d2043 100644 --- a/enterprise/enterprise_hooks/openai_moderation.py +++ b/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 diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index d6fc5c25509..cad1a069d56 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/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]: diff --git a/tests/integration/security/test_tool_text_guardrail_consumers.py b/tests/integration/security/test_tool_text_guardrail_consumers.py index 302cf54ecaa..e753af51484 100644 --- a/tests/integration/security/test_tool_text_guardrail_consumers.py +++ b/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)) diff --git a/tests/unit/proxy/guardrails/test_guardrail_coverage.py b/tests/unit/proxy/guardrails/test_guardrail_coverage.py index 6418f474f56..fc9146d6acd 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_coverage.py +++ b/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 + ]