mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 10:21:52 +00:00
feat(ui): added UI for Zscaler AI Guard (#21077)
* fix: allow Management keys to access user/daily/activity and team/daily/activity * feat(ui): added UI for Zscaler AI Guard * feat(ui): addressed UI comment * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: naaa760 <neh6a683@gmail.com> Co-authored-by: yuneng-jiang <yuneng.jiang@gmail.com> Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
This commit is contained in:
datzscaler
co-authored by
naaa760
yuneng-jiang
Krish Dholakia
greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
naaa760
yuneng-jiang
Krish Dholakia
greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
parent
1bc90d2db7
commit
f74fdfbb61
@@ -514,6 +514,8 @@ class LiteLLMRoutes(enum.Enum):
|
||||
"/user/delete",
|
||||
"/user/info",
|
||||
"/user/list",
|
||||
"/user/daily/activity",
|
||||
"/user/daily/activity/aggregated",
|
||||
# team
|
||||
"/team/new",
|
||||
"/team/update",
|
||||
@@ -526,6 +528,7 @@ class LiteLLMRoutes(enum.Enum):
|
||||
"/team/available",
|
||||
"/team/permissions_list",
|
||||
"/team/permissions_update",
|
||||
"/team/daily/activity",
|
||||
# model
|
||||
"/model/new",
|
||||
"/model/update",
|
||||
|
||||
@@ -21,6 +21,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
||||
GUARDRAIL_TIMEOUT = 5
|
||||
|
||||
@@ -334,3 +335,11 @@ class ZscalerAIGuard(CustomGuardrail):
|
||||
user_facing_error = self._create_user_facing_error(f"{str(e)})")
|
||||
# This exception will be caught by the proxy and returned to the user
|
||||
raise HTTPException(status_code=500, detail=user_facing_error)
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> Optional[type["GuardrailConfigModel"]]:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import (
|
||||
ZscalerAIGuardConfigModel,
|
||||
)
|
||||
|
||||
return ZscalerAIGuardConfigModel
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import Field, model_validator
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.types.guardrails import GuardrailParamUITypes
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class ZscalerAIGuardConfigModel(GuardrailConfigModel):
|
||||
api_key: Optional[str] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"API key for Zscaler AI Guard authentication. "
|
||||
"If not provided, falls back to ZSCALER_AI_GUARD_API_KEY environment variable."
|
||||
),
|
||||
)
|
||||
|
||||
api_base: Optional[str] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Zscaler AI Guard API endpoint. Determines policy resolution behavior:\n"
|
||||
"• /execute-policy (default) - Requires explicit policy_id in configuration\n"
|
||||
"• /resolve-and-execute-policy - Infers policy from user-api-key-alias header\n"
|
||||
"Default: https://api.us1.zseclipse.net/v1/detection/execute-policy\n"
|
||||
"Falls back to ZSCALER_AI_GUARD_URL environment variable."
|
||||
),
|
||||
json_schema_extra={
|
||||
"examples": [
|
||||
"https://api.us1.zseclipse.net/v1/detection/execute-policy",
|
||||
"https://api.us1.zseclipse.net/v1/detection/resolve-and-execute-policy",
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
policy_id: Optional[int] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Global policy ID for Zscaler AI Guard. Required when using /execute-policy endpoint.\n\n"
|
||||
"Set to 0 or leave empty when using /resolve-and-execute-policy with dynamic policy resolution.\n"
|
||||
"Falls back to ZSCALER_AI_GUARD_POLICY_ID environment variable."
|
||||
),
|
||||
json_schema_extra={
|
||||
"ui_hint": "conditional_required",
|
||||
"condition": "Required when api_base ends with /execute-policy",
|
||||
},
|
||||
)
|
||||
|
||||
send_user_api_key_alias: Optional[bool] = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"Send user API key alias in request headers as 'user-api-key-alias'. "
|
||||
"CRITICAL when using /resolve-and-execute-policy endpoint - the policy is inferred from this value. "
|
||||
"Also useful for tracking/auditing with /execute-policy endpoint."
|
||||
),
|
||||
json_schema_extra={
|
||||
"ui_type": GuardrailParamUITypes.BOOL,
|
||||
"ui_hint": "recommended_when",
|
||||
"condition": "Recommended when api_base ends with /resolve-and-execute-policy",
|
||||
},
|
||||
)
|
||||
|
||||
send_user_api_key_user_id: Optional[bool] = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"Send user API key user_id in request headers as 'user-api-key-user-id'. "
|
||||
"Enables user-level tracking and analytics in Zscaler AI Guard."
|
||||
),
|
||||
json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL},
|
||||
)
|
||||
|
||||
send_user_api_key_team_id: Optional[bool] = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"Send user API key team_id in request headers as 'user-api-key-team-id'. "
|
||||
"Enables team-level tracking and analytics in Zscaler AI Guard."
|
||||
),
|
||||
json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL},
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_endpoint_configuration(self) -> "ZscalerAIGuardConfigModel":
|
||||
"""
|
||||
Validate configuration consistency between api_base and other fields.
|
||||
Provides warnings but doesn't block (since env vars might provide values).
|
||||
"""
|
||||
import os
|
||||
|
||||
# Resolve actual api_base value (including env fallback)
|
||||
api_base = self.api_base or os.getenv(
|
||||
"ZSCALER_AI_GUARD_URL",
|
||||
"https://api.us1.zseclipse.net/v1/detection/execute-policy",
|
||||
)
|
||||
|
||||
# Resolve actual policy_id value
|
||||
policy_id = self.policy_id
|
||||
if policy_id is None:
|
||||
env_policy = os.getenv("ZSCALER_AI_GUARD_POLICY_ID")
|
||||
if env_policy:
|
||||
try:
|
||||
policy_id = int(env_policy)
|
||||
except ValueError:
|
||||
verbose_proxy_logger.warning(
|
||||
f"ZSCALER_AI_GUARD_POLICY_ID env var is not a valid integer: {env_policy}"
|
||||
)
|
||||
|
||||
# Check for configuration issues
|
||||
is_resolve_policy = api_base.endswith("/resolve-and-execute-policy")
|
||||
is_execute_policy = api_base.endswith("/execute-policy") and not is_resolve_policy
|
||||
|
||||
# Scenario A: execute-policy without policy_id
|
||||
if is_execute_policy and (policy_id is None or policy_id < 1):
|
||||
verbose_proxy_logger.warning(
|
||||
"Using /execute-policy endpoint without a valid policy_id. "
|
||||
"Ensure ZSCALER_AI_GUARD_POLICY_ID environment variable is set, "
|
||||
"or provide policy_id via request/key/team metadata."
|
||||
)
|
||||
|
||||
# Scenario B: resolve-and-execute-policy without user_api_key_alias
|
||||
if is_resolve_policy and not self.send_user_api_key_alias:
|
||||
verbose_proxy_logger.warning(
|
||||
"Using /resolve-and-execute-policy endpoint without send_user_api_key_alias=true. "
|
||||
"The endpoint requires user-api-key-alias header to resolve the policy. "
|
||||
"Set send_user_api_key_alias to true or ensure the header is sent via other means."
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "Zscaler AI Guard"
|
||||
@@ -0,0 +1,5 @@
|
||||
<svg width="50" height="41" viewBox="0 0 35 22" fill="none" xmlns="http://www.w3.org/2000/svg">
|
||||
<path d="M34.9305 11.1952C35.4985 14.7139 32.4864 16.6962 29.4313 16.9252C27.4867 20.976 20.6324 23.7107 14.6971 20.762C12.1552 21.1535 10.764 20.5092 9.6896 19.2752C11.861 16.4425 18.4366 11.3459 26.47 13.9755C30.7569 15.3849 31.8715 12.0843 30.7961 10.8513C26.7505 6.19284 17.6134 10.3865 17.2879 10.7353C20.8729 5.16699 33.7065 3.60491 34.9305 11.1952ZM22.6298 4.80521C22.6522 4.79728 19.7125 3.74467 15.6034 5.54462C15.4367 5.4832 15.2735 5.41239 15.1146 5.33251C19.0701 2.66529 22.5359 1.53437 25.5001 1.97445C23.7042 -0.11988 15.0833 -1.65817 10.1432 3.36208C4.03683 2.14096 -0.233535 7.38224 0.00989967 12.1765C0.253334 16.9708 5.41043 19.9502 8.34533 19.1424C8.41578 19.1335 8.48704 19.1335 8.55748 19.1424C9.21348 15.9766 11.8658 8.61029 22.6298 4.80521Z" fill="#2160E1"/>
|
||||
</svg>
|
||||
|
||||
<!--65 41, 50 32-->
|
||||
|
After Width: | Height: | Size: 906 B |
@@ -104,6 +104,7 @@ export const shouldRenderContentFilterConfigSettings = (provider: string | null)
|
||||
const asset_logos_folder = "../ui/assets/logos/";
|
||||
|
||||
export const guardrailLogoMap: Record<string, string> = {
|
||||
"Zscaler AI Guard": `${asset_logos_folder}zscaler.svg`,
|
||||
"Presidio PII": `${asset_logos_folder}presidio.png`,
|
||||
"Bedrock Guardrail": `${asset_logos_folder}bedrock.svg`,
|
||||
Lakera: `${asset_logos_folder}lakeraai.jpeg`,
|
||||
|
||||
Reference in New Issue
Block a user