From 5dc651bad9beeac3599f6338a766ecfc2973c970 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:21:30 -0700 Subject: [PATCH] fix(guardrails): confine tool-traffic coverage to hide-secrets Restore iter_message_text to its main behavior, so banned keywords, Azure Content Safety, IBM detectors and the OpenAI and Google moderation hooks read exactly what they read on main. Only walk_user_text, which hide-secrets alone calls, still reaches system, instructions and tool payloads. --- .../google_text_moderation.py | 64 ++-- .../enterprise_hooks/openai_moderation.py | 26 +- litellm/proxy/guardrails/_content_utils.py | 88 +---- .../test_tool_text_guardrail_consumers.py | 313 ------------------ .../proxy/guardrails/test_content_utils.py | 77 +---- .../guardrails/test_guardrail_coverage.py | 312 +---------------- 6 files changed, 57 insertions(+), 823 deletions(-) delete mode 100644 tests/integration/security/test_tool_text_guardrail_consumers.py diff --git a/enterprise/enterprise_hooks/google_text_moderation.py b/enterprise/enterprise_hooks/google_text_moderation.py index 1841e507aa2..b3ef5a3d496 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 split_message_text +from litellm.proxy.guardrails._content_utils import iter_message_text from litellm.types.utils import CallTypesLiteral @@ -95,42 +95,36 @@ class ENTERPRISE_GoogleTextModeration(CustomLogger): - Calls Google's Text Moderation API - Rejects request if it fails safety check """ - # 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) + # 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}" + }, + ) # 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 74ceaf5d5eb..d8536ad1ac3 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 split_message_text +from litellm.proxy.guardrails._content_utils import iter_message_text from litellm.types.utils import CallTypesLiteral @@ -37,24 +37,22 @@ class ENTERPRISE_OpenAI_Moderation(CustomLogger): user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral, ): - # 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 = tuple(text for text in ("".join(message_text), "".join(added_text)) if text) or ("",) + # Covers multimodal list content + Responses-API input. + text = "".join(iter_message_text(data)) from litellm.proxy.proxy_server import llm_router if llm_router is None: return - for text in texts: - moderation_response = await llm_router.amoderation( - model=self.model_name, input=text - ) + 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"} - ) + 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"} + ) + pass _ENTERPRISE_OpenAI_Moderation = ENTERPRISE_OpenAI_Moderation diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index cad1a069d56..665d68c0a1a 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -11,7 +11,6 @@ every text fragment. import json import math import re -from collections import Counter from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass from typing import Any, Final @@ -156,13 +155,6 @@ def iter_tool_json_text(value: object) -> Iterator[str]: yield text -def _tool_payload_text(value: object) -> Iterator[str]: - """Yield a tool payload's string values as one newline-joined fragment, like one text part.""" - text: Final = "\n".join(iter_tool_json_text(value)) - if text: - yield text - - @dataclass(frozen=True, slots=True) class _ToolFrame: source: object @@ -210,20 +202,6 @@ def _tool_json(value: object) -> JsonValue: return map_tool_json_text(value, lambda text: text) -def _iter_native_content_text(content: object) -> Iterator[str]: - parts: Final = _object_list(content) - if parts is None: - yield from _iter_text_parts_in_content(content) - return - for part in parts: - yield from _iter_text_parts_in_content([part]) - fields = _text_object(part) - if fields is not None and fields.get("type") == "tool_result": - yield from _iter_text_parts_in_content(fields.get("content")) - elif fields is not None and fields.get("type") == "tool_use": - yield from _tool_payload_text(fields.get("input")) - - # Responses-API item types whose ``output`` field carries user/tool text # that guardrails should inspect. ``function_call_output`` is the # built-in shape; ``custom_tool_call_output`` is the custom-tool @@ -309,59 +287,15 @@ def _iter_inspection_messages(data: Mapping[str, object]) -> Iterator[object]: def iter_message_text(data: Mapping[str, object]) -> Iterator[str]: - """Yield every text fragment from ``messages``, ``input``, ``system``, ``instructions`` AND tool payloads. + """Yield every text fragment from ``messages`` AND ``input``. - Message text parts come first, in conversation order, followed by the - system, instructions and tool text. Walks every role (user, assistant, - system, …) — guardrails inspect the entire conversation, not just user turns. + Walks every role (user, assistant, system, …) — guardrails inspect + the entire conversation, not just user turns. """ - 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]: - pending[text] -= 1 - else: - yield text - - -def _iter_message_parts_text(data: Mapping[str, object]) -> Iterator[str]: for message in _iter_inspection_messages(data): - fields = _text_object(message) - if fields is not None: - yield from _iter_text_parts_in_content(fields.get("content")) - - -def _iter_request_text(data: Mapping[str, object]) -> Iterator[str]: - yield from _iter_native_content_text(data.get("system")) - yield from _iter_native_content_text(data.get("instructions")) - for message in _iter_inspection_messages(data): - fields = _text_object(message) - if fields is None: + if not isinstance(message, dict): continue - yield from _iter_native_content_text(fields.get("content")) - calls = _object_list(fields.get("tool_calls")) - if calls is not None: - for call in calls: - call_fields = _text_object(call) - function = _text_object(call_fields.get("function")) if call_fields is not None else None - if function is not None: - yield from _iter_arguments_text(function.get("arguments")) - items: Final = _object_list(data.get("input")) if isinstance(data.get("input"), list) else None - for item in items or (): - call = _call_item(item) - if call is not None: - yield from _iter_arguments_text(call[1]) + yield from _iter_text_parts_in_content(message.get("content")) def _call_item(item: object) -> tuple[str, object] | None: @@ -371,18 +305,6 @@ def _call_item(item: object) -> tuple[str, object] | None: return (field, fields.get(field)) if fields is not None and field is not None and field in fields else None -def _iter_arguments_text(arguments: object) -> Iterator[str]: - if isinstance(arguments, str) and arguments: - try: - parsed: Final = _tool_arguments(arguments) - except (ValueError, RecursionError): - yield arguments - else: - yield from _tool_payload_text(parsed) - elif isinstance(arguments, dict): - yield from _tool_payload_text(arguments) - - def walk_user_text( data: dict[str, object], visit: Callable[[str], str], diff --git a/tests/integration/security/test_tool_text_guardrail_consumers.py b/tests/integration/security/test_tool_text_guardrail_consumers.py deleted file mode 100644 index e753af51484..00000000000 --- a/tests/integration/security/test_tool_text_guardrail_consumers.py +++ /dev/null @@ -1,313 +0,0 @@ -from __future__ import annotations - -import json -from collections.abc import Iterator -from functools import reduce -from typing import Final - -import pytest -from integration._support.wire import Reply, Request -from integration.security.test_hide_secrets_tool_traffic import Rig, _provider, tool_rig -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" -AZURE_REQUEST: Final = TypeAdapter(dict[str, JsonValue]) - - -def _azure(request: Request) -> Reply: - text: Final = AZURE_REQUEST.validate_json(request.body)["text"] - assert isinstance(text, str) - if len(text) > AZURE_LIMIT: - error: Final = { - "code": "InvalidRequestBody", - "message": f"The length of given text {len(text)} exceeds the limit", - } - return Reply(status=400, body=json.dumps({"error": error}).encode()) - severity: Final = 6 if AZURE_FLAGGED in text else 0 - categories: Final = ("Hate", "SelfHarm", "Sexual", "Violence") - analysis: Final = [{"category": name, "severity": severity if name == "Violence" else 0} for name in categories] - return Reply(body=json.dumps({"blocklistsMatch": [], "categoriesAnalysis": analysis}).encode()) - - -def _vendors(request: Request) -> Reply: - if request.method != "POST": - return Reply() - body: Final = json.loads(request.body) - if request.target.endswith("/moderations"): - flagged: Final = FLAGGED in body["input"] and len(body["input"]) <= MODERATION_LIMIT - return Reply( - body=json.dumps( - { - "id": "modr-integration", - "model": body["model"], - "results": [ - { - "flagged": flagged, - "categories": {"violence": flagged}, - "category_scores": {"violence": 0.99 if flagged else 0.0}, - } - ], - } - ).encode() - ) - if "/api/v1/text/contents" in request.target: - return Reply(body=json.dumps([[] for _ in body["contents"]]).encode()) - if request.target.endswith("/api/v2/text/detection/content"): - return Reply(body=b'{"detections": []}') - if AZURE_ANALYZE in request.target: - return _azure(request) - return _provider(request) - - -def _consumers(configuration: dict[str, JsonValue]) -> None: - models: Final = configuration["model_list"] - settings: Final = configuration["litellm_settings"] - assert isinstance(models, list) and isinstance(settings, dict) - vendor: Final = next(model for model in models if isinstance(model, dict) and model["model_name"] == "gpt") - assert isinstance(vendor, dict) and isinstance(vendor["litellm_params"], dict) - base_url: Final = vendor["litellm_params"]["api_base"] - models.append( - { - "model_name": "omni-moderation-latest", - "litellm_params": { - "model": "openai/omni-moderation-latest", - "api_key": "synthetic-provider-key", - "api_base": base_url, - }, - } - ) - settings["banned_keywords_list"] = [KEYWORD] - settings["callbacks"] = ["banned_keywords", "openai_moderations", "azure_content_safety"] - settings["azure_content_safety_params"] = {"endpoint": base_url, "api_key": "synthetic-azure-key"} - configuration["guardrails"] = [ - { - "guardrail_name": f"ibm-{name}", - "litellm_params": { - "guardrail": "ibm_guardrails", - "mode": "pre_call", - "default_on": True, - "auth_token": "synthetic-detector-token", - "base_url": base_url, - "detector_id": "tool-text", - "is_detector_server": detector_server, - }, - } - for name, detector_server in (("detector", True), ("orchestrator", False)) - ] - - -@pytest.fixture(scope="module") -def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: - with tool_rig(tmp_path_factory.mktemp("tool-text-consumers"), _vendors, configure=_consumers) as value: - yield value - - -def _chat(arguments: str, *, stream: bool = False, prompt: str = "say hi") -> dict[str, JsonValue]: - return { - "model": "gpt", - "stream": stream, - "cache": {"no-cache": True}, - "messages": [ - {"role": "user", "content": "write notes"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - {"id": "call_1", "type": "function", "function": {"name": "write_file", "arguments": arguments}} - ], - }, - {"role": "tool", "tool_call_id": "call_1", "content": "written"}, - {"role": "user", "content": prompt}, - ], - } - - -def _messages(result: str) -> dict[str, JsonValue]: - return { - "model": "claude", - "max_tokens": 32, - "cache": {"no-cache": True}, - "system": [{"type": "text", "text": "system note"}], - "messages": [ - {"role": "user", "content": "read notes"}, - { - "role": "assistant", - "content": [ - { - "type": "tool_use", - "id": "tool_1", - "name": "read_file", - "input": {"path": "notes.txt", "options": [7, {"mode": "tool input"}]}, - } - ], - }, - { - "role": "user", - "content": [ - {"type": "tool_result", "tool_use_id": "tool_1", "content": [{"type": "text", "text": result}]}, - {"type": "text", "text": "summarize"}, - ], - }, - ], - } - - -def _provider_calls(sent: tuple[Request, ...]) -> tuple[Request, ...]: - return tuple( - request for request in sent if request.target.endswith(("/chat/completions", "/messages", "/responses")) - ) - - -def _analyzed(sent: tuple[Request, ...]) -> tuple[str, ...]: - return tuple(json.loads(request.body)["text"] for request in sent if AZURE_ANALYZE in request.target) - - -@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) -def test_keyword_in_tool_call_arguments_is_blocked(rig: Rig, stream: bool) -> None: - rig.wire.drain() - response: Final = rig.proxy.request( - "POST", "/v1/chat/completions", _chat(json.dumps({"path": f"notes/{KEYWORD}.txt"}), stream=stream) - ) - assert response.status_code == 400, response.text - assert f"Keyword={KEYWORD}" in response.text - assert _provider_calls(rig.wire.drain()) == () - - -def test_tool_call_arguments_without_keyword_reach_provider(rig: Rig) -> None: - rig.wire.drain() - arguments: Final = json.dumps({"path": "notes/plain.txt"}) - response: Final = rig.proxy.request("POST", "/v1/chat/completions", _chat(arguments)) - assert response.status_code == 200, response.text - sent: Final = _provider_calls(rig.wire.drain()) - assert len(sent) == 1 - assert json.loads(sent[0].body)["messages"][1]["tool_calls"][0]["function"]["arguments"] == arguments - - -def test_deep_tool_result_chain_cannot_hide_a_later_keyword(rig: Rig) -> None: - rig.wire.drain() - chain: Final = reduce( - lambda content, _: [{"type": "tool_result", "tool_use_id": "tool_1", "content": content}], - range(1000), - [{"type": "text", "text": "inner"}], - ) - response: Final = rig.proxy.request( - "POST", - "/v1/chat/completions", - { - "model": "gpt", - "cache": {"no-cache": True}, - "messages": [{"role": "assistant", "content": chain}, {"role": "user", "content": f"say {KEYWORD}"}], - }, - ) - assert response.status_code == 400, response.text - assert f"Keyword={KEYWORD}" in response.text - assert _provider_calls(rig.wire.drain()) == () - - -@pytest.mark.parametrize( - ("path", "body", "message", "added"), - ( - pytest.param( - "/v1/messages", - _messages("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",), - id="chat", - ), - ), -) -def test_read_only_vendors_receive_tool_and_system_text( - 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 - sent: Final = rig.wire.drain() - 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(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)) - assert response.status_code == 403, response.text - assert "Violated content safety policy" in response.text - - -def test_azure_content_safety_checks_message_text_before_tool_text(rig: Rig) -> None: - rig.wire.drain() - response: Final = rig.proxy.request( - "POST", "/v1/chat/completions", _chat(json.dumps({"path": "n.txt", "content": "argument text"})) - ) - assert response.status_code == 200, response.text - sent: Final = rig.wire.drain() - assert _analyzed(sent) == ("write notes", "written", "say hi", "n.txt\nargument text", "accepted") - assert len(_provider_calls(sent)) == 1 - - -@pytest.mark.parametrize( - ("path", "body", "analyzed"), - ( - pytest.param( - "/v1/chat/completions", - _chat(json.dumps({"content": "x" * (AZURE_LIMIT + 1)}), prompt=f"say {AZURE_FLAGGED}"), - ("write notes", "written", f"say {AZURE_FLAGGED}"), - id="chat-tool-arguments", - ), - pytest.param( - "/v1/responses", - { - "model": "gpt", - "cache": {"no-cache": True}, - "instructions": "x" * (AZURE_LIMIT + 1), - "input": f"say {AZURE_FLAGGED}", - }, - (f"say {AZURE_FLAGGED}",), - id="responses-instructions", - ), - ), -) -def test_azure_content_safety_blocks_message_text_after_oversized_tool_or_system_text( - rig: Rig, path: str, body: dict[str, JsonValue], analyzed: tuple[str, ...] -) -> None: - rig.wire.drain() - response: Final = rig.proxy.request("POST", path, body) - assert response.status_code == 400, response.text - assert "Violated content safety policy" in response.text - sent: Final = rig.wire.drain() - assert _analyzed(sent) == analyzed - assert _provider_calls(sent) == () diff --git a/tests/unit/proxy/guardrails/test_content_utils.py b/tests/unit/proxy/guardrails/test_content_utils.py index cf0f2db2c95..99fe15f37d7 100644 --- a/tests/unit/proxy/guardrails/test_content_utils.py +++ b/tests/unit/proxy/guardrails/test_content_utils.py @@ -720,18 +720,16 @@ def test_walk_user_text_redacts_custom_tool_call_output(): ), ids=("function-call", "function-call-malformed", "custom-tool-call"), ) -def test_responses_tool_call_items_are_read_and_masked(item, field, masked): +def test_responses_tool_call_items_are_masked(item, field, masked): data = {"input": [{"role": "user", "content": "go"}, item]} - assert any("secret" in text and "secret_tool_name" not in text for text in iter_message_text(data)) assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 2 assert data["input"][1][field] == masked assert data["input"][1].get("name", "secret_tool_name") == "secret_tool_name" @pytest.mark.parametrize("instructions", ("secret rules", [{"type": "input_text", "text": "secret rules"}])) -def test_responses_instructions_are_read_and_masked(instructions): +def test_responses_instructions_are_masked(instructions): data = {"instructions": instructions, "input": "hi"} - assert "secret rules" in tuple(iter_message_text(data)) assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 2 masked = data["instructions"] if isinstance(instructions, str) else data["instructions"][0]["text"] assert masked == "[MASKED] rules" @@ -811,7 +809,6 @@ def test_native_tool_payload_rewrites_values_preserving_structure(tool_input): }, ], } - assert "secret" in "\n".join(iter_message_text(data)) visited = walk_user_text(data, lambda value: value.replace("secret", "[MASKED]")) assert visited == 3 block = data["messages"][0]["content"][0] @@ -846,7 +843,6 @@ def _tool_call(arguments: JsonValue) -> dict[str, JsonValue]: ) def test_native_tool_call_arguments_are_masked_as_json(arguments: JsonValue) -> None: data = _tool_call(arguments) - assert "secret" in "\n".join(iter_message_text(data)) assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1 function = data["messages"][0]["tool_calls"][0]["function"] masked = json.loads(function["arguments"]) if isinstance(arguments, str) else function["arguments"] @@ -861,15 +857,13 @@ def test_native_tool_call_arguments_are_masked_as_json(arguments: JsonValue) -> ) def test_native_tool_call_arguments_that_are_not_clean_json_are_masked_as_raw_text(arguments: str) -> None: data = _tool_call(arguments) - assert arguments in tuple(iter_message_text(data)) assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1 assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == arguments.replace("secret", "[MASKED]") -def test_native_tool_arguments_unchanged_bytes_and_decoded_reader(): +def test_native_tool_arguments_keep_their_bytes_when_nothing_is_masked(): arguments = '{ "content" : "caf\\u00e9", "count": 3 }' data = {"messages": [{"tool_calls": [{"function": {"arguments": arguments}}]}]} - assert "café" in "\n".join(iter_message_text(data)) assert walk_user_text(data, lambda text: text) == 1 assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == arguments @@ -886,8 +880,9 @@ def test_native_tool_arguments_unchanged_bytes_and_decoded_reader(): {"messages": [{"tool_calls": [{"function": {"arguments": {"field": "payload"}}}]}]}, ), ) -def test_native_reader_exposes_each_tool_and_system_slot(data): - assert tuple(iter_message_text(data)) == ("payload",) +def test_native_writer_visits_each_tool_and_system_slot(data): + assert walk_user_text(data, lambda text: text.replace("payload", "[MASKED]")) == 1 + assert "payload" not in json.dumps(data) @pytest.mark.parametrize( @@ -900,7 +895,7 @@ def test_native_reader_exposes_each_tool_and_system_slot(data): def test_native_tool_envelope_text_cannot_shadow_its_payload(part): data = {"messages": [{"role": "assistant", "content": [part]}]} walk_user_text(data, lambda text: text.replace("secret", "[REDACTED]")) - assert list(iter_message_text(data)) == ["safe", "[REDACTED]"] + assert "secret" not in json.dumps(data) assert data["messages"][0]["content"][0]["text"] == "safe" @@ -917,11 +912,10 @@ def _deep_tool_leaf(value: JsonValue, depth: int) -> JsonValue: @pytest.mark.parametrize("depth", (255, 1000)) -def test_deep_tool_input_reader_and_writer_preserve_values(depth: int) -> None: +def test_deep_tool_input_writer_preserves_values(depth: int) -> None: leaf = {"secret-key": "secret", "empty": "", "enabled": False, "count": 7, "ratio": 1.5, "none": None} payload = _deep_tool_value(depth, leaf) data = {"messages": [{"content": [{"type": "tool_use", "id": "tool", "name": "read", "input": payload}]}]} - assert tuple(iter_message_text(data)) == ("secret",) assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1 result = _deep_tool_leaf(data["messages"][0]["content"][0]["input"], depth) assert isinstance(result, dict) @@ -930,7 +924,6 @@ def test_deep_tool_input_reader_and_writer_preserve_values(depth: int) -> None: assert result["enabled"] is False and result["none"] is None assert type(result["count"]) is int and type(result["ratio"]) is float assert leaf["secret-key"] == "secret" - assert tuple(iter_message_text(data)) == ("[MASKED]",) @pytest.mark.parametrize("depth", (255, 1000)) @@ -945,7 +938,6 @@ def test_deep_tool_json_direct_map_and_iteration(depth: int) -> None: def test_deep_tool_arguments_decode_values_and_preserve_no_hit_bytes() -> None: arguments = '{ "nested" : ' * 255 + '{"caf\\u00e9":"caf\\u00e9 secret","empty":"","count":7}' + " }" * 255 data = {"messages": [{"tool_calls": [{"function": {"name": "secret_tool", "arguments": arguments}}]}]} - assert tuple(iter_message_text(data)) == ("café secret",) assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 1 function = data["messages"][0]["tool_calls"][0]["function"] decoded = json.loads(function["arguments"]) @@ -971,11 +963,6 @@ def test_deep_tool_arguments_parser_limit_keeps_raw_masking_fallback() -> None: break assert arguments, "JSON decoder did not reach its recursion guard" data = {"messages": [{"tool_calls": [{"function": {"arguments": arguments}}]}]} - try: - texts = tuple(iter_message_text(data)) - except RecursionError: - texts = () - assert texts == (arguments,) try: visited = walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) except RecursionError: @@ -1002,49 +989,17 @@ def test_tool_json_graph_repeated_shared_containers_keep_order_and_empty_leaves( assert rewritten["empty"] == [] and shared["secret-key"] == "secret" -def test_nested_tool_result_chain_leaves_later_messages_inspected() -> None: +def test_nested_tool_result_chain_leaves_later_messages_masked() -> None: chain = reduce( lambda content, _: [{"type": "tool_result", "tool_use_id": "tool", "content": content}], range(1500), [{"type": "text", "text": "inner"}], ) data = {"messages": [{"role": "assistant", "content": chain}, {"role": "user", "content": "secret"}]} - assert "secret" in tuple(iter_message_text(data)) walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) assert data["messages"][1]["content"] == "[MASKED]" -def test_message_text_comes_before_system_instructions_and_tool_text() -> None: - data = { - "system": "system", - "instructions": "instructions", - "messages": [ - {"role": "user", "content": "first"}, - { - "role": "assistant", - "content": [{"type": "tool_use", "input": {"path": "tool input"}}], - "tool_calls": [{"function": {"arguments": '{"path":"arguments"}'}}], - }, - { - "role": "user", - "content": [{"type": "tool_result", "content": "first"}, {"type": "text", "text": "second"}], - }, - ], - "input": [{"type": "function_call", "arguments": '{"path":"call"}'}, {"role": "user", "content": "third"}], - } - assert tuple(iter_message_text(data)) == ( - "first", - "second", - "third", - "system", - "instructions", - "tool input", - "arguments", - "first", - "call", - ) - - @pytest.mark.parametrize( "arguments", ('{"token":"secret","limit":1e400}', '{"token":"secret","limit":-Infinity}', '{"token":"secret","ratio":NaN}'), @@ -1092,17 +1047,3 @@ def test_tool_visitor_that_changes_more_than_strings_writes_its_payload_as_argum assert walk_user_text(data, lambda text: text, lambda payload: rewritten) == 1 assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == json.dumps(rewritten) - -@pytest.mark.parametrize( - "data", - ( - { - "messages": [ - {"content": [{"type": "tool_use", "input": {"a": "first", "b": [7, {"c": "second"}, "third"]}}]} - ] - }, - {"messages": [{"tool_calls": [{"function": {"arguments": '{"a":"first","b":[7,{"c":"second"},"third"]}'}}]}]}, - ), -) -def test_native_reader_yields_one_fragment_per_tool_payload(data) -> None: - assert tuple(iter_message_text(data)) == ("first\nsecond\nthird",) diff --git a/tests/unit/proxy/guardrails/test_guardrail_coverage.py b/tests/unit/proxy/guardrails/test_guardrail_coverage.py index 5cb46491d0e..1feae213eba 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/unit/proxy/guardrails/test_guardrail_coverage.py @@ -12,21 +12,13 @@ inspection by the relevant guardrail hook: not just ``choices[0]``. """ -import json -import sys -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, call, patch +from typing import Any, Dict, Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import HTTPException from httpx import Request, Response -from pydantic import JsonValue import litellm -from enterprise.enterprise_hooks.google_text_moderation import ENTERPRISE_GoogleTextModeration from litellm import DualCache from litellm.constants import DEFAULT_OPENAI_MODERATIONS_MODEL from litellm.proxy._types import UserAPIKeyAuth @@ -502,31 +494,6 @@ def test_banned_keywords_blocks_responses_api_input(monkeypatch): asyncio.run(_run()) -def test_banned_keywords_still_checks_messages_after_a_deep_tool_result_chain(monkeypatch): - monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False) - from enterprise.enterprise_hooks.banned_keywords import ENTERPRISE_BannedKeywords - - chain: Final = reduce( - lambda content, _: [{"type": "tool_result", "tool_use_id": "tool", "content": content}], - range(1500), - [{"type": "text", "text": "inner"}], - ) - data: Final = {"messages": [{"role": "assistant", "content": chain}, {"role": "user", "content": "forbidden"}]} - - import asyncio - - with pytest.raises(HTTPException) as exc: - asyncio.run( - ENTERPRISE_BannedKeywords().async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"), - cache=DualCache(), - data=data, - call_type="acompletion", - ) - ) - assert "forbidden" in str(exc.value.detail).lower() - - @pytest.mark.parametrize("call_type", ["completion", "acompletion", "aresponses"]) def test_banned_keywords_fires_on_text_content_call_types(monkeypatch, call_type): """Locks the call-type gate to the runtime ``route_type`` values the @@ -656,58 +623,6 @@ async def test_azure_content_safety_pre_call_fires_on_runtime_call_types( assert ("scan me", "input") in seen -@pytest.mark.asyncio -@pytest.mark.parametrize( - "call_type, data", - [ - ("aresponses", {"instructions": "x" * 10001, "input": [{"role": "user", "content": "flagged text"}]}), - ( - "acompletion", - { - "messages": [ - { - "role": "assistant", - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": {"name": "write_file", "arguments": json.dumps({"content": "x" * 10001})}, - } - ], - }, - {"role": "user", "content": "flagged text"}, - ] - }, - ), - ], - ids=("responses-instructions", "chat-tool-arguments"), -) -async def test_azure_content_safety_checks_message_text_before_oversized_tool_or_system_text( - user_api_key: UserAPIKeyAuth, call_type: str, data: dict[str, JsonValue] -) -> None: - from litellm.proxy.hooks.azure_content_safety import PROXY_AzureContentSafety - - guard: Final = PROXY_AzureContentSafety.__new__(PROXY_AzureContentSafety) - seen: Final[list[str]] = [] - - async def fake_test_violation(content: str, source: str | None = None) -> None: - seen.append(content) - if len(content) > 10000: - raise RuntimeError("(InvalidRequestBody) The length of given text exceeds the limit 10000.") - if "flagged" in content: - raise HTTPException(status_code=400, detail={"error": "Violated content safety policy"}) - - guard.test_violation = fake_test_violation - with pytest.raises(HTTPException): - await guard.async_pre_call_hook( - user_api_key_dict=user_api_key, - cache=DualCache(), - data=data, - call_type=call_type, - ) - assert seen == ["flagged text"] - - @pytest.mark.asyncio async def test_azure_content_safety_post_call_checks_all_choices(user_api_key): """Krrish blocker: ``n>1`` responses must not bypass Azure Content Safety @@ -851,93 +766,6 @@ 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( - ("data", "inputs"), - [ - ({"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "x"}}]}]}, [""]), - ( - { - "system": "be brief", - "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "x"}}]}], - }, - ["be brief"], - ), - ], - ids=("no-text", "system-only"), -) -async def test_openai_moderation_makes_one_call_when_one_side_has_no_text( - monkeypatch: pytest.MonkeyPatch, user_api_key: UserAPIKeyAuth, data: dict[str, JsonValue], inputs: list[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) - return SimpleNamespace(results=[SimpleNamespace(flagged=False)]) - - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", SimpleNamespace(amoderation=amoderation)) - assert await ENTERPRISE_OpenAI_Moderation().async_moderation_hook(data, user_api_key, "acompletion") is None - assert seen == inputs - - @pytest.mark.asyncio @pytest.mark.parametrize( ("configured_after_init", "expected_model"), @@ -1077,139 +905,3 @@ async def test_google_text_moderation_inspects_multimodal_content(user_api_key): ) assert seen_documents == ["hello world"] - - -class _GoogleDocument(NamedTuple): - content: str - type_: object - - -class _GoogleRequest(NamedTuple): - document: _GoogleDocument - - -@pytest.mark.asyncio -@pytest.mark.parametrize("blocked", (False, True), ids=("allow", "block")) -@pytest.mark.parametrize( - ("data", "documents"), - ( - pytest.param( - {"messages": [{"role": "user", "content": "ordinary blocked fragment"}]}, - ("ordinary blocked fragment",), - id="ordinary-text-control", - ), - pytest.param( - { - "system": [{"type": "text", "text": "system blocked fragment"}], - "messages": [{"role": "user", "content": "ordinary control"}], - }, - ("ordinary control", "system blocked fragment"), - id="native-system", - ), - pytest.param( - { - "messages": [ - {"role": "user", "content": "ordinary control"}, - { - "role": "assistant", - "content": [ - { - "type": "tool_use", - "id": "ignored-id", - "name": "ignored-name", - "input": { - "first": "input first ", - "nested": [7, False, None, {"value": "input blocked fragment"}], - }, - } - ], - }, - ] - }, - ("ordinary control", "input first \ninput blocked fragment"), - id="native-tool-input", - ), - pytest.param( - { - "messages": [ - {"role": "user", "content": "ordinary control"}, - { - "role": "user", - "content": [ - { - "type": "tool_result", - "tool_use_id": "ignored-id", - "content": [ - {"type": "text", "text": "result first "}, - {"type": "image", "source": {"data": "ignored-image"}}, - {"type": "text", "text": "result blocked fragment"}, - ], - } - ], - }, - ] - }, - ("ordinary control", "result first result blocked fragment"), - id="native-tool-result", - ), - pytest.param( - { - "messages": [ - {"role": "user", "content": "ordinary control"}, - { - "role": "assistant", - "tool_calls": [ - { - "id": "ignored-id", - "type": "function", - "function": { - "name": "ignored-name", - "arguments": json.dumps( - { - "first": "argument first ", - "nested": [7, False, None, {"value": "argument blocked fragment"}], - } - ), - }, - } - ], - }, - ] - }, - ("ordinary control", "argument first \nargument blocked fragment"), - id="chat-tool-arguments", - ), - ), -) -async def test_google_text_moderation_inspects_tool_and_system_payloads( - user_api_key: UserAPIKeyAuth, - data: Mapping[str, JsonValue], - documents: tuple[str, ...], - blocked: bool, -) -> None: - sdk: Final = MagicMock() - sdk.types.Document.side_effect = _GoogleDocument - sdk.ModerateTextRequest.side_effect = _GoogleRequest - - def moderate(*, request: _GoogleRequest) -> SimpleNamespace: - flagged: Final = blocked and "blocked fragment" in request.document.content - return SimpleNamespace(moderation_categories=[SimpleNamespace(name="Toxic", confidence=float(flagged))]) - - client: Final = sdk.LanguageServiceClient.return_value - client.moderate_text.side_effect = moderate - cloud: Final = SimpleNamespace(language_v1=sdk) - request_data: Final = dict(data) - with patch.dict(sys.modules, {"google.cloud": cloud, "google.cloud.language_v1": sdk}): - guard: Final = ENTERPRISE_GoogleTextModeration() - if blocked: - with pytest.raises(HTTPException, match="Violated content safety policy") as rejection: - await guard.async_moderation_hook(request_data, user_api_key, "acompletion") - assert rejection.value.status_code == 400 - else: - assert await guard.async_moderation_hook(request_data, user_api_key, "acompletion") is request_data - sdk.LanguageServiceClient.assert_called_once_with() - 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 - ]