From a8d3c0504ed9e14313069e3bb0f7db9e69a380f2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 13 Feb 2026 18:08:07 -0800 Subject: [PATCH] Add test guardrails for pipeline E2E testing --- .../pipeline_test_guardrails.py | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 litellm/proxy/example_config_yaml/pipeline_test_guardrails.py diff --git a/litellm/proxy/example_config_yaml/pipeline_test_guardrails.py b/litellm/proxy/example_config_yaml/pipeline_test_guardrails.py new file mode 100644 index 00000000000..539a520fcef --- /dev/null +++ b/litellm/proxy/example_config_yaml/pipeline_test_guardrails.py @@ -0,0 +1,69 @@ +""" +Test guardrails for pipeline E2E testing. + +- StrictFilter: blocks any message containing "bad" (case-insensitive) +- PermissiveFilter: always passes (simulates an advanced guardrail that is more lenient) +""" + +from typing import Optional, Union + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral + + +class StrictFilter(CustomGuardrail): + """Blocks any message containing the word 'bad'.""" + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> Optional[Union[Exception, str, dict]]: + for msg in data.get("messages", []): + content = msg.get("content", "") + if isinstance(content, str) and "bad" in content.lower(): + verbose_proxy_logger.info("StrictFilter: BLOCKED - found 'bad'") + raise HTTPException( + status_code=400, + detail="StrictFilter: content contains forbidden word 'bad'", + ) + verbose_proxy_logger.info("StrictFilter: PASSED") + return data + + +class PermissiveFilter(CustomGuardrail): + """Always passes - simulates a lenient advanced guardrail.""" + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> Optional[Union[Exception, str, dict]]: + verbose_proxy_logger.info("PermissiveFilter: PASSED (always passes)") + return data + + +class AlwaysBlockFilter(CustomGuardrail): + """Always blocks - for testing full escalation->block path.""" + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> Optional[Union[Exception, str, dict]]: + verbose_proxy_logger.info("AlwaysBlockFilter: BLOCKED") + raise HTTPException( + status_code=400, + detail="AlwaysBlockFilter: all content blocked", + )