diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 23c1202ea07..61651659556 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -141,6 +141,7 @@ from litellm.proxy.utils import normalize_route_for_root_path from litellm.repositories.team_repository import TeamRepository from litellm.secret_managers.main import get_secret_str from litellm.types import utils as types_utils +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.litellm_params import ProxyRequestState, wire_names from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -863,11 +864,11 @@ async def _inherit_auth_guardrails( configured_guardrails: Mapping[str, bool] | None, inherit: bool, ) -> Mapping[str, bool] | Sequence[object] | None: - """Return the guardrails a pass-through request runs. + """Return the guardrails a pass-through request runs after the provider responds. The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team, - project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` and returned, or None - when there are none. + project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` for the pre-call hook, + and returned only when at least one of them runs post_call, so the response is relayed untouched otherwise. """ if configured_guardrails or not inherit or not parsed_body: return configured_guardrails @@ -879,7 +880,18 @@ async def _inherit_auth_guardrails( user_api_key_dict=user_api_key_dict, ) merged: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") - return tuple(_GUARDRAIL_ENTRIES.validate_python(merged)) if isinstance(merged, list) and merged else None + if not isinstance(merged, list) or not merged or not _named_guardrail_runs_post_call(parsed_body): + return None + return tuple(_GUARDRAIL_ENTRIES.validate_python(merged)) + + +def _named_guardrail_runs_post_call(data: dict[str, object]) -> bool: + return any( + guardrail.default_on is not True + and guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) + for guardrail in litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) + if isinstance(guardrail, CustomGuardrail) + ) def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | None) -> None: diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index e92ebcb5594..09b20bf6653 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -18,6 +18,7 @@ from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter GUARDRAIL: Final = "passthrough-auth-guardrail" +GLOBAL_POST_CALL: Final = "passthrough-global-post-call" BLOCK_MARKER: Final = "block-this-prompt" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" @@ -43,8 +44,17 @@ def _verdict(request: Request) -> Reply: return Reply(body=json.dumps({"action": action, "blocked_reason": "synthetic block"}).encode()) +_OPENAI_BODY: Final = json.dumps( + _JSON.validate_json(chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False).body), indent=2 +).encode() + + def _openai_reply(_request: Request) -> Reply: - return chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False) + return Reply(body=_OPENAI_BODY) + + +def _vendor_down(_request: Request) -> Reply: + return Reply(status=500, body=b'{"error": "guardrail vendor down"}') def _anthropic_reply(_request: Request) -> Reply: @@ -66,6 +76,7 @@ def _gemini_reply(_request: Request) -> Reply: class AuthGuardrailRig: owned: OwnedProxy vendor: Wire + post_call_vendor: Wire upstreams: Mapping[str, Wire] @@ -74,6 +85,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I directory: Final = tmp_path_factory.mktemp("passthrough-auth-guardrails") with ExitStack() as stack: vendor: Final = stack.enter_context(wire_server(_verdict)) + post_call_vendor: Final = stack.enter_context(wire_server(_vendor_down)) upstreams: Final = MappingProxyType( { OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), @@ -97,7 +109,17 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "api_base": vendor.url, "api_key": "synthetic-guardrail-key", }, - } + }, + { + "guardrail_name": GLOBAL_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": True, + "api_base": post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, ], } ) @@ -111,7 +133,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "GEMINI_API_KEY": "synthetic-gemini-key", } owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) - yield AuthGuardrailRig(owned, vendor, upstreams) + yield AuthGuardrailRig(owned, vendor, post_call_vendor, upstreams) def _key_with_guardrail(scenario: Scenario, level: Level) -> str: @@ -132,6 +154,7 @@ def _send(gateway: Gateway, route: str, body: Mapping[str, JsonValue], key: str) def _drain(rig: AuthGuardrailRig) -> None: rig.vendor.drain() + rig.post_call_vendor.drain() for upstream in rig.upstreams.values(): upstream.drain() @@ -191,3 +214,17 @@ def test_request_without_a_body_is_not_sent_to_the_guardrail(rig: AuthGuardrailR assert response.status_code == 200, response.text assert rig.vendor.drain() == () assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + +def test_pre_call_guardrail_relays_the_provider_reply_untouched(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "key") + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _route_body(OPENAI_ROUTE, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert response.content == _OPENAI_BODY + assert len(rig.vendor.drain()) == 1 + assert rig.post_call_vendor.drain() == () diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 294ad45d795..c70e66329ad 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -11360,8 +11360,25 @@ class TestInheritAuthGuardrails: @staticmethod async def _run( - body: dict[str, object], *, inherit: bool, user_api_key_dict: UserAPIKeyAuth + body: dict[str, object], + *, + inherit: bool, + user_api_key_dict: UserAPIKeyAuth, + guardrails: tuple[CustomGuardrail, ...] = (), ) -> tuple[list[dict[str, object]], AsyncMock]: + _, pre_call_data, post_call = await TestInheritAuthGuardrails._run_with_response( + body, inherit=inherit, user_api_key_dict=user_api_key_dict, guardrails=guardrails + ) + return pre_call_data, post_call + + @staticmethod + async def _run_with_response( + body: dict[str, object], + *, + inherit: bool, + user_api_key_dict: UserAPIKeyAuth, + guardrails: tuple[CustomGuardrail, ...] = (), + ) -> tuple[Response, list[dict[str, object]], AsyncMock]: pre_call_data: Final[list[dict[str, object]]] = [] async def _record_pre_call(**kwargs: object) -> object: @@ -11375,7 +11392,7 @@ class TestInheritAuthGuardrails: upstream: Final = httpx.Response( status_code=200, headers={"content-type": "application/json"}, - content=b'{"id": "chatcmpl-1"}', + content=b'{\n "id":"chatcmpl-1"\n}', request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions"), ) request: Final = MagicMock() @@ -11403,7 +11420,8 @@ class TestInheritAuthGuardrails: ) stack.enter_context(patch(f"{module}.read_request_body", new_callable=AsyncMock, return_value=body)) stack.enter_context(patch(f"{module}.safe_get_request_headers", return_value={})) - await pass_through_request( + stack.enter_context(patch.object(litellm, "callbacks", list(guardrails))) + response: Final = await pass_through_request( request=request, target="https://api.openai.com/v1/chat/completions", custom_headers={}, @@ -11411,7 +11429,11 @@ class TestInheritAuthGuardrails: stream=False, inherit_auth_guardrails=inherit, ) - return pre_call_data, proxy_logging.post_call_success_hook + return response, pre_call_data, proxy_logging.post_call_success_hook + + @staticmethod + def _guardrail(name: str, event_hook: GuardrailEventHooks, *, default_on: bool = False) -> CustomGuardrail: + return CustomGuardrail(guardrail_name=name, event_hook=event_hook, default_on=default_on) @staticmethod def _chat_body(**extra: object) -> dict[str, object]: @@ -11426,7 +11448,15 @@ class TestInheritAuthGuardrails: project_metadata={"guardrails": ["project-guard"]}, ) - pre_call, post_call = await self._run(self._chat_body(), inherit=True, user_api_key_dict=auth) + pre_call, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=( + self._guardrail("key-guard", GuardrailEventHooks.pre_call), + self._guardrail("team-guard", GuardrailEventHooks.post_call), + ), + ) assert sorted(pre_call[0]["metadata"]["guardrails"]) == ["key-guard", "project-guard", "team-guard"] post_call.assert_awaited_once() @@ -11436,6 +11466,24 @@ class TestInheritAuthGuardrails: "team-guard", ] + @pytest.mark.asyncio + async def test_pre_call_only_guardrails_relay_the_response_untouched(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + response, pre_call, post_call = await self._run_with_response( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=( + self._guardrail("key-guard", GuardrailEventHooks.pre_call), + self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True), + ), + ) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard"] + post_call.assert_not_awaited() + assert response.body == b'{\n "id":"chatcmpl-1"\n}' + @pytest.mark.asyncio async def test_request_body_guardrails_are_added_to_the_key_guardrails(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]})