diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..5e1d792718e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -69,6 +69,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions" # field to suppress a guardrail on the direct-SDK path that never reaches the # proxy's metadata sanitizer. _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16) +_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16) _GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422}) @@ -824,13 +825,13 @@ class CustomGuardrail(CustomLogger): return False - def _pre_call_marker(self) -> str | None: + def _pre_call_marker(self, *, raised: bool = False) -> str | None: name: Final = self.guardrail_name if not name: return None - return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" + return f"{_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN}:{name}" - def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None: + def mark_pre_call_hook_ran(self, data: dict[str, object], *, raised: bool = False) -> None: """ Record that this guardrail's ``async_pre_call_hook`` already ran for this request, so the deployment-level hook does not run it a second time. @@ -839,8 +840,11 @@ class CustomGuardrail(CustomLogger): router later spreads a deployment's model-level ``guardrails`` into the top-level request kwargs, which would otherwise re-trigger the same hook from ``async_pre_call_deployment_hook``. + + With ``raised``, records instead that the hook raised as a policy pipeline + step the pipeline then moved past. The deployment-level hook still runs it. """ - marker: Final = self._pre_call_marker() + marker: Final = self._pre_call_marker(raised=raised) if marker is None: return for meta_key in ("metadata", "litellm_metadata"): @@ -855,8 +859,12 @@ class CustomGuardrail(CustomLogger): return data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} - def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool: - marker: Final = self._pre_call_marker() + def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool: + """Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided.""" + return self._pre_call_hook_already_ran(data) or self._pre_call_hook_already_ran(data, raised=True) + + def _pre_call_hook_already_ran(self, data: Mapping[str, object], *, raised: bool = False) -> bool: + marker: Final = self._pre_call_marker(raised=raised) if marker is None: return False for meta_key in ("metadata", "litellm_metadata"): diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 4ad29bd3478..7f804286282 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -329,6 +329,32 @@ def _prepare_hook_input( return hook_input, scans_raw_request +async def _run_pre_call_hook( + target: CustomLogger, + callback: CustomGuardrail, + data: dict[str, object], # mutable-ok: the request payload the run is recorded on + hook_input: dict[str, object], # mutable-ok: same request-payload shape the hooks mutate + user_api_key_dict: "UserAPIKeyAuth", + call_type: str, +) -> dict[str, object] | None: + try: + response: Final = await target.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=None, + data=hook_input, + call_type=call_type, + ) + except Exception: + callback.mark_pre_call_hook_ran(data, raised=True) + raise + callback.mark_pre_call_hook_ran(data) + if not isinstance(response, dict): + return None + payload: Final = cast(dict[str, object], response) # cast-ok: a pre_call hook returns the request payload + callback.mark_pre_call_hook_ran(payload) + return payload + + def _release_original_chunks( guardrail_name: str, reason: str, @@ -489,8 +515,8 @@ class PipelineExecutor: live, one that drops or adds a tool call, or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the originals and the step passes, so the - client gets the stream the merge base sent, and the guardrail stays out of the - applied-guardrails header since its output never reached the client. The response an + client gets the stream the merge base sent. The guardrail is named in the + applied-guardrails header either way, since it ran. The response an earlier step's translation stored under ``request_data["response"]`` is dropped first, so this step's hook sees the stream as the steps before it left it.""" scanner: Final = ( @@ -525,12 +551,12 @@ class PipelineExecutor: except UndeliverableStreamRewrite as undeliverable: _release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals) return + finally: + if not callback.records_own_guardrail_information: + add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) discard_reason: Final = observer.discard_reason(rewrite_undeliverable_reason) if discard_reason is not None: _release_original_chunks(step.guardrail, discard_reason, streaming_chunks, originals) - return - if not callback.records_own_guardrail_information: - add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) @staticmethod async def _run_step( @@ -582,16 +608,7 @@ class PipelineExecutor: try: if mode == "pre_call": - response = await target.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=None, - data=hook_input, - call_type=call_type, - ) - if isinstance(callback, CustomGuardrail): - callback.mark_pre_call_hook_ran(data) - if isinstance(response, dict): - callback.mark_pre_call_hook_ran(response) + response = await _run_pre_call_hook(target, callback, data, hook_input, user_api_key_dict, call_type) elif mode == "post_call" and streaming_chunks is not None: if endpoint_translation is None: return ( @@ -612,6 +629,7 @@ class PipelineExecutor: ) response = None elif mode == "post_call": + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail) response = await target.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, data=data, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ce7e6afde24..9a960d30d3d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -706,6 +706,62 @@ def pipeline_managed_guardrail_names( ) +def _ran_or_runs_outside_pipelines( + callback: CustomGuardrail | None, + data: Mapping[str, object], + pipeline_modes: frozenset[str], +) -> bool: + return callback is not None and ( + callback.pre_call_hook_ran_or_raised(data) + or any( + callback._event_hook_is_event_type(hook) # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks + for hook in GuardrailEventHooks + if hook.value not in pipeline_modes + ) + ) + + +# The A2A JSON-RPC methods answered with a stream, as both protocol versions spell them +_A2A_STREAMING_METHODS: Final = frozenset( + ("message/stream", "tasks/resubscribe", "SendStreamingMessage", "SubscribeToTask") +) + + +def _post_call_steps_named_before_headers( + data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> frozenset[str]: + """ + Post_call pipeline steps the executor names, as it runs them, before the response + headers go out: every step of a whole response, and of a streamed one the steps of + the pipelines that hold the stream until they finish. + """ + if data.get("stream") is not True and data.get("method") not in _A2A_STREAMING_METHODS: + return pipeline_managed_guardrail_names(data, "post_call") + gated: Final = _stream_gated_pipelines(data, user_api_key_dict) + return frozenset() if _pipelines_stream_live(gated) else _pipeline_step_guardrail_names(gated) + + +def _pipeline_steps_that_have_not_run(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: + """ + Guardrails only policy pipelines run whose pre_call hook has not run for this request: + pre_call steps a pipeline skipped and the post_call steps the executor names when it + runs them. A stream runs a pipeline it cannot hold as guardrails on their own, and one + whose steps all stream live after the headers go out, so those steps are named up front + like any other guardrail. + """ + pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") + post_call: Final = _post_call_steps_named_before_headers(data, user_api_key_dict) + return frozenset( + name + for name in pre_call | post_call + if not _ran_or_runs_outside_pipelines( + PipelineExecutor.find_guardrail_callback(name), + data, + frozenset(mode for mode, names in (("pre_call", pre_call), ("post_call", post_call)) if name in names), + ) + ) + + def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]: resolved: Final = tuple( litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -927,19 +983,23 @@ def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseT return None if resolved is None else resolved[1] +def _stream_gated_pipelines( + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> tuple[tuple[str, "GuardrailPipeline"], ...]: + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if translation is None: + return () + return tuple( + (policy_name, pipeline) + for policy_name, pipeline in _post_call_pipelines(request_data) + if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) + ) + + def stream_gated_guardrail_names( request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth ) -> frozenset[str]: - translation: Final = _streaming_pipeline_translation(user_api_key_dict) - if translation is None: - return frozenset() - return _pipeline_step_guardrail_names( - tuple( - (policy_name, pipeline) - for policy_name, pipeline in _post_call_pipelines(request_data) - if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) - ) - ) + return _pipeline_step_guardrail_names(_stream_gated_pipelines(request_data, user_api_key_dict)) def _streamable_post_call_pipelines( @@ -2639,7 +2699,7 @@ class ProxyLogging: data.pop("prompt_version", None) data.pop("prompt_environment", None) - def _process_guardrail_metadata(self, data: dict) -> None: + def _process_guardrail_metadata(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> None: """Process guardrails from metadata and add to applied_guardrails.""" from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -2664,8 +2724,16 @@ class ProxyLogging: if not isinstance(applied_guardrails, list): applied_guardrails = [] + not_run: Final = _pipeline_steps_that_have_not_run( + cast(dict[str, object], data), # cast-ok: request payload + user_api_key_dict, + ) for guardrail_name in guardrails_in_metadata: - if isinstance(guardrail_name, str) and guardrail_name not in applied_guardrails: + if ( + isinstance(guardrail_name, str) + and guardrail_name not in applied_guardrails + and guardrail_name not in not_run + ): add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name) async def _maybe_execute_pipelines( @@ -2952,7 +3020,7 @@ class ProxyLogging: and (guardrails_only or not caps.has_pre_call_override) ): if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, user_api_key_dict) return data parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( @@ -3056,17 +3124,17 @@ class ProxyLogging: ) if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, user_api_key_dict) return data except SensitiveDataRouteException as e: data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict) if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, user_api_key_dict) return data except Exception: if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, user_api_key_dict) raise async def _run_parallel_pre_call_guardrails( diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py new file mode 100644 index 00000000000..5d08aa325e7 --- /dev/null +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml + +from tests.integration._support.client import Gateway, gateway_from_environment +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, Wire, wire_server + +_HEADER: Final = "x-litellm-applied-guardrails" +_VENDOR_PATH: Final = "/beta/litellm_basic_guardrail_api" +_PROMPT_SHIELD_PATH: Final = "/contentsafety/text:shieldPrompt" +_PROMPT_SHIELD: Final = "prompt-shield-step" +_REQUEST_FIRST: Final = "request-first-step" +_REQUEST_SECOND: Final = "request-second-step" +_REQUEST_BLOCKER: Final = "request-blocking-step" +_RESPONSE_FIRST: Final = "response-first-step" +_RESPONSE_SECOND: Final = "response-second-step" +_RESPONSE_BLOCKER: Final = "response-blocking-step" +_LIVE_FIRST: Final = "live-first-step" +_LIVE_SECOND: Final = "live-second-step" +_UNGATED_STEP: Final = "ungated-step" +_STOPS_AFTER_FIRST: Final = "stops-after-first" +_RUNS_EVERY_STEP: Final = "runs-every-step" +_BLOCKS_AT_FIRST: Final = "blocks-at-first" +_RESPONSE_STOPS_AFTER_FIRST: Final = "response-stops-after-first" +_SHIELDS_THEN_CHECKS: Final = "shields-then-checks" +_MOVES_PAST_FIRST: Final = "moves-past-first" +_RESPONSE_MOVES_PAST_FIRST: Final = "response-moves-past-first" +_SCANS_AFTER_RELEASE: Final = "scans-after-release" +_CANNOT_HOLD_STREAM: Final = "cannot-hold-stream" + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + vendor: Wire + upstream: Wire + model: str + + +def _completion_chunk(identity: str, delta: dict[str, str], finish_reason: str | None) -> bytes: + chunk: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(chunk)}\n\n".encode() + + +def _completion(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + identity: Final = f"chatcmpl-{uuid.uuid4().hex}" + if json.loads(request.body).get("stream") is True: + return Reply( + content_type="text/event-stream", + chunks=( + _completion_chunk(identity, {"role": "assistant", "content": "guarded"}, None), + _completion_chunk(identity, {}, "stop"), + b"data: [DONE]\n\n", + ), + ) + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "guarded"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +def _vendor(request: Request) -> Reply: + if request.target.startswith(f"/{_PROMPT_SHIELD}{_PROMPT_SHIELD_PATH}?"): + return Reply( + body=json.dumps({"userPromptAnalysis": {"attackDetected": False}, "documentsAnalysis": []}).encode() + ) + if request.target.startswith(f"/{_UNGATED_STEP}/"): + return Reply(body=json.dumps({"finalAction": "NONE", "appliedPolicies": []}).encode()) + assert request.target.endswith(_VENDOR_PATH), request.target + if request.target.startswith((f"/{_REQUEST_BLOCKER}/", f"/{_RESPONSE_BLOCKER}/")): + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "first step refused"}).encode()) + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def _guardrail(name: str, mode: str, vendor: Wire, **params: object) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": f"{vendor.url}/{name}", + "api_key": "synthetic-guardrail-key", + **params, + }, + } + + +def _policy(mode: str, steps: tuple[tuple[str, str, str], ...]) -> dict[str, object]: + return { + "guardrails": {"add": [guardrail for guardrail, _, _ in steps]}, + "pipeline": { + "mode": mode, + "steps": [ + {"guardrail": guardrail, "on_pass": on_pass, "on_fail": on_fail} + for guardrail, on_pass, on_fail in steps + ], + }, + } + + +def _config(directory: Path, vendor: Wire) -> Path: + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["guardrails"] = [ + *(_guardrail(name, "pre_call", vendor) for name in (_REQUEST_FIRST, _REQUEST_SECOND, _REQUEST_BLOCKER)), + *(_guardrail(name, "post_call", vendor) for name in (_RESPONSE_FIRST, _RESPONSE_SECOND, _RESPONSE_BLOCKER)), + *( + _guardrail(name, "post_call", vendor, streaming_buffer_until_moderated=False) + for name in (_LIVE_FIRST, _LIVE_SECOND) + ), + { + "guardrail_name": _UNGATED_STEP, + "litellm_params": { + "guardrail": "dynamoai", + "mode": "post_call", + "default_on": False, + "api_base": f"{vendor.url}/{_UNGATED_STEP}", + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": _PROMPT_SHIELD, + "litellm_params": { + "guardrail": "azure/prompt_shield", + "mode": "pre_call", + "default_on": False, + "api_base": f"{vendor.url}/{_PROMPT_SHIELD}", + "api_key": "synthetic-guardrail-key", + }, + }, + ] + configuration["policies"] = { + _STOPS_AFTER_FIRST: _policy( + "pre_call", ((_REQUEST_FIRST, "allow", "block"), (_REQUEST_SECOND, "next", "block")) + ), + _RUNS_EVERY_STEP: _policy("pre_call", ((_REQUEST_FIRST, "next", "block"), (_REQUEST_SECOND, "next", "block"))), + _BLOCKS_AT_FIRST: _policy( + "pre_call", ((_REQUEST_BLOCKER, "next", "block"), (_REQUEST_SECOND, "next", "block")) + ), + _RESPONSE_STOPS_AFTER_FIRST: _policy( + "post_call", ((_RESPONSE_FIRST, "allow", "block"), (_RESPONSE_SECOND, "next", "block")) + ), + _SHIELDS_THEN_CHECKS: _policy( + "pre_call", ((_PROMPT_SHIELD, "next", "block"), (_REQUEST_FIRST, "next", "block")) + ), + _MOVES_PAST_FIRST: _policy( + "pre_call", ((_REQUEST_BLOCKER, "next", "next"), (_REQUEST_SECOND, "next", "block")) + ), + _RESPONSE_MOVES_PAST_FIRST: _policy( + "post_call", ((_RESPONSE_BLOCKER, "next", "next"), (_RESPONSE_SECOND, "next", "block")) + ), + _SCANS_AFTER_RELEASE: _policy("post_call", ((_LIVE_FIRST, "next", "next"), (_LIVE_SECOND, "next", "next"))), + _CANNOT_HOLD_STREAM: _policy( + "post_call", ((_UNGATED_STEP, "next", "block"), (_RESPONSE_SECOND, "next", "block")) + ), + } + configuration["policy_attachments"] = [{"policy": policy, "tags": [policy]} for policy in configuration["policies"]] + path: Final = directory / "guardrail-pipeline-applied-header.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("guardrail-pipeline-applied-header") + with ( + wire_server(_vendor) as vendor, + wire_server(_completion) as upstream, + gateway_from_environment() as shared, + owned_proxy(shared, directory, {}, config=_config(directory, vendor)) as owned, + owned.scenario() as scenario, + ): + yield _Rig(owned, vendor, upstream, scenario.model(api_base=f"{upstream.url}/v1", api_key="sk-fixture")) + + +def _ask_as_team_tagged(rig: _Rig, tag: str, *, stream: bool = False) -> httpx.Response: + rig.vendor.drain() + rig.upstream.drain() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [tag]}) + key: Final = scenario.key(team_id=team) + return rig.gateway.request( + "POST", + "/v1/chat/completions", + { + "model": rig.model, + "messages": [{"role": "user", "content": f"hello {uuid.uuid4().hex}"}], + "stream": stream, + }, + key=key, + ) + + +def _steps_that_ran(rig: _Rig) -> tuple[str, ...]: + return tuple(request.target.split("/")[1] for request in rig.vendor.drain()) + + +def test_a_request_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert len(rig.upstream.drain()) == 1 + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +def test_a_streamed_request_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert len(rig.upstream.drain()) == 1 + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +def test_a_request_pipeline_that_runs_every_step_names_every_step_in_order(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RUNS_EVERY_STEP) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST, _REQUEST_SECOND) + assert response.headers.get(_HEADER) == f"{_REQUEST_FIRST},{_REQUEST_SECOND}", dict(response.headers) + + +def test_a_request_pipeline_names_a_prompt_shield_step_that_ran_before_a_generic_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _SHIELDS_THEN_CHECKS) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_PROMPT_SHIELD, _REQUEST_FIRST) + assert set(response.headers.get(_HEADER, "").split(",")) == {_PROMPT_SHIELD, _REQUEST_FIRST}, dict(response.headers) + + +def test_a_request_pipeline_that_blocks_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _BLOCKS_AT_FIRST) + assert response.status_code == 400, response.text + assert _steps_that_ran(rig) == (_REQUEST_BLOCKER,) + assert rig.upstream.drain() == () + assert response.headers.get(_HEADER) == _REQUEST_BLOCKER, dict(response.headers) + + +def test_a_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + +def test_a_streamed_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + +def test_a_request_pipeline_names_a_step_that_flagged_the_request_before_moving_on(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _MOVES_PAST_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_BLOCKER, _REQUEST_SECOND) + assert len(rig.upstream.drain()) == 1 + assert set(response.headers.get(_HEADER, "").split(",")) == {_REQUEST_BLOCKER, _REQUEST_SECOND}, dict( + response.headers + ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +def test_a_response_pipeline_names_a_step_that_flagged_the_response_before_moving_on(rig: _Rig, stream: bool) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_MOVES_PAST_FIRST, stream=stream) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_BLOCKER, _RESPONSE_SECOND) + assert set(response.headers.get(_HEADER, "").split(",")) == {_RESPONSE_BLOCKER, _RESPONSE_SECOND}, dict( + response.headers + ) + + +def test_a_stream_released_before_its_pipeline_scans_it_names_the_steps_in_its_headers(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _SCANS_AFTER_RELEASE, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_LIVE_FIRST, _LIVE_SECOND) + assert set(response.headers.get(_HEADER, "").split(",")) == {_LIVE_FIRST, _LIVE_SECOND}, dict(response.headers) + + +def test_a_stream_its_pipeline_cannot_hold_names_the_steps_that_run_on_their_own(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _CANNOT_HOLD_STREAM, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _RESPONSE_SECOND in _steps_that_ran(rig) + assert set(response.headers.get(_HEADER, "").split(",")) == {_UNGATED_STEP, _RESPONSE_SECOND}, dict( + response.headers + ) diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index c42878bd4f7..eaee18c3929 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -1270,7 +1270,7 @@ def _assert_passed_with_discard_warning(result, caplog, reason): assert len(discard_warnings) == 1 assert reason in discard_warnings[0] assert "text rewrites included" in discard_warnings[0] - assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] @pytest.mark.asyncio @@ -1433,6 +1433,16 @@ async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_he assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] +@pytest.mark.asyncio +async def test_streaming_step_that_flags_the_stream_is_named_when_the_pipeline_moves_on(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert [(step.outcome, step.action_taken) for step in result.step_results] == [("fail", "next")] + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + + @pytest.mark.asyncio async def test_streaming_step_records_guardrail_information_once_on_block(monkeypatch): monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()]) 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 acc1463b072..a96c515228f 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -30,12 +30,16 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + get_logging_caching_headers, +) from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler from litellm.proxy.policy_engine.pipeline_executor import recorded_guardrail_information from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.utils import CallTypes from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, @@ -275,7 +279,7 @@ def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypat monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"]}} - proxy_logging._process_guardrail_metadata(data) + proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth()) snapshot = { "call_count": len(calls), "first_name": calls[0]["name"], @@ -300,7 +304,7 @@ def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeyp monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}} - proxy_logging._process_guardrail_metadata(data) + proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth()) assert calls == ["g2"] @@ -312,12 +316,12 @@ def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypat "add_guardrail_to_applied_guardrails_header", MagicMock(side_effect=AssertionError("should not be called")), ) - proxy_logging._process_guardrail_metadata({}) + proxy_logging._process_guardrail_metadata({}, UserAPIKeyAuth()) def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): with pytest.raises(AttributeError): - proxy_logging._process_guardrail_metadata(None) # type: ignore[arg-type] + proxy_logging._process_guardrail_metadata(None, UserAPIKeyAuth()) # type: ignore[arg-type] # --------------------------------------------------------------------------- @@ -1160,6 +1164,233 @@ async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipe assert seen["count"] == 1 +class _PipelineStepGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str, event_hook: str | list[str]) -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=False) + self.calls = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + return data + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.calls += 1 + + +def _two_step_pipeline_data(mode: str, first_on_pass: str, first_on_fail: str = "block") -> dict[str, Any]: + pipeline = GuardrailPipeline( + mode=mode, + steps=[ + PipelineStep(guardrail="first", on_pass=first_on_pass, on_fail=first_on_fail), + PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + ], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second"], + "_guardrail_pipelines": [("governance", pipeline)], + "_pipeline_managed_guardrails": {"first", "second"}, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first_on_pass", "expected_calls", "expected_header"), + [("allow", (1, 0), ["first"]), ("next", (1, 1), ["first", "second"])], +) +async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( + proxy_logging, make_user_api_key_auth, monkeypatch, first_on_pass, expected_calls, expected_header +): + first = _PipelineStepGuardrail("first", "pre_call") + second = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", first_on_pass), + call_type="completion", + ) + + assert (first.calls, second.calls) == expected_calls + assert result["metadata"]["applied_guardrails"] == expected_header + + +@pytest.mark.asyncio +async def test_pre_call_hook_names_a_skipped_step_whose_guardrail_also_runs_outside_the_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + first = _PipelineStepGuardrail("first", "pre_call") + second = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "allow"), + call_type="completion", + ) + + assert (first.calls, second.calls) == (1, 0) + assert result["metadata"]["applied_guardrails"] == ["first", "second"] + + +class _RaisingPipelineStepGuardrail(_PipelineStepGuardrail): + def __init__(self, guardrail_name: str, error: Exception) -> None: + super().__init__(guardrail_name, "pre_call") + self.error = error + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + raise self.error + + +_PIPELINE_STEP_ERRORS: Final = ( + pytest.param(HTTPException(status_code=400, detail="flagged"), id="flagged"), + pytest.param(RuntimeError("vendor unreachable"), id="errored"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) +async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_on( + proxy_logging, make_user_api_key_auth, monkeypatch, error +): + first = _RaisingPipelineStepGuardrail("first", error) + second = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), + call_type="completion", + ) + + assert (first.calls, second.calls) == (1, 1) + assert result["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) +async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( + proxy_logging, make_user_api_key_auth, monkeypatch, error +): + first = _RaisingPipelineStepGuardrail("first", error) + second = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), + call_type="completion", + ) + + with pytest.raises(type(error)): + await first.async_pre_call_deployment_hook( + kwargs={**result, "guardrails": ["first"]}, call_type=CallTypes.completion + ) + await second.async_pre_call_deployment_hook( + kwargs={**result, "guardrails": ["second"]}, call_type=CallTypes.completion + ) + + assert (first.calls, second.calls) == (2, 1) + + +class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + data["litellm_metadata"] = {} + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + return data + + +@pytest.mark.asyncio +async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_later_step_opened( + proxy_logging, make_user_api_key_auth, monkeypatch +): + first = _PipelineStepGuardrail("first", "pre_call") + second = _OpensLitellmMetadataGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "next"), + call_type="completion", + ) + + assert (first.calls, second.calls) == (1, 1) + assert "applied_guardrails" not in result["metadata"] + assert result["litellm_metadata"]["applied_guardrails"] == ["second", "first"] + assert set(get_logging_caching_headers(result)["x-litellm-applied-guardrails"].split(",")) == {"first", "second"} + + +@pytest.mark.asyncio +async def test_post_call_pipeline_steps_are_named_only_once_they_run( + proxy_logging, make_user_api_key_auth, monkeypatch +): + first = _PipelineStepGuardrail("first", "post_call") + second = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + data = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("post_call", "allow"), + call_type="completion", + ) + assert "applied_guardrails" not in data["metadata"] + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 0) + assert data["metadata"]["applied_guardrails"] == ["first"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("asks_for", "request_route", "streams_live", "expected_header"), + [ + pytest.param({"stream": False}, "/custom/stream", False, None, id="whole-response"), + pytest.param({"stream": True}, "/v1/chat/completions", False, None, id="stream-held-for-the-pipeline"), + pytest.param({"stream": True}, "/v1/chat/completions", True, ["first", "second"], id="stream-released-live"), + pytest.param( + {"stream": True}, "/custom/stream", False, ["first", "second"], id="stream-runs-steps-on-their-own" + ), + pytest.param({"method": "message/send"}, "/a2a/agent", False, None, id="a2a-whole-response"), + pytest.param({"method": "message/stream"}, "/a2a/agent", False, ["first", "second"], id="a2a-stream"), + pytest.param({"method": "SubscribeToTask"}, "/a2a/agent", False, ["first", "second"], id="a2a-v1-resubscribe"), + ], +) +async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_after_its_headers( + proxy_logging, make_user_api_key_auth, monkeypatch, asks_for, request_route, streams_live, expected_header +): + guardrails = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] + for guardrail in guardrails: + guardrail.streaming_buffer_until_moderated = not streams_live + monkeypatch.setattr(litellm, "callbacks", guardrails) + data = _two_step_pipeline_data("post_call", "next") + data["metadata"]["_guardrail_pipelines"] = [ + ( + "governance", + GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="next") for name in ("first", "second")], + ), + ) + ] + data.update(asks_for) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route=request_route), + data=data, + call_type="completion", + ) + + assert result["metadata"].get("applied_guardrails") == expected_header + + @pytest.mark.asyncio async def test_post_call_pipeline_replacement_response_reaches_caller( proxy_logging, make_user_api_key_auth, monkeypatch diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index b7fc8c6dbeb..615e39c495d 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -170,7 +170,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox proxy_logging.slack_alerting_instance = MagicMock(alerting=None) invoked = {} - def fake_process(d): + def fake_process(d, user_api_key_dict): invoked["data"] = d proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]