tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-21 12:20:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(guardrails): Zscaler AI Guard bug fixes and support during post-call (#20801)

Browse Source 
* fix(guardrails): fixed post-call issue with Zscaler guardrail and invalid headers. Added unittests

* fix(guardrails): Addressed greptil comments. Make policyid hanlding more clear

* fix(guardrails): Address greptile comment

* Update litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
This commit is contained in:
datzscaler
2026-02-13 11:28:14 -08:00
committed by GitHub
co-authored by greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
parent 066e694f5e
commit aab8edde67
3 changed files with 161 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py
+4
View File
@@ -14,6 +14,10 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
_zscaler_ai_guard_callback = ZscalerAIGuard(
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
policy_id=litellm_params.policy_id,
send_user_api_key_alias=litellm_params.send_user_api_key_alias,
send_user_api_key_user_id=litellm_params.send_user_api_key_user_id,
send_user_api_key_team_id=litellm_params.send_user_api_key_team_id,
guardrail_name=guardrail.get("guardrail_name", ""),
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py
+60 -28
View File
@@ -32,9 +32,9 @@ class ZscalerAIGuard(CustomGuardrail):
api_key: Optional[str] = None,
api_base: Optional[str] = None,
policy_id: Optional[int] = None,
send_user_api_key_alias: Optional[bool] = False,
send_user_api_key_user_id: Optional[bool] = False,
send_user_api_key_team_id: Optional[bool] = False,
send_user_api_key_alias: Optional[bool] = None,
send_user_api_key_user_id: Optional[bool] = None,
send_user_api_key_team_id: Optional[bool] = None,
**kwargs,
):
self.optional_params = kwargs
@@ -42,15 +42,15 @@ class ZscalerAIGuard(CustomGuardrail):
"ZSCALER_AI_GUARD_URL",
"https://api.us1.zseclipse.net/v1/detection/execute-policy",
)
self.policy_id = policy_id or int(os.getenv("ZSCALER_AI_GUARD_POLICY_ID", -1))
self.policy_id = policy_id if policy_id is not None else int(os.getenv("ZSCALER_AI_GUARD_POLICY_ID", -1))
self.api_key = api_key or os.getenv("ZSCALER_AI_GUARD_API_KEY")
self.send_user_api_key_alias = send_user_api_key_alias or os.getenv(
self.send_user_api_key_alias = send_user_api_key_alias if send_user_api_key_alias is not None else os.getenv(
"SEND_USER_API_KEY_ALIAS", "False"
).lower() in ("true", "1")
self.send_user_api_key_user_id = send_user_api_key_user_id or os.getenv(
self.send_user_api_key_user_id = send_user_api_key_user_id if send_user_api_key_user_id is not None else os.getenv(
"SEND_USER_API_KEY_USER_ID", "False"
).lower() in ("true", "1,")
self.send_user_api_key_team_id = send_user_api_key_team_id or os.getenv(
).lower() in ("true", "1")
self.send_user_api_key_team_id = send_user_api_key_team_id if send_user_api_key_team_id is not None else os.getenv(
"SEND_USER_API_KEY_TEAM_ID", "False"
).lower() in ("true", "1")
@@ -60,19 +60,48 @@ class ZscalerAIGuard(CustomGuardrail):
send_user_api_key_team_id:{self.send_user_api_key_team_id}"""
)
super().__init__(default_on=True)
super().__init__(**kwargs)
verbose_proxy_logger.debug("ZscalerAIGuard Initializing ...")
def _get_stripped_metadata_value(
self, request_data: Optional[dict], key: str
) -> Optional[str]:
@staticmethod
def _resolve_metadata_value(request_data: Optional[dict], key: str) -> Optional[str]:
"""
Resolve metadata value from request_data, checking both metadata locations.
During pre-call: metadata is at request_data["metadata"][key]
During post-call: metadata is at request_data["litellm_metadata"][key]
(set by transform_user_api_key_dict_to_metadata which prefixes keys with 'user_api_key_')
Also handles key name mapping for UserAPIKeyAuth fields:
- key_alias -> user_api_key_key_alias (in litellm_metadata)
- user_id -> user_api_key_user_id
- team_id -> user_api_key_team_id
"""
if request_data is None:
return "N/A"
value = request_data.get("metadata", {}).get(key, "N/A")
if value is not None:
return str(value).strip()
return "N/A"
return None
# Check litellm_metadata first (set during post-call by guardrail framework)
litellm_metadata = request_data.get("litellm_metadata", {})
if litellm_metadata:
value = litellm_metadata.get(key)
if value is not None:
return str(value).strip()
# Handle key_alias -> user_api_key_key_alias mapping
# transform_user_api_key_dict_to_metadata prefixes "key_alias" -> "user_api_key_key_alias"
if key == "user_api_key_alias":
value = litellm_metadata.get("user_api_key_key_alias")
if value is not None:
return str(value).strip()
# Then check regular metadata (set during pre-call by proxy_server)
metadata = request_data.get("metadata", {})
if metadata:
value = metadata.get(key)
if value is not None:
return str(value).strip()
return None
@log_guardrail_information
async def apply_guardrail(
@@ -128,17 +157,17 @@ class ZscalerAIGuard(CustomGuardrail):
kwargs = {}
if self.send_user_api_key_alias:
kwargs["user_api_key_alias"] = self._get_stripped_metadata_value(
kwargs["user_api_key_alias"] = self._resolve_metadata_value(
request_data, "user_api_key_alias"
)
) or "N/A"
if self.send_user_api_key_team_id:
kwargs["user_api_key_team_id"] = self._get_stripped_metadata_value(
kwargs["user_api_key_team_id"] = self._resolve_metadata_value(
request_data, "user_api_key_team_id"
)
) or "N/A"
if self.send_user_api_key_user_id:
kwargs["user_api_key_user_id"] = self._get_stripped_metadata_value(
kwargs["user_api_key_user_id"] = self._resolve_metadata_value(
request_data, "user_api_key_user_id"
)
) or "N/A"
verbose_proxy_logger.debug(f"inside apply_guardrail kwargs: {kwargs}")
zscaler_ai_guard_result = None
@@ -155,6 +184,7 @@ class ZscalerAIGuard(CustomGuardrail):
content=concatenated_text,
**kwargs,
)
verbose_proxy_logger.debug(f"response from zscaler ai guards: {zscaler_ai_guard_result}")
if (
zscaler_ai_guard_result
and zscaler_ai_guard_result.get("action") == "BLOCK"
@@ -217,7 +247,7 @@ class ZscalerAIGuard(CustomGuardrail):
extra_headers.update({"user-api-key-team-id": user_api_key_team_id})
if self.send_user_api_key_user_id:
user_api_key_user_id = kwargs.get("user-api-key-user-id", "N/A")
user_api_key_user_id = kwargs.get("user_api_key_user_id", "N/A")
extra_headers.update({"user-api-key-user-id": user_api_key_user_id})
verbose_proxy_logger.debug(f"extra_headers: {extra_headers}")
@@ -320,11 +350,14 @@ class ZscalerAIGuard(CustomGuardrail):
extra_headers = self._prepare_headers(api_key, **kwargs)
data = {
"policyId": policy_id,
"direction": direction,
"content": content,
}
# Only include policyId when explicitly configured (policy_id >= 1)
# When policy_id is None, 0, or -1 (default), use resolve-and-execute-policy which infers
# the policy from headers (e.g., user-api-key-alias)
if policy_id is not None and policy_id >= 1:
data["policyId"] = policy_id
try:
response = await self._send_request(
zscaler_ai_guard_url, extra_headers, data
@@ -332,8 +365,7 @@ class ZscalerAIGuard(CustomGuardrail):
return self._handle_response(response, direction)
except Exception as e:
verbose_proxy_logger.error(f"{e}. Blocking request.")
user_facing_error = self._create_user_facing_error(f"{str(e)})")
# This exception will be caught by the proxy and returned to the user
user_facing_error = self._create_user_facing_error(f"{str(e)}")
raise HTTPException(status_code=500, detail=user_facing_error)
@staticmethod
tests/guardrails_tests/test_zscaler_ai_guard.py
+97
View File
@@ -1,3 +1,15 @@
"""
Unit tests for Zscaler AI Guard guardrail
Tests covering:
- API call handling (ALLOW, BLOCK, errors)
- Policy ID precedence (metadata > user_api_key > team > init)
- Boolean config parameter handling (send_user_api_key_*)
- Metadata resolution (pre-call vs post-call)
- Header preparation with kwargs
- resolve-and-execute-policy endpoint (policyId omission)
"""
import pytest
from unittest.mock import AsyncMock, Mock, patch
from fastapi import HTTPException
@@ -244,3 +256,88 @@ async def test_policy_id_zero_from_request_metadata(mock_api_call):
await guardrail.apply_guardrail(inputs, request_data, "request")
mock_api_call.assert_called_once()
assert mock_api_call.call_args.kwargs["policy_id"] == 0
@pytest.mark.asyncio
async def test_should_use_config_send_user_api_key_alias_when_true():
"""Test that send_user_api_key_alias=True from config is used (not overridden by env)"""
guardrail = ZscalerAIGuard(
api_key="test_key",
send_user_api_key_alias=True,
)
assert guardrail.send_user_api_key_alias is True
@pytest.mark.asyncio
async def test_should_preserve_policy_id_zero_in_init():
"""Test that policy_id=0 is preserved (not treated as falsy and overridden by env)"""
guardrail = ZscalerAIGuard(
api_key="test_key",
policy_id=0,
)
assert guardrail.policy_id == 0
@pytest.mark.asyncio
async def test_should_resolve_from_litellm_metadata_during_post_call():
"""Test that user_api_key_alias is resolved from litellm_metadata during post-call"""
request_data = {
"litellm_metadata": {
"user_api_key_alias": "test-alias-post-call"
}
}
result = ZscalerAIGuard._resolve_metadata_value(request_data, "user_api_key_alias")
assert result == "test-alias-post-call"
@pytest.mark.asyncio
async def test_should_resolve_user_api_key_key_alias_mapping():
"""Test key_alias -> user_api_key_key_alias mapping in litellm_metadata"""
request_data = {
"litellm_metadata": {
"user_api_key_key_alias": "test-key-alias"
}
}
result = ZscalerAIGuard._resolve_metadata_value(request_data, "user_api_key_alias")
assert result == "test-key-alias"
@pytest.mark.asyncio
async def test_should_include_user_api_key_alias_header():
"""Test that user-api-key-alias header is included when send_user_api_key_alias is True"""
guardrail = ZscalerAIGuard(
api_key="test_key",
send_user_api_key_alias=True,
)
headers = guardrail._prepare_headers("test_key", user_api_key_alias="test-alias")
assert headers.get("user-api-key-alias") == "test-alias"
@pytest.mark.asyncio
async def test_should_omit_policy_id_when_zero_or_negative():
"""Test that policyId is omitted from request body when policy_id <= 0 (for resolve-and-execute-policy)"""
guardrail = ZscalerAIGuard(
api_key="test_key",
policy_id=-1,
)
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {
"statusCode": 200,
"action": "ALLOW",
}
with patch.object(guardrail, "_send_request", new_callable=AsyncMock) as mock_send:
mock_send.return_value = mock_response
await guardrail.make_zscaler_ai_guard_api_call(
zscaler_ai_guard_url="http://example.com",
api_key="test_key",
policy_id=-1,
direction="OUT",
content="test content",
)
call_args = mock_send.call_args
data = call_args[0][2] # Third positional arg is data
assert "policyId" not in data