diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py index 4f44a4adeed..c49b334760b 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py @@ -61,9 +61,11 @@ async def test_apply_guardrail_endpoint_returns_correct_response( assert response.response_text == "Redacted text: [REDACTED] and [REDACTED]" # Verify the guardrail was called with correct parameters + # the endpoint always forwards the proxy's own metadata, which this + # fixture leaves empty, so the guardrail can read the calling key mock_guardrail.apply_guardrail.assert_called_once_with( inputs={"texts": ["Test text with PII"]}, - request_data={}, + request_data={"metadata": {}}, input_type="request", ) @@ -198,5 +200,7 @@ async def test_apply_guardrail_endpoint_without_optional_params(mock_proxy_loggi # Verify the guardrail was called with correct parameters mock_guardrail.apply_guardrail.assert_called_once_with( - inputs={"texts": ["Test text"]}, request_data={}, input_type="request" + inputs={"texts": ["Test text"]}, + request_data={"metadata": {}}, + input_type="request", ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py index 14338a58248..19b579cebc4 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py @@ -2181,11 +2181,12 @@ class TestXecGuardApplyGuardrailKeyTargeting: } resp = _make_response({"decision": "SAFE", "trace_id": "tr"}) with patch.object(gr.async_handler, "post", return_value=resp) as post: - await gr.apply_guardrail( + out = await gr.apply_guardrail( inputs={"texts": ["hi"]}, request_data=data, input_type="request", ) + assert out == {"texts": ["hi"]}, "a SAFE verdict hands the inputs back untouched" post.assert_called_once() assert post.call_args.kwargs["json"]["policy_names"] == ["Config_Level_Policy"] @@ -2272,24 +2273,28 @@ class TestXecGuardLoggingHookKeyTargeting: gr = _extension_guardrail(apply_to_aliases=["prod"]) kwargs = self._kwargs({"user_api_key_alias": "dev"}) with patch.object(gr.async_handler, "post") as post: - await gr.async_logging_hook( + out_kwargs, _ = await gr.async_logging_hook( kwargs=kwargs, result=_build_model_response("some answer"), call_type="acompletion", ) post.assert_not_called() + assert out_kwargs is kwargs + assert "guardrail_information" not in kwargs["standard_logging_object"] @pytest.mark.asyncio async def test_top_level_metadata_is_also_honoured(self): gr = _extension_guardrail(except_aliases=["internal"]) kwargs = self._kwargs({"user_api_key_alias": "internal"}, nested=False) with patch.object(gr.async_handler, "post") as post: - await gr.async_logging_hook( + out_kwargs, _ = await gr.async_logging_hook( kwargs=kwargs, result=_build_model_response("some answer"), call_type="acompletion", ) post.assert_not_called() + assert out_kwargs is kwargs + assert "guardrail_information" not in kwargs["standard_logging_object"] @pytest.mark.asyncio async def test_targeted_key_is_still_scanned(self): @@ -2828,11 +2833,12 @@ class TestXecGuardScanMeta: with patch.object( gr.async_handler, "post", side_effect=[scan_ok, grounding_ok] ) as mock_post: - await gr.apply_guardrail( + out = await gr.apply_guardrail( inputs={"texts": ["text"]}, request_data=data, input_type="response", ) + assert out["texts"] == ["text"], "two SAFE verdicts leave the response alone" assert "meta" in mock_post.call_args_list[0].kwargs["json"] assert "meta" not in mock_post.call_args_list[1].kwargs["json"] @@ -3023,6 +3029,8 @@ class TestXecGuardScanMetaAutoFields: result=_build_model_response("some answer"), call_type="acompletion", ) + info_list = kwargs["standard_logging_object"]["guardrail_information"] + assert info_list[0]["guardrail_response"]["trace_id"] == "meta-auto-log" assert mock_post.call_args.kwargs["json"]["meta"]["data"] == { "key_id": "abc123", "key_alias": "team-alpha",