ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add team policy mapping for zguard (#20608)

Browse Source 
* support policy mapping on team key level

* update document

* update document

* address comments

* update document

* add unit test for new feature

* add more test case
This commit is contained in:
jwang-gif 2026-02-07 22:44:17 -08:00 committed by GitHub
parent 7f93ff9e83
commit c9df996b77
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 196 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
docs/my-website/docs/proxy/guardrails/zscaler_ai_guard.md
View File

@ -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": <the custom 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}
```

60
litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py
View File

@ -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)

129
tests/guardrails_tests/test_zscaler_ai_guard.py
View File

@ -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"]
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