Merge branch 'litellm_lakera_v2_gemini_contents' into litellm_passthrough_applied_guardrails_header

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -147,6 +147,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 (
@@ -922,11 +923,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
@@ -938,7 +939,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:

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() == ()

tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py

@@ -11361,8 +11361,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:
@@ -11376,7 +11393,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()
@@ -11404,7 +11421,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={},
@@ -11412,7 +11430,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]:
@@ -11427,7 +11449,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()
@@ -11437,6 +11467,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"]})