diff --git a/docs/my-website/docs/proxy/guardrails/zscaler_ai_guard.md b/docs/my-website/docs/proxy/guardrails/zscaler_ai_guard.md index 94f31c3bfd..2e62600423 100644 --- a/docs/my-website/docs/proxy/guardrails/zscaler_ai_guard.md +++ b/docs/my-website/docs/proxy/guardrails/zscaler_ai_guard.md @@ -100,7 +100,7 @@ In cases where encounter other errors when apply Zscaler AI Guard, return exampl } } ``` -## 6. Sending User Information to Zscaler AI Guard for Analysis (Optional) +## 6. Sending User Information to Zscaler AI Guard (Optional) If you need to send end-user information to Zscaler AI Guard for analysis, you can set the configuration in the environment variables to True and include the relevant information in custom_headers on Zscaler AI Guard. - To send user_api_key_alias: @@ -133,4 +133,30 @@ curl -i http://localhost:8165/v1/chat/completions \ "zguard_policy_id": } }' +``` + +## 8. Set Custom Zscaler AI Guard Policy on Litellm Team OR Key Metadata (Optional) +In addition to setting `zguard_policy_id` in a request or the configuration file, you can also set it in the metadata for LiteLLM Team or Key. The `zguard_policy_id` is determined using the following order of precedence: request, Key, Team, config file. This logic is illustrated below: +``` +user_api_key_metadata = metadata.get("user_api_key_metadata", {}) or {} +team_metadata = metadata.get("team_metadata", {}) or {} +policy_id = ( + metadata.get("zguard_policy_id") + if "zguard_policy_id" in metadata + else ( + user_api_key_metadata.get("zguard_policy_id") + if "zguard_policy_id" in user_api_key_metadata + else ( + team_metadata.get("zguard_policy_id") + if "zguard_policy_id" in team_metadata + else self.policy_id + ) + ) + ) +``` +You can leverage this feature to apply multiple policies configured on the Zscaler AI Guard (ZGuard) to traffic from different applications. (Note: It is recommended to map policies using either Team or Key metadata, but not a mix of both.) + +Example set in Team/Key Metadata, you can set From UI: +``` +{"zguard_policy_id": 100} ``` \ No newline at end of file diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index d62bbb0b45..c60752d795 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -92,14 +92,34 @@ class ZscalerAIGuard(CustomGuardrail): Raises: Exception: If content is blocked by Zscaler AI Guard """ + texts = inputs.get("texts", []) try: verbose_proxy_logger.debug(f"ZscalerAIGuard: Checking {len(texts)} text(s)") + metadata = request_data.get("metadata", {}) - custom_policy_id = request_data.get("metadata", {}).get( - "zguard_policy_id", self.policy_id + user_api_key_metadata = metadata.get("user_api_key_metadata", {}) or {} + team_metadata = metadata.get("team_metadata", {}) or {} + + # Precedence for policy_id: + # 1. metadata.zguard_policy_id # request level + # 2. user_api_key_metadata.zguard_policy_id # Key level + # 3. team_metadata.zguard_policy_id # Team level + # 4. self.policy_id (from environment) # Global + policy_id = ( + metadata.get("zguard_policy_id") + if "zguard_policy_id" in metadata + else ( + user_api_key_metadata.get("zguard_policy_id") + if "zguard_policy_id" in user_api_key_metadata + else ( + team_metadata.get("zguard_policy_id") + if "zguard_policy_id" in team_metadata + else self.policy_id + ) + ) ) - verbose_proxy_logger.debug(f"custom_policy_id: {custom_policy_id}") + verbose_proxy_logger.info(f"policy_id applied: {policy_id}") kwargs = {} if self.send_user_api_key_alias: @@ -116,27 +136,29 @@ class ZscalerAIGuard(CustomGuardrail): ) verbose_proxy_logger.debug(f"inside apply_guardrail kwargs: {kwargs}") - # Check each text (Zscaler processes one at a time) - for text in texts: + zscaler_ai_guard_result = None + direction = "OUT" if input_type == "response" else "IN" + verbose_proxy_logger.debug(f"direction: {direction}") + # Concatenate all texts and send to Zscaler AI Guard + if texts: + concatenated_text = " ".join(texts) zscaler_ai_guard_result = await self.make_zscaler_ai_guard_api_call( zscaler_ai_guard_url=self.zscaler_ai_guard_url, api_key=self.api_key, - policy_id=self.policy_id, - direction="IN", - content=text, + policy_id=policy_id, + direction=direction, + content=concatenated_text, **kwargs, ) - - if ( - zscaler_ai_guard_result - and zscaler_ai_guard_result.get("action") == "BLOCK" - ): - blocking_info = zscaler_ai_guard_result.get( - "zscaler_ai_guard_response" - ) - error_message = f"Content blocked by Zscaler AI Guard: {self.extract_blocking_info(blocking_info)}" - raise Exception(error_message) - + if ( + zscaler_ai_guard_result + and zscaler_ai_guard_result.get("action") == "BLOCK" + ): + blocking_info = zscaler_ai_guard_result.get( + "zscaler_ai_guard_response" + ) + error_message = f"Content blocked by Zscaler AI Guard: {self.extract_blocking_info(blocking_info)}" + raise Exception(error_message) except Exception as e: verbose_proxy_logger.error( "ZscalerAIGuard: Failed to apply guardrail: %s", str(e) diff --git a/tests/guardrails_tests/test_zscaler_ai_guard.py b/tests/guardrails_tests/test_zscaler_ai_guard.py index cf70af510c..9d519c17f6 100644 --- a/tests/guardrails_tests/test_zscaler_ai_guard.py +++ b/tests/guardrails_tests/test_zscaler_ai_guard.py @@ -116,4 +116,131 @@ def test_extract_blocking_info(): blocking_info = guardrail.extract_blocking_info(response) assert blocking_info["transactionId"] == "12345" - assert blocking_info["blockingDetectors"] == ["detector1"] \ No newline at end of file + assert blocking_info["blockingDetectors"] == ["detector1"] + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call", + new_callable=AsyncMock, +) +async def test_apply_guardrail_text_concatenation(mock_api_call): + """ + Test that `apply_guardrail` correctly concatenates texts. + """ + guardrail = ZscalerAIGuard(policy_id=100) + inputs = {"texts": ["Hello", "world"]} + request_data = {} + + await guardrail.apply_guardrail(inputs, request_data, "request") + + mock_api_call.assert_called_once() + call_args = mock_api_call.call_args + assert call_args.kwargs["content"] == "Hello world" + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call", + new_callable=AsyncMock, +) +async def test_policy_id_from_request_metadata(mock_api_call): + """ + Test policy_id is picked from request metadata (highest precedence). + """ + guardrail = ZscalerAIGuard(policy_id=100) + inputs = {"texts": ["test"]} + request_data = { + "metadata": { + "zguard_policy_id": 1, + "user_api_key_metadata": {"zguard_policy_id": 2}, + "team_metadata": {"zguard_policy_id": 3}, + } + } + + await guardrail.apply_guardrail(inputs, request_data, "request") + + mock_api_call.assert_called_once() + assert mock_api_call.call_args.kwargs["policy_id"] == 1 + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call", + new_callable=AsyncMock, +) +async def test_policy_id_from_user_api_key_metadata(mock_api_call): + """ + Test policy_id is picked from user_api_key_metadata (2nd precedence). + """ + guardrail = ZscalerAIGuard(policy_id=100) + inputs = {"texts": ["test"]} + request_data = { + "metadata": { + "user_api_key_metadata": {"zguard_policy_id": 2}, + "team_metadata": {"zguard_policy_id": 3}, + } + } + + await guardrail.apply_guardrail(inputs, request_data, "request") + + mock_api_call.assert_called_once() + assert mock_api_call.call_args.kwargs["policy_id"] == 2 + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call", + new_callable=AsyncMock, +) +async def test_policy_id_from_team_metadata(mock_api_call): + """ + Test policy_id is picked from team_metadata (3rd precedence). + """ + guardrail = ZscalerAIGuard(policy_id=100) + inputs = {"texts": ["test"]} + request_data = {"metadata": {"team_metadata": {"zguard_policy_id": 3}}} + + await guardrail.apply_guardrail(inputs, request_data, "request") + + mock_api_call.assert_called_once() + assert mock_api_call.call_args.kwargs["policy_id"] == 3 + + +@pytest.mark.asyncio +@patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call", + new_callable=AsyncMock, +) +async def test_policy_id_from_init(mock_api_call): + """ + Test policy_id is picked from guardrail initialization (lowest precedence). + """ + guardrail = ZscalerAIGuard(policy_id=100) + inputs = {"texts": ["test"]} + request_data = {"metadata": {}} + + await guardrail.apply_guardrail(inputs, request_data, "request") + + mock_api_call.assert_called_once() + assert mock_api_call.call_args.kwargs["policy_id"] == 100 + +@pytest.mark.asyncio +@patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call", + new_callable=AsyncMock, +) +async def test_policy_id_zero_from_request_metadata(mock_api_call): + """ + Test policy_id=0 is correctly picked. Make sure pick exact policy_id which users set + """ + guardrail = ZscalerAIGuard(policy_id=100) + inputs = {"texts": ["test"]} + request_data = { + "metadata": { + "zguard_policy_id": 0, + } + } + await guardrail.apply_guardrail(inputs, request_data, "request") + mock_api_call.assert_called_once() + assert mock_api_call.call_args.kwargs["policy_id"] == 0