From dcf0f624006bdbd619f3247487b47c97f455fa9a Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 29 Oct 2026 22:14:54 -0700 Subject: [PATCH] fix(guardrails): keep naming pipeline steps that ran at the deployment level, in a later pipeline, or on another request --- litellm/integrations/custom_guardrail.py | 98 +++++++-- litellm/proxy/utils.py | 49 +++-- .../test_guardrail_pipeline_applied_header.py | 52 ++++- .../proxy_logging/test_guardrail_pipeline.py | 188 ++++++++++++++++++ .../utils/proxy_logging/test_pre_call_hook.py | 5 +- 5 files changed, 351 insertions(+), 41 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 5e1d792718e..312fd2a542c 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -3,7 +3,7 @@ import copy import hashlib import os import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args @@ -70,6 +70,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions" # proxy's metadata sanitizer. _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16) _PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16) +_DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16) _GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422}) @@ -219,6 +220,25 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = ) +def _str_collection(value: object) -> frozenset[str]: + if not isinstance(value, (list, tuple, set, frozenset)): + return frozenset() + items: Final = cast("Iterable[object]", value) # cast-ok: a metadata collection of guardrail names + return frozenset(item for item in items if isinstance(item, str)) + + +def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> bool: + if marker is None: + return False + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if isinstance(executed, list) and marker in executed: + return True + return False + + class CustomGuardrail(CustomLogger): # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. use_native_during_call_hook: ClassVar[bool] = False @@ -825,11 +845,14 @@ class CustomGuardrail(CustomLogger): return False - def _pre_call_marker(self, *, raised: bool = False) -> str | None: + def _marker(self, token: str) -> str | None: name: Final = self.guardrail_name if not name: return None - return f"{_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN}:{name}" + return f"{token}:{name}" + + def _pre_call_marker(self, *, raised: bool = False) -> str | None: + return self._marker(_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN) def mark_pre_call_hook_ran(self, data: dict[str, object], *, raised: bool = False) -> None: """ @@ -864,16 +887,47 @@ class CustomGuardrail(CustomLogger): 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"): - meta = data.get(meta_key) - if isinstance(meta, dict): - executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) - if isinstance(executed, list) and marker in executed: - return True - return False + return _request_carries_marker(data, self._pre_call_marker(raised=raised)) + + def deployment_hook_ran(self, data: Mapping[str, object]) -> bool: + """Whether a deployment-level hook ran this guardrail as a policy pipeline step of this request.""" + return _request_carries_marker(data, self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN)) + + def _record_deployment_hook_run( + self, + data: dict[str, object], # mutable-ok: the request payload the run is recorded on + ) -> None: + """ + Record that a deployment-level hook ran this guardrail when it is a policy pipeline step of + this request, and name it in the applied guardrails when the request's metadata lists it. + """ + name: Final = self.guardrail_name + marker: Final = self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN) + buckets: Final = tuple( + cast("dict[str, object]", bucket) # cast-ok: metadata buckets are str-keyed + for bucket in (data.get("metadata"), data.get("litellm_metadata")) + if isinstance(bucket, dict) + ) + policy_state: Final = next( + (bucket for bucket in buckets if name in _str_collection(bucket.get("_pipeline_managed_guardrails"))), + None, + ) + if marker is None or policy_state is None: + return + recorded: Final = policy_state.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if not isinstance(recorded, list): + policy_state[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] + elif marker not in recorded: + cast("list[object]", recorded).append(marker) # cast-ok: markers are only compared as strings + requested: Final = next((bucket.get("guardrails") for bucket in buckets if "guardrails" in bucket), None) + if name not in _str_collection(requested): + 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 @classmethod def uses_apply_guardrail_interface(cls) -> bool: @@ -922,12 +976,15 @@ class CustomGuardrail(CustomLogger): target: Final = self._deployment_hook_target() if target is not self: kwargs["guardrail_to_apply"] = self - result: Final = await target.async_pre_call_hook( - user_api_key_dict=_user_api_key_auth_from_request(kwargs), - cache=dc, - data=kwargs, - call_type="completion" if call_type == CallTypes.completion else "acompletion", - ) + try: + result: Final = await target.async_pre_call_hook( + user_api_key_dict=_user_api_key_auth_from_request(kwargs), + cache=dc, + data=kwargs, + call_type="completion" if call_type == CallTypes.completion else "acompletion", + ) + finally: + self._record_deployment_hook_run(kwargs) if result is not None and isinstance(result, dict): result_messages: Final = result.get("messages") @@ -1011,6 +1068,9 @@ class CustomGuardrail(CustomLogger): finally: if target is not self: request_data.pop("guardrail_to_apply", None) + self._record_deployment_hook_run( + cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed + ) if not self._is_valid_response_type(result): return None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5fd0e9d1f36..48ee6ba35d1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -713,6 +713,7 @@ def _ran_or_runs_outside_pipelines( ) -> bool: return callback is not None and ( callback.pre_call_hook_ran_or_raised(data) + or callback.deployment_hook_ran(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 @@ -766,13 +767,25 @@ def _pipelines_settled_by( return pipelines[index:] if result.terminal_action != "allow" else pipelines[index : index + 1] +def _steps_later_pipelines_still_run( + pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult +) -> frozenset[str]: + """The steps of the later response pipelines, which still run when this pipeline's verdict lets the request on.""" + if result.terminal_action != "allow": + return frozenset() + later: Final = tuple( + (policy_name, pipeline) for policy_name, pipeline in pipelines[index + 1 :] if pipeline.mode == "post_call" + ) + return _pipeline_step_guardrail_names(later) + + def _drop_response_steps_that_did_not_run( data: dict[str, object], # mutable-ok: the request's applied guardrails are trimmed in place pipelines: Sequence[tuple[str, "GuardrailPipeline"]], - steps_that_ran: frozenset[str], + steps_to_keep: frozenset[str], ) -> None: """ - Drop from the applied guardrails the steps of these response pipelines that no pipeline ran, + Drop from the applied guardrails the steps of these response pipelines outside ``steps_to_keep``, unless their guardrail also runs outside pipelines. Every step is named up front, so a log built before the pipeline finished, such as a cache hit's, still lists the steps that ran. The list is trimmed in place, since a route's logging object may hold a @@ -781,7 +794,7 @@ def _drop_response_steps_that_did_not_run( pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") post_call: Final = _pipeline_step_guardrail_names(pipelines) not_run: Final = frozenset( - name for name in post_call - steps_that_ran if _runs_only_in_pipelines(name, data, pre_call, post_call) + name for name in post_call - steps_to_keep if _runs_only_in_pipelines(name, data, pre_call, post_call) ) if not not_run: return @@ -2735,8 +2748,12 @@ class ProxyLogging: data.pop("prompt_version", None) data.pop("prompt_environment", None) - def _process_guardrail_metadata(self, data: dict) -> None: - """Process guardrails from metadata and add to applied_guardrails.""" + def _process_guardrail_metadata(self, data: dict, *, skip_guardrails: bool = False) -> None: + """ + Process guardrails from metadata and add to applied_guardrails, leaving out the pre_call + policy pipeline steps that have not run. With ``skip_guardrails`` the pipelines run on + another request, so every step is added. + """ from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) @@ -2760,8 +2777,10 @@ class ProxyLogging: if not isinstance(applied_guardrails, list): applied_guardrails = [] - not_run: Final = _request_steps_that_have_not_run( - cast(dict[str, object], data) # cast-ok: request payload + not_run: Final = ( + frozenset[str]() + if skip_guardrails + else _request_steps_that_have_not_run(cast(dict[str, object], data)) # cast-ok: request payload ) for guardrail_name in guardrails_in_metadata: if ( @@ -2822,7 +2841,9 @@ class ProxyLogging: steps_that_ran = steps_that_ran | _steps_that_ran(result) if log_waits: _drop_response_steps_that_did_not_run( - request, _pipelines_settled_by(pipelines, index, result), steps_that_ran + request, + _pipelines_settled_by(pipelines, index, result), + steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result), ) data = self._handle_pipeline_result( @@ -3064,7 +3085,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, skip_guardrails=skip_guardrails) return data parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( @@ -3168,17 +3189,17 @@ class ProxyLogging: ) if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) 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, skip_guardrails=skip_guardrails) return data except Exception: if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) raise async def _run_parallel_pre_call_guardrails( @@ -4757,7 +4778,9 @@ class ProxyLogging: steps_that_ran = steps_that_ran | _steps_that_ran(result) if _stream_log_waits_for_its_pipelines(request_data): _drop_response_steps_that_did_not_run( - request_data, _pipelines_settled_by(pipelines, index, result), steps_that_ran + request_data, + _pipelines_settled_by(pipelines, index, result), + steps_that_ran | _steps_later_pipelines_still_run(pipelines, index, result), ) try: ProxyLogging._handle_pipeline_result( diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index bebef167f92..b5129d03ee5 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -45,6 +45,7 @@ _RUNS_EVERY_STEP: Final = "runs-every-step" _BLOCKS_AT_FIRST: Final = "blocks-at-first" _RESPONSE_STOPS_AFTER_FIRST: Final = "response-stops-after-first" _RESPONSE_RUNS_EVERY_STEP: Final = "response-runs-every-step" +_RESPONSE_STOPS_AFTER_SECOND: Final = "response-stops-after-second" _SHIELDS_THEN_CHECKS: Final = "shields-then-checks" _MOVES_PAST_FIRST: Final = "moves-past-first" _RESPONSE_MOVES_PAST_FIRST: Final = "response-moves-past-first" @@ -418,6 +419,9 @@ def _config(directory: Path, vendor: Wire, upstream: Wire) -> Path: _RESPONSE_RUNS_EVERY_STEP: _policy( "post_call", ((_RESPONSE_FIRST, "next", "block"), (_RESPONSE_SECOND, "next", "block")) ), + _RESPONSE_STOPS_AFTER_SECOND: _policy( + "post_call", ((_RESPONSE_SECOND, "allow", "block"), (_RESPONSE_FIRST, "next", "block")) + ), _SHIELDS_THEN_CHECKS: _policy( "pre_call", ((_PROMPT_SHIELD, "next", "block"), (_REQUEST_FIRST, "next", "block")) ), @@ -695,6 +699,35 @@ def test_the_spend_log_and_the_trace_name_only_the_steps_that_ran( ] * len(spans), spans +@pytest.mark.parametrize( + ("tag", "model_guardrail", "expected"), + [ + pytest.param(_STOPS_AFTER_FIRST, _REQUEST_SECOND, [_REQUEST_FIRST, _REQUEST_SECOND], id="request-pipeline"), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, _RESPONSE_SECOND, [_RESPONSE_FIRST, _RESPONSE_SECOND], id="response-pipeline" + ), + ], +) +def test_a_model_level_guardrail_the_pipeline_skipped_is_named_everywhere_once_the_model_runs_it( + rig: _Rig, tag: str, model_guardrail: str, expected: list[str] +) -> None: + since: Final = recorded_spans(rig.traces)[0] + with rig.gateway.scenario() as scenario: + model: Final = scenario.model( + api_base=f"{rig.upstream.url}/v1", api_key="sk-fixture", guardrails=[model_guardrail] + ) + team: Final = scenario.team(metadata={"tags": [tag]}) + key: Final = scenario.key(team_id=team) + rig.vendor.drain() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _body("/v1/chat/completions", model, False), key=key + ) + header, rows, spans = _applied_in_header_spend_log_and_trace(rig, since, team, response) + assert response.status_code == 200, response.text + assert sorted(_steps_that_ran(rig)) == expected + assert (header, rows, spans) == (expected, [expected], [expected] * len(spans)) + + @pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) def test_a_response_pipeline_names_a_judge_step_that_records_its_own_results(rig: _Rig, stream: bool) -> None: team, response = _ask_as_new_team(rig, _JUDGES_THEN_CHECKS, stream=stream) @@ -1005,17 +1038,22 @@ def agent(rig: _Rig) -> Iterator[str]: @pytest.mark.parametrize( - ("tag", "expected"), + ("tags", "expected"), [ - pytest.param(_SCANS_AFTER_RELEASE, [_LIVE_FIRST, _LIVE_SECOND], id="every-step-passes"), + pytest.param((_SCANS_AFTER_RELEASE,), [_LIVE_FIRST, _LIVE_SECOND], id="every-step-passes"), pytest.param( - _RESPONSE_MOVES_PAST_FIRST, [_RESPONSE_BLOCKER, _RESPONSE_SECOND], id="first-step-flags-and-moves-on" + (_RESPONSE_MOVES_PAST_FIRST,), [_RESPONSE_BLOCKER, _RESPONSE_SECOND], id="first-step-flags-and-moves-on" + ), + pytest.param((_RESPONSE_STOPS_AFTER_FIRST,), [_RESPONSE_FIRST], id="first-step-allows"), + pytest.param( + (_RESPONSE_STOPS_AFTER_FIRST, _RESPONSE_STOPS_AFTER_SECOND), + [_RESPONSE_FIRST, _RESPONSE_SECOND], + id="two-policies-each-run-the-step-the-other-skipped", ), - pytest.param(_RESPONSE_STOPS_AFTER_FIRST, [_RESPONSE_FIRST], id="first-step-allows"), ], ) def test_an_agent_message_names_every_response_step_that_ran_in_the_spend_log_and_the_trace( - rig: _Rig, agent: str, tag: str, expected: list[str] + rig: _Rig, agent: str, tags: tuple[str, ...], expected: list[str] ) -> None: since: Final = recorded_spans(rig.traces)[0] rig.vendor.drain() @@ -1026,7 +1064,7 @@ def test_an_agent_message_names_every_response_step_that_ran_in_the_spend_log_an "parts": [{"kind": "text", "text": f"hello {uuid.uuid4().hex}"}], } with rig.gateway.scenario() as scenario: - team: Final = scenario.team(metadata={"tags": [tag]}) + team: Final = scenario.team(metadata={"tags": list(tags)}) response: Final = rig.gateway.request( "POST", f"/a2a/{agent}", @@ -1035,7 +1073,7 @@ def test_an_agent_message_names_every_response_step_that_ran_in_the_spend_log_an headers={"Accept": "application/json"}, ) assert response.status_code == 200, response.text - assert list(_steps_that_ran(rig)) == expected + assert sorted(_steps_that_ran(rig)) == expected rows: Final = eventually( lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), lambda found: len(found) == 1, 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 ac314667347..39f4e9a10ef 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1672,6 +1672,194 @@ async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_ assert data["metadata"]["applied_guardrails"] == ["first", "second"] +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_step_a_later_pipeline_runs_after_an_earlier_one_opens_litellm_metadata( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + stops_before_second: Final = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="first", on_pass="allow"), PipelineStep(guardrail="second", on_pass="next")], + ) + stops_before_first: Final = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="second", on_pass="allow"), PipelineStep(guardrail="first", on_pass="next")], + ) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second"], + "_guardrail_pipelines": [ + ("stops-before-second", stops_before_second), + ("stops-before-first", stops_before_first), + ], + "_pipeline_managed_guardrails": {"first", "second"}, + }, + "litellm_logging_obj": _logging_obj(log_parked=True), + }, + call_type="completion", + ) + 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, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_pre_call_hook_names_every_pipeline_step_when_the_pipelines_run_on_another_request( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result: Final = 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", + skip_guardrails=True, + ) + + assert (first.calls, second.calls) == (0, 0) + assert result["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_deployment_hook_names_a_skipped_pre_call_step_it_ran( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + data: Final = 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", + ) + + await second.async_pre_call_deployment_hook( + kwargs={**data, "guardrails": ["second"]}, call_type=CallTypes.completion + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + assert second.deployment_hook_ran(data) + assert not first.deployment_hook_ran(data) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) +async def test_deployment_hook_names_a_skipped_pre_call_step_that_raised( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + error: Exception, +) -> None: + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _RaisingPipelineStepGuardrail("second", error) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + data: Final = 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", + ) + + with pytest.raises(type(error)): + await second.async_pre_call_deployment_hook( + kwargs={**data, "guardrails": ["second"]}, call_type=CallTypes.completion + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_skipped_step_a_deployment_hook_ran( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={**_two_step_pipeline_data("post_call", "allow"), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + + await second.async_post_call_success_deployment_hook( + request_data={**data, "guardrails": ["second"]}, + response=litellm.ModelResponse(), + call_type=CallTypes.acompletion, + ) + 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, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +async def test_deployment_hook_leaves_a_request_without_policy_pipelines_untouched() -> None: + guardrail: Final = _PipelineStepGuardrail("second", "pre_call") + kwargs: Final[dict[str, object]] = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["second"], + "metadata": {"guardrails": ["second"]}, + } + + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) + + assert guardrail.calls == 1 + assert kwargs["metadata"] == {"guardrails": ["second"]} + + +@pytest.mark.asyncio +async def test_deployment_hook_names_a_pipeline_step_it_ran_before_any_guardrail_was_named() -> None: + guardrail: Final = _PipelineStepGuardrail("second", "pre_call") + metadata: Final[dict[str, object]] = { + "guardrails": ["second"], + "_pipeline_managed_guardrails": frozenset({"second"}), + } + kwargs: Final[dict[str, object]] = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["second"], + "metadata": metadata, + } + + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) + + assert guardrail.calls == 1 + assert metadata["applied_guardrails"] == ["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") + kwargs: Final[dict[str, object]] = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "guardrails": ["second"], + "metadata": {"_pipeline_managed_guardrails": frozenset({"second"})}, + } + + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) + + assert guardrail.deployment_hook_ran(kwargs) + assert not _PipelineStepGuardrail("first", "pre_call").deployment_hook_ran(kwargs) + assert "applied_guardrails" not in kwargs["metadata"] + + @pytest.mark.asyncio @pytest.mark.parametrize( ("log_waits", "expected_after_stream"), 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..dc1267ec289 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,8 +170,9 @@ 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, *, skip_guardrails): invoked["data"] = d + invoked["skip_guardrails"] = skip_guardrails proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment] out = await proxy_logging.pre_call_hook( @@ -180,7 +181,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox call_type="completion", ) assert out is data - assert invoked["data"] is data + assert invoked == {"data": data, "skip_guardrails": False} @pytest.mark.asyncio