diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index ed330cf1930..1b1b2302236 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -411,7 +411,7 @@ class PipelineExecutor: duration = time.perf_counter() - start_time - action = _pipeline_action_for_outcome(step, outcome) + action = _pipeline_action_for_step(step, outcome, original_exception) step_result = PipelineStepResult( guardrail_name=step.guardrail, @@ -635,7 +635,7 @@ class PipelineExecutor: return ("pass", response if isinstance(response, dict) else None, None, None) except Exception as e: - if CustomGuardrail._is_guardrail_intervention(e) or isinstance(e, UnappliableRequestRewrite): + if CustomGuardrail._is_guardrail_intervention(e): error_msg: Final = _extract_error_message(e) return ("fail", None, error_msg, e) else: @@ -749,13 +749,24 @@ def _carry_working_guardrail_information( _append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing]) +def _pipeline_action_for_step(step: PipelineStep, outcome: str, original_exception: Exception | None) -> str: + """ + The action a step takes: block when its guardrail returned a rewrite the endpoint + cannot apply, whatever on_fail and on_error say, else the action configured for + the outcome. + """ + if isinstance(original_exception, UnappliableRequestRewrite): + return "block" + return _pipeline_action_for_outcome(step, outcome) + + def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. - skip -> next (stream_scope mismatch; do not apply on_pass/on_fail) - pass -> on_pass - - fail -> on_fail (content/policy intervention, or a rewrite the endpoint cannot apply) + - fail -> on_fail (content/policy intervention) - error -> on_error if set, else on_fail (backward compatible) """ if outcome == "skip": diff --git a/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py index ee93245b6d0..002b671fcb9 100644 --- a/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py +++ b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py @@ -28,7 +28,12 @@ from tests.integration._support.wire import Reply, Request, Wire, wire_server _MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" _MODEL: Final = "converse-guarded" -_PIPELINED_MODEL: Final = "converse-pipelined" +_PIPELINED_ACTIONS: Final = { + "converse-pipelined-error-allow": {"on_fail": "block", "on_error": "allow"}, + "converse-pipelined-fail-next": {"on_fail": "next", "on_error": "block"}, + "converse-pipelined-fail-allow": {"on_fail": "allow", "on_error": "block"}, + "converse-pipelined-allow-all": {"on_fail": "allow", "on_error": "allow"}, +} _CARD: Final = "4111 1111 1111 1111" _CARD_PATTERN: Final = re.compile(r"\b(?:\d{4} ){3}\d{4}\b") _MASK: Final = "****" @@ -520,27 +525,23 @@ def _config(directory: Path, guardrails: tuple[dict[str, JsonValue], ...], bedro "aws_region_name": "us-east-1", }, } - for model_name in (_MODEL, _PIPELINED_MODEL) + for model_name in (_MODEL, *_PIPELINED_ACTIONS) ), ] configuration["guardrails"] = list(guardrails) configuration["policies"] = { - "converse-pipeline": { + model_name: { "guardrails": {"add": [_generic(_Vendor.DROPS_A_ROW)]}, "pipeline": { "mode": "pre_call", - "steps": [ - { - "guardrail": _generic(_Vendor.DROPS_A_ROW), - "on_pass": "next", - "on_fail": "block", - "on_error": "allow", - } - ], + "steps": [{"guardrail": _generic(_Vendor.DROPS_A_ROW), "on_pass": "next", **actions}], }, } + for model_name, actions in _PIPELINED_ACTIONS.items() } - configuration["policy_attachments"] = [{"policy": "converse-pipeline", "models": [_PIPELINED_MODEL]}] + configuration["policy_attachments"] = [ + {"policy": model_name, "models": [model_name]} for model_name in _PIPELINED_ACTIONS + ] path: Final = directory / "converse-guardrails.yaml" path.write_text(yaml.safe_dump(configuration)) return path @@ -969,9 +970,10 @@ def test_missing_structured_row_is_refused_before_bedrock(rig: _Rig) -> None: assert _UNAPPLIABLE in _refused(rig, _converse(rig, _generic(_Vendor.DROPS_A_ROW), body)) -def test_pipeline_step_allowing_errors_still_refuses_a_rewrite_it_cannot_apply(rig: _Rig) -> None: +@pytest.mark.parametrize("model_name", tuple(_PIPELINED_ACTIONS)) +def test_pipeline_step_refuses_a_rewrite_it_cannot_apply_whatever_its_actions(rig: _Rig, model_name: str) -> None: body: Final = {"messages": [_user({"text": "Hello."}), _assistant({"text": "Hi."}), _user({"text": f"{_CARD}"})]} - response: Final = rig.gateway.request("POST", f"/bedrock/model/{_PIPELINED_MODEL}/converse", body) + response: Final = rig.gateway.request("POST", f"/bedrock/model/{model_name}/converse", body) assert _UNAPPLIABLE in _refused(rig, response) diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index ea902aa2d38..c42878bd4f7 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -816,15 +816,23 @@ async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content(monke @pytest.mark.asyncio -async def test_unappliable_rewrite_follows_on_fail_not_on_error_allow(monkeypatch): +@pytest.mark.parametrize( + "on_fail, on_error", + [("block", "allow"), ("next", "block"), ("allow", "block"), ("allow", "allow"), ("next", "next"), ("next", None)], +) +async def test_unappliable_rewrite_blocks_whatever_on_fail_and_on_error_say(monkeypatch, on_fail, on_error): """ - A rewrite the endpoint cannot apply uses on_fail, so on_error: allow does not send the original request on. + A rewrite the endpoint cannot apply blocks the pipeline, so no on_fail or on_error action sends the original request on. """ rewriter = UnappliableRewriteGuardrail("rewriter") - monkeypatch.setattr(litellm, "callbacks", [rewriter]) + later = AlwaysPassGuardrail("later-step") + monkeypatch.setattr(litellm, "callbacks", [rewriter, later]) result = await PipelineExecutor.execute_steps( - steps=[PipelineStep(guardrail="rewriter", on_pass="allow", on_fail="block", on_error="allow")], + steps=[ + PipelineStep(guardrail="rewriter", on_pass="next", on_fail=on_fail, on_error=on_error), + PipelineStep(guardrail="later-step", on_pass="allow", on_fail="block"), + ], mode="pre_call", data={"messages": [{"role": "user", "content": "card 4111 1111 1111 1111"}]}, user_api_key_dict=MagicMock(), @@ -833,8 +841,9 @@ async def test_unappliable_rewrite_follows_on_fail_not_on_error_allow(monkeypatc ) assert rewriter.calls == 1 + assert later.calls == 0 assert result.terminal_action == "block" - assert result.step_results[0].outcome == "fail" + assert result.step_results[0].action_taken == "block" assert isinstance(result.original_exception, UnappliableRequestRewrite)