From b1b00e4bdc1384a9cc6d2bfb67e461f5356c068c Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 03:50:15 +0000 Subject: [PATCH 01/11] chore(guardrails): cover multimodal + Responses-API content shapes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Several guardrail hooks short-circuit when ``message.content`` is a list or when the request uses the Responses-API ``input`` field instead of ``messages``. Centralise the content-walking logic in a shared helper and update the affected hooks so list-format and Responses-API payloads no longer skip inspection. Also: Aim's ``async_post_call_success_hook`` now inspects every choice (via ``asyncio.gather``) instead of only ``choices[0]`` — the prior behaviour let ``n>1`` callers hide content in subsequent completions. Hooks updated to use the new helper: - aim, lakera_ai_v2, lasso (post a synthesised messages list to a remote guardrail service) - azure_content_safety, ibm_detector, banned_keywords, openai_moderation, google_text_moderation (iterate text fragments locally) - secret_detection (walk-and-rewrite to redact in place) Drive-by fix: the legacy ``data["prompt"]`` list-handling path in secret_detection rebound the loop variable instead of mutating the list, leaving secrets unredacted on text-completion calls; corrected to index back into the list. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../enterprise_hooks/banned_keywords.py | 9 +- .../google_text_moderation.py | 9 +- .../enterprise_hooks/openai_moderation.py | 8 +- .../enterprise_callbacks/secret_detection.py | 64 +-- litellm/proxy/guardrails/_content_utils.py | 160 +++++++ .../guardrails/guardrail_hooks/aim/aim.py | 36 +- .../ibm_guardrails/ibm_detector.py | 173 +++---- .../guardrail_hooks/lakera_ai_v2.py | 18 +- .../guardrails/guardrail_hooks/lasso/lasso.py | 4 +- litellm/proxy/hooks/azure_content_safety.py | 9 +- .../proxy/guardrails/test_content_utils.py | 201 ++++++++ .../guardrails/test_guardrail_coverage.py | 437 ++++++++++++++++++ 12 files changed, 948 insertions(+), 180 deletions(-) create mode 100644 litellm/proxy/guardrails/_content_utils.py create mode 100644 tests/test_litellm/proxy/guardrails/test_content_utils.py create mode 100644 tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py diff --git a/enterprise/enterprise_hooks/banned_keywords.py b/enterprise/enterprise_hooks/banned_keywords.py index 4df138939a..aafbb8cf40 100644 --- a/enterprise/enterprise_hooks/banned_keywords.py +++ b/enterprise/enterprise_hooks/banned_keywords.py @@ -11,6 +11,7 @@ from typing import Literal import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails._content_utils import iter_user_text from litellm.integrations.custom_logger import CustomLogger from litellm._logging import verbose_proxy_logger from fastapi import HTTPException @@ -73,10 +74,10 @@ class _ENTERPRISE_BannedKeywords(CustomLogger): - check if user id part of blocked list """ self.print_verbose("Inside Banned Keyword List Pre-Call Hook") - if call_type == "completion" and "messages" in data: - for m in data["messages"]: - if "content" in m and isinstance(m["content"], str): - self.test_violation(test_str=m["content"]) + if call_type == "completion": + # Covers multimodal list content + Responses-API input. + for text in iter_user_text(data): + self.test_violation(test_str=text) except HTTPException as e: raise e diff --git a/enterprise/enterprise_hooks/google_text_moderation.py b/enterprise/enterprise_hooks/google_text_moderation.py index 1f26d52adf..75829c41ca 100644 --- a/enterprise/enterprise_hooks/google_text_moderation.py +++ b/enterprise/enterprise_hooks/google_text_moderation.py @@ -12,6 +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_user_text from litellm.types.utils import CallTypesLiteral @@ -94,11 +95,9 @@ class _ENTERPRISE_GoogleTextModeration(CustomLogger): - Calls Google's Text Moderation API - Rejects request if it fails safety check """ - if "messages" in data and isinstance(data["messages"], list): - text = "" - for m in data["messages"]: # assume messages is a list - if "content" in m and isinstance(m["content"], str): - text += m["content"] + # Covers multimodal list content + Responses-API input. + text = "".join(iter_user_text(data)) + if text: document = self.language_document(content=text, type_=self.document_type) request = self.moderate_text_request( diff --git a/enterprise/enterprise_hooks/openai_moderation.py b/enterprise/enterprise_hooks/openai_moderation.py index a1db9818e5..a98d5b30d7 100644 --- a/enterprise/enterprise_hooks/openai_moderation.py +++ b/enterprise/enterprise_hooks/openai_moderation.py @@ -19,6 +19,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_user_text from litellm.types.utils import CallTypesLiteral @@ -37,11 +38,8 @@ class _ENTERPRISE_OpenAI_Moderation(CustomLogger): user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral, ): - text = "" - if "messages" in data and isinstance(data["messages"], list): - for m in data["messages"]: # assume messages is a list - if "content" in m and isinstance(m["content"], str): - text += m["content"] + # Covers multimodal list content + Responses-API input. + text = "".join(iter_user_text(data)) from litellm.proxy.proxy_server import llm_router diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index 8a7a82df68..f441ce71ab 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -18,6 +18,7 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails._content_utils import walk_user_text GUARDRAIL_NAME = "hide_secrets" @@ -473,23 +474,19 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): if await self.should_run_check(user_api_key_dict) is False: return - if "messages" in data and isinstance(data["messages"], list): - for message in data["messages"]: - if "content" in message and isinstance(message["content"], str): - detected_secrets = self.scan_message_for_secrets(message["content"]) + # Covers multimodal list content + Responses-API input. + def _redact_message_text(text: str) -> str: + detected_secrets = self.scan_message_for_secrets(text) + for secret in detected_secrets: + text = text.replace(secret["value"], "[REDACTED]") + if detected_secrets: + secret_types = [secret["type"] for secret in detected_secrets] + verbose_proxy_logger.warning( + f"Detected and redacted secrets in message: {secret_types}" + ) + return text - for secret in detected_secrets: - message["content"] = message["content"].replace( - secret["value"], "[REDACTED]" - ) - - if len(detected_secrets) > 0: - secret_types = [secret["type"] for secret in detected_secrets] - verbose_proxy_logger.warning( - f"Detected and redacted secrets in message: {secret_types}" - ) - else: - verbose_proxy_logger.debug("No secrets detected on input.") + walk_user_text(data, _redact_message_text) if "prompt" in data: if isinstance(data["prompt"], str): @@ -504,11 +501,15 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): f"Detected and redacted secrets in prompt: {secret_types}" ) elif isinstance(data["prompt"], list): - for item in data["prompt"]: + # Index back into the list — assigning to ``item`` would only + # rebind the loop variable and leave ``data["prompt"]`` + # carrying the unredacted secret. + for idx, item in enumerate(data["prompt"]): if isinstance(item, str): detected_secrets = self.scan_message_for_secrets(item) for secret in detected_secrets: item = item.replace(secret["value"], "[REDACTED]") + data["prompt"][idx] = item if len(detected_secrets) > 0: secret_types = [ secret["type"] for secret in detected_secrets @@ -517,31 +518,6 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): f"Detected and redacted secrets in prompt: {secret_types}" ) - if "input" in data: - if isinstance(data["input"], str): - detected_secrets = self.scan_message_for_secrets(data["input"]) - for secret in detected_secrets: - data["input"] = data["input"].replace(secret["value"], "[REDACTED]") - if len(detected_secrets) > 0: - secret_types = [secret["type"] for secret in detected_secrets] - verbose_proxy_logger.warning( - f"Detected and redacted secrets in input: {secret_types}" - ) - elif isinstance(data["input"], list): - _input_in_request = data["input"] - for idx, item in enumerate(_input_in_request): - if isinstance(item, str): - detected_secrets = self.scan_message_for_secrets(item) - for secret in detected_secrets: - _input_in_request[idx] = item.replace( - secret["value"], "[REDACTED]" - ) - if len(detected_secrets) > 0: - secret_types = [ - secret["type"] for secret in detected_secrets - ] - verbose_proxy_logger.warning( - f"Detected and redacted secrets in input: {secret_types}" - ) - verbose_proxy_logger.debug("Data after redacting input %s", data) + # ``data["input"]`` (Responses API and embeddings/moderation) is + # already covered by ``walk_user_text`` above. return diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py new file mode 100644 index 0000000000..97a022768e --- /dev/null +++ b/litellm/proxy/guardrails/_content_utils.py @@ -0,0 +1,160 @@ +""" +Shared helpers for guardrail hooks: extract user-supplied text from a +request body regardless of whether it uses Chat Completions ``messages``, +Responses-API ``input``, or multimodal list-format ``content`` parts. + +Hooks that only check ``data["messages"]`` for string content silently +skip the other shapes — these helpers normalise that so every hook sees +every text fragment. +""" + +from typing import Any, Callable, Dict, Iterator, List + + +def _iter_text_parts_in_content(content: Any) -> Iterator[str]: + """Yield text fragments from a ``message.content`` value (string or + multimodal list). Non-text parts (images, audio, …) are skipped.""" + if isinstance(content, str): + if content: + yield content + elif isinstance(content, list): + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") == "text": + text = part.get("text") + if isinstance(text, str) and text: + yield text + + +def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]: + """Coerce a Responses-API ``data["input"]`` value into chat-style messages.""" + if isinstance(input_value, str): + return [{"role": "user", "content": input_value}] + if isinstance(input_value, list): + if input_value and all( + isinstance(item, dict) and "role" in item for item in input_value + ): + return list(input_value) + if input_value and all(isinstance(item, str) for item in input_value): + return [{"role": "user", "content": item} for item in input_value] + return [{"role": "user", "content": input_value}] + return [] + + +def _resolve_messages(data: Dict[str, Any]) -> List[Dict[str, Any]]: + """Return the messages list to inspect, falling back to ``input``.""" + messages = data.get("messages") + if isinstance(messages, list) and messages: + return messages + return _coerce_input_to_messages(data.get("input")) + + +def _iter_inspection_messages(data: Dict[str, Any]) -> Iterator[Dict[str, Any]]: + """Yield every message-like dict, walking ``messages`` AND ``input``.""" + messages = data.get("messages") + if isinstance(messages, list): + yield from messages + yield from _coerce_input_to_messages(data.get("input")) + + +def iter_user_text(data: Dict[str, Any]) -> Iterator[str]: + """Yield every user-supplied text fragment from ``messages`` and ``input``.""" + for message in _iter_inspection_messages(data): + if not isinstance(message, dict): + continue + yield from _iter_text_parts_in_content(message.get("content")) + + +def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: + """Rewrite every user-supplied text fragment in place via ``visit``. + + Mutates ``data["messages"]`` and ``data["input"]``. Returns the number + of fragments visited so callers can short-circuit when nothing was + inspected. + """ + visited = 0 + + def _rewrite_content(content: Any) -> Any: + nonlocal visited + if isinstance(content, str): + if content: + visited += 1 + return visit(content) + return content + if isinstance(content, list): + new_parts = [] + for part in content: + if ( + isinstance(part, dict) + and part.get("type") == "text" + and isinstance(part.get("text"), str) + and part["text"] + ): + visited += 1 + new_parts.append({**part, "text": visit(part["text"])}) + else: + new_parts.append(part) + return new_parts + return content + + messages = data.get("messages") + if isinstance(messages, list): + for message in messages: + if isinstance(message, dict) and "content" in message: + message["content"] = _rewrite_content(message["content"]) + + input_value = data.get("input") + if isinstance(input_value, str): + if input_value: + visited += 1 + data["input"] = visit(input_value) + return visited + if isinstance(input_value, list): + # List of full messages: rewrite each message's content. + if input_value and all( + isinstance(item, dict) and "role" in item for item in input_value + ): + for item in input_value: + if "content" in item: + item["content"] = _rewrite_content(item["content"]) + return visited + # List of content parts or strings: rewrite in place. + for idx, item in enumerate(input_value): + if isinstance(item, str) and item: + visited += 1 + input_value[idx] = visit(item) + elif ( + isinstance(item, dict) + and item.get("type") == "text" + and isinstance(item.get("text"), str) + and item["text"] + ): + visited += 1 + input_value[idx] = {**item, "text": visit(item["text"])} + return visited + + return visited + + +def build_inspection_messages(data: Dict[str, Any]) -> List[Dict[str, str]]: + """Synthesize a chat-style messages list for posting to a guardrail API. + + Each returned message has a plain-string ``content`` — multimodal text + parts are joined with newlines and Responses-API ``input`` is lifted + into synthetic messages. Messages with no inspectable text are dropped. + + Hooks that POST ``{"messages": [...]}`` to an external service should + call this instead of ``data.get("messages", [])`` so the Responses API + and multimodal content are covered. + """ + flattened: List[Dict[str, str]] = [] + for message in _iter_inspection_messages(data): + if not isinstance(message, dict): + continue + text = "\n".join(_iter_text_parts_in_content(message.get("content"))) + if not text: + continue + role = message.get("role", "user") or "user" + flattened.append({"role": role, "content": text}) + return flattened diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 1ae87e99c9..a386617428 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -22,6 +22,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails._content_utils import build_inspection_messages from litellm.types.utils import ( CallTypesLiteral, Choices, @@ -101,10 +102,11 @@ class AimGuardrail(CustomGuardrail): user_email=user_email, litellm_call_id=call_id, ) + # Covers multimodal list content + Responses-API input. response = await self.async_handler.post( f"{self.api_base}/fw/v1/analyze", headers=headers, - json={"messages": data.get("messages", [])}, + json={"messages": build_inspection_messages(data)}, ) response.raise_for_status() res = response.json() @@ -162,7 +164,7 @@ class AimGuardrail(CustomGuardrail): litellm_call_id=call_id, ), json={ - "messages": request_data.get("messages", []) + "messages": build_inspection_messages(request_data) + [{"role": "assistant", "content": output}] }, ) @@ -233,15 +235,27 @@ class AimGuardrail(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, response: Union[Any, ModelResponse, EmbeddingResponse, ImageResponse], ) -> Any: - if ( - isinstance(response, ModelResponse) - and response.choices - and isinstance(response.choices[0], Choices) - ): - content = response.choices[0].message.content or "" - aim_output_guardrail_result = await self.call_aim_guardrail_on_output( - data, content, hook="output", key_alias=user_api_key_dict.key_alias + if not (isinstance(response, ModelResponse) and response.choices): + return response + # Inspect every choice — when ``n>1`` the additional completions + # used to bypass Aim entirely because the hook only inspected + # ``choices[0]``. Run inspections concurrently so multi-completion + # responses don't pay an n× latency penalty. + choices_to_inspect = [c for c in response.choices if isinstance(c, Choices)] + if not choices_to_inspect: + return response + results = await asyncio.gather( + *( + self.call_aim_guardrail_on_output( + data, + choice.message.content or "", + hook="output", + key_alias=user_api_key_dict.key_alias, + ) + for choice in choices_to_inspect ) + ) + for choice, aim_output_guardrail_result in zip(choices_to_inspect, results): if aim_output_guardrail_result and aim_output_guardrail_result.get( "detection_message" ): @@ -252,7 +266,7 @@ class AimGuardrail(CustomGuardrail): if aim_output_guardrail_result and aim_output_guardrail_result.get( "redacted_output" ): - response.choices[0].message.content = aim_output_guardrail_result.get( + choice.message.content = aim_output_guardrail_result.get( "redacted_output" ) return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 2fc0521364..435bb383df 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails._content_utils import iter_user_text from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( IBMDetectorDetection, @@ -463,65 +464,53 @@ class IBMGuardrailDetector(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data - _messages = data.get("messages") - if _messages: - contents_to_check: List[str] = [] - for message in _messages: - _content = message.get("content") - if isinstance(_content, str): - contents_to_check.append(_content) + # Covers multimodal list content + Responses-API input. + contents_to_check: List[str] = list(iter_user_text(data)) + if contents_to_check: + if self.is_detector_server: + # Call detector server with all contents at once + result = await self._call_detector_server( + contents=contents_to_check, + request_data=data, + event_type=GuardrailEventHooks.pre_call, + ) - if contents_to_check: - if self.is_detector_server: - # Call detector server with all contents at once - result = await self._call_detector_server( - contents=contents_to_check, + verbose_proxy_logger.debug( + "IBM Detector Server async_pre_call_hook result: %s", result + ) + + # Check if any detections were found + has_violations = False + for message_detections in result: + filtered = self._filter_detections_by_threshold(message_detections) + if filtered: + has_violations = True + break + + if has_violations and self.block_on_detection: + error_message = self._create_error_message_detector_server(result) + raise ValueError(error_message) + + else: + # Call orchestrator for each content separately + for content in contents_to_check: + orchestrator_result = await self._call_orchestrator( + content=content, request_data=data, event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( - "IBM Detector Server async_pre_call_hook result: %s", result + "IBM Orchestrator async_pre_call_hook result: %s", + orchestrator_result, ) - # Check if any detections were found - has_violations = False - for message_detections in result: - filtered = self._filter_detections_by_threshold( - message_detections - ) - if filtered: - has_violations = True - break - - if has_violations and self.block_on_detection: - error_message = self._create_error_message_detector_server( - result - ) - raise ValueError(error_message) - - else: - # Call orchestrator for each content separately - for content in contents_to_check: - orchestrator_result = await self._call_orchestrator( - content=content, - request_data=data, - event_type=GuardrailEventHooks.pre_call, - ) - - verbose_proxy_logger.debug( - "IBM Orchestrator async_pre_call_hook result: %s", - orchestrator_result, - ) - - filtered = self._filter_detections_by_threshold( + filtered = self._filter_detections_by_threshold(orchestrator_result) + if filtered and self.block_on_detection: + error_message = self._create_error_message_orchestrator( orchestrator_result ) - if filtered and self.block_on_detection: - error_message = self._create_error_message_orchestrator( - orchestrator_result - ) - raise ValueError(error_message) + raise ValueError(error_message) # Add guardrail to applied guardrails header add_guardrail_to_applied_guardrails_header( @@ -550,65 +539,53 @@ class IBMGuardrailDetector(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return - _messages = data.get("messages") - if _messages: - contents_to_check: List[str] = [] - for message in _messages: - _content = message.get("content") - if isinstance(_content, str): - contents_to_check.append(_content) + # Covers multimodal list content + Responses-API input. + contents_to_check: List[str] = list(iter_user_text(data)) + if contents_to_check: + if self.is_detector_server: + # Call detector server with all contents at once + result = await self._call_detector_server( + contents=contents_to_check, + request_data=data, + event_type=GuardrailEventHooks.during_call, + ) - if contents_to_check: - if self.is_detector_server: - # Call detector server with all contents at once - result = await self._call_detector_server( - contents=contents_to_check, + verbose_proxy_logger.debug( + "IBM Detector Server async_moderation_hook result: %s", result + ) + + # Check if any detections were found + has_violations = False + for message_detections in result: + filtered = self._filter_detections_by_threshold(message_detections) + if filtered: + has_violations = True + break + + if has_violations and self.block_on_detection: + error_message = self._create_error_message_detector_server(result) + raise ValueError(error_message) + + else: + # Call orchestrator for each content separately + for content in contents_to_check: + orchestrator_result = await self._call_orchestrator( + content=content, request_data=data, event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( - "IBM Detector Server async_moderation_hook result: %s", result + "IBM Orchestrator async_moderation_hook result: %s", + orchestrator_result, ) - # Check if any detections were found - has_violations = False - for message_detections in result: - filtered = self._filter_detections_by_threshold( - message_detections - ) - if filtered: - has_violations = True - break - - if has_violations and self.block_on_detection: - error_message = self._create_error_message_detector_server( - result - ) - raise ValueError(error_message) - - else: - # Call orchestrator for each content separately - for content in contents_to_check: - orchestrator_result = await self._call_orchestrator( - content=content, - request_data=data, - event_type=GuardrailEventHooks.during_call, - ) - - verbose_proxy_logger.debug( - "IBM Orchestrator async_moderation_hook result: %s", - orchestrator_result, - ) - - filtered = self._filter_detections_by_threshold( + filtered = self._filter_detections_by_threshold(orchestrator_result) + if filtered and self.block_on_detection: + error_message = self._create_error_message_orchestrator( orchestrator_result ) - if filtered and self.block_on_detection: - error_message = self._create_error_message_orchestrator( - orchestrator_result - ) - raise ValueError(error_message) + raise ValueError(error_message) # Add guardrail to applied guardrails header add_guardrail_to_applied_guardrails_header( diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 7aa26435ba..167ad4f790 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -13,6 +13,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails._content_utils import build_inspection_messages from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import AllMessageValues @@ -214,10 +215,11 @@ class LakeraAIGuardrail(CustomGuardrail): ) return data - new_messages: Optional[List[AllMessageValues]] = data.get("messages") - if new_messages is None: + # Covers multimodal list content + Responses-API input. + new_messages = build_inspection_messages(data) + if not new_messages: verbose_proxy_logger.warning( - "Lakera AI: not running guardrail. No messages in data" + "Lakera AI: not running guardrail. No inspectable text in data" ) return data @@ -225,7 +227,7 @@ class LakeraAIGuardrail(CustomGuardrail): ########## 1. Make the Lakera AI v2 guard API request ########## ######################################################### lakera_guardrail_response, masked_entity_count = await self.call_v2_guard( - messages=new_messages, + messages=new_messages, # type: ignore[arg-type] request_data=data, event_type=GuardrailEventHooks.pre_call, ) @@ -280,10 +282,10 @@ class LakeraAIGuardrail(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return - new_messages: Optional[List[AllMessageValues]] = data.get("messages") - if new_messages is None: + new_messages = build_inspection_messages(data) + if not new_messages: verbose_proxy_logger.warning( - "Lakera AI: not running guardrail. No messages in data" + "Lakera AI: not running guardrail. No inspectable text in data" ) return @@ -291,7 +293,7 @@ class LakeraAIGuardrail(CustomGuardrail): ########## 1. Make the Lakera AI v2 guard API request ########## ######################################################### lakera_guardrail_response, masked_entity_count = await self.call_v2_guard( - messages=new_messages, + messages=new_messages, # type: ignore[arg-type] request_data=data, event_type=GuardrailEventHooks.during_call, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index a72f3e4c3f..a7719124b2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -50,6 +50,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails._content_utils import build_inspection_messages from litellm.types.guardrails import GuardrailEventHooks import litellm @@ -366,7 +367,8 @@ class LassoGuardrail(CustomGuardrail): LassoGuardrailAPIError: If the Lasso API call fails HTTPException: If blocking violations are detected """ - messages: List[Dict[str, str]] = data.get("messages", []) + # Covers multimodal list content + Responses-API input. + messages: List[Dict[str, str]] = build_inspection_messages(data) if not messages: return data diff --git a/litellm/proxy/hooks/azure_content_safety.py b/litellm/proxy/hooks/azure_content_safety.py index b35d671117..e66ada4208 100644 --- a/litellm/proxy/hooks/azure_content_safety.py +++ b/litellm/proxy/hooks/azure_content_safety.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails._content_utils import iter_user_text class _PROXY_AzureContentSafety( @@ -118,10 +119,10 @@ class _PROXY_AzureContentSafety( ): verbose_proxy_logger.debug("Inside Azure Content-Safety Pre-Call Hook") try: - if call_type == "completion" and "messages" in data: - for m in data["messages"]: - if "content" in m and isinstance(m["content"], str): - await self.test_violation(content=m["content"], source="input") + if call_type == "completion": + # Covers multimodal list content + Responses-API input. + for text in iter_user_text(data): + await self.test_violation(content=text, source="input") except HTTPException as e: raise e diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py new file mode 100644 index 0000000000..2491363330 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -0,0 +1,201 @@ +"""Tests for the shared guardrail content extraction helpers.""" + +from litellm.proxy.guardrails._content_utils import ( + build_inspection_messages, + iter_user_text, + walk_user_text, +) + + +# ── iter_user_text ──────────────────────────────────────────────────────────── + + +def test_iter_user_text_string_messages(): + data = { + "messages": [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + ] + } + assert list(iter_user_text(data)) == ["hello", "hi"] + + +def test_iter_user_text_multimodal_list_content(): + """VERIA-11: list-format content must be inspected, not silently skipped.""" + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "AWS_KEY=AKIA..."}, + {"type": "image_url", "image_url": {"url": "..."}}, + {"type": "text", "text": "more secrets"}, + ], + } + ] + } + assert list(iter_user_text(data)) == ["AWS_KEY=AKIA...", "more secrets"] + + +def test_iter_user_text_responses_api_string_input(): + """fniVO9-F: Responses-API ``input`` must be inspectable when ``messages`` absent.""" + data = {"input": "tell me a secret"} + assert list(iter_user_text(data)) == ["tell me a secret"] + + +def test_iter_user_text_responses_api_list_input_messages(): + data = { + "input": [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ] + } + assert list(iter_user_text(data)) == ["first", "second"] + + +def test_iter_user_text_responses_api_list_input_content_parts(): + data = { + "input": [ + {"type": "text", "text": "alpha"}, + {"type": "image_url", "image_url": {"url": "..."}}, + {"type": "text", "text": "beta"}, + ] + } + assert list(iter_user_text(data)) == ["alpha", "beta"] + + +def test_iter_user_text_walks_messages_and_input_independently(): + """When both are present (rare), every fragment from either field is + inspected — a stricter guarantee than "first one wins".""" + data = { + "messages": [{"role": "user", "content": "msg-content"}], + "input": "input-content", + } + assert list(iter_user_text(data)) == ["msg-content", "input-content"] + + +def test_iter_user_text_empty_data(): + assert list(iter_user_text({})) == [] + assert list(iter_user_text({"messages": []})) == [] + assert list(iter_user_text({"input": ""})) == [] + + +# ── walk_user_text ──────────────────────────────────────────────────────────── + + +def test_walk_user_text_redacts_string_messages_in_place(): + data = { + "messages": [ + {"role": "user", "content": "leak: AKIAEXAMPLE"}, + {"role": "assistant", "content": "ok"}, + ] + } + visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]")) + assert visited == 2 + assert data["messages"][0]["content"] == "leak: [REDACTED]" + assert data["messages"][1]["content"] == "ok" + + +def test_walk_user_text_redacts_multimodal_text_parts(): + """VERIA-11: list-content text parts must be mutable for in-place redaction.""" + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "AKIAEXAMPLE here"}, + {"type": "image_url", "image_url": {"url": "..."}}, + {"type": "text", "text": "no secret"}, + ], + } + ] + } + visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]")) + assert visited == 2 + parts = data["messages"][0]["content"] + assert parts[0] == {"type": "text", "text": "[REDACTED] here"} + # Non-text part must be left untouched. + assert parts[1] == {"type": "image_url", "image_url": {"url": "..."}} + assert parts[2] == {"type": "text", "text": "no secret"} + + +def test_walk_user_text_redacts_responses_api_string_input(): + data = {"input": "leak AKIAEXAMPLE"} + visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]")) + assert visited == 1 + assert data["input"] == "leak [REDACTED]" + + +def test_walk_user_text_redacts_responses_api_list_input(): + data = { + "input": [ + {"type": "text", "text": "AKIAEXAMPLE"}, + {"type": "image_url", "image_url": {"url": "..."}}, + ] + } + visited = walk_user_text(data, lambda s: f"[redacted]{s}[/]") + assert visited == 1 + assert data["input"][0] == {"type": "text", "text": "[redacted]AKIAEXAMPLE[/]"} + assert data["input"][1] == {"type": "image_url", "image_url": {"url": "..."}} + + +# ── build_inspection_messages ───────────────────────────────────────────────── + + +def test_build_inspection_messages_chat_completion_passthrough(): + data = { + "messages": [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hi"}, + ] + } + assert build_inspection_messages(data) == [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hi"}, + ] + + +def test_build_inspection_messages_joins_multimodal_text_parts(): + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "first part"}, + {"type": "image_url", "image_url": {"url": "..."}}, + {"type": "text", "text": "second part"}, + ], + } + ] + } + assert build_inspection_messages(data) == [ + {"role": "user", "content": "first part\nsecond part"} + ] + + +def test_build_inspection_messages_lifts_responses_api_input(): + """fniVO9-F: ``input`` must be visible to hooks that POST messages to a remote API.""" + data = {"input": "responses-api content"} + assert build_inspection_messages(data) == [ + {"role": "user", "content": "responses-api content"} + ] + + +def test_build_inspection_messages_drops_messages_with_no_text(): + data = { + "messages": [ + {"role": "user", "content": ""}, + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "..."}}], + }, + {"role": "user", "content": "kept"}, + ] + } + assert build_inspection_messages(data) == [{"role": "user", "content": "kept"}] + + +def test_build_inspection_messages_empty_data(): + assert build_inspection_messages({}) == [] + assert build_inspection_messages({"messages": []}) == [] + assert build_inspection_messages({"input": ""}) == [] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py new file mode 100644 index 0000000000..4cfcbf3c6d --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -0,0 +1,437 @@ +""" +Regression tests for guardrail-coverage gaps. + +Each test confirms that a previously-bypassable input shape now triggers +inspection by the relevant guardrail hook: + +- VERIA-11: multimodal list-format ``content`` is inspected (no longer + silently skipped because of an ``isinstance(content, str)`` check). +- fniVO9-F: Responses-API ``data["input"]`` is inspected (no longer + silently skipped because the hook only looked at ``data["messages"]``). +- yVS0wMDO: Aim's post-call hook inspects every choice when ``n>1``, + not just ``choices[0]``. +""" + +from typing import Any, Dict +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from httpx import Request, Response + +from litellm import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import Choices, Message, ModelResponse + + +@pytest.fixture +def user_api_key(): + return UserAPIKeyAuth(api_key="hashed", user_id="u", key_alias=None) + + +# ── Aim ─────────────────────────────────────────────────────────────────────── + + +def _aim_no_action_response() -> Response: + return Response( + status_code=200, + json={"required_action": None}, + request=Request("POST", "https://api.aim.security/fw/v1/analyze"), + ) + + +@pytest.mark.asyncio +async def test_aim_inspects_multimodal_list_content(user_api_key, monkeypatch): + monkeypatch.setenv("AIM_API_KEY", "hs-aim-key") + from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail + + guard = AimGuardrail() + sent_payload: Dict[str, Any] = {} + + async def capture(url, headers, json): + sent_payload.update(json) + return _aim_no_action_response() + + with patch.object(guard.async_handler, "post", side_effect=capture): + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data={ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "secret payload"}, + {"type": "image_url", "image_url": {"url": "..."}}, + ], + } + ] + }, + call_type="acompletion", + ) + + # The multimodal text part must be visible to Aim. + assert sent_payload["messages"] == [{"role": "user", "content": "secret payload"}] + + +@pytest.mark.asyncio +async def test_aim_inspects_responses_api_input(user_api_key, monkeypatch): + monkeypatch.setenv("AIM_API_KEY", "hs-aim-key") + from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail + + guard = AimGuardrail() + sent_payload: Dict[str, Any] = {} + + async def capture(url, headers, json): + sent_payload.update(json) + return _aim_no_action_response() + + with patch.object(guard.async_handler, "post", side_effect=capture): + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data={"input": "responses-api content"}, + call_type="acompletion", + ) + + assert sent_payload["messages"] == [ + {"role": "user", "content": "responses-api content"} + ] + + +@pytest.mark.asyncio +async def test_aim_post_call_inspects_all_choices(user_api_key, monkeypatch): + """yVS0wMDO: ``n>1`` no longer bypasses Aim by hiding violations in + ``choices[1+]``.""" + monkeypatch.setenv("AIM_API_KEY", "hs-aim-key") + from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail + + guard = AimGuardrail() + inspected_outputs = [] + + async def capture(request_data, output, hook, key_alias): + inspected_outputs.append(output) + return {"redacted_output": output} + + response = ModelResponse( + choices=[ + Choices(index=0, message=Message(role="assistant", content="first")), + Choices(index=1, message=Message(role="assistant", content="second")), + Choices(index=2, message=Message(role="assistant", content="third")), + ] + ) + + with patch.object(guard, "call_aim_guardrail_on_output", side_effect=capture): + await guard.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hi"}]}, + user_api_key_dict=user_api_key, + response=response, + ) + + # ``asyncio.gather`` is used for parallelism, so order of inspection is + # not guaranteed. + assert sorted(inspected_outputs) == ["first", "second", "third"] + + +# ── Lakera v2 ───────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_lakera_v2_inspects_responses_api_input(user_api_key, monkeypatch): + monkeypatch.setenv("LAKERA_API_KEY", "lk-test") + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( + LakeraAIGuardrail, + ) + + guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="monitor") + + seen_messages = [] + + async def fake_call_v2_guard(messages, request_data, event_type): + seen_messages.append(messages) + return {"flagged": False}, {} + + with patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard): + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data={"input": "responses-api content"}, + call_type="responses", + ) + + assert seen_messages == [[{"role": "user", "content": "responses-api content"}]] + + +@pytest.mark.asyncio +async def test_lakera_v2_inspects_multimodal_list_content(user_api_key, monkeypatch): + monkeypatch.setenv("LAKERA_API_KEY", "lk-test") + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( + LakeraAIGuardrail, + ) + + guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="monitor") + seen_messages = [] + + async def fake_call_v2_guard(messages, request_data, event_type): + seen_messages.append(messages) + return {"flagged": False}, {} + + with patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard): + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data={ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "AKIAEXAMPLE"}, + {"type": "image_url", "image_url": {"url": "..."}}, + ], + } + ] + }, + call_type="acompletion", + ) + + assert seen_messages == [[{"role": "user", "content": "AKIAEXAMPLE"}]] + + +# ── Lasso ───────────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_lasso_inspects_responses_api_input(user_api_key, monkeypatch): + monkeypatch.setenv("LASSO_API_KEY", "ls-test") + from litellm.proxy.guardrails.guardrail_hooks.lasso.lasso import LassoGuardrail + + guard = LassoGuardrail(lasso_api_key="ls-test") + + seen_messages = [] + + async def fake_handle_classification(data, cache, message_type, messages): + seen_messages.append(messages) + return data + + with patch.object( + guard, "_handle_classification", side_effect=fake_handle_classification + ): + await guard._run_lasso_guardrail( + data={"input": "responses-api content"}, + cache=DualCache(), + message_type="PROMPT", + ) + + assert seen_messages == [[{"role": "user", "content": "responses-api content"}]] + + +# ── Banned Keywords ─────────────────────────────────────────────────────────── + + +def test_banned_keywords_blocks_multimodal_content(monkeypatch): + """VERIA-11: a banned word hidden in a multimodal text part is now caught.""" + monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False) + from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords + from fastapi import HTTPException + + guard = _ENTERPRISE_BannedKeywords() + + async def _run(): + await guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"), + cache=DualCache(), + data={ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "forbidden word here"}, + {"type": "image_url", "image_url": {"url": "..."}}, + ], + } + ] + }, + call_type="completion", + ) + + import asyncio + + with pytest.raises(HTTPException) as exc: + asyncio.run(_run()) + assert "forbidden" in str(exc.value.detail).lower() + + +def test_banned_keywords_blocks_responses_api_input(monkeypatch): + monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False) + from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords + from fastapi import HTTPException + + guard = _ENTERPRISE_BannedKeywords() + + async def _run(): + await guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"), + cache=DualCache(), + data={"input": "this contains forbidden content"}, + call_type="completion", + ) + + import asyncio + + with pytest.raises(HTTPException): + asyncio.run(_run()) + + +# ── Secret Detection ────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_secret_detection_redacts_multimodal_text_parts(user_api_key): + from enterprise.litellm_enterprise.enterprise_callbacks.secret_detection import ( + _ENTERPRISE_SecretDetection, + ) + + guard = _ENTERPRISE_SecretDetection() + data = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "AKIAIOSFODNN7EXAMPLE is the key", + }, + {"type": "image_url", "image_url": {"url": "..."}}, + ], + } + ] + } + + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data=data, + call_type="completion", + ) + + parts = data["messages"][0]["content"] + assert "AKIAIOSFODNN7EXAMPLE" not in parts[0]["text"] + assert "[REDACTED]" in parts[0]["text"] + # Non-text part is preserved untouched. + assert parts[1] == {"type": "image_url", "image_url": {"url": "..."}} + + +@pytest.mark.asyncio +async def test_secret_detection_redacts_responses_api_input(user_api_key): + from enterprise.litellm_enterprise.enterprise_callbacks.secret_detection import ( + _ENTERPRISE_SecretDetection, + ) + + guard = _ENTERPRISE_SecretDetection() + data = {"input": "leak: AKIAIOSFODNN7EXAMPLE"} + + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data=data, + call_type="moderation", + ) + + assert "AKIAIOSFODNN7EXAMPLE" not in data["input"] + assert "[REDACTED]" in data["input"] + + +# ── OpenAI Moderation ───────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_openai_moderation_inspects_multimodal_content(monkeypatch, user_api_key): + """The aggregated text passed to ``llm_router.amoderation`` must include + list-format text parts and Responses-API input — without this, multimodal + content silently passed moderation.""" + from enterprise.enterprise_hooks.openai_moderation import ( + _ENTERPRISE_OpenAI_Moderation, + ) + + guard = _ENTERPRISE_OpenAI_Moderation() + + seen_inputs = [] + + class FakeModeration: + results = [type("R", (), {"flagged": False})()] + + async def fake_amoderation(model, input): + seen_inputs.append(input) + return FakeModeration() + + fake_router = MagicMock() + fake_router.amoderation = AsyncMock(side_effect=fake_amoderation) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", fake_router, raising=False + ) + + await guard.async_moderation_hook( + data={ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "alpha "}, + {"type": "image_url", "image_url": {"url": "..."}}, + {"type": "text", "text": "beta"}, + ], + } + ] + }, + user_api_key_dict=user_api_key, + call_type="acompletion", + ) + + assert seen_inputs == ["alpha beta"] + + +# ── Google Text Moderation ──────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_google_text_moderation_inspects_multimodal_content(user_api_key): + """The text passed to Google's moderation client must include list-format + text parts.""" + from enterprise.enterprise_hooks.google_text_moderation import ( + _ENTERPRISE_GoogleTextModeration, + ) + + guard = _ENTERPRISE_GoogleTextModeration.__new__(_ENTERPRISE_GoogleTextModeration) + seen_documents = [] + + def fake_language_document(content, type_): + seen_documents.append(content) + return MagicMock() + + fake_response = MagicMock() + fake_response.moderation_categories = [] + + guard.language_document = fake_language_document + guard.moderate_text_request = MagicMock(return_value=MagicMock()) + guard.document_type = MagicMock() + guard.client = MagicMock() + guard.client.moderate_text = MagicMock(return_value=fake_response) + + await guard.async_moderation_hook( + data={ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello "}, + {"type": "image_url", "image_url": {"url": "..."}}, + {"type": "text", "text": "world"}, + ], + } + ] + }, + user_api_key_dict=user_api_key, + call_type="acompletion", + ) + + assert seen_documents == ["hello world"] From d7677085712958f5d14e6b0ae5072600663afb4b Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 04:01:27 +0000 Subject: [PATCH 02/11] fix(guardrails): add mypy type-ignore for lakera_v2 mask call CI mypy flagged the two ``_mask_pii_in_messages(messages=new_messages, ...)`` call sites because ``build_inspection_messages`` returns ``List[Dict[str, str]]`` while ``_mask_pii_in_messages`` declares ``List[AllMessageValues]``. The runtime payload matches; widening the return type to a Union of TypedDicts is a larger refactor than is warranted by this hardening change, so add the localised ``# type: ignore`` to match the two we already added on the ``call_v2_guard`` calls above. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 167ad4f790..d3f19966be 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -239,7 +239,7 @@ class LakeraAIGuardrail(CustomGuardrail): # If only PII violations exist, mask the PII if self._is_only_pii_violation(lakera_guardrail_response): data["messages"] = self._mask_pii_in_messages( - messages=new_messages, + messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, masked_entity_count=masked_entity_count, ) @@ -305,7 +305,7 @@ class LakeraAIGuardrail(CustomGuardrail): # If only PII violations exist, mask the PII if self._is_only_pii_violation(lakera_guardrail_response): data["messages"] = self._mask_pii_in_messages( - messages=new_messages, + messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, masked_entity_count=masked_entity_count, ) From 7514bb4740647876216364fba632274d06cf13c0 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 04:10:04 +0000 Subject: [PATCH 03/11] fix(guardrails): close mixed-list gap, drop dead code, rename helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Greptile P2 follow-ups on _content_utils.py: - Drop unreachable ``_resolve_messages``. The new ``_iter_inspection_messages`` walks ``messages`` AND ``input`` independently; leaving the old fallback-only variant around invited a future maintainer to wire it back up and silently narrow coverage. - Rename ``iter_user_text`` → ``iter_message_text``. The helper walks every role (user, assistant, system); the old name implied user-turn content only. Callers and tests updated. - Close mixed-list coverage gap. When ``data["input"]`` was a list mixing content-part dicts and bare strings, ``iter_message_text`` and ``build_inspection_messages`` only saw the dict parts while ``walk_user_text`` already inspected both. ``_iter_text_parts_in_content`` now treats bare strings inside a content list as text fragments, so read and write helpers agree on coverage. Adds two regression tests for the mixed-list shape. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../enterprise_hooks/banned_keywords.py | 4 +- .../google_text_moderation.py | 4 +- .../enterprise_hooks/openai_moderation.py | 4 +- litellm/proxy/guardrails/_content_utils.py | 42 ++++++----- .../ibm_guardrails/ibm_detector.py | 6 +- litellm/proxy/hooks/azure_content_safety.py | 4 +- .../proxy/guardrails/test_content_utils.py | 72 ++++++++++++++----- 7 files changed, 89 insertions(+), 47 deletions(-) diff --git a/enterprise/enterprise_hooks/banned_keywords.py b/enterprise/enterprise_hooks/banned_keywords.py index aafbb8cf40..a86b59daf5 100644 --- a/enterprise/enterprise_hooks/banned_keywords.py +++ b/enterprise/enterprise_hooks/banned_keywords.py @@ -11,7 +11,7 @@ from typing import Literal import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails._content_utils import iter_user_text +from litellm.proxy.guardrails._content_utils import iter_message_text from litellm.integrations.custom_logger import CustomLogger from litellm._logging import verbose_proxy_logger from fastapi import HTTPException @@ -76,7 +76,7 @@ class _ENTERPRISE_BannedKeywords(CustomLogger): self.print_verbose("Inside Banned Keyword List Pre-Call Hook") if call_type == "completion": # Covers multimodal list content + Responses-API input. - for text in iter_user_text(data): + for text in iter_message_text(data): self.test_violation(test_str=text) except HTTPException as e: diff --git a/enterprise/enterprise_hooks/google_text_moderation.py b/enterprise/enterprise_hooks/google_text_moderation.py index 75829c41ca..5b2d71c5cc 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_user_text +from litellm.proxy.guardrails._content_utils import iter_message_text from litellm.types.utils import CallTypesLiteral @@ -96,7 +96,7 @@ class _ENTERPRISE_GoogleTextModeration(CustomLogger): - Rejects request if it fails safety check """ # Covers multimodal list content + Responses-API input. - text = "".join(iter_user_text(data)) + text = "".join(iter_message_text(data)) if text: document = self.language_document(content=text, type_=self.document_type) diff --git a/enterprise/enterprise_hooks/openai_moderation.py b/enterprise/enterprise_hooks/openai_moderation.py index a98d5b30d7..2162370804 100644 --- a/enterprise/enterprise_hooks/openai_moderation.py +++ b/enterprise/enterprise_hooks/openai_moderation.py @@ -19,7 +19,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_user_text +from litellm.proxy.guardrails._content_utils import iter_message_text from litellm.types.utils import CallTypesLiteral @@ -39,7 +39,7 @@ class _ENTERPRISE_OpenAI_Moderation(CustomLogger): call_type: CallTypesLiteral, ): # Covers multimodal list content + Responses-API input. - text = "".join(iter_user_text(data)) + text = "".join(iter_message_text(data)) from litellm.proxy.proxy_server import llm_router diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 97a022768e..576fefe85a 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -1,7 +1,7 @@ """ -Shared helpers for guardrail hooks: extract user-supplied text from a -request body regardless of whether it uses Chat Completions ``messages``, -Responses-API ``input``, or multimodal list-format ``content`` parts. +Shared helpers for guardrail hooks: extract text from a request body +regardless of whether it uses Chat Completions ``messages``, Responses-API +``input``, or multimodal list-format ``content`` parts. Hooks that only check ``data["messages"]`` for string content silently skip the other shapes — these helpers normalise that so every hook sees @@ -19,6 +19,12 @@ def _iter_text_parts_in_content(content: Any) -> Iterator[str]: yield content elif isinstance(content, list): for part in content: + if isinstance(part, str): + # A bare string in a content/input list is itself a text + # fragment (Responses-API mixed-list shape). + if part: + yield part + continue if not isinstance(part, dict): continue if part.get("type") == "text": @@ -36,20 +42,13 @@ def _coerce_input_to_messages(input_value: Any) -> List[Dict[str, Any]]: isinstance(item, dict) and "role" in item for item in input_value ): return list(input_value) - if input_value and all(isinstance(item, str) for item in input_value): - return [{"role": "user", "content": item} for item in input_value] + # Mixed lists (content-part dicts + bare strings) and pure + # string/dict lists all become a single user message; the content + # iterator below handles each element type uniformly. return [{"role": "user", "content": input_value}] return [] -def _resolve_messages(data: Dict[str, Any]) -> List[Dict[str, Any]]: - """Return the messages list to inspect, falling back to ``input``.""" - messages = data.get("messages") - if isinstance(messages, list) and messages: - return messages - return _coerce_input_to_messages(data.get("input")) - - def _iter_inspection_messages(data: Dict[str, Any]) -> Iterator[Dict[str, Any]]: """Yield every message-like dict, walking ``messages`` AND ``input``.""" messages = data.get("messages") @@ -58,8 +57,12 @@ def _iter_inspection_messages(data: Dict[str, Any]) -> Iterator[Dict[str, Any]]: yield from _coerce_input_to_messages(data.get("input")) -def iter_user_text(data: Dict[str, Any]) -> Iterator[str]: - """Yield every user-supplied text fragment from ``messages`` and ``input``.""" +def iter_message_text(data: Dict[str, Any]) -> Iterator[str]: + """Yield every text fragment from ``messages`` AND ``input``. + + Walks every role (user, assistant, system, …) — guardrails inspect + the entire conversation, not just user turns. + """ for message in _iter_inspection_messages(data): if not isinstance(message, dict): continue @@ -67,7 +70,7 @@ def iter_user_text(data: Dict[str, Any]) -> Iterator[str]: def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: - """Rewrite every user-supplied text fragment in place via ``visit``. + """Rewrite every text fragment in place via ``visit``. Mutates ``data["messages"]`` and ``data["input"]``. Returns the number of fragments visited so callers can short-circuit when nothing was @@ -85,7 +88,10 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: if isinstance(content, list): new_parts = [] for part in content: - if ( + if isinstance(part, str) and part: + visited += 1 + new_parts.append(visit(part)) + elif ( isinstance(part, dict) and part.get("type") == "text" and isinstance(part.get("text"), str) @@ -119,7 +125,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: if "content" in item: item["content"] = _rewrite_content(item["content"]) return visited - # List of content parts or strings: rewrite in place. + # List of content parts and/or bare strings: rewrite in place. for idx, item in enumerate(input_value): if isinstance(item, str) and item: visited += 1 diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 435bb383df..36d70a37c7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -20,7 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails._content_utils import iter_user_text +from litellm.proxy.guardrails._content_utils import iter_message_text from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( IBMDetectorDetection, @@ -465,7 +465,7 @@ class IBMGuardrailDetector(CustomGuardrail): return data # Covers multimodal list content + Responses-API input. - contents_to_check: List[str] = list(iter_user_text(data)) + contents_to_check: List[str] = list(iter_message_text(data)) if contents_to_check: if self.is_detector_server: # Call detector server with all contents at once @@ -540,7 +540,7 @@ class IBMGuardrailDetector(CustomGuardrail): return # Covers multimodal list content + Responses-API input. - contents_to_check: List[str] = list(iter_user_text(data)) + contents_to_check: List[str] = list(iter_message_text(data)) if contents_to_check: if self.is_detector_server: # Call detector server with all contents at once diff --git a/litellm/proxy/hooks/azure_content_safety.py b/litellm/proxy/hooks/azure_content_safety.py index e66ada4208..a21f0566cc 100644 --- a/litellm/proxy/hooks/azure_content_safety.py +++ b/litellm/proxy/hooks/azure_content_safety.py @@ -8,7 +8,7 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails._content_utils import iter_user_text +from litellm.proxy.guardrails._content_utils import iter_message_text class _PROXY_AzureContentSafety( @@ -121,7 +121,7 @@ class _PROXY_AzureContentSafety( try: if call_type == "completion": # Covers multimodal list content + Responses-API input. - for text in iter_user_text(data): + for text in iter_message_text(data): await self.test_violation(content=text, source="input") except HTTPException as e: diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 2491363330..e26dc7aa2d 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -2,25 +2,25 @@ from litellm.proxy.guardrails._content_utils import ( build_inspection_messages, - iter_user_text, + iter_message_text, walk_user_text, ) -# ── iter_user_text ──────────────────────────────────────────────────────────── +# ── iter_message_text ──────────────────────────────────────────────────────────── -def test_iter_user_text_string_messages(): +def test_iter_message_text_string_messages(): data = { "messages": [ {"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}, ] } - assert list(iter_user_text(data)) == ["hello", "hi"] + assert list(iter_message_text(data)) == ["hello", "hi"] -def test_iter_user_text_multimodal_list_content(): +def test_iter_message_text_multimodal_list_content(): """VERIA-11: list-format content must be inspected, not silently skipped.""" data = { "messages": [ @@ -34,26 +34,26 @@ def test_iter_user_text_multimodal_list_content(): } ] } - assert list(iter_user_text(data)) == ["AWS_KEY=AKIA...", "more secrets"] + assert list(iter_message_text(data)) == ["AWS_KEY=AKIA...", "more secrets"] -def test_iter_user_text_responses_api_string_input(): +def test_iter_message_text_responses_api_string_input(): """fniVO9-F: Responses-API ``input`` must be inspectable when ``messages`` absent.""" data = {"input": "tell me a secret"} - assert list(iter_user_text(data)) == ["tell me a secret"] + assert list(iter_message_text(data)) == ["tell me a secret"] -def test_iter_user_text_responses_api_list_input_messages(): +def test_iter_message_text_responses_api_list_input_messages(): data = { "input": [ {"role": "user", "content": "first"}, {"role": "user", "content": "second"}, ] } - assert list(iter_user_text(data)) == ["first", "second"] + assert list(iter_message_text(data)) == ["first", "second"] -def test_iter_user_text_responses_api_list_input_content_parts(): +def test_iter_message_text_responses_api_list_input_content_parts(): data = { "input": [ {"type": "text", "text": "alpha"}, @@ -61,23 +61,42 @@ def test_iter_user_text_responses_api_list_input_content_parts(): {"type": "text", "text": "beta"}, ] } - assert list(iter_user_text(data)) == ["alpha", "beta"] + assert list(iter_message_text(data)) == ["alpha", "beta"] -def test_iter_user_text_walks_messages_and_input_independently(): +def test_iter_message_text_responses_api_list_input_mixed_dicts_and_strings(): + """Greptile P2: mixed-list ``input`` with content-part dicts AND bare + strings must yield every text fragment — read helpers used to truncate + the bare strings.""" + data = { + "input": [ + {"type": "text", "text": "from-dict"}, + "from-bare-string", + {"type": "image_url", "image_url": {"url": "..."}}, + "another-bare-string", + ] + } + assert list(iter_message_text(data)) == [ + "from-dict", + "from-bare-string", + "another-bare-string", + ] + + +def test_iter_message_text_walks_messages_and_input_independently(): """When both are present (rare), every fragment from either field is inspected — a stricter guarantee than "first one wins".""" data = { "messages": [{"role": "user", "content": "msg-content"}], "input": "input-content", } - assert list(iter_user_text(data)) == ["msg-content", "input-content"] + assert list(iter_message_text(data)) == ["msg-content", "input-content"] -def test_iter_user_text_empty_data(): - assert list(iter_user_text({})) == [] - assert list(iter_user_text({"messages": []})) == [] - assert list(iter_user_text({"input": ""})) == [] +def test_iter_message_text_empty_data(): + assert list(iter_message_text({})) == [] + assert list(iter_message_text({"messages": []})) == [] + assert list(iter_message_text({"input": ""})) == [] # ── walk_user_text ──────────────────────────────────────────────────────────── @@ -139,6 +158,23 @@ def test_walk_user_text_redacts_responses_api_list_input(): assert data["input"][1] == {"type": "image_url", "image_url": {"url": "..."}} +def test_walk_user_text_redacts_mixed_list_input(): + """Read and write helpers must agree on coverage — bare strings inside + a mixed ``input`` list are inspected by both.""" + data = { + "input": [ + {"type": "text", "text": "secret-one"}, + "secret-two", + {"type": "image_url", "image_url": {"url": "..."}}, + ] + } + visited = walk_user_text(data, lambda s: f"<{s}>") + assert visited == 2 + assert data["input"][0] == {"type": "text", "text": ""} + assert data["input"][1] == "" + assert data["input"][2] == {"type": "image_url", "image_url": {"url": "..."}} + + # ── build_inspection_messages ───────────────────────────────────────────────── From 9fcf2347505c10037d8a6ee65c2f951d612a0f77 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 04:22:57 +0000 Subject: [PATCH 04/11] fix(guardrails): degrade Lakera v2 mask mode to block on multimodal input MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mask-in-place uses the offsets that Lakera returns for the inspection payload. ``build_inspection_messages`` flattens multimodal content into joined text before sending to Lakera, so the offsets refer to the flattened representation. Writing those offsets back via ``_mask_pii_in_messages`` and overwriting ``data["messages"]`` would silently strip image/audio parts from the original request — that is a real functional regression for Lakera + mask mode + multimodal input. Detect multimodal input (any list-format ``content`` or non-string ``data["input"]``) up front and skip the mask-in-place branch in that case. The hook then falls into the standard block-on-detect path so PII is still blocked but the multimodal payload is never silently rewritten. Per-part masking that preserves multimodal structure is the right long-term fix; tracking that as a follow-up. Also: add ``has_non_string_content`` to ``_content_utils`` (with tests) and a regression test that asserts multimodal+PII raises an HTTPException instead of returning a flattened request body. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/guardrails/_content_utils.py | 24 ++++++++- .../guardrail_hooks/lakera_ai_v2.py | 36 +++++++++---- .../proxy/guardrails/test_content_utils.py | 28 ++++++++++ .../guardrails/test_guardrail_coverage.py | 52 +++++++++++++++++++ 4 files changed, 130 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 576fefe85a..6c681c322f 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -86,7 +86,7 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: return visit(content) return content if isinstance(content, list): - new_parts = [] + new_parts: List[Any] = [] for part in content: if isinstance(part, str) and part: visited += 1 @@ -143,6 +143,28 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: return visited +def has_non_string_content(data: Dict[str, Any]) -> bool: + """Return True if any inspected content is not a plain string. + + Used by hooks whose mask/redact path operates on string offsets and + therefore cannot preserve multimodal non-text parts. Such hooks should + degrade to block-on-detect when this returns True so image/audio parts + are not silently stripped during in-place masking. + """ + messages = data.get("messages") + if isinstance(messages, list): + for message in messages: + if isinstance(message, dict) and not isinstance( + message.get("content"), str + ): + if message.get("content") is not None: + return True + input_value = data.get("input") + if input_value is not None and not isinstance(input_value, str): + return True + return False + + def build_inspection_messages(data: Dict[str, Any]) -> List[Dict[str, str]]: """Synthesize a chat-style messages list for posting to a guardrail API. diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index d3f19966be..4d6e6a302b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -13,7 +13,10 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails._content_utils import build_inspection_messages +from litellm.proxy.guardrails._content_utils import ( + build_inspection_messages, + has_non_string_content, +) from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import AllMessageValues @@ -223,6 +226,13 @@ class LakeraAIGuardrail(CustomGuardrail): ) return data + # Mask-in-place uses offsets returned by Lakera and can only + # preserve non-text parts (images, audio, …) when the original + # content is a plain string. For multimodal/Responses-API input + # we degrade to block-on-detect so we never silently strip image + # parts while attempting to redact text. + is_multimodal_input = has_non_string_content(data) + ######################################################### ########## 1. Make the Lakera AI v2 guard API request ########## ######################################################### @@ -236,8 +246,11 @@ class LakeraAIGuardrail(CustomGuardrail): ########## 2. Handle flagged content ########## ######################################################### if lakera_guardrail_response.get("flagged") is True: - # If only PII violations exist, mask the PII - if self._is_only_pii_violation(lakera_guardrail_response): + # If only PII violations exist, mask the PII (string input only). + if ( + self._is_only_pii_violation(lakera_guardrail_response) + and not is_multimodal_input + ): data["messages"] = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, @@ -254,7 +267,9 @@ class LakeraAIGuardrail(CustomGuardrail): ) # Log violation but continue elif self.on_flagged == "block": - # If there are other violations or not set to mask PII, raise exception + # Either non-PII violations, or PII on multimodal input + # (which cannot be masked in place without dropping + # image/audio parts) — raise the standard block error. raise self._get_http_exception_for_blocked_guardrail( lakera_guardrail_response ) @@ -289,6 +304,10 @@ class LakeraAIGuardrail(CustomGuardrail): ) return + # See ``async_pre_call_hook`` — multimodal input degrades to + # block-on-detect because mask-in-place would drop image parts. + is_multimodal_input = has_non_string_content(data) + ######################################################### ########## 1. Make the Lakera AI v2 guard API request ########## ######################################################### @@ -302,8 +321,10 @@ class LakeraAIGuardrail(CustomGuardrail): ########## 2. Handle flagged content ########## ######################################################### if lakera_guardrail_response.get("flagged") is True: - # If only PII violations exist, mask the PII - if self._is_only_pii_violation(lakera_guardrail_response): + if ( + self._is_only_pii_violation(lakera_guardrail_response) + and not is_multimodal_input + ): data["messages"] = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, @@ -313,14 +334,11 @@ class LakeraAIGuardrail(CustomGuardrail): "Lakera AI: Masked PII in messages instead of blocking request" ) else: - # Check on_flagged setting if self.on_flagged == "monitor": verbose_proxy_logger.warning( "Lakera Guardrail: Monitoring mode - violation detected but allowing request" ) - # Log violation but continue elif self.on_flagged == "block": - # If there are other violations or not set to mask PII, raise exception raise self._get_http_exception_for_blocked_guardrail( lakera_guardrail_response ) diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index e26dc7aa2d..23683abd1c 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -2,6 +2,7 @@ from litellm.proxy.guardrails._content_utils import ( build_inspection_messages, + has_non_string_content, iter_message_text, walk_user_text, ) @@ -235,3 +236,30 @@ def test_build_inspection_messages_empty_data(): assert build_inspection_messages({}) == [] assert build_inspection_messages({"messages": []}) == [] assert build_inspection_messages({"input": ""}) == [] + + +# ── has_non_string_content ──────────────────────────────────────────────────── + + +def test_has_non_string_content_string_messages(): + data = {"messages": [{"role": "user", "content": "hello"}]} + assert has_non_string_content(data) is False + + +def test_has_non_string_content_multimodal_messages(): + data = {"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]} + assert has_non_string_content(data) is True + + +def test_has_non_string_content_responses_api_string_input(): + assert has_non_string_content({"input": "plain string"}) is False + + +def test_has_non_string_content_responses_api_list_input(): + assert has_non_string_content({"input": ["a", "b"]}) is True + + +def test_has_non_string_content_empty_data(): + assert has_non_string_content({}) is False + assert has_non_string_content({"messages": []}) is False + assert has_non_string_content({"input": ""}) is False diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py index 4cfcbf3c6d..a4fb995ed5 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -161,6 +161,58 @@ async def test_lakera_v2_inspects_responses_api_input(user_api_key, monkeypatch) assert seen_messages == [[{"role": "user", "content": "responses-api content"}]] +@pytest.mark.asyncio +async def test_lakera_v2_multimodal_pii_degrades_to_block(user_api_key, monkeypatch): + """Mask-in-place uses Lakera offsets and cannot preserve image/audio + parts of multimodal input. When PII is detected on a multimodal + request, the hook must raise the block exception instead of silently + flattening ``data["messages"]`` to text-only.""" + monkeypatch.setenv("LAKERA_API_KEY", "lk-test") + from fastapi import HTTPException + + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( + LakeraAIGuardrail, + ) + + guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="block") + + async def fake_call_v2_guard(messages, request_data, event_type): + return ( + { + "flagged": True, + "payload": [{"detector_type": "pii/email", "start": 0, "end": 5}], + }, + {"EMAIL": 1}, + ) + + with ( + patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard), + patch.object(guard, "_is_only_pii_violation", return_value=True), + patch.object( + guard, + "_get_http_exception_for_blocked_guardrail", + return_value=HTTPException(status_code=400, detail="blocked"), + ), + ): + with pytest.raises(HTTPException): + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data={ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "leak"}, + {"type": "image_url", "image_url": {"url": "..."}}, + ], + } + ] + }, + call_type="acompletion", + ) + + @pytest.mark.asyncio async def test_lakera_v2_inspects_multimodal_list_content(user_api_key, monkeypatch): monkeypatch.setenv("LAKERA_API_KEY", "lk-test") From 40817caa4a1008149dd132c09d8edfa5f66f835b Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 04:28:22 +0000 Subject: [PATCH 05/11] fix(guardrails): degrade Lasso/Aim mask paths to block on multimodal Two more in-place rewrite paths exhibit the same regression as Lakera v2: overwriting ``data["messages"]`` with text-only redacted versions silently strips image/audio parts from multimodal requests. - ``LassoGuardrail._run_lasso_guardrail``: when ``mask=True`` AND input is multimodal/Responses-API list, fall back to the classify endpoint (which raises on BLOCK actions but never overwrites the payload). - ``AimGuardrail._anonymize_request``: when input is multimodal, raise the standard 400 instead of replacing ``data["messages"]`` with the text-only ``redacted_chat`` from Aim. The error message tells the user to either send plain string content or rely on block-mode. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../guardrails/guardrail_hooks/aim/aim.py | 19 +++++++- .../guardrails/guardrail_hooks/lasso/lasso.py | 17 ++++--- .../guardrails/test_guardrail_coverage.py | 48 +++++++++++++++++++ 3 files changed, 77 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index a386617428..749bd4acf2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -22,7 +22,10 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails._content_utils import build_inspection_messages +from litellm.proxy.guardrails._content_utils import ( + build_inspection_messages, + has_non_string_content, +) from litellm.types.utils import ( CallTypesLiteral, Choices, @@ -139,6 +142,20 @@ class AimGuardrail(CustomGuardrail): redacted_chat = res.get("redacted_chat") if not redacted_chat: return data + # Aim returns text-only redacted messages. Overwriting + # ``data["messages"]`` with that would silently strip image/audio + # parts from a multimodal request — degrade to block so the + # multimodal payload is never silently rewritten. + if has_non_string_content(data): + raise HTTPException( + status_code=400, + detail=( + "Aim: anonymize action requested for multimodal input " + "but mask-in-place would drop non-text parts. Send the " + "request with plain string content to use anonymize, " + "or rely on block-mode policies." + ), + ) data["messages"] = [ { "role": message["role"], diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index a7719124b2..0469798108 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -50,7 +50,10 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails._content_utils import build_inspection_messages +from litellm.proxy.guardrails._content_utils import ( + build_inspection_messages, + has_non_string_content, +) from litellm.types.guardrails import GuardrailEventHooks import litellm @@ -372,12 +375,14 @@ class LassoGuardrail(CustomGuardrail): if not messages: return data - if self.mask: + # Lasso's classifix endpoint returns masked text that we copy back + # into ``data["messages"]``. For multimodal/Responses-API input we + # would silently strip image/audio parts, so fall back to the + # classify endpoint (which still raises on BLOCK actions) and + # leave the original payload intact. + if self.mask and not has_non_string_content(data): return await self._handle_masking(data, cache, message_type, messages) - else: - return await self._handle_classification( - data, cache, message_type, messages - ) + return await self._handle_classification(data, cache, message_type, messages) async def _handle_classification( self, diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py index a4fb995ed5..df321b8074 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -251,6 +251,54 @@ async def test_lakera_v2_inspects_multimodal_list_content(user_api_key, monkeypa # ── Lasso ───────────────────────────────────────────────────────────────────── +@pytest.mark.asyncio +async def test_lasso_multimodal_falls_back_to_classify(user_api_key, monkeypatch): + """Lasso's classifix (mask) endpoint returns text that overwrites + ``data["messages"]``. For multimodal input that would silently strip + image parts — the hook must use the classify endpoint instead and + leave the original payload intact.""" + monkeypatch.setenv("LASSO_API_KEY", "ls-test") + from litellm.proxy.guardrails.guardrail_hooks.lasso.lasso import LassoGuardrail + + guard = LassoGuardrail(lasso_api_key="ls-test", mask=True) + + masking_called = False + classify_called = False + + async def fake_masking(data, cache, message_type, messages): + nonlocal masking_called + masking_called = True + return data + + async def fake_classification(data, cache, message_type, messages): + nonlocal classify_called + classify_called = True + return data + + with ( + patch.object(guard, "_handle_masking", side_effect=fake_masking), + patch.object(guard, "_handle_classification", side_effect=fake_classification), + ): + await guard._run_lasso_guardrail( + data={ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image_url", "image_url": {"url": "..."}}, + ], + } + ] + }, + cache=DualCache(), + message_type="PROMPT", + ) + + assert classify_called is True + assert masking_called is False + + @pytest.mark.asyncio async def test_lasso_inspects_responses_api_input(user_api_key, monkeypatch): monkeypatch.setenv("LASSO_API_KEY", "ls-test") From 5397ac4562e853e41be0b7a8d0bb056536e2b79b Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 04:41:57 +0000 Subject: [PATCH 06/11] fix(guardrails): redact ``data["input"]`` for Responses-API mask paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Greptile P1: Aim's ``_anonymize_request`` and Lakera v2's mask-PII path both wrote redacted content only to ``data["messages"]``. The Responses API backend reads ``data["input"]``, so when a request arrived via ``/v1/responses`` with a plain string ``input`` the hook would update ``messages`` (which the backend ignores) and leave ``input`` carrying the original unredacted text. Net effect: anonymize/mask silently passed PII through to the LLM. Add ``apply_redacted_messages_back`` to ``_content_utils`` — it writes the redacted messages back to ``data["messages"]`` AND, when present, re-flattens the redacted content into ``data["input"]``. Aim and Lakera v2 now route their mask writeback through this helper. List ``input`` (multimodal) is still handled by the upstream block-on-multimodal guard. Adds unit tests for the helper and regression tests asserting ``data["input"]`` is redacted for both hooks on Responses-API string input. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/guardrails/_content_utils.py | 24 ++++++ .../guardrails/guardrail_hooks/aim/aim.py | 7 +- .../guardrail_hooks/lakera_ai_v2.py | 13 +++- .../proxy/guardrails/test_content_utils.py | 38 ++++++++++ .../guardrails/test_guardrail_coverage.py | 75 +++++++++++++++++++ 5 files changed, 154 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 6c681c322f..8fdd8b23cb 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -143,6 +143,30 @@ def walk_user_text(data: Dict[str, Any], visit: Callable[[str], str]) -> int: return visited +def apply_redacted_messages_back( + data: Dict[str, Any], redacted_messages: List[Dict[str, Any]] +) -> None: + """Write redacted messages back to whichever field(s) the caller used. + + Mask/anonymize paths take a synthesised messages list (from + :func:`build_inspection_messages`), get a redacted version back from a + third-party guardrail, and need to rewrite the request body. Writing + only to ``data["messages"]`` leaves the Responses-API ``data["input"]`` + field untouched, so the unredacted text still reaches the LLM. + + This helper updates both fields when both are present. + """ + if "messages" in data: + data["messages"] = redacted_messages + if isinstance(data.get("input"), str): + text_parts: List[str] = [] + for msg in redacted_messages: + if not isinstance(msg, dict): + continue + text_parts.extend(_iter_text_parts_in_content(msg.get("content"))) + data["input"] = "\n".join(text_parts) + + def has_non_string_content(data: Dict[str, Any]) -> bool: """Return True if any inspected content is not a plain string. diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 749bd4acf2..03a73f9fc9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -23,6 +23,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import ( + apply_redacted_messages_back, build_inspection_messages, has_non_string_content, ) @@ -156,13 +157,17 @@ class AimGuardrail(CustomGuardrail): "or rely on block-mode policies." ), ) - data["messages"] = [ + redacted_messages = [ { "role": message["role"], "content": message["content"], } for message in redacted_chat["all_redacted_messages"] ] + # Write back to ``messages`` AND ``input``. The Responses-API + # backend reads ``input``; writing only to ``messages`` would let + # unredacted text reach the LLM for ``/v1/responses`` calls. + apply_redacted_messages_back(data, redacted_messages) return data async def call_aim_guardrail_on_output( diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 4d6e6a302b..c6af4d3c42 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -14,6 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import ( + apply_redacted_messages_back, build_inspection_messages, has_non_string_content, ) @@ -251,11 +252,15 @@ class LakeraAIGuardrail(CustomGuardrail): self._is_only_pii_violation(lakera_guardrail_response) and not is_multimodal_input ): - data["messages"] = self._mask_pii_in_messages( + redacted_messages = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, masked_entity_count=masked_entity_count, ) + # Write back to ``messages`` AND ``input``. The Responses-API + # backend reads ``input``; writing only to ``messages`` + # would let unredacted PII reach the LLM for /v1/responses. + apply_redacted_messages_back(data, list(redacted_messages)) # type: ignore[arg-type] verbose_proxy_logger.debug( "Lakera AI: Masked PII in messages instead of blocking request" ) @@ -325,11 +330,15 @@ class LakeraAIGuardrail(CustomGuardrail): self._is_only_pii_violation(lakera_guardrail_response) and not is_multimodal_input ): - data["messages"] = self._mask_pii_in_messages( + redacted_messages = self._mask_pii_in_messages( messages=new_messages, # type: ignore[arg-type] lakera_response=lakera_guardrail_response, masked_entity_count=masked_entity_count, ) + # Write back to ``messages`` AND ``input``. The Responses-API + # backend reads ``input``; writing only to ``messages`` + # would let unredacted PII reach the LLM for /v1/responses. + apply_redacted_messages_back(data, list(redacted_messages)) # type: ignore[arg-type] verbose_proxy_logger.debug( "Lakera AI: Masked PII in messages instead of blocking request" ) diff --git a/tests/test_litellm/proxy/guardrails/test_content_utils.py b/tests/test_litellm/proxy/guardrails/test_content_utils.py index 23683abd1c..099fca78a6 100644 --- a/tests/test_litellm/proxy/guardrails/test_content_utils.py +++ b/tests/test_litellm/proxy/guardrails/test_content_utils.py @@ -1,6 +1,7 @@ """Tests for the shared guardrail content extraction helpers.""" from litellm.proxy.guardrails._content_utils import ( + apply_redacted_messages_back, build_inspection_messages, has_non_string_content, iter_message_text, @@ -263,3 +264,40 @@ def test_has_non_string_content_empty_data(): assert has_non_string_content({}) is False assert has_non_string_content({"messages": []}) is False assert has_non_string_content({"input": ""}) is False + + +# ── apply_redacted_messages_back ────────────────────────────────────────────── + + +def test_apply_redacted_messages_back_chat_completion(): + data = {"messages": [{"role": "user", "content": "secret"}]} + apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) + assert data["messages"] == [{"role": "user", "content": "[REDACTED]"}] + assert "input" not in data + + +def test_apply_redacted_messages_back_responses_api_string_input(): + """A Responses-API request reads ``data["input"]``; writing only to + ``messages`` would let unredacted text reach the LLM.""" + data = {"input": "secret payload"} + apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) + assert data["input"] == "[REDACTED]" + + +def test_apply_redacted_messages_back_both_fields(): + """Defensive: when both fields are present, both are updated.""" + data = { + "messages": [{"role": "user", "content": "old"}], + "input": "old", + } + apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) + assert data["messages"] == [{"role": "user", "content": "[REDACTED]"}] + assert data["input"] == "[REDACTED]" + + +def test_apply_redacted_messages_back_skips_input_when_not_string(): + """List ``input`` (multimodal Responses-API) is left alone — the + multimodal-degrades-to-block guard runs upstream.""" + data = {"input": [{"type": "text", "text": "leak"}]} + apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) + assert data["input"] == [{"type": "text", "text": "leak"}] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py index df321b8074..04bd57fa32 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -161,6 +161,81 @@ async def test_lakera_v2_inspects_responses_api_input(user_api_key, monkeypatch) assert seen_messages == [[{"role": "user", "content": "responses-api content"}]] +@pytest.mark.asyncio +async def test_lakera_v2_responses_api_input_redacted_writeback( + user_api_key, monkeypatch +): + """Greptile P1: when input arrives via Responses-API ``data["input"]`` + (string) and Lakera flags PII, the redacted content must be written + back to ``data["input"]`` — the Responses-API backend reads from + ``input``, so writing only to ``messages`` would let unredacted PII + reach the LLM.""" + monkeypatch.setenv("LAKERA_API_KEY", "lk-test") + from litellm.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( + LakeraAIGuardrail, + ) + + guard = LakeraAIGuardrail(api_key="lk-test", on_flagged="block") + + async def fake_call_v2_guard(messages, request_data, event_type): + return ({"flagged": True, "payload": []}, {"EMAIL": 1}) + + def fake_mask(messages, lakera_response, masked_entity_count): + return [{"role": "user", "content": "[REDACTED EMAIL]"}] + + with ( + patch.object(guard, "call_v2_guard", side_effect=fake_call_v2_guard), + patch.object(guard, "_is_only_pii_violation", return_value=True), + patch.object(guard, "_mask_pii_in_messages", side_effect=fake_mask), + ): + data = {"input": "user@example.com leaked"} + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data=data, + call_type="responses", + ) + + assert data["input"] == "[REDACTED EMAIL]" + + +@pytest.mark.asyncio +async def test_aim_responses_api_input_anonymize_writeback(user_api_key, monkeypatch): + """Greptile P1: Aim's anonymize action must redact ``data["input"]`` + for Responses-API requests, not just ``data["messages"]``.""" + monkeypatch.setenv("AIM_API_KEY", "hs-aim-key") + from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail + + guard = AimGuardrail() + + aim_response_body = { + "required_action": {"action_type": "anonymize_action"}, + "redacted_chat": { + "all_redacted_messages": [ + {"role": "user", "content": "[REDACTED] anonymised"} + ] + }, + } + + async def capture(url, headers, json): + return Response( + status_code=200, + json=aim_response_body, + request=Request("POST", "https://api.aim.security/fw/v1/analyze"), + ) + + with patch.object(guard.async_handler, "post", side_effect=capture): + data = {"input": "user@example.com leaked"} + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data=data, + call_type="responses", + ) + + assert data["input"] == "[REDACTED] anonymised" + + @pytest.mark.asyncio async def test_lakera_v2_multimodal_pii_degrades_to_block(user_api_key, monkeypatch): """Mask-in-place uses Lakera offsets and cannot preserve image/audio From 41eeaaf6309ba7e983d35705ab04e076085efccc Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 04:58:00 +0000 Subject: [PATCH 07/11] fix(guardrails): handle Aim multi-choice gather exceptions cleanly Greptile P2: ``asyncio.gather`` without ``return_exceptions=True`` lets the first failing call propagate immediately, leaving the other in-flight inspections running until they complete on their own. Pass ``return_exceptions=True`` so every inspection finishes, then re-raise the first exception encountered while iterating results. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/guardrails/guardrail_hooks/aim/aim.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 03a73f9fc9..5b5f91195e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -266,6 +266,9 @@ class AimGuardrail(CustomGuardrail): choices_to_inspect = [c for c in response.choices if isinstance(c, Choices)] if not choices_to_inspect: return response + # ``return_exceptions=True`` lets every inspection finish even if + # one fails — without it, the first exception would propagate and + # leave the remaining tasks running in the background. results = await asyncio.gather( *( self.call_aim_guardrail_on_output( @@ -275,9 +278,12 @@ class AimGuardrail(CustomGuardrail): key_alias=user_api_key_dict.key_alias, ) for choice in choices_to_inspect - ) + ), + return_exceptions=True, ) for choice, aim_output_guardrail_result in zip(choices_to_inspect, results): + if isinstance(aim_output_guardrail_result, BaseException): + raise aim_output_guardrail_result if aim_output_guardrail_result and aim_output_guardrail_result.get( "detection_message" ): From f34b968a5f680751cb6df2c71962f2eab8242cbc Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 14:08:53 -0700 Subject: [PATCH 08/11] chore: retrigger PR checks From f909fa79b50d8428e5e4d44fd1a329264655bdc9 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 19:26:40 -0700 Subject: [PATCH 09/11] fix(guardrails): close post-call coverage gaps --- .../enterprise_hooks/banned_keywords.py | 15 ++- .../guardrails/guardrail_hooks/lasso/lasso.py | 6 +- litellm/proxy/hooks/azure_content_safety.py | 16 ++-- .../guardrails/test_guardrail_coverage.py | 93 +++++++++++++++++++ 4 files changed, 117 insertions(+), 13 deletions(-) diff --git a/enterprise/enterprise_hooks/banned_keywords.py b/enterprise/enterprise_hooks/banned_keywords.py index a86b59daf5..621fde011d 100644 --- a/enterprise/enterprise_hooks/banned_keywords.py +++ b/enterprise/enterprise_hooks/banned_keywords.py @@ -94,11 +94,16 @@ class _ENTERPRISE_BannedKeywords(CustomLogger): user_api_key_dict: UserAPIKeyAuth, response, ): - if isinstance(response, litellm.ModelResponse) and isinstance( - response.choices[0], litellm.utils.Choices - ): - for word in self.banned_keywords_list: - self.test_violation(test_str=response.choices[0].message.content or "") + if not isinstance(response, litellm.ModelResponse): + return + + for choice in response.choices: + if not isinstance(choice, litellm.utils.Choices): + continue + message = getattr(choice, "message", None) + content = getattr(message, "content", None) + if isinstance(content, str): + self.test_violation(test_str=content) async def async_post_call_streaming_hook( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 0469798108..e64b69efcc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -51,6 +51,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import ( + apply_redacted_messages_back, build_inspection_messages, has_non_string_content, ) @@ -420,8 +421,9 @@ class LassoGuardrail(CustomGuardrail): self._process_lasso_response(response) # Apply masking to messages if violations detected and masked messages are available - if response.get("violations_detected") and response.get("messages"): - data["messages"] = response["messages"] + redacted_messages = response.get("messages") + if response.get("violations_detected") and redacted_messages: + apply_redacted_messages_back(data, list(redacted_messages)) self._log_masking_applied(message_type, dict(response)) return data diff --git a/litellm/proxy/hooks/azure_content_safety.py b/litellm/proxy/hooks/azure_content_safety.py index a21f0566cc..844bbc11af 100644 --- a/litellm/proxy/hooks/azure_content_safety.py +++ b/litellm/proxy/hooks/azure_content_safety.py @@ -141,12 +141,16 @@ class _PROXY_AzureContentSafety( response, ): verbose_proxy_logger.debug("Inside Azure Content-Safety Post-Call Hook") - if isinstance(response, litellm.ModelResponse) and isinstance( - response.choices[0], litellm.utils.Choices - ): - await self.test_violation( - content=response.choices[0].message.content or "", source="output" - ) + if not isinstance(response, litellm.ModelResponse): + return + + for choice in response.choices: + if not isinstance(choice, litellm.utils.Choices): + continue + message = getattr(choice, "message", None) + content = getattr(message, "content", None) + if isinstance(content, str): + await self.test_violation(content=content, source="output") # async def async_post_call_streaming_hook( # self, diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py index 04bd57fa32..329f78bf50 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -399,6 +399,36 @@ async def test_lasso_inspects_responses_api_input(user_api_key, monkeypatch): assert seen_messages == [[{"role": "user", "content": "responses-api content"}]] +@pytest.mark.asyncio +async def test_lasso_masking_writes_back_responses_api_input(user_api_key, monkeypatch): + """Krrish blocker: Lasso classifix masking must update ``data["input"]`` + for Responses-API requests, not only ``data["messages"]``.""" + monkeypatch.setenv("LASSO_API_KEY", "ls-test") + from litellm.proxy.guardrails.guardrail_hooks.lasso.lasso import LassoGuardrail + + guard = LassoGuardrail(lasso_api_key="ls-test", mask=True) + lasso_response = { + "violations_detected": True, + "deputies": {"pii": True}, + "findings": {"pii": [{"action": "AUTO_MASKING"}]}, + "messages": [{"role": "user", "content": "[REDACTED]"}], + } + + async def fake_call_lasso_api(headers, payload, api_url=None): + return lasso_response + + data = {"input": "user@example.com leaked"} + + with patch.object(guard, "_call_lasso_api", side_effect=fake_call_lasso_api): + await guard._run_lasso_guardrail( + data=data, + cache=DualCache(), + message_type="PROMPT", + ) + + assert data["input"] == "[REDACTED]" + + # ── Banned Keywords ─────────────────────────────────────────────────────────── @@ -456,6 +486,69 @@ def test_banned_keywords_blocks_responses_api_input(monkeypatch): asyncio.run(_run()) +@pytest.mark.asyncio +async def test_banned_keywords_post_call_checks_all_choices(monkeypatch, user_api_key): + """Krrish blocker: ``n>1`` responses must not bypass post-call checks by + placing the banned text in ``choices[1+]``.""" + monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False) + from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords + from fastapi import HTTPException + + guard = _ENTERPRISE_BannedKeywords() + response = ModelResponse( + choices=[ + Choices(index=0, message=Message(role="assistant", content="clean")), + Choices(index=1, message=Message(role="assistant", content="forbidden")), + ] + ) + + with pytest.raises(HTTPException) as exc: + await guard.async_post_call_success_hook( + data={}, + user_api_key_dict=user_api_key, + response=response, + ) + + assert "forbidden" in str(exc.value.detail).lower() + + +# ── Azure Content Safety ────────────────────────────────────────────────────── + + +@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 + by placing the unsafe text in ``choices[1+]``.""" + from fastapi import HTTPException + from litellm.proxy.hooks.azure_content_safety import _PROXY_AzureContentSafety + + guard = _PROXY_AzureContentSafety.__new__(_PROXY_AzureContentSafety) + seen_outputs = [] + + async def fake_test_violation(content, source=None): + seen_outputs.append((content, source)) + if "unsafe" in content: + raise HTTPException(status_code=400, detail={"error": "unsafe"}) + + guard.test_violation = fake_test_violation + response = ModelResponse( + choices=[ + Choices(index=0, message=Message(role="assistant", content="clean")), + Choices(index=1, message=Message(role="assistant", content="unsafe")), + Choices(index=2, message=Message(role="assistant", content="later")), + ] + ) + + with pytest.raises(HTTPException): + await guard.async_post_call_success_hook( + data={}, + user_api_key_dict=user_api_key, + response=response, + ) + + assert seen_outputs == [("clean", "output"), ("unsafe", "output")] + + # ── Secret Detection ────────────────────────────────────────────────────────── From f4e7dde2d8ff150c969fea80c7fd663482458a89 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 19:38:04 -0700 Subject: [PATCH 10/11] fix(guardrails): cover multi-choice output variants --- .../guardrail_hooks/azure/text_moderation.py | 32 ++-- .../litellm_content_filter/content_filter.py | 145 ++++++++++-------- .../semantic_guard/semantic_guard.py | 25 ++- .../guardrail_hooks/xecguard/xecguard.py | 31 ++-- tests/guardrails_tests/test_semantic_guard.py | 19 +++ .../azure/test_azure_text_moderation.py | 46 ++++++ .../content_filter/test_content_filter.py | 65 ++++++++ .../guardrail_hooks/test_xecguard.py | 21 ++- 8 files changed, 296 insertions(+), 88 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index a465c5428d..55b446eb52 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -254,15 +254,16 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr ) -> Any: from litellm.types.utils import Choices, ModelResponse - if ( - isinstance(response, ModelResponse) - and response.choices - and isinstance(response.choices[0], Choices) - ): - content = response.choices[0].message.content or "" - await self.async_make_request( - text=content, - ) + if isinstance(response, ModelResponse) and response.choices: + for choice in response.choices: + if not isinstance(choice, Choices): + continue + content = _message_content_to_text(choice.message.content) + if not content: + continue + await self.async_make_request( + text=content, + ) return response async def async_post_call_streaming_hook( @@ -279,3 +280,16 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr error_returned = json.dumps({"error": e.detail}) return f"data: {error_returned}\n\n" + + +def _message_content_to_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + text_parts = [ + item.get("text") + for item in content + if isinstance(item, dict) and isinstance(item.get("text"), str) + ] + return "\n".join(part for part in text_parts if part) + return "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 3ef99e4ad6..d6d2e01494 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -1873,8 +1873,9 @@ class ContentFilterGuardrail(CustomGuardrail): and the UI Request Lifecycle panel. Mirrors apply_guardrail's finally-block contract. """ - accumulated_full_text = "" - yielded_masked_text_len = 0 + accumulated_text_by_choice: Dict[int, str] = {} + yielded_masked_text_len_by_choice: Dict[int, int] = {} + latest_detections_by_choice: Dict[int, List[ContentFilterDetection]] = {} buffer_size = 50 # Increased buffer to catch patterns split across many chunks start_time = datetime.now() @@ -1890,79 +1891,90 @@ class ContentFilterGuardrail(CustomGuardrail): try: async for item in response: if isinstance(item, ModelResponseStream) and item.choices: - delta_content = "" - is_final = False for choice in item.choices: - if hasattr(choice, "delta") and choice.delta: - content = getattr(choice.delta, "content", None) - if content and isinstance(content, str): - delta_content += content - if getattr(choice, "finish_reason", None): - is_final = True + if not (hasattr(choice, "delta") and choice.delta): + continue - accumulated_full_text += delta_content + choice_index = getattr(choice, "index", 0) + if not isinstance(choice_index, int): + choice_index = 0 - # Check for blocking or apply masking - # Add a space at the end if it's the final chunk to trigger word boundaries (\b) - text_to_check = accumulated_full_text - if is_final: - text_to_check += " " + content = getattr(choice.delta, "content", None) + is_final = bool(getattr(choice, "finish_reason", None)) + if isinstance(content, str) and content: + accumulated_text_by_choice[choice_index] = ( + accumulated_text_by_choice.get(choice_index, "") + + content + ) + elif not is_final: + continue - try: - # Reset before each scan: _filter_single_text scans the - # whole accumulated buffer every chunk, so previous-chunk - # matches are guaranteed to be re-found. Keeping only the - # latest scan's detections avoids N× duplication in the - # final log row. BLOCK still records correctly because - # handlers append to detections before raising. - detections.clear() - masked_text = self._filter_single_text( - text_to_check, detections=detections + text_to_check = accumulated_text_by_choice.get(choice_index, "") + if not text_to_check: + continue + + # Add a space at the end if it's the final chunk to trigger word boundaries (\b) + text_to_scan = text_to_check + (" " if is_final else "") + choice_detections: List[ContentFilterDetection] = [] + + try: + # _filter_single_text scans the whole accumulated + # choice buffer every chunk, so previous-chunk + # matches are guaranteed to be re-found. Keeping + # only each choice's latest scan avoids duplicate + # detections in the final log row. + masked_text = self._filter_single_text( + text_to_scan, detections=choice_detections + ) + if is_final and masked_text.endswith(" "): + masked_text = masked_text[:-1] + latest_detections_by_choice[choice_index] = ( + choice_detections + ) + except HTTPException: + latest_detections_by_choice[choice_index] = ( + choice_detections + ) + raise + except Exception as e: + verbose_proxy_logger.error( + f"ContentFilterGuardrail: Error in masking: {e}" + ) + masked_text = text_to_scan # Fallback to current text + + # Determine how much can be safely yielded + if is_final: + safe_to_yield_len = len(masked_text) + else: + safe_to_yield_len = max(0, len(masked_text) - buffer_size) + + yielded_masked_text_len = yielded_masked_text_len_by_choice.get( + choice_index, 0 ) - if is_final and masked_text.endswith(" "): - masked_text = masked_text[:-1] - except HTTPException: - raise - except Exception as e: - verbose_proxy_logger.error( - f"ContentFilterGuardrail: Error in masking: {e}" - ) - masked_text = text_to_check # Fallback to current text + if safe_to_yield_len > yielded_masked_text_len: + new_masked_content = masked_text[ + yielded_masked_text_len:safe_to_yield_len + ] + choice.delta.content = new_masked_content + yielded_masked_text_len_by_choice[choice_index] = ( + safe_to_yield_len + ) + else: + # Hold content by yielding empty content on this choice + # while preserving chunk metadata and other choices. + choice.delta.content = "" - # Determine how much can be safely yielded - if is_final: - safe_to_yield_len = len(masked_text) - else: - safe_to_yield_len = max(0, len(masked_text) - buffer_size) - - if safe_to_yield_len > yielded_masked_text_len: - new_masked_content = masked_text[ - yielded_masked_text_len:safe_to_yield_len - ] - # Modify the chunk to contain only the new masked content - if ( - item.choices - and hasattr(item.choices[0], "delta") - and item.choices[0].delta - ): - item.choices[0].delta.content = new_masked_content - yielded_masked_text_len = safe_to_yield_len - yield item - else: - # Hold content by yielding empty content chunk (keeps metadata/structure) - if ( - item.choices - and hasattr(item.choices[0], "delta") - and item.choices[0].delta - ): - item.choices[0].delta.content = "" - yield item + yield item else: # Not a ModelResponseStream or no choices - yield as is yield item # Any remaining content (should have been handled by is_final, but just in case) - if yielded_masked_text_len < len(accumulated_full_text): + if any( + yielded_masked_text_len_by_choice.get(choice_index, 0) + < len(accumulated_text) + for choice_index, accumulated_text in accumulated_text_by_choice.items() + ): # We already reached the end of the generator pass except HTTPException: @@ -1973,6 +1985,11 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str = str(e) raise e finally: + detections = [ + detection + for choice_detections in latest_detections_by_choice.values() + for detection in choice_detections + ] self._count_masked_entities(detections, masked_entity_count) self._log_guardrail_information( request_data=request_data, diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py index be48991500..21c01da029 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py @@ -187,11 +187,28 @@ def _extract_user_text(messages: List) -> str: def _extract_response_text(response: Any) -> str: - """Extract text from LLM response object.""" + """Extract text from every LLM response choice.""" if hasattr(response, "choices") and response.choices: - choice = response.choices[0] - if hasattr(choice, "message") and choice.message: - return choice.message.content or "" + text_parts: List[str] = [] + for choice in response.choices: + if hasattr(choice, "message") and choice.message: + text = _content_to_text(choice.message.content) + if text: + text_parts.append(text) + return "\n".join(text_parts) + return "" + + +def _content_to_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + text_parts = [ + block.get("text") + for block in content + if isinstance(block, dict) and isinstance(block.get("text"), str) + ] + return " ".join(part for part in text_parts if part) return "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py index 5c374540e2..5acfc5403c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py @@ -480,21 +480,32 @@ class XecGuardGuardrail(CustomGuardrail): choices = response.get("choices") if not choices: return None - first = choices[0] - if hasattr(first, "message"): - message = first.message - elif isinstance(first, dict): - message = first.get("message") + text_parts: List[str] = [] + for choice in choices: + content = XecGuardGuardrail._extract_choice_content(choice) + text = XecGuardGuardrail._content_to_text(content) + if text: + text_parts.append(text) + return "\n".join(text_parts) or None + + @staticmethod + def _extract_choice_content(choice: Any) -> Any: + if hasattr(choice, "message"): + message = choice.message + elif isinstance(choice, dict): + message = choice.get("message") else: return None if message is None: return None if hasattr(message, "content"): - content = message.content - elif isinstance(message, dict): - content = message.get("content") - else: - return None + return message.content + if isinstance(message, dict): + return message.get("content") + return None + + @staticmethod + def _content_to_text(content: Any) -> Optional[str]: if isinstance(content, str) and content: return content if isinstance(content, list): diff --git a/tests/guardrails_tests/test_semantic_guard.py b/tests/guardrails_tests/test_semantic_guard.py index 44d8e417f5..a7e6230d02 100644 --- a/tests/guardrails_tests/test_semantic_guard.py +++ b/tests/guardrails_tests/test_semantic_guard.py @@ -171,6 +171,25 @@ class TestHelperFunctions: mock_response.choices[0].message.content = "Hello from LLM" assert _extract_response_text(mock_response) == "Hello from LLM" + def test_extract_response_text_combines_all_choices(self): + from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.semantic_guard import ( + _extract_response_text, + ) + + first_choice = MagicMock() + first_choice.message.content = "first response" + second_choice = MagicMock() + second_choice.message.content = [ + {"type": "text", "text": "second"}, + {"type": "text", "text": "response"}, + ] + mock_response = MagicMock() + mock_response.choices = [first_choice, second_choice] + + assert ( + _extract_response_text(mock_response) == "first response\nsecond response" + ) + def test_extract_response_text_empty(self): from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.semantic_guard import ( _extract_response_text, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py index 5e3cb76f44..94c7cefaeb 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py @@ -232,6 +232,52 @@ async def test_azure_text_moderation_guardrail_post_call_success_hook(): assert mock_async_make_request.call_args.kwargs["text"] == "Hello world" +@pytest.mark.asyncio +async def test_azure_text_moderation_guardrail_post_call_checks_all_choices(): + azure_text_moderation_guardrail = AzureContentSafetyTextModerationGuardrail( + guardrail_name="azure_text_moderation", + api_key="azure_text_moderation_api_key", + api_base="azure_text_moderation_api_base", + ) + with patch.object( + azure_text_moderation_guardrail, "async_make_request" + ) as mock_async_make_request: + mock_async_make_request.side_effect = [ + { + "blocklistsMatch": [], + "categoriesAnalysis": [{"category": "Hate", "severity": 0}], + }, + HTTPException( + status_code=400, + detail={"error": "blocked second choice"}, + ), + ] + + with pytest.raises(HTTPException): + await azure_text_moderation_guardrail.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth( + api_key="azure_text_moderation_api_key" + ), + response=ModelResponse( + choices=[ + Choices( + index=0, + message=Message(content="safe response"), + ), + Choices( + index=1, + message=Message(content="unsafe response"), + ), + ] + ), + ) + + assert [ + call.kwargs["text"] for call in mock_async_make_request.call_args_list + ] == ["safe response", "unsafe response"] + + @pytest.mark.asyncio async def test_azure_text_moderation_guardrail_post_call_streaming_hook(): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index f7377926ef..fb952d4b18 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -453,6 +453,71 @@ class TestContentFilterGuardrail: assert "[EMAIL_REDACTED]" in full_content assert "Contact me at [EMAIL_REDACTED] for info" in full_content + @pytest.mark.asyncio + async def test_streaming_hook_mask_checks_all_choices(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + patterns = [ + ContentFilterPattern( + pattern_type="prebuilt", + pattern_name="email", + action=ContentFilterAction.MASK, + ), + ] + + guardrail = ContentFilterGuardrail( + guardrail_name="test-streaming-mask-all-choices", + patterns=patterns, + event_hook=GuardrailEventHooks.during_call, + ) + + async def mock_stream(): + yield ModelResponseStream( + id="chunk1", + choices=[ + StreamingChoices( + delta=Delta(content="Contact first@ex"), + index=0, + ), + StreamingChoices( + delta=Delta(content="Email second@ex"), + index=1, + ), + ], + model="gpt-4", + ) + yield ModelResponseStream( + id="chunk2", + choices=[ + StreamingChoices( + delta=Delta(content="ample.com for help"), + index=0, + finish_reason="stop", + ), + StreamingChoices( + delta=Delta(content="ample.com for support"), + index=1, + finish_reason="stop", + ), + ], + model="gpt-4", + ) + + content_by_choice = {0: "", 1: ""} + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), + response=mock_stream(), + request_data={}, + ): + for choice in chunk.choices: + if choice.delta.content: + content_by_choice[choice.index] += choice.delta.content + + assert "first@example.com" not in content_by_choice[0] + assert "second@example.com" not in content_by_choice[1] + assert content_by_choice[0] == "Contact [EMAIL_REDACTED] for help" + assert content_by_choice[1] == "Email [EMAIL_REDACTED] for support" + @pytest.mark.asyncio async def test_streaming_hook_block(self): """ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py index b663544238..60c595c64e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py @@ -6,7 +6,6 @@ branch coverage. Network calls are always mocked; the companion live suite lives in ``test_xecguard_live.py``. """ -import asyncio import os from unittest.mock import MagicMock, patch @@ -1196,6 +1195,26 @@ class TestXecGuardMessageAssembly: is None ) + def test_extract_assistant_text_combines_all_choices(self, xecguard_guardrail): + assert ( + xecguard_guardrail._extract_assistant_text_from_response( + { + "choices": [ + {"message": {"content": "first response"}}, + { + "message": { + "content": [ + {"type": "text", "text": "second"}, + {"type": "text", "text": "response"}, + ] + } + }, + ] + } + ) + == "first response\nsecond\nresponse" + ) + def test_synthesize_user_inputs_not_dict(self, xecguard_guardrail): assert xecguard_guardrail._synthesize_user_from_inputs("not-dict") is None From abbefccad41b319a2077d5330c7cfdae1964119b Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Mon, 4 May 2026 21:27:24 +0000 Subject: [PATCH 11/11] fix(guardrails): align banned_keywords + azure_content_safety call_type gates with runtime route_type MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The hooks gated on ``call_type == "completion"`` but the proxy ingress passes ``route_type`` straight through as ``call_type`` — ``"acompletion"`` for /v1/chat/completions and ``"aresponses"`` for /v1/responses. Tests passed because they used the literal sync ``"completion"`` value, masking the gap. Switch both hooks to ``is_text_content_call_type`` (matches the canonical runtime values: completion / acompletion / aresponses) and update existing tests to assert against runtime values, plus parametrize a regression test that pins the gate. --- .../enterprise_hooks/banned_keywords.py | 8 +- litellm/proxy/guardrails/_content_utils.py | 26 +++- litellm/proxy/hooks/azure_content_safety.py | 8 +- .../guardrails/test_guardrail_coverage.py | 112 +++++++++++++++++- 4 files changed, 144 insertions(+), 10 deletions(-) diff --git a/enterprise/enterprise_hooks/banned_keywords.py b/enterprise/enterprise_hooks/banned_keywords.py index 621fde011d..47421c9605 100644 --- a/enterprise/enterprise_hooks/banned_keywords.py +++ b/enterprise/enterprise_hooks/banned_keywords.py @@ -11,7 +11,10 @@ from typing import Literal import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails._content_utils import iter_message_text +from litellm.proxy.guardrails._content_utils import ( + is_text_content_call_type, + iter_message_text, +) from litellm.integrations.custom_logger import CustomLogger from litellm._logging import verbose_proxy_logger from fastapi import HTTPException @@ -74,8 +77,7 @@ class _ENTERPRISE_BannedKeywords(CustomLogger): - check if user id part of blocked list """ self.print_verbose("Inside Banned Keyword List Pre-Call Hook") - if call_type == "completion": - # Covers multimodal list content + Responses-API input. + if is_text_content_call_type(call_type): for text in iter_message_text(data): self.test_violation(test_str=text) diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 8fdd8b23cb..7cad1352a7 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -8,7 +8,31 @@ skip the other shapes — these helpers normalise that so every hook sees every text fragment. """ -from typing import Any, Callable, Dict, Iterator, List +from typing import Any, Callable, Dict, FrozenSet, Iterator, List + + +# Call types whose body carries free-form chat / prompt text that +# text-content guardrails (banned keywords, content moderation, secret +# detection, …) should inspect. The proxy ingress passes ``route_type`` +# straight through as ``call_type``, so the literal values here are +# what the guardrail dispatcher actually receives: +# +# /v1/chat/completions -> "acompletion" +# /v1/responses -> "aresponses" +# +# ``"completion"`` is included for SDK / internal callers that invoke +# ``pre_call_hook`` directly with the sync name. Embedding, moderation, +# audio, and transcription endpoints are deliberately excluded — text +# guardrails on those paths are a separate scope. +TEXT_CONTENT_CALL_TYPES: FrozenSet[str] = frozenset( + {"completion", "acompletion", "aresponses"} +) + + +def is_text_content_call_type(call_type: str) -> bool: + """Return True if ``call_type`` carries free-form text that text + guardrails should inspect (Chat Completions or Responses API).""" + return call_type in TEXT_CONTENT_CALL_TYPES def _iter_text_parts_in_content(content: Any) -> Iterator[str]: diff --git a/litellm/proxy/hooks/azure_content_safety.py b/litellm/proxy/hooks/azure_content_safety.py index 844bbc11af..ddcd154059 100644 --- a/litellm/proxy/hooks/azure_content_safety.py +++ b/litellm/proxy/hooks/azure_content_safety.py @@ -8,7 +8,10 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache 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 ( + is_text_content_call_type, + iter_message_text, +) class _PROXY_AzureContentSafety( @@ -119,8 +122,7 @@ class _PROXY_AzureContentSafety( ): verbose_proxy_logger.debug("Inside Azure Content-Safety Pre-Call Hook") try: - if call_type == "completion": - # Covers multimodal list content + Responses-API input. + if is_text_content_call_type(call_type): for text in iter_message_text(data): await self.test_violation(content=text, source="input") diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py index 329f78bf50..6def548b93 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_coverage.py @@ -433,7 +433,13 @@ async def test_lasso_masking_writes_back_responses_api_input(user_api_key, monke def test_banned_keywords_blocks_multimodal_content(monkeypatch): - """VERIA-11: a banned word hidden in a multimodal text part is now caught.""" + """VERIA-11: a banned word hidden in a multimodal text part is now caught. + + Uses ``acompletion`` — the value the proxy ingress actually passes + for ``/v1/chat/completions``. Asserting against the literal sync + ``"completion"`` would pass even if the hook's call-type gate were + misaligned with the runtime, so the test wouldn't catch regressions. + """ monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False) from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords from fastapi import HTTPException @@ -455,7 +461,7 @@ def test_banned_keywords_blocks_multimodal_content(monkeypatch): } ] }, - call_type="completion", + call_type="acompletion", ) import asyncio @@ -477,7 +483,7 @@ def test_banned_keywords_blocks_responses_api_input(monkeypatch): user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"), cache=DualCache(), data={"input": "this contains forbidden content"}, - call_type="completion", + call_type="aresponses", ) import asyncio @@ -486,6 +492,59 @@ def test_banned_keywords_blocks_responses_api_input(monkeypatch): asyncio.run(_run()) +@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 + proxy actually emits — pinning a regression where the hook had + ``call_type == "completion"`` and silently no-op'd both + ``acompletion`` (chat completions) and ``aresponses`` (Responses API). + """ + monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False) + from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords + from fastapi import HTTPException + + guard = _ENTERPRISE_BannedKeywords() + + import asyncio + + with pytest.raises(HTTPException): + asyncio.run( + guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"), + cache=DualCache(), + data={ + "messages": [{"role": "user", "content": "forbidden text"}], + "input": "forbidden text", + }, + call_type=call_type, + ) + ) + + +def test_banned_keywords_skips_non_text_call_types(monkeypatch): + """Embedding / moderation / audio paths don't carry chat text and + aren't in the text-guardrail scope. They must not trigger the hook + even when the request body otherwise looks like a chat payload. + """ + monkeypatch.setattr("litellm.banned_keywords_list", ["forbidden"], raising=False) + from enterprise.enterprise_hooks.banned_keywords import _ENTERPRISE_BannedKeywords + + guard = _ENTERPRISE_BannedKeywords() + + import asyncio + + for call_type in ("aembedding", "amoderation", "aspeech", "atranscription"): + # Should return without raising, even though the data carries the banned word. + asyncio.run( + guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_id="u"), + cache=DualCache(), + data={"input": "forbidden text"}, + call_type=call_type, + ) + ) + + @pytest.mark.asyncio async def test_banned_keywords_post_call_checks_all_choices(monkeypatch, user_api_key): """Krrish blocker: ``n>1`` responses must not bypass post-call checks by @@ -515,6 +574,53 @@ async def test_banned_keywords_post_call_checks_all_choices(monkeypatch, user_ap # ── Azure Content Safety ────────────────────────────────────────────────────── +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type, data", + [ + ( + "acompletion", + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "scan me"}, + {"type": "image_url", "image_url": {"url": "..."}}, + ], + } + ] + }, + ), + ("aresponses", {"input": "scan me"}), + ], +) +async def test_azure_content_safety_pre_call_fires_on_runtime_call_types( + user_api_key, call_type, data +): + """The proxy ingress passes ``route_type`` straight through as + ``call_type`` — ``acompletion`` for chat completions and + ``aresponses`` for the Responses API. The hook must inspect text + fragments under both, not only the literal ``"completion"`` string + used by some SDK callers.""" + from litellm.proxy.hooks.azure_content_safety import _PROXY_AzureContentSafety + + guard = _PROXY_AzureContentSafety.__new__(_PROXY_AzureContentSafety) + seen = [] + + async def fake_test_violation(content, source=None): + seen.append((content, source)) + + guard.test_violation = fake_test_violation + await guard.async_pre_call_hook( + user_api_key_dict=user_api_key, + cache=DualCache(), + data=data, + call_type=call_type, + ) + assert ("scan me", "input") in seen + + @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