From b45f9124bcb612caa18f927293b9ddde7230cbd3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:12:19 -0700 Subject: [PATCH] fix(guardrails): list guardrails that scanned decision and extraction text, and scan every part the model reads Decision Model, Aim, Cato Networks, and Lasso add themselves to x-litellm-applied-guardrails when they scan a request, so the decisions routes keep listing them. Aim, Cato Networks, and Pangea no longer call their API when a request has no text for them. The systemone translation scans keys and numbers in state plus criteria descriptions, and decisions scans choice and level descriptions. Model Armor honors its mask flags and fail_on_error in apply_guardrail, and Pillar, DynamoAI, and IBM Guardrails list themselves only when they had text to scan. --- .../guardrail_translation/base_translation.py | 2 + .../guardrail_translation/handler.py | 145 +++++++++++++++-- .../guardrails/guardrail_hooks/aim/aim.py | 9 +- .../cato_networks/cato_networks.py | 9 +- .../decision_model/decision_model.py | 4 + .../guardrail_hooks/dynamoai/dynamoai.py | 9 +- .../ibm_guardrails/ibm_detector.py | 9 +- .../guardrails/guardrail_hooks/lasso/lasso.py | 4 + .../model_armor/model_armor.py | 40 +++-- .../guardrail_hooks/pangea/pangea.py | 2 + .../guardrail_hooks/pillar/pillar.py | 8 +- .../unified_guardrail/unified_guardrail.py | 5 +- .../test_decisions_guardrail_handler.py | 146 +++++++++++++++--- .../guardrail_hooks/dynamoai/test_dynamoai.py | 70 +++++++++ .../ibm_guardrails/test_ibm_detector.py | 85 ++++++++++ .../guardrails/guardrail_hooks/test_aim.py | 33 +++- .../guardrail_hooks/test_cato_networks.py | 32 ++++ .../guardrail_hooks/test_decision_model.py | 23 +++ .../guardrails/guardrail_hooks/test_lasso.py | 32 ++++ .../guardrail_hooks/test_model_armor.py | 128 +++++++++++++++ .../guardrails/guardrail_hooks/test_pangea.py | 37 +++++ .../guardrails/test_pillar_guardrails.py | 50 ++++++ .../test_proxy_logging_hook_detection.py | 15 ++ 23 files changed, 832 insertions(+), 65 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 104c82812e9..6b044a87521 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -74,6 +74,8 @@ class BaseTranslation(ABC): scans_responses: ClassVar[bool] = True + lists_guardrail_when_called: ClassVar[bool] = False + def post_call_hook_response(self, response: object) -> object: """The ``response`` this endpoint's non-streaming post-call hooks receive, derived from the object the translation stores under ``request_data["response"]`` while scanning an diff --git a/litellm/llms/openai/decisions/guardrail_translation/handler.py b/litellm/llms/openai/decisions/guardrail_translation/handler.py index 0244b44f688..76032483949 100644 --- a/litellm/llms/openai/decisions/guardrail_translation/handler.py +++ b/litellm/llms/openai/decisions/guardrail_translation/handler.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping +import json +from collections.abc import Iterator, Mapping from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final @@ -6,7 +7,7 @@ from typing import TYPE_CHECKING, ClassVar, Final from pydantic import JsonValue, TypeAdapter from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation -from litellm.proxy.guardrails._content_utils import iter_tool_json_text, map_tool_json_text +from litellm.llms.base_llm.guardrail_translation.utils import unappliable_request_rewrite from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -15,6 +16,7 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_DESCRIBED_ITEM_KEYS: Final = ("choices", "levels") _QUESTION_LIST_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, JsonValue], ...]]] = TypeAdapter( tuple[Mapping[str, JsonValue], ...] ) @@ -73,20 +75,123 @@ def _guardrailed_input(decision_input: JsonValue, guardrailed: Mapping[str, str] return [_guardrailed_message(message, guardrailed) for message in decision_input] -def _guardrailed_json(value: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue: - return map_tool_json_text(value, lambda text: guardrailed.get(text, text)) +def _json_text_forms(value: JsonValue) -> Iterator[str]: + match value: + case bool() | None: + return + case str(): + yield value + case int() | float(): + yield json.dumps(value) + case list(): + yield from chain.from_iterable(map(_json_text_forms, value)) + case dict(): + for key, item in value.items(): + yield key + yield from _json_text_forms(item) + + +def _json_texts(value: JsonValue) -> tuple[str, ...]: + return tuple(text for text in _json_text_forms(value) if text) + + +def _guardrailed_number(value: float, guardrailed: Mapping[str, str]) -> JsonValue: + text: Final = json.dumps(value) + rewritten: Final = guardrailed.get(text, text) + return value if rewritten == text else rewritten + + +def _guardrailed_json(value: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None) -> JsonValue: + match value: + case bool() | None: + return value + case str(): + return guardrailed.get(value, value) + case int() | float(): + return _guardrailed_number(value, guardrailed) + case list(): + return [_guardrailed_json(item, guardrailed, guardrail_name) for item in value] + case dict(): + keys: Final = tuple(guardrailed.get(key, key) for key in value) + if len(frozenset(keys)) != len(keys): + raise unappliable_request_rewrite(guardrail_name) + return { + key: _guardrailed_json(item, guardrailed, guardrail_name) + for key, item in zip(keys, value.values(), strict=True) + } def _instruction_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: - return tuple(iter_tool_json_text(question.get("instructions"))) + return _json_texts(question.get("instructions")) + + +def _described_items(question: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + return tuple(chain.from_iterable(_list_items(question.get(key)) for key in _DESCRIBED_ITEM_KEYS)) + + +def _list_items(value: JsonValue) -> tuple[JsonValue, ...]: + return tuple(value) if isinstance(value, list) else () + + +def _item_description(item: JsonValue) -> tuple[str, ...]: + description: Final = item.get("description") if isinstance(item, dict) else None + return (description,) if isinstance(description, str) and description else () + + +def _decision_question_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: + descriptions: Final = chain.from_iterable(map(_item_description, _described_items(question))) + return (*_instruction_texts(question), *descriptions) + + +def _guardrailed_item(item: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue: + if not isinstance(item, dict) or not isinstance(description := item.get("description"), str): + return item + return {**item, "description": guardrailed.get(description, description)} + + +def _with_guardrailed_decision_question( + question: Mapping[str, JsonValue], guardrailed: Mapping[str, str], guardrail_name: str | None +) -> Mapping[str, JsonValue]: + described: Final = { + key: [_guardrailed_item(item, guardrailed) for item in _list_items(question[key])] + for key in _DESCRIBED_ITEM_KEYS + if key in question + } + return {**_with_guardrailed_instructions(question, guardrailed, guardrail_name), **described} + + +def _criteria_texts(criteria: JsonValue) -> tuple[str, ...]: + descriptions: Final = tuple(criteria.values()) if isinstance(criteria, dict) else (criteria,) + return tuple(chain.from_iterable(map(_json_texts, descriptions))) + + +def _systemone_question_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: + return (*_instruction_texts(question), *_criteria_texts(question.get("criteria"))) + + +def _guardrailed_criteria(criteria: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None) -> JsonValue: + if not isinstance(criteria, dict): + return _guardrailed_json(criteria, guardrailed, guardrail_name) + return { + label: _guardrailed_json(description, guardrailed, guardrail_name) for label, description in criteria.items() + } + + +def _with_guardrailed_systemone_question( + question: Mapping[str, JsonValue], guardrailed: Mapping[str, str], guardrail_name: str | None +) -> Mapping[str, JsonValue]: + with_instructions: Final = _with_guardrailed_instructions(question, guardrailed, guardrail_name) + if "criteria" not in question: + return with_instructions + return {**with_instructions, "criteria": _guardrailed_criteria(question["criteria"], guardrailed, guardrail_name)} def _with_guardrailed_instructions( - question: Mapping[str, JsonValue], guardrailed: Mapping[str, str] + question: Mapping[str, JsonValue], guardrailed: Mapping[str, str], guardrail_name: str | None ) -> Mapping[str, JsonValue]: if "instructions" not in question: return question - return {**question, "instructions": _guardrailed_json(question["instructions"], guardrailed)} + return {**question, "instructions": _guardrailed_json(question["instructions"], guardrailed, guardrail_name)} async def _guardrailed_texts( @@ -103,6 +208,14 @@ async def _guardrailed_texts( input_type="request", logging_obj=litellm_logging_obj, ) + if not guardrail_to_apply.records_own_guardrail_information: + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict + ) + + add_guardrail_to_applied_guardrails_header( + request_data=request_data, guardrail_name=guardrail_to_apply.guardrail_name + ) returned: Final = result.get("texts", []) return MappingProxyType( {text: returned[index] if index < len(returned) else text for index, text in enumerate(unique_texts)} @@ -111,6 +224,7 @@ async def _guardrailed_texts( class _DecisionAnswersTranslation(BaseTranslation): scans_responses: ClassVar[bool] = False + lists_guardrail_when_called: ClassVar[bool] = True async def process_output_response( self, @@ -132,14 +246,17 @@ class DecisionsHandler(_DecisionAnswersTranslation): ) -> Mapping[str, object]: decision_input: Final = _JSON_ADAPTER.validate_python(data.get("input")) questions: Final = _QUESTION_LIST_ADAPTER.validate_python(data.get("questions") or ()) - texts: Final = (*_input_texts(decision_input), *chain.from_iterable(map(_instruction_texts, questions))) + texts: Final = (*_input_texts(decision_input), *chain.from_iterable(map(_decision_question_texts, questions))) if not texts: return data guardrailed: Final = await _guardrailed_texts(texts, guardrail_to_apply, data, litellm_logging_obj) return { **data, "input": _guardrailed_input(decision_input, guardrailed), - "questions": [_with_guardrailed_instructions(question, guardrailed) for question in questions], + "questions": [ + _with_guardrailed_decision_question(question, guardrailed, guardrail_to_apply.guardrail_name) + for question in questions + ], } @@ -152,17 +269,15 @@ class SystemOneHandler(_DecisionAnswersTranslation): ) -> Mapping[str, object]: state: Final = _JSON_ADAPTER.validate_python(data.get("state")) questions: Final = _QUESTION_MAP_ADAPTER.validate_python(data.get("questions") or {}) - texts: Final = ( - *iter_tool_json_text(state), - *chain.from_iterable(map(_instruction_texts, questions.values())), - ) + texts: Final = (*_json_texts(state), *chain.from_iterable(map(_systemone_question_texts, questions.values()))) if not texts: return data guardrailed: Final = await _guardrailed_texts(texts, guardrail_to_apply, data, litellm_logging_obj) return { **data, - "state": _guardrailed_json(state, guardrailed), + "state": _guardrailed_json(state, guardrailed, guardrail_to_apply.guardrail_name), "questions": { - name: _with_guardrailed_instructions(question, guardrailed) for name, question in questions.items() + name: _with_guardrailed_systemone_question(question, guardrailed, guardrail_to_apply.guardrail_name) + for name, question in questions.items() }, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 61117fbc55e..b564a1df041 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -23,6 +23,9 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict +) from litellm.proxy.guardrails._content_utils import ( apply_redacted_messages_back, build_inspection_messages, @@ -169,6 +172,10 @@ class AimGuardrail(CustomGuardrail): return data async def call_aim_guardrail(self, data: dict, hook: str, key_alias: str | None) -> dict: + messages: Final = self._build_aim_inspection_messages(data) # pyright: ignore[reportUnknownMemberType] # hook request data is an untyped dict + if not messages: + return data # pyright: ignore[reportUnknownVariableType] # hook request data is an untyped dict + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) user_email: Final = data.get("metadata", {}).get("headers", {}).get("x-aim-user-email") call_id: Final = data.get("litellm_call_id") headers: Final = self._build_aim_headers( @@ -180,7 +187,7 @@ class AimGuardrail(CustomGuardrail): response: Final = await self.async_handler.post( f"{self.api_base}/fw/v1/analyze", headers=headers, - json={"messages": self._build_aim_inspection_messages(data)}, + json={"messages": messages}, timeout=self.timeout, ) response.raise_for_status() diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index 19b9bccd8aa..be7d967836f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -30,6 +30,9 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict +) from litellm.proxy.guardrails._content_utils import ( apply_redacted_messages_back, build_inspection_messages, @@ -295,6 +298,10 @@ class CatoNetworksGuardrail(CustomGuardrail): key_alias: str | None, user_email: str | None = None, ) -> dict: + messages: Final = self._inspection_messages(data) # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType] # hook request data is an untyped dict + if not messages: + return data # pyright: ignore[reportUnknownVariableType] # hook request data is an untyped dict + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) call_id: Final = data.get("litellm_call_id") headers: Final = self._build_cato_headers( hook=hook, @@ -305,7 +312,7 @@ class CatoNetworksGuardrail(CustomGuardrail): response: Final = await self.async_handler.post( f"{self.api_base}/fw/v1/analyze", headers=headers, - json={"messages": self._inspection_messages(data)}, + json={"messages": messages}, timeout=self.timeout, ) response.raise_for_status() diff --git a/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py b/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py index f3bd2b3f889..11db4690cde 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py +++ b/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py @@ -20,6 +20,9 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped upstream ) from litellm.litellm_core_utils.llm_judge import attached_auth, caller_team_for, default_router_provider, judge_target +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict +) from litellm.types.decisions import MAX_DECISION_QUESTIONS, OpenAIDecisionResponse, OpenAIPredicateAnswer from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.llms.openai import ChatCompletionAssistantToolCall, ChatCompletionToolCallChunk @@ -347,6 +350,7 @@ class DecisionModelGuardrail(CustomGuardrail): status: Final = "guardrail_intervened" if intervened else ("guardrail_flagged" if flagged else "success") self._log(request_data, verdicts, status, start_time) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) if intervened: raise HTTPException( status_code=400, diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index d507df24faa..6974c33b0bf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -328,8 +328,7 @@ class DynamoAIGuardrails(CustomGuardrail): error_message: Final = self._create_error_message(processed_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -373,8 +372,7 @@ class DynamoAIGuardrails(CustomGuardrail): error_message: Final = self._create_error_message(processed_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -441,8 +439,7 @@ class DynamoAIGuardrails(CustomGuardrail): error_message: Final = self._create_error_message(processed_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) async def async_post_call_streaming_iterator_hook( self, 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 83dd90dafb4..ff4ae15b6f1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -476,8 +476,7 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -545,8 +544,7 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -641,8 +639,7 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) async def async_post_call_streaming_iterator_hook( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 97b35769f63..ca06b3c1675 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -45,6 +45,9 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict +) from litellm.proxy.guardrails._content_utils import ( build_inspection_messages, has_non_string_content, @@ -409,6 +412,7 @@ class LassoGuardrail(CustomGuardrail): messages: Final[Sequence[Mapping[str, object]]] = [*expanded_messages, *input_messages] if not messages: return data + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) # Lasso's classifix endpoint returns masked text that we copy back # into ``data["messages"]``. For multimodal/Responses-API input we diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index b9871b1e276..04f495506a7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,6 +1,8 @@ +import asyncio import time from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence from enum import Enum, auto +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Protocol, runtime_checkable import httpx @@ -291,7 +293,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return _redact_scanned_content(armor_response) return armor_response - def _raise_if_fail_closed(self, e: ModelArmorAPIError) -> None: + def _raise_if_fail_closed(self, e: ModelArmorAPIError | httpx.HTTPError) -> None: if self.optional_params.get("fail_on_error", True): raise e from None @@ -1302,17 +1304,23 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: - content: Final = "\n".join(text for text in inputs.get("texts") or () if text) + texts: Final = tuple(inputs.get("texts") or ()) + content: Final = "\n".join(text for text in texts if text) if not content: return inputs source: Final[Literal["user_prompt", "model_response"]] = ( "user_prompt" if input_type == "request" else "model_response" ) + masking: Final = self.mask_request_content if input_type == "request" else self.mask_response_content + scanned: Final = tuple(dict.fromkeys(text for text in texts if text)) if masking else (content,) start_time: Final = time.time() try: - armor_response: Final = await self.make_model_armor_request( - content=content, source=source, request_data=request_data + armor_responses: Final = await asyncio.gather( + *( + self.make_model_armor_request(content=text, source=source, request_data=request_data) + for text in scanned + ) ) except (ModelArmorAPIError, httpx.HTTPError) as e: error_end_time: Final = time.time() @@ -1325,12 +1333,19 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=error_end_time, duration=error_end_time - start_time, ) + self._raise_if_fail_closed(e) return inputs - flagged: Final = self._should_block_content(armor_response, allow_sanitization=False) + flagged: Final = tuple( + response for response in armor_responses if self._should_block_content(response, allow_sanitization=masking) + ) end_time: Final = time.time() self.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response=self._build_logging_response(armor_response), + guardrail_json_response=( + tuple(self._build_logging_response(response) for response in armor_responses) + if masking + else self._build_logging_response(armor_responses[0]) + ), request_data=request_data, guardrail_status="guardrail_flagged" if flagged else "success", guardrail_provider="model_armor", @@ -1338,15 +1353,22 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=end_time, duration=end_time - start_time, ) - if flagged and not self._event_hook_is_event_type(GuardrailEventHooks.logging_only): + if self._event_hook_is_event_type(GuardrailEventHooks.logging_only): + return inputs + if flagged: raise HTTPException( status_code=400, detail=self._build_block_error_detail( "Response blocked by Model Armor" if input_type == "response" else "Content blocked by Model Armor", - armor_response, + flagged[0], ), ) - return inputs + if not masking: + return inputs + sanitized: Final = MappingProxyType( + {text: self._get_sanitized_content(response) or text for text, response in zip(scanned, armor_responses)} + ) + return {**inputs, "texts": [sanitized.get(text, text) for text in texts]} @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 2194238cea0..808bb68acf3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -167,6 +167,8 @@ class PangeaHandler(CustomGuardrail): messages = transformer.get_messages() else: messages = data.get("messages") + if not messages: + return ai_guard_payload: Final = { "debug": False, diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 3f025cfc8b3..1983be794b3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -351,8 +351,8 @@ class PillarGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook") result: Final = await self.run_pillar_guardrail(data, user_api_key_dict) - # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if data.get("messages"): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return result @@ -394,8 +394,8 @@ class PillarGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook") result: Final = await self.run_pillar_guardrail(data, user_api_key_dict) - # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if data.get("messages"): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return result diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index a147580a1f9..3aff4779235 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -284,7 +284,8 @@ class UnifiedLLMGuardrails(CustomLogger): except ValueError: return data # handle unmapped call types - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + translation: Final = mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation(translation) _ensure_litellm_metadata(data, user_api_key_dict) @@ -296,7 +297,7 @@ class UnifiedLLMGuardrails(CustomLogger): litellm_logging_obj=data.get("litellm_logging_obj"), ) - if not guardrail_to_apply.records_own_guardrail_information: + if not guardrail_to_apply.records_own_guardrail_information and not translation.lists_guardrail_when_called: add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=guardrail_to_apply.guardrail_name ) diff --git a/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py b/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py index a2993ac1cea..b84769dc8b9 100644 --- a/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py +++ b/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py @@ -5,6 +5,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms import get_guardrail_translation_mapping +from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite from litellm.llms.openai.decisions.guardrail_translation.handler import DecisionsHandler, SystemOneHandler from litellm.types.decisions import OpenAIDecisionResponse from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -27,6 +28,22 @@ class _TaggingGuardrail(CustomGuardrail): return GenericGuardrailAPIInputs(texts=[f"{text} [GUARDRAILED]" for text in texts]) +class _MasksSecrets(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + secrets: Final = {"AKIAIOSFODNN7EXAMPLE": "", "AKIAI44QH8DHBEXAMPLE": ""} + return GenericGuardrailAPIInputs( + texts=[ + "" if text == "4111111111111111" else secrets.get(text, text) for text in inputs.get("texts", []) + ] + ) + + class _ReturnsNoTexts(CustomGuardrail): async def apply_guardrail( self, @@ -52,7 +69,7 @@ def test_decision_call_types_resolve_to_their_handler(call_type: CallTypes, hand @pytest.mark.asyncio -async def test_decision_message_text_parts_and_question_instructions_are_guardrailed() -> None: +async def test_decision_message_text_parts_and_question_text_are_guardrailed() -> None: guardrail: Final = _TaggingGuardrail() image: Final = {"type": "input_image", "image_url": "data:image/png;base64,AA=="} data: Final = { @@ -64,12 +81,29 @@ async def test_decision_message_text_parts_and_question_instructions_are_guardra "questions": [ {"type": "predicate", "name": "urgent", "instructions": "Is it urgent?"}, {"type": "box", "name": "logo", "instructions": "ticket one"}, + { + "type": "choice", + "name": "team", + "instructions": "Which team?", + "choices": [{"value": "billing", "description": "Payment issues"}, {"value": True}], + }, + { + "type": "score", + "name": "mood", + "instructions": "How upset?", + "levels": [{"label": "calm"}, {"label": "angry", "description": "Uses caps"}], + }, ], } result: Final = await DecisionsHandler().process_input_messages(data, guardrail) - assert guardrail.seen == [("request", ("ticket one", "ticket two", "Is it urgent?"))] + assert guardrail.seen == [ + ( + "request", + ("ticket one", "ticket two", "Is it urgent?", "Which team?", "Payment issues", "How upset?", "Uses caps"), + ) + ] assert {key: value for key, value in result.items() if key != "metadata"} == { "model": "judge", "input": [ @@ -79,48 +113,111 @@ async def test_decision_message_text_parts_and_question_instructions_are_guardra "questions": [ {"type": "predicate", "name": "urgent", "instructions": "Is it urgent? [GUARDRAILED]"}, {"type": "box", "name": "logo", "instructions": "ticket one [GUARDRAILED]"}, + { + "type": "choice", + "name": "team", + "instructions": "Which team? [GUARDRAILED]", + "choices": [{"value": "billing", "description": "Payment issues [GUARDRAILED]"}, {"value": True}], + }, + { + "type": "score", + "name": "mood", + "instructions": "How upset? [GUARDRAILED]", + "levels": [{"label": "calm"}, {"label": "angry", "description": "Uses caps [GUARDRAILED]"}], + }, ], } assert data["input"][0]["content"][0]["text"] == "ticket one" @pytest.mark.asyncio -async def test_systemone_state_leaves_and_question_instructions_are_guardrailed() -> None: +async def test_systemone_state_keys_scalars_and_question_text_are_guardrailed() -> None: guardrail: Final = _TaggingGuardrail() - criteria_question: Final = {"type": "noul", "criteria": {"yes": "refund"}} data: Final = { "model": "judge", "state": {"subject": "refund", "lines": ["order 7", "", 3, True], "meta": {"id": None}}, "questions": { "refund": {"type": "noul", "instructions": {"ask": "Does the customer want a refund?"}}, - "routed": criteria_question, + "routed": {"type": "noul", "criteria": {"true": "The customer asks for money back", "false": None}}, + "tier": {"type": "score", "criteria": ["calm", {"level": "angry"}]}, }, } result: Final = await SystemOneHandler().process_input_messages(data, guardrail) - assert guardrail.seen == [("request", ("refund", "order 7", "Does the customer want a refund?"))] + assert guardrail.seen == [ + ( + "request", + ( + "subject", + "refund", + "lines", + "order 7", + "3", + "meta", + "id", + "ask", + "Does the customer want a refund?", + "The customer asks for money back", + "calm", + "level", + "angry", + ), + ) + ] assert {key: value for key, value in result.items() if key != "metadata"} == { "model": "judge", "state": { - "subject": "refund [GUARDRAILED]", - "lines": ["order 7 [GUARDRAILED]", "", 3, True], - "meta": {"id": None}, + "subject [GUARDRAILED]": "refund [GUARDRAILED]", + "lines [GUARDRAILED]": ["order 7 [GUARDRAILED]", "", "3 [GUARDRAILED]", True], + "meta [GUARDRAILED]": {"id [GUARDRAILED]": None}, }, "questions": { - "refund": {"type": "noul", "instructions": {"ask": "Does the customer want a refund? [GUARDRAILED]"}}, - "routed": criteria_question, + "refund": { + "type": "noul", + "instructions": {"ask [GUARDRAILED]": "Does the customer want a refund? [GUARDRAILED]"}, + }, + "routed": { + "type": "noul", + "criteria": {"true": "The customer asks for money back [GUARDRAILED]", "false": None}, + }, + "tier": {"type": "score", "criteria": ["calm [GUARDRAILED]", {"level [GUARDRAILED]": "angry [GUARDRAILED]"}]}, }, } +@pytest.mark.asyncio +async def test_systemone_secret_key_and_card_number_are_masked_before_the_provider() -> None: + data: Final = { + "model": "judge", + "state": {"AKIAIOSFODNN7EXAMPLE": "", "card": 4111111111111111, "items": 2}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + result: Final = await SystemOneHandler().process_input_messages(data, _MasksSecrets(guardrail_name="mask")) + + assert result["state"] == {"": "", "card": "", "items": 2} + + +@pytest.mark.asyncio +async def test_systemone_rejects_masked_keys_that_collide() -> None: + data: Final = { + "model": "judge", + "state": {"AKIAIOSFODNN7EXAMPLE": "first", "AKIAI44QH8DHBEXAMPLE": "second"}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + with pytest.raises(UnappliableRequestRewrite): + await SystemOneHandler().process_input_messages(data, _MasksSecrets(guardrail_name="mask")) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("handler", "data"), ( ( SystemOneHandler, - {"model": "judge", "state": {"count": 3}, "questions": {"q": {"type": "noul", "criteria": {}}}}, + {"model": "judge", "state": {}, "questions": {"q": {"type": "noul", "criteria": {}}}}, ), ( DecisionsHandler, @@ -144,14 +241,23 @@ async def test_request_without_text_is_returned_unscanned( @pytest.mark.asyncio -async def test_guardrail_that_returns_no_texts_leaves_the_request_text_as_sent() -> None: - data: Final = { - "model": "judge", - "input": "ticket", - "questions": [{"type": "predicate", "instructions": "Is it urgent?"}], - } - - result: Final = await DecisionsHandler().process_input_messages(data, _ReturnsNoTexts(guardrail_name="noop")) +@pytest.mark.parametrize( + ("handler", "data"), + ( + ( + DecisionsHandler, + {"model": "judge", "input": "ticket", "questions": [{"type": "predicate", "instructions": "Is it urgent?"}]}, + ), + ( + SystemOneHandler, + {"model": "judge", "state": {"total": 3, "paid": 2.5}, "questions": {"q": {"type": "noul", "criteria": {}}}}, + ), + ), +) +async def test_guardrail_that_returns_no_texts_leaves_the_request_text_as_sent( + handler: type[DecisionsHandler] | type[SystemOneHandler], data: dict[str, object] +) -> None: + result: Final = await handler().process_input_messages(data, _ReturnsNoTexts(guardrail_name="noop")) assert result == data diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py index 80bde77afdc..7208136f69a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py @@ -203,6 +203,76 @@ async def test_post_call_hook_skips_dynamoai_when_the_reply_has_no_text(respx_mo assert route.called is False + +_INPUT_HOOKS = [ + pytest.param( + "pre_call", + lambda guardrail, data: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion" + ), + id="pre_call", + ), + pytest.param( + "during_call", + lambda guardrail, data: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + ), + id="during_call", + ), +] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) +@pytest.mark.parametrize( + ("request_fields", "applied"), + [ + pytest.param({"messages": [{"role": "user", "content": "hi"}]}, ["test-dynamoai"], id="messages"), + pytest.param({"state": "hi"}, None, id="no-messages"), + ], +) +async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( + respx_mock: respx.MockRouter, event_hook, run_hook, request_fields, applied +) -> None: + respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []}) + guardrail = DynamoAIGuardrails( + guardrail_name="test-dynamoai", + api_key="test-api-key", + api_base="https://api.dynamo.ai", + event_hook=event_hook, + default_on=True, + ) + metadata: dict[str, object] = {} + + await run_hook(guardrail, {**request_fields, "metadata": metadata}) + + assert metadata.get("applied_guardrails") == applied + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("response", "applied"), + [ + pytest.param(_assistant_reply("Hello, how are you?"), ["test-dynamoai"], id="chat-reply"), + pytest.param({"answers": {"aws_key": {"type": "text", "value": None}}}, None, id="no-choices"), + ], +) +async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text( + respx_mock: respx.MockRouter, response, applied +) -> None: + respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []}) + metadata: dict[str, object] = {} + + await _post_call_guardrail().async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hi"}], "metadata": metadata}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert metadata.get("applied_guardrails") == applied + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py index 4a591c08bf3..912021e4bcd 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py @@ -3,6 +3,8 @@ import json import pytest import respx +import litellm + from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.ibm_guardrails.ibm_detector import IBMGuardrailDetector @@ -122,3 +124,86 @@ async def test_orchestrator_detection_blocks_the_request_and_names_the_detector( ) assert "- HAP (detector: hap-detector, score: 0.970)\n Text: 'hello'" in str(blocked.value) + + +_INPUT_HOOKS = [ + pytest.param( + "pre_call", + lambda guardrail, data: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion" + ), + id="pre_call", + ), + pytest.param( + "during_call", + lambda guardrail, data: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + ), + id="during_call", + ), +] +_REQUEST_TEXT_CASES = [ + pytest.param({"messages": [{"role": "user", "content": "hello there"}]}, ["ibm-guard"], id="messages"), + pytest.param({"state": "hello there"}, None, id="no-messages"), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) +@pytest.mark.parametrize(("request_fields", "applied"), _REQUEST_TEXT_CASES) +async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( + respx_mock: respx.MockRouter, event_hook, run_hook, request_fields, applied +) -> None: + respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]]) + guardrail = IBMGuardrailDetector( + guardrail_name="ibm-guard", + auth_token="ibm-token", + base_url=_BASE_URL, + detector_id="hap-detector", + is_detector_server=True, + event_hook=event_hook, + default_on=True, + ) + metadata: dict[str, object] = {} + + await run_hook(guardrail, {**request_fields, "metadata": metadata}) + + assert metadata.get("applied_guardrails") == applied + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "applied"), + [ + pytest.param( + litellm.ModelResponse( + choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}] + ), + ["ibm-guard"], + id="chat-reply", + ), + pytest.param({"answers": {"aws_key": {"type": "text", "value": None}}}, None, id="no-choices"), + ], +) +async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text( + respx_mock: respx.MockRouter, response, applied +) -> None: + respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]]) + guardrail = IBMGuardrailDetector( + guardrail_name="ibm-guard", + auth_token="ibm-token", + base_url=_BASE_URL, + detector_id="hap-detector", + is_detector_server=True, + event_hook="post_call", + default_on=True, + ) + metadata: dict[str, object] = {} + + await guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hello"}], "metadata": metadata}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py index 276fcb86e4d..8c7f0d451e7 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py @@ -316,7 +316,7 @@ async def test_aim_anonymize_action_blocks_malformed_redacted_messages( ) assert exc_info.value.code == "400" - assert data == request_body + assert {key: value for key, value in data.items() if key != "metadata"} == request_body _OUTPUT_REQUEST = { @@ -815,3 +815,34 @@ async def test_aim_streaming_block_stops_after_verified_chunks(monkeypatch: pyte {"choices": [{"delta": {"content": "A"}}]}, {"done": True}, ] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("call_type", "request_fields", "aim_calls", "applied"), + [ + pytest.param( + "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["aim"], id="decisions" + ), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + ], +) +async def test_aim_calls_aim_and_lists_itself_only_when_the_request_has_text( + respx_mock: respx.MockRouter, call_type: str, request_fields: dict, aim_calls: int, applied: list[str] | None +): + analyze = respx_mock.post("http://aim.test/fw/v1/analyze").respond( + json={"analysis_result": {"analysis_time_ms": 1, "policy_drill_down": {}}, "required_action": None} + ) + guardrail = AimGuardrail(api_key="hs-aim-key", api_base="http://aim.test", guardrail_name="aim") + metadata: dict[str, object] = {} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data={**request_fields, "metadata": metadata}, + call_type=call_type, + ) + + assert analyze.call_count == aim_calls + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py index 349030c6c75..66a8f5a85a8 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py @@ -4,6 +4,7 @@ import ssl from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx from fastapi.exceptions import HTTPException from httpx import Request, Response from websockets.exceptions import ConnectionClosed @@ -2785,3 +2786,34 @@ async def test_cato_still_inspects_every_conversational_call_type(call_type: str ) mock_post.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("call_type", "request_fields", "cato_calls", "applied"), + [ + pytest.param( + "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["cato"], id="decisions" + ), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + ], +) +async def test_cato_calls_cato_and_lists_itself_only_when_the_request_has_text( + respx_mock, call_type, request_fields, cato_calls, applied +): + analyze = respx_mock.post("http://cato.test/fw/v1/analyze").respond( + json={"analysis_result": {}, "required_action": None} + ) + guardrail = CatoNetworksGuardrail(api_key="cato-key", api_base="http://cato.test", guardrail_name="cato") + metadata = {} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data={**request_fields, "metadata": metadata}, + call_type=call_type, + ) + + assert analyze.call_count == cato_calls + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py index 660cc42a07d..ae6ee15d64a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py @@ -1307,3 +1307,26 @@ async def test_tool_call_in_both_tool_calls_and_structured_messages_is_screened_ await guardrail.apply_guardrail(inputs, _request_data(), "request") assert [call.kwargs["input"] for call in router.adecisions.await_args_list] == ['send_email({"body": "hello"})'] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("texts", "unreachable", "applied"), + [ + pytest.param(["hello"], False, ["test_decision_model"], id="scanned"), + pytest.param([], False, None, id="no-text"), + pytest.param(["hello"], True, None, id="unreachable-fail-open"), + ], +) +async def test_decision_model_lists_itself_only_when_it_returned_verdicts( + texts: list[str], unreachable: bool, applied: list[str] | None +): + router: Final = _decision_router(probability=0.1) + if unreachable: + router.adecisions.side_effect = httpx.ConnectError("unreachable") + guardrail: Final = _make_guardrail(router_provider=lambda: router, unreachable_fallback="fail_open") + request_data: Final = _request_data() + + await guardrail.apply_guardrail({"texts": texts}, request_data, "request") + + assert cast(dict[str, object], request_data["metadata"]).get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py index b8a219ba863..dec3510d6ab 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -1549,3 +1549,35 @@ async def test_api_error_handling_legacy_messages( assert "Failed to verify request safety with Lasso API" in str(excinfo.value) assert error_message in str(excinfo.value) assert route.called + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("call_type", "request_fields", "applied"), + [ + pytest.param( + "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, ["lasso"], id="decisions" + ), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, None, id="extraction"), + ], +) +async def test_lasso_lists_itself_only_when_it_classified_request_text( + respx_mock: respx.MockRouter, call_type: str, request_fields: dict[str, object], applied: list[str] | None +) -> None: + respx_mock.post("https://server.lasso.security/gateway/v3/classify").respond( + json={"deputies": {}, "findings": {}, "violations_detected": False} + ) + guardrail: Final = LassoGuardrail( + lasso_api_key="test-api-key", guardrail_name="lasso", event_hook="pre_call", default_on=True + ) + metadata: Final[dict[str, object]] = {} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data={**request_fields, "metadata": metadata}, + call_type=call_type, + ) + + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index 5aa3c7dfc17..7cc553e2aee 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -5291,6 +5291,134 @@ async def test_logging_only_flagged_prompt_still_scans_response(): assert len(flagged) == 2 +class _RedactsPhoneNumbers(AsyncHTTPHandler): + def __init__(self): + self.scanned: list[str] = [] + + async def post( + self, + url: str, + json: dict[str, object] | None = None, + headers: dict[str, str] | None = None, + **kwargs: object, + ) -> httpx.Response: + payload = cast(dict[str, dict[str, str]], json) + text = (payload.get("userPromptData") or payload["modelResponseData"])["text"] + self.scanned.append(text) + redacted = text.replace("555-123-4567", "[REDACTED]") + result = ( + _clean_armor_response() + if redacted == text + else { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": {"matchState": "MATCH_FOUND", "data": {"text": redacted}} + } + } + }, + } + } + ) + return httpx.Response(200, json=result, request=httpx.Request("POST", url)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ("request", "response")) +async def test_apply_guardrail_masks_each_text_when_masking_is_on(input_type): + handler = _RedactsPhoneNumbers() + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-mask", + event_hook=GuardrailEventHooks.pre_call, + mask_request_content=input_type == "request", + mask_response_content=input_type == "response", + async_handler=handler, + access_token_provider=_async_token_provider, + ) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["call me at 555-123-4567", "refund order 4411"]}, + request_data={"metadata": {}}, + input_type=input_type, + ) + + assert result["texts"] == ["call me at [REDACTED]", "refund order 4411"] + assert sorted(handler.scanned) == ["call me at 555-123-4567", "refund order 4411"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_a_masking_finding_when_masking_is_off(): + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-block", + event_hook=GuardrailEventHooks.pre_call, + async_handler=_RedactsPhoneNumbers(), + access_token_provider=_async_token_provider, + ) + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["call me at 555-123-4567", "refund order 4411"]}, + request_data={"metadata": {}}, + input_type="request", + ) + + assert exc_info.value.status_code == 400 + + + +@pytest.mark.asyncio +async def test_apply_guardrail_returns_clean_text_unchanged_when_masking_is_off(): + inputs = {"texts": ["refund order 4411"]} + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-clean", + event_hook=GuardrailEventHooks.pre_call, + async_handler=_RedactsPhoneNumbers(), + access_token_provider=_async_token_provider, + ) + + assert await guardrail.apply_guardrail(inputs=inputs, request_data={"metadata": {}}, input_type="request") is inputs + +class _UnavailableArmor(AsyncHTTPHandler): + def __init__(self): + pass + + async def post(self, url: str, **kwargs: object) -> httpx.Response: + return httpx.Response(503, text="unavailable", request=httpx.Request("POST", url)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_on_error", (True, False)) +async def test_apply_guardrail_honors_fail_on_error_when_model_armor_is_down(fail_on_error): + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-down", + event_hook=GuardrailEventHooks.pre_call, + fail_on_error=fail_on_error, + async_handler=_UnavailableArmor(), + access_token_provider=_async_token_provider, + ) + inputs = {"texts": ["refund order 4411"]} + + if fail_on_error: + with pytest.raises(ModelArmorAPIError): + await guardrail.apply_guardrail(inputs=inputs, request_data={"metadata": {}}, input_type="request") + else: + assert await guardrail.apply_guardrail(inputs=inputs, request_data={"metadata": {}}, input_type="request") == inputs + + @pytest.mark.asyncio async def test_apply_guardrail_raises_on_flagged_when_not_logging_only(): """The /guardrails/apply_guardrail endpoint calls apply_guardrail directly; a diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py index 4e31a3c2055..1a61dc06131 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py @@ -2,8 +2,11 @@ from unittest.mock import AsyncMock, patch import httpx import pytest +import respx from fastapi import HTTPException +from litellm import DualCache +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.pangea import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.pangea.pangea import ( PangeaGuardrailMissingSecrets, @@ -368,3 +371,37 @@ async def test_pangea_ai_guard_response_transformed(pangea_guardrail): ) assert response.choices[0]["message"]["content"] == "Yes, here is an SSN: " + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("call_type", "request_fields", "pangea_calls", "applied"), + [ + pytest.param( + "acompletion", {"messages": [{"role": "user", "content": "hi"}]}, 1, ["pangea-ai-guard"], id="chat" + ), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + pytest.param("adecisions", {"input": "Is this a refund?"}, 0, None, id="decisions"), + ], +) +async def test_pangea_calls_pangea_and_lists_itself_only_when_the_request_has_messages( + respx_mock, call_type, request_fields, pangea_calls, applied +): + guard = respx_mock.post("https://ai-guard.aws.us.pangea.cloud/v1beta/guard").respond( + json={"result": {"blocked": False, "transformed": False}} + ) + guardrail = PangeaHandler( + mode="pre_call", guardrail_name="pangea-ai-guard", api_key="pts_pangeatokenid", default_on=True + ) + metadata = {} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data={**request_fields, "metadata": metadata}, + call_type=call_type, + ) + + assert guard.call_count == pangea_calls + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/test_pillar_guardrails.py b/tests/unit/proxy/guardrails/test_pillar_guardrails.py index c6be433399a..49c21fa60f4 100644 --- a/tests/unit/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/unit/proxy/guardrails/test_pillar_guardrails.py @@ -17,6 +17,7 @@ import json from urllib.parse import unquote import pytest +import respx from fastapi.exceptions import HTTPException from httpx import Request, Response @@ -1513,3 +1514,52 @@ async def test_mcp_call_masking( if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +_INPUT_HOOKS = [ + pytest.param( + "pre_call", + lambda guardrail, data: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion" + ), + id="pre_call", + ), + pytest.param( + "during_call", + lambda guardrail, data: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + ), + id="during_call", + ), +] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) +@pytest.mark.parametrize( + ("request_fields", "applied"), + [ + pytest.param({"messages": [{"role": "user", "content": "hi"}]}, ["pillar-test"], id="messages"), + pytest.param({"state": "hi"}, None, id="no-messages"), + ], +) +async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( + respx_mock: respx.MockRouter, event_hook, run_hook, request_fields, applied +): + respx_mock.post(url__startswith="https://api.pillar.security/api/v1/protect").respond( + json={"session_id": "s", "flagged": False, "scanners": {}, "evidence": []} + ) + guardrail = PillarGuardrail( + guardrail_name="pillar-test", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + on_flagged_action="block", + event_hook=event_hook, + default_on=True, + ) + metadata: dict[str, object] = {} + + await run_hook(guardrail, {**request_fields, "metadata": metadata}) + + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index d3d00aa0794..17e481db49a 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -1134,6 +1134,21 @@ async def test_pre_call_hook_scans_decision_and_extraction_text_with_a_native_ho assert _applied_guardrails(result) == ["keeps_native"] +@pytest.mark.asyncio +async def test_pre_call_hook_does_not_list_a_guardrail_for_a_systemone_request_without_text(monkeypatch): + guardrail = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + result = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={"model": "m", "state": {}, "questions": {"routed": {"type": "choice", "criteria": {"refund": None}}}}, + call_type="asystemone", + ) + + assert guardrail.scanned == [] + assert _applied_guardrails(result) == [] + + @pytest.mark.asyncio async def test_during_call_hook_scans_extraction_state_with_a_native_hook_guardrail(monkeypatch): guardrail = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.during_call, default_on=True)