From 744dadabd8dcb1bc1921b227d2eddc65495f4ee1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:29:56 -0700 Subject: [PATCH] fix(guardrails): keep chat headers as before and check escaped systemone entries Vendor guardrails with their own hooks list themselves in the applied guardrails header as before on chat, Responses and Messages, and only when they scanned text on extraction and decision routes Model Armor lists itself only once its API answered, so a call it let through with fail_on_error false is no longer listed A systemone entry whose JSON escaping splits a value is scanned in its decoded form too, and a rewrite where the two forms disagree is rejected --- litellm/integrations/custom_guardrail.py | 7 +++ .../guardrail_translation/handler.py | 27 ++++++++-- .../guardrails/guardrail_hooks/aim/aim.py | 17 ++++-- .../cato_networks/cato_networks.py | 8 ++- .../guardrail_hooks/dynamoai/dynamoai.py | 11 +++- .../ibm_guardrails/ibm_detector.py | 11 +++- .../guardrails/guardrail_hooks/lakera_ai.py | 4 +- .../guardrails/guardrail_hooks/lasso/lasso.py | 10 ++-- .../model_armor/model_armor.py | 6 +++ .../guardrail_hooks/pillar/pillar.py | 5 +- .../unified_guardrail/unified_guardrail.py | 10 +++- .../test_decisions_guardrail_handler.py | 31 +++++++++++ .../guardrail_hooks/dynamoai/test_dynamoai.py | 45 +++++++++++----- .../ibm_guardrails/test_ibm_detector.py | 44 ++++++++++----- .../guardrails/guardrail_hooks/test_aim.py | 1 + .../guardrail_hooks/test_cato_networks.py | 1 + .../guardrail_hooks/test_lakera_ai.py | 19 +++++++ .../guardrails/guardrail_hooks/test_lasso.py | 35 +++++++----- .../guardrails/test_pillar_guardrails.py | 21 ++++---- .../test_proxy_logging_hook_detection.py | 54 +++++++++++++++++++ 20 files changed, 299 insertions(+), 68 deletions(-) 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,