diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index fc76779f634..b063670c346 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4255,10 +4255,20 @@ def _apply_resolved_guardrails_to_metadata( return existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) - existing_guardrails_list: Final = existing_guardrails if isinstance(existing_guardrails, list) else [] + existing_guardrails_list: Final = ( + _REQUEST_ENTRIES.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else () + ) - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = list(dict.fromkeys((*existing_guardrails_list, *resolved_guardrails))) + selected_names: Final = frozenset( + itertools.chain.from_iterable( + _REQUEST_FIELDS.validate_python(entry).keys() if isinstance(entry, dict) else (entry,) + for entry in existing_guardrails_list + ) + ) + combined: Final = [ + *existing_guardrails_list, + *(name for name in dict.fromkeys(resolved_guardrails) if name not in selected_names), + ] data[metadata_variable_name]["guardrails"] = combined verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 11dbc1b0ca2..63e5315ecd1 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterErrors +from litellm.types.utils import CallTypesLiteral from tests._master_key import MASTER_KEY @@ -1983,15 +1984,23 @@ class _FrameGuardrail(CustomGuardrail): self.rewrites = rewrites self.blocks = blocks - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: return self._scan(data) - async def async_moderation_hook(self, data, user_api_key_dict, call_type): + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: self._scan(data) - def _scan(self, data): + def _scan(self, data: dict[str, object]) -> dict[str, object]: self.inputs.append(data["input"]) - self.configs.append(data["litellm_metadata"].get("guardrail_config")) + self.configs.append(cast(dict[str, object], data["litellm_metadata"]).get("guardrail_config")) if self.blocks: raise HTTPException(status_code=400, detail={"error": f"Content blocked by {self.guardrail_name}"}) return {**data, "input": f"{data['input']} [{self.guardrail_name}]"} if self.rewrites else data diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 81a196af676..e82d2d66e3c 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5292,6 +5292,42 @@ async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_ assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"] +@pytest.mark.asyncio +async def test_add_guardrails_from_policy_engine_keeps_configured_guardrails_and_skips_ones_already_selected(): + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"guardrails": [configured, "audit"]}, + } + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [PolicyAttachment(policy="baseline", scope="*")] + attachment_registry._initialized = True + + try: + await add_guardrails_from_policy_engine( + data=data, + metadata_variable_name="metadata", + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_applies_inherited_parent_guardrail_when_child_condition_misses(): from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry