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 9380a539ecd..e8f0dc68d83 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 @@ -4,6 +4,7 @@ # # +-------------------------------------------------------------+ import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException @@ -105,6 +106,11 @@ class ZscalerAIGuard(CustomGuardrail): super().update_in_memory_litellm_params(litellm_params) self.timeout = self._resolve_timeout(litellm_params.timeout) + @staticmethod + def _resolve_request_user(request_data: Mapping[str, object]) -> str | None: + request_user: Final = request_data.get("user") + return request_user if isinstance(request_user, str) and request_user else None + @staticmethod def _resolve_metadata_value(request_data: dict | None, key: str) -> str | None: """ @@ -207,6 +213,8 @@ class ZscalerAIGuard(CustomGuardrail): ) verbose_proxy_logger.debug("inside apply_guardrail kwargs: %s", kwargs) + user_for_payload: Final = self._resolve_request_user(request_data) + zscaler_ai_guard_result = None direction: Final = "OUT" if input_type == "response" else "IN" verbose_proxy_logger.debug("direction: %s", direction) @@ -219,6 +227,7 @@ class ZscalerAIGuard(CustomGuardrail): policy_id=policy_id, direction=direction, content=concatenated_text, + user=user_for_payload, **kwargs, ) verbose_proxy_logger.debug("response from zscaler ai guards: %s", zscaler_ai_guard_result) @@ -357,7 +366,14 @@ class ZscalerAIGuard(CustomGuardrail): raise HTTPException(status_code=response.status_code, detail=user_facing_error) async def make_zscaler_ai_guard_api_call( - self, zscaler_ai_guard_url, api_key, policy_id, direction, content, **kwargs + self, + zscaler_ai_guard_url: str, + api_key: str | None, + policy_id, + direction: str, + content: str, + user: str | None = None, + **kwargs, ): """ Makes an API call to the Zscaler AI Guard service and handles retries, errors, and response parsing. @@ -374,6 +390,8 @@ class ZscalerAIGuard(CustomGuardrail): # the policy from headers (e.g., user-api-key-alias) if policy_id is not None and policy_id >= 1: data["policyId"] = policy_id + if user: + data["user"] = user try: response: Final = await self._send_request(zscaler_ai_guard_url, extra_headers, data) return self._handle_response(response, direction) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py index 7445db47185..bccc531793f 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py @@ -10,11 +10,15 @@ Tests covering: - resolve-and-execute-policy endpoint (policyId omission) """ +import asyncio import importlib import os +from typing import Final, Literal from unittest.mock import AsyncMock, Mock, patch import pytest +from aiohttp import web +from aiohttp.test_utils import TestServer from fastapi import HTTPException import litellm @@ -561,3 +565,47 @@ def setup_and_teardown(): if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() yield + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["request", "response"]) +@pytest.mark.parametrize( + ("request_data", "expected_body_user"), + [ + ({"user": "alice@example.com"}, "alice@example.com"), + ({}, None), + ({"user": None}, None), + ({"user": ""}, None), + ({"user": 12345}, None), + ({"user": ["alice@example.com"]}, None), + ], +) +async def test_apply_guardrail_sends_request_user_in_execute_policy_body( + request_data: dict[str, object], + expected_body_user: str | None, + input_type: Literal["request", "response"], +) -> None: + received: Final = asyncio.Queue[object]() + + async def execute_policy(request: web.Request) -> web.Response: + await received.put(await request.json()) + return web.json_response({"statusCode": 200, "action": "ALLOW"}) + + app: Final = web.Application() + app.router.add_post("/v1/detection/execute-policy", execute_policy) + async with TestServer(app) as server: + guardrail: Final = ZscalerAIGuard( + api_key="test_key", + api_base=str(server.make_url("/v1/detection/execute-policy")), + policy_id=1, + ) + await guardrail.apply_guardrail({"texts": ["hello"]}, request_data, input_type) + + body: Final = received.get_nowait() + assert received.empty() + assert body == { + "direction": "OUT" if input_type == "response" else "IN", + "content": "hello", + "policyId": 1, + **({"user": expected_body_user} if expected_body_user is not None else {}), + }