diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 4418ac20895..7fa2721e5b3 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -14,6 +14,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, @@ -85,6 +86,11 @@ APPLY_GUARDRAIL_ONLY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) +def is_apply_guardrail_only_route(route: str | None) -> bool: + call_types: Final = get_call_types_for_route(route) if route else None + return bool(call_types) and call_types[0] in APPLY_GUARDRAIL_ONLY_CALL_TYPES + + def runs_through_apply_guardrail(guardrail: "CustomGuardrail", call_type: str | None) -> bool: return "apply_guardrail" in type(guardrail).__dict__ and ( not guardrail.use_native_lifecycle_hooks or call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES @@ -245,6 +251,7 @@ class CustomGuardrail(CustomLogger): use_native_lifecycle_hooks: ClassVar[bool] = False records_own_guardrail_information: ClassVar[bool] = False + adds_itself_to_applied_guardrails: ClassVar[bool] = False logging_only_scope: LoggingOnlyScope | None logging_only_continue_on_input_failure: bool diff --git a/litellm/llms/openai/decisions/guardrail_translation/handler.py b/litellm/llms/openai/decisions/guardrail_translation/handler.py index 58f13129186..5b04bde9572 100644 --- a/litellm/llms/openai/decisions/guardrail_translation/handler.py +++ b/litellm/llms/openai/decisions/guardrail_translation/handler.py @@ -84,12 +84,18 @@ def _entry_text(key: str, value: str | float | None) -> str: return json.dumps({key: value}, ensure_ascii=False)[1:-1] +def _escaped_texts(text: str, key: str, item: JsonValue) -> tuple[str, ...]: + return tuple(form for form in (key, item) if isinstance(form, str) and form not in text) + + def _entry_text_forms(key: str, item: JsonValue) -> Iterator[str]: if isinstance(item, (list, dict)): yield key yield from _json_text_forms(item) else: - yield _entry_text(key, item) + text: Final = _entry_text(key, item) + yield text + yield from _escaped_texts(text, key, item) def _json_text_forms(value: JsonValue) -> Iterator[str]: @@ -163,6 +169,10 @@ def _rewritten_entry(rewritten: str, guardrail_name: str | None) -> tuple[str, J return _KEY_ADAPTER.validate_json(key.group(1)), _rewritten_value(rewritten[key.end() :]) +def _decoded_text(text: str, form: str, guardrailed: Mapping[str, str]) -> str: + return form if form in text else guardrailed.get(form, form) + + def _guardrailed_entry( key: str, item: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None ) -> tuple[str, JsonValue]: @@ -170,7 +180,16 @@ def _guardrailed_entry( return guardrailed.get(key, key), _guardrailed_json(item, guardrailed, guardrail_name) text: Final = _entry_text(key, item) rewritten: Final = guardrailed.get(text, text) - return (key, item) if rewritten == text else _rewritten_entry(rewritten, guardrail_name) + decoded: Final = ( + _decoded_text(text, key, guardrailed), + _decoded_text(text, item, guardrailed) if isinstance(item, str) else item, + ) + if rewritten == text: + return decoded + entry: Final = _rewritten_entry(rewritten, guardrail_name) + if decoded not in ((key, item), entry): + raise unappliable_request_rewrite(guardrail_name) + return entry def _instruction_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: @@ -260,7 +279,9 @@ async def _guardrailed_texts( input_type="request", logging_obj=litellm_logging_obj, ) - if not guardrail_to_apply.records_own_guardrail_information: + if not ( + guardrail_to_apply.records_own_guardrail_information or guardrail_to_apply.adds_itself_to_applied_guardrails + ): 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 ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index b564a1df041..9629a1e5e93 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -17,7 +17,7 @@ from websockets.asyncio.client import ClientConnection, connect from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -155,7 +155,9 @@ class AimGuardrail(CustomGuardrail): if is_non_conversational_call_type(call_type) and not self.inspect_embeddings: verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type) return data - return await self.call_aim_guardrail(data, hook="pre_call", key_alias=user_api_key_dict.key_alias) + return await self.call_aim_guardrail( + data, hook="pre_call", key_alias=user_api_key_dict.key_alias, call_type=call_type + ) async def async_moderation_hook( self, @@ -168,14 +170,19 @@ class AimGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type) return data - await self.call_aim_guardrail(data, hook="moderation", key_alias=user_api_key_dict.key_alias) + await self.call_aim_guardrail( + data, hook="moderation", key_alias=user_api_key_dict.key_alias, call_type=call_type + ) return data - async def call_aim_guardrail(self, data: dict, hook: str, key_alias: str | None) -> dict: + async def call_aim_guardrail( + self, data: dict, hook: str, key_alias: str | None, call_type: CallTypesLiteral | None = 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) + if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + 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( 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 be7d967836f..a80575c3cc1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -22,7 +22,7 @@ from websockets.exceptions import ConnectionClosed from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -176,6 +176,7 @@ class CatoNetworksGuardrail(CustomGuardrail): hook="pre_call", key_alias=user_api_key_dict.key_alias, user_email=self._resolve_cato_user_email(user_api_key_dict), + call_type=call_type, ) async def async_moderation_hook( @@ -193,6 +194,7 @@ class CatoNetworksGuardrail(CustomGuardrail): hook="moderation", key_alias=user_api_key_dict.key_alias, user_email=self._resolve_cato_user_email(user_api_key_dict), + call_type=call_type, ) @classmethod @@ -297,11 +299,13 @@ class CatoNetworksGuardrail(CustomGuardrail): hook: str, key_alias: str | None, user_email: str | None = None, + call_type: CallTypesLiteral | 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) + if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + 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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index 6974c33b0bf..0bb6454f35f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -15,7 +15,11 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, + CustomGuardrail, + is_apply_guardrail_only_route, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -328,6 +332,7 @@ class DynamoAIGuardrails(CustomGuardrail): error_message: Final = self._create_error_message(processed_result) raise ValueError(error_message) + if _messages or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -372,6 +377,7 @@ class DynamoAIGuardrails(CustomGuardrail): error_message: Final = self._create_error_message(processed_result) raise ValueError(error_message) + if _messages or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -439,6 +445,9 @@ class DynamoAIGuardrails(CustomGuardrail): error_message: Final = self._create_error_message(processed_result) raise ValueError(error_message) + if isinstance(response, litellm.ModelResponse) or not is_apply_guardrail_only_route( + user_api_key_dict.request_route + ): add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) async def async_post_call_streaming_iterator_hook( 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 ff4ae15b6f1..e7ff6ec1039 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -15,7 +15,11 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, + CustomGuardrail, + is_apply_guardrail_only_route, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -476,6 +480,7 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) + if contents_to_check or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -544,6 +549,7 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) + if contents_to_check or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -639,6 +645,9 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) + if isinstance(response, litellm.ModelResponse) or not is_apply_guardrail_only_route( + user_api_key_dict.request_route + ): add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) async def async_post_call_streaming_iterator_hook( diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index d27ad711fdb..a4bf5165eb1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -19,6 +19,7 @@ from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail, log_guardrail_information, ) @@ -292,7 +293,8 @@ class lakeraAI_Moderation(CustomGuardrail): } """ self._check_response_flagged(response=response.json()) - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @log_guardrail_information async def async_pre_call_hook( diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index ca06b3c1675..e8d1584430b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -34,6 +34,7 @@ from litellm import DualCache from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, GUARDRAIL_SESSIONS_TARGET, CustomGuardrail, log_guardrail_information, @@ -207,7 +208,7 @@ class LassoGuardrail(CustomGuardrail): # The conversation_id is being stored in the cache so it can be used by the post_call hook self._get_or_generate_conversation_id(data, global_cache) - return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT") + return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT", call_type=call_type) @log_guardrail_information async def async_moderation_hook( @@ -235,7 +236,7 @@ class LassoGuardrail(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data - return await self._run_lasso_guardrail(data, cache, message_type="PROMPT") + return await self._run_lasso_guardrail(data, cache, message_type="PROMPT", call_type=call_type) @log_guardrail_information async def async_post_call_success_hook( @@ -375,6 +376,7 @@ class LassoGuardrail(CustomGuardrail): data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT", + call_type: str | None = None, ): """ Run the Lasso guardrail with the specified message type. @@ -395,6 +397,7 @@ class LassoGuardrail(CustomGuardrail): data: The request data containing messages cache: The cache instance for storing conversation_id (optional for post-call) message_type: Either "PROMPT" for input or "COMPLETION" for output + call_type: The request's call type, which decides whether Lasso lists itself as applied Raises: LassoGuardrailAPIError: If the Lasso API call fails @@ -412,7 +415,8 @@ 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) + if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + 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 b50b8c740ad..5eb8e3df6fd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -185,6 +185,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): """ use_native_lifecycle_hooks: ClassVar[bool] = True + adds_itself_to_applied_guardrails: ClassVar[bool] = True @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -1328,6 +1329,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): self._raise_if_fail_closed(e) return inputs + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) flagged: Final = self._should_block_content(armor_response, allow_sanitization=False) end_time: Final = time.time() self.add_standard_logging_guardrail_information_to_request_data( diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 1983be794b3..2444eee6abb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -22,6 +22,7 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail, log_guardrail_information, ) @@ -351,7 +352,7 @@ class PillarGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook") result: Final = await self.run_pillar_guardrail(data, user_api_key_dict) - if data.get("messages"): + if data.get("messages") or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return result @@ -394,7 +395,7 @@ 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) - if data.get("messages"): + if data.get("messages") or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: 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 60a0a1f3813..74bbe510e71 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -306,7 +306,11 @@ class UnifiedLLMGuardrails(CustomLogger): litellm_logging_obj=data.get("litellm_logging_obj"), ) - if not guardrail_to_apply.records_own_guardrail_information and not translation.lists_guardrail_when_called: + if not ( + guardrail_to_apply.records_own_guardrail_information + or guardrail_to_apply.adds_itself_to_applied_guardrails + or translation.lists_guardrail_when_called + ): add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=guardrail_to_apply.guardrail_name ) @@ -447,7 +451,9 @@ class UnifiedLLMGuardrails(CustomLogger): if e.original_response is None: e.original_response = response raise - if not guardrail_to_apply.records_own_guardrail_information: + if not ( + guardrail_to_apply.records_own_guardrail_information or guardrail_to_apply.adds_itself_to_applied_guardrails + ): 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 620739721cd..080f231f5e7 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 @@ -10,6 +10,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging 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.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.types.guardrails import ContentFilterAction, ContentFilterPattern from litellm.types.decisions import OpenAIDecisionResponse from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -212,6 +214,35 @@ async def test_systemone_secret_key_and_card_number_are_masked_before_the_provid assert result["state"] == {"": "", "card": "", "items": 2} +@pytest.mark.asyncio +async def test_systemone_masks_a_card_number_that_json_escaping_splits() -> None: + guardrail: Final = ContentFilterGuardrail( + guardrail_name="cards", + patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="visa", action=ContentFilterAction.MASK)], + ) + data: Final = { + "model": "judge", + "state": {"card": "4111\t1111\t1111\t1111", "4111\n1111\n1111\n1111": "on file", "note": "first\nsecond"}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + result: Final = await SystemOneHandler().process_input_messages(data, guardrail) + + assert result["state"] == {"card": "[VISA_REDACTED]", "[VISA_REDACTED]": "on file", "note": "first\nsecond"} + + +@pytest.mark.asyncio +async def test_systemone_rejects_an_escaped_entry_whose_two_scanned_forms_disagree() -> None: + data: Final = { + "model": "judge", + "state": {"note": "first\nsecond"}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + with pytest.raises(UnappliableRequestRewrite): + await SystemOneHandler().process_input_messages(data, _TaggingGuardrail()) + + @pytest.mark.asyncio async def test_systemone_keeps_every_state_entry_when_masked_keys_collide() -> None: data: Final = { 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 1ad0890f0db..1f691540145 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py @@ -15,6 +15,7 @@ import respx import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral from litellm.proxy.guardrails.guardrail_hooks.dynamoai import DynamoAIGuardrails from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -209,15 +210,15 @@ async def test_post_call_hook_skips_dynamoai_when_the_reply_has_no_text(respx_mo _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" + lambda guardrail, data, call_type: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type ), id="pre_call", ), pytest.param( "during_call", - lambda guardrail, data: guardrail.async_moderation_hook( - data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + lambda guardrail, data, call_type: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type ), id="during_call", ), @@ -228,17 +229,19 @@ _INPUT_HOOKS = [ @pytest.mark.usefixtures("httpx_transport") @pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) @pytest.mark.parametrize( - ("request_fields", "applied"), + ("request_fields", "call_type", "applied"), [ - pytest.param({"messages": [{"role": "user", "content": "hi"}]}, ["test-dynamoai"], id="messages"), - pytest.param({"state": "hi"}, None, id="no-messages"), + pytest.param({"messages": [{"role": "user", "content": "hi"}]}, "completion", ["test-dynamoai"], id="messages"), + pytest.param({"state": "hi"}, "aextraction", None, id="no-messages-on-extraction"), + pytest.param({"state": "hi"}, "completion", ["test-dynamoai"], id="no-messages-on-chat"), ], ) async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( respx_mock: respx.MockRouter, event_hook: str, - run_hook: Callable[[DynamoAIGuardrails, Mapping[str, object]], Awaitable[object]], + run_hook: Callable[[DynamoAIGuardrails, Mapping[str, object], CallTypesLiteral], Awaitable[object]], request_fields: Mapping[str, object], + call_type: CallTypesLiteral, applied: Sequence[str] | None, ) -> None: respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []}) @@ -251,7 +254,7 @@ async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( ) metadata: Final[dict[str, object]] = {} - await run_hook(guardrail, {**request_fields, "metadata": metadata}) + await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type) assert metadata.get("applied_guardrails") == applied @@ -259,21 +262,35 @@ async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( @pytest.mark.asyncio @pytest.mark.usefixtures("httpx_transport") @pytest.mark.parametrize( - ("response", "applied"), + ("response", "request_route", "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"), + pytest.param(_assistant_reply("Hello, how are you?"), "/v1/chat/completions", ["test-dynamoai"], id="chat-reply"), + pytest.param( + {"answers": {"aws_key": {"type": "text", "value": None}}}, + "/v1/extraction", + None, + id="no-choices-on-extraction", + ), + pytest.param( + {"answers": {"aws_key": {"type": "text", "value": None}}}, + "/v1/responses", + ["test-dynamoai"], + id="no-choices-on-responses", + ), ], ) async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text( - respx_mock: respx.MockRouter, response: litellm.ModelResponse | Mapping[str, object], applied: Sequence[str] | None + respx_mock: respx.MockRouter, + response: litellm.ModelResponse | Mapping[str, object], + request_route: str, + applied: Sequence[str] | None, ) -> None: respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []}) metadata: Final[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(), + user_api_key_dict=UserAPIKeyAuth(request_route=request_route), response=response, ) 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 ba9a8b370fa..1396c0599ac 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 @@ -9,6 +9,7 @@ import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral from litellm.proxy.guardrails.guardrail_hooks.ibm_guardrails.ibm_detector import IBMGuardrailDetector pytestmark = pytest.mark.usefixtures("httpx_transport") @@ -131,33 +132,35 @@ async def test_orchestrator_detection_blocks_the_request_and_names_the_detector( _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" + lambda guardrail, data, call_type: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type ), id="pre_call", ), pytest.param( "during_call", - lambda guardrail, data: guardrail.async_moderation_hook( - data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + lambda guardrail, data, call_type: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type ), 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.param({"messages": [{"role": "user", "content": "hello there"}]}, "completion", ["ibm-guard"], id="messages"), + pytest.param({"state": "hello there"}, "aextraction", None, id="no-messages-on-extraction"), + pytest.param({"state": "hello there"}, "completion", ["ibm-guard"], id="no-messages-on-chat"), ] @pytest.mark.asyncio @pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) -@pytest.mark.parametrize(("request_fields", "applied"), _REQUEST_TEXT_CASES) +@pytest.mark.parametrize(("request_fields", "call_type", "applied"), _REQUEST_TEXT_CASES) async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( respx_mock: respx.MockRouter, event_hook: str, - run_hook: Callable[[IBMGuardrailDetector, Mapping[str, object]], Awaitable[object]], + run_hook: Callable[[IBMGuardrailDetector, Mapping[str, object], CallTypesLiteral], Awaitable[object]], request_fields: Mapping[str, object], + call_type: CallTypesLiteral, applied: Sequence[str] | None, ) -> None: respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]]) @@ -172,27 +175,42 @@ async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( ) metadata: Final[dict[str, object]] = {} - await run_hook(guardrail, {**request_fields, "metadata": metadata}) + await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type) assert metadata.get("applied_guardrails") == applied @pytest.mark.asyncio @pytest.mark.parametrize( - ("response", "applied"), + ("response", "request_route", "applied"), [ pytest.param( litellm.ModelResponse( choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}] ), + "/v1/chat/completions", ["ibm-guard"], id="chat-reply", ), - pytest.param({"answers": {"aws_key": {"type": "text", "value": None}}}, None, id="no-choices"), + pytest.param( + {"answers": {"aws_key": {"type": "text", "value": None}}}, + "/v1/extraction", + None, + id="no-choices-on-extraction", + ), + pytest.param( + {"answers": {"aws_key": {"type": "text", "value": None}}}, + "/v1/responses", + ["ibm-guard"], + id="no-choices-on-responses", + ), ], ) async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text( - respx_mock: respx.MockRouter, response: litellm.ModelResponse | Mapping[str, object], applied: Sequence[str] | None + respx_mock: respx.MockRouter, + response: litellm.ModelResponse | Mapping[str, object], + request_route: str, + applied: Sequence[str] | None, ) -> None: respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]]) guardrail: Final = IBMGuardrailDetector( @@ -208,7 +226,7 @@ async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text( await guardrail.async_post_call_success_hook( data={"messages": [{"role": "user", "content": "hello"}], "metadata": metadata}, - user_api_key_dict=UserAPIKeyAuth(), + user_api_key_dict=UserAPIKeyAuth(request_route=request_route), response=response, ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py index 52bb6a9b6ae..20b9f0cc726 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py @@ -829,6 +829,7 @@ async def test_aim_streaming_block_stops_after_verified_chunks(monkeypatch: pyte "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["aim"], id="decisions" ), pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"), ], ) async def test_aim_calls_aim_and_lists_itself_only_when_the_request_has_text( 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 fca978cb4c8..8eba3360320 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py @@ -2800,6 +2800,7 @@ async def test_cato_still_inspects_every_conversational_call_type(call_type: str "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["cato"], id="decisions" ), pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"), ], ) async def test_cato_calls_cato_and_lists_itself_only_when_the_request_has_text( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py index bf57529b9a0..e1902480cfd 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -83,3 +83,22 @@ async def test_clean_scan_of_string_input_lists_the_guardrail() -> None: ) assert _applied_guardrails(data) == ("lakera-v1",) + + +@pytest.mark.asyncio +async def test_clean_scan_of_chat_messages_keeps_the_header_as_before() -> None: + scans: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + scans.append(request) + return httpx.Response(200, json={"results": [{"flagged": False}]}) + + guardrail: Final = _lakera(httpx.MockTransport(respond)) + data: Final[dict[str, object]] = {"model": "m", "messages": [{"role": "user", "content": "is this a refund?"}]} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), cache=DualCache(), data=data, call_type="completion" + ) + + assert len(scans) == 1 + assert _applied_guardrails(data) == () diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py index dec3510d6ab..cd62b9fbe89 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -1554,30 +1554,41 @@ async def test_api_error_handling_legacy_messages( @pytest.mark.asyncio @pytest.mark.usefixtures("httpx_transport") @pytest.mark.parametrize( - ("call_type", "request_fields", "applied"), + ("call_type", "request_fields", "lasso_calls", "applied"), [ pytest.param( - "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, ["lasso"], id="decisions" + "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["lasso"], id="decisions" ), - pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, None, id="extraction"), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"), ], ) +@pytest.mark.parametrize("event_hook", ["pre_call", "during_call"]) 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 + respx_mock: respx.MockRouter, + event_hook: str, + call_type: str, + request_fields: dict[str, object], + lasso_calls: int, + applied: list[str] | None, ) -> None: - respx_mock.post("https://server.lasso.security/gateway/v3/classify").respond( + classify: Final = 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 + lasso_api_key="test-api-key", guardrail_name="lasso", event_hook=event_hook, default_on=True ) metadata: Final[dict[str, object]] = {} + data: Final = {**request_fields, "metadata": metadata} - await guardrail.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(), - cache=DualCache(), - data={**request_fields, "metadata": metadata}, - call_type=call_type, - ) + if event_hook == "pre_call": + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type + ) + else: + await guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type, cache=DualCache() + ) + assert classify.call_count == lasso_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 9d650e11634..0a8dadd5d11 100644 --- a/tests/unit/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/unit/proxy/guardrails/test_pillar_guardrails.py @@ -27,6 +27,7 @@ from httpx import Request, Response import litellm from litellm import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers from litellm.proxy.guardrails.guardrail_hooks.pillar import ( PillarGuardrail, @@ -1521,15 +1522,15 @@ if __name__ == "__main__": _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" + lambda guardrail, data, call_type: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type ), id="pre_call", ), pytest.param( "during_call", - lambda guardrail, data: guardrail.async_moderation_hook( - data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + lambda guardrail, data, call_type: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type ), id="during_call", ), @@ -1540,17 +1541,19 @@ _INPUT_HOOKS = [ @pytest.mark.usefixtures("httpx_transport") @pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) @pytest.mark.parametrize( - ("request_fields", "applied"), + ("request_fields", "call_type", "applied"), [ - pytest.param({"messages": [{"role": "user", "content": "hi"}]}, ["pillar-test"], id="messages"), - pytest.param({"state": "hi"}, None, id="no-messages"), + pytest.param({"messages": [{"role": "user", "content": "hi"}]}, "completion", ["pillar-test"], id="messages"), + pytest.param({"state": "hi"}, "aextraction", None, id="no-messages-on-extraction"), + pytest.param({"state": "hi"}, "completion", ["pillar-test"], id="no-messages-on-chat"), ], ) async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( respx_mock: respx.MockRouter, event_hook: str, - run_hook: Callable[[PillarGuardrail, Mapping[str, object]], Awaitable[object]], + run_hook: Callable[[PillarGuardrail, Mapping[str, object], CallTypesLiteral], Awaitable[object]], request_fields: Mapping[str, object], + call_type: CallTypesLiteral, applied: Sequence[str] | None, ) -> None: respx_mock.post(url__startswith="https://api.pillar.security/api/v1/protect").respond( @@ -1566,6 +1569,6 @@ async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( ) metadata: Final[dict[str, object]] = {} - await run_hook(guardrail, {**request_fields, "metadata": metadata}) + await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type) 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 2735f46f50f..70ac3d5c73c 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from itertools import chain from typing import Final, Literal +import httpx import pytest from fastapi import HTTPException from pydantic import BaseModel, ConfigDict @@ -13,7 +14,9 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ModelArmorGuardrail from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ToolPermissionGuardrail from litellm.proxy.utils import ProxyLogging from litellm.types.decisions import OpenAIDecisionResponse @@ -1199,6 +1202,57 @@ async def test_pre_call_hook_keeps_request_tags_where_the_router_reads_them( assert _applied_guardrails(result) == ("keeps_native",) +class _ModelArmorAnswering(AsyncHTTPHandler): + def __init__(self, status_code: int) -> None: + self.status_code = status_code + + async def post(self, url: str, **kwargs: object) -> httpx.Response: + clean: Final = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}} + return httpx.Response(self.status_code, json=clean, request=httpx.Request("POST", url)) + + +async def _model_armor_token() -> tuple[str, str]: + return ("test-token", "test-project") + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("status_code", "applied"), ((200, ("model-armor",)), (503, ()))) +@pytest.mark.parametrize( + ("call_type", "request_fields"), + ( + ("aextraction", {"state": "refund order 4411", "fields": {"total": {"type": "number"}}}), + ("asystemone", {"state": {"ticket": "refund order 4411"}, "questions": {"q": {"type": "noul"}}}), + ("adecisions", {"input": "refund order 4411", "questions": [{"type": "predicate", "instructions": "Refund?"}]}), + ), +) +async def test_pre_call_hook_lists_model_armor_only_when_it_answered( + monkeypatch: pytest.MonkeyPatch, + call_type: CallTypesLiteral, + request_fields: Mapping[str, object], + status_code: int, + applied: tuple[str, ...], +) -> None: + guardrail: Final = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + fail_on_error=False, + async_handler=_ModelArmorAnswering(status_code), + access_token_provider=_model_armor_token, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={"model": "m", **request_fields}, + call_type=call_type, + ) + + assert _applied_guardrails(result) == applied + + @pytest.mark.asyncio async def test_pre_call_hook_does_not_list_a_guardrail_for_a_systemone_request_without_text( monkeypatch: pytest.MonkeyPatch,