diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/__init__.py index 747ef3cc911..f7c05fb8c45 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/__init__.py @@ -24,6 +24,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" **{ **litellm_params.model_dump(), "guardrail_name": guardrail_name, + "event_hook": litellm_params.mode, "default_on": litellm_params.default_on or False, } ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index f75f6db82a0..41b3c1368a5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -293,7 +293,7 @@ async def async_pre_call_hook( return None @log_guardrail_information - async def async_post_call_hook( + async def async_post_call_success_hook( self, data: Dict[str, Any], user_api_key_dict: UserAPIKeyAuth, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index e4c9d056dd9..12d84f9530e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -280,7 +280,7 @@ async def test_safe_response_allowed( mock_response = {"action": "allow", "category": "benign"} with patch.object(handler, "_call_panw_api", return_value=mock_response): - result = await handler.async_post_call_hook( + result = await handler.async_post_call_success_hook( data=request_data, user_api_key_dict=user_api_key_dict, response=safe_response, @@ -299,7 +299,7 @@ async def test_harmful_response_blocked( with patch.object(handler, "_call_panw_api", return_value=mock_response): with pytest.raises(HTTPException) as exc_info: - await handler.async_post_call_hook( + await handler.async_post_call_success_hook( data=request_data, user_api_key_dict=user_api_key_dict, response=harmful_response,