diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 312fd2a542c..6825d2fd2a0 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -922,12 +922,18 @@ class CustomGuardrail(CustomLogger): requested: Final = next((bucket.get("guardrails") for bucket in buckets if "guardrails" in bucket), None) if name not in _str_collection(requested): return + named_lists: Final = tuple( + cast("list[object]", bucket["applied_guardrails"]) # cast-ok: entries are only compared with names + for bucket in buckets + if isinstance(bucket.get("applied_guardrails"), list) + ) + for applied in named_lists: + if name not in applied: + applied.append(name) + if named_lists: + return metadata_bucket: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str keys - applied: Final = metadata_bucket.get("applied_guardrails") - if not isinstance(applied, list): - metadata_bucket["applied_guardrails"] = [name] - elif name not in applied: - cast("list[object]", applied).append(name) # cast-ok: entries are only compared with names + metadata_bucket["applied_guardrails"] = [name] @classmethod def uses_apply_guardrail_interface(cls) -> bool: diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 39f4e9a10ef..b02b5aecb8a 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1843,6 +1843,32 @@ async def test_deployment_hook_names_a_pipeline_step_it_ran_before_any_guardrail assert metadata["applied_guardrails"] == ["second"] +@pytest.mark.asyncio +@pytest.mark.parametrize("named_bucket", ["metadata", "litellm_metadata"]) +async def test_deployment_hook_names_a_pipeline_step_next_to_the_steps_already_named_when_both_buckets_exist( + named_bucket: str, +) -> None: + guardrail: Final = _PipelineStepGuardrail("second", "pre_call") + metadata: Final[dict[str, object]] = { + "guardrails": ["first", "second"], + "_pipeline_managed_guardrails": frozenset({"first", "second"}), + } + litellm_metadata: Final[dict[str, object]] = {} + {"metadata": metadata, "litellm_metadata": litellm_metadata}[named_bucket]["applied_guardrails"] = ["first"] + kwargs: Final[dict[str, object]] = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["second"], + "metadata": metadata, + "litellm_metadata": litellm_metadata, + } + + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) + + assert guardrail.calls == 1 + assert sorted(get_logging_caching_headers(kwargs)["x-litellm-applied-guardrails"].split(",")) == ["first", "second"] + + @pytest.mark.asyncio async def test_deployment_hook_records_a_pipeline_step_it_ran_without_naming_one_the_request_did_not_list() -> None: guardrail: Final = _PipelineStepGuardrail("second", "pre_call")