From 93ec0c48bc2bce64e092ef6f1eeb4f0e08d47b35 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 17 May 2025 16:27:57 -0700 Subject: [PATCH] [UI QA Guardrails] - Minor UI Fixes (#10920) * fix: presidio params on ui * move mode and always on to the top * explain each mode * show switch for default on / off for guards * allow multiple modes for guardrails * fix: presidio output parsing with stream * Revert "git commit fix linting error on cryptography bindings" This reverts commit 2615f04afd809f74d70aec2ad6591d517d052b7a. --- .circleci/config.yml | 2 +- .../proxy/guardrails/guardrail_endpoints.py | 17 +++- .../guardrails/guardrail_hooks/presidio.py | 94 ++++++++++++++++++- litellm/types/guardrails.py | 23 ++++- .../guardrails/add_guardrail_form.tsx | 73 +++++++++++--- .../guardrails/guardrail_provider_fields.tsx | 8 ++ 6 files changed, 193 insertions(+), 24 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 559df6b917..59616f1990 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -143,7 +143,7 @@ jobs: command: | cd litellm python -m pip install types-requests types-setuptools types-redis types-PyYAML - if ! python -m mypy . --ignore-missing-imports --exclude '(^|/)cryptography/'; then + if ! python -m mypy . --ignore-missing-imports; then echo "mypy detected errors" exit 1 fi diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index b9a7d85f20..054eb61d96 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -16,6 +16,7 @@ from litellm.types.guardrails import ( Guardrail, GuardrailEventHooks, GuardrailInfoResponse, + GuardrailParamUITypes, GuardrailUIAddGuardrailSettings, LakeraV2GuardrailConfigModel, ListGuardrailsResponse, @@ -23,7 +24,7 @@ from litellm.types.guardrails import ( PatchGuardrailRequest, PiiAction, PiiEntityType, - PresidioConfigModel, + PresidioPresidioConfigModelUserInterface, SupportedGuardrailIntegrations, ) @@ -671,11 +672,17 @@ def _get_fields_from_model(model_class: Type[BaseModel]) -> List[Dict[str, Any]] # Check if this field is in the required_fields class variable required = field.is_required() + field_type: Optional[GuardrailParamUITypes] = None + field_json_schema_extra = getattr(field, "json_schema_extra", {}) + if field_json_schema_extra and "ui_type" in field_json_schema_extra: + field_type = field_json_schema_extra["ui_type"] + fields.append( { "param": field_name, "description": description, "required": required, + "type": field_type.value if field_type else None, } ) return fields @@ -718,6 +725,12 @@ async def get_provider_specific_params(): "param": "presidio_anonymizer_api_base", "description": "Base URL for the Presidio anonymizer API", "required": true + }, + { + "param": "output_parse_pii", + "description": "Whether to parse PII in model outputs", + "required": false, + "type": "bool" } ] } @@ -725,7 +738,7 @@ async def get_provider_specific_params(): """ # Get fields from the models bedrock_fields = _get_fields_from_model(BedrockGuardrailConfigModel) - presidio_fields = _get_fields_from_model(PresidioConfigModel) + presidio_fields = _get_fields_from_model(PresidioPresidioConfigModelUserInterface) lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel) # Return the provider-specific parameters diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 8052dfd3a0..3c5ca3d9c2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -12,7 +12,17 @@ import asyncio import json import uuid from datetime import datetime -from typing import Any, Dict, List, Literal, Optional, Tuple, Union, cast +from typing import ( + Any, + AsyncGenerator, + Dict, + List, + Literal, + Optional, + Tuple, + Union, + cast, +) import aiohttp @@ -39,6 +49,7 @@ from litellm.utils import ( EmbeddingResponse, ImageResponse, ModelResponse, + ModelResponseStream, StreamingChoices, ) @@ -536,6 +547,87 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ].message.content.replace(key, value) return response + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + request_data: dict, + ) -> AsyncGenerator[ModelResponseStream, None]: + """ + Process streaming response chunks to unmask PII tokens when needed. + + If PII processing is enabled, this collects all chunks, applies PII unmasking, + and returns a reconstructed stream. Otherwise, it passes through the original stream. + """ + # If PII unmasking not needed, just pass through the original stream + if not (self.output_parse_pii and self.pii_tokens): + async for chunk in response: + yield chunk + return + + # Import here to avoid circular imports + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + from litellm.types.utils import Choices, Message + + try: + # Collect all chunks to process them together + collected_content = "" + last_chunk = None + + async for chunk in response: + last_chunk = chunk + + # Extract content safely with proper attribute checks + if ( + hasattr(chunk, "choices") + and chunk.choices + and hasattr(chunk.choices[0], "delta") + and hasattr(chunk.choices[0].delta, "content") + and isinstance(chunk.choices[0].delta.content, str) + ): + collected_content += chunk.choices[0].delta.content + + # No need to proceed if we didn't capture a valid chunk + if not last_chunk: + async for chunk in response: + yield chunk + return + + # Apply PII unmasking to the complete content + for token, original_text in self.pii_tokens.items(): + collected_content = collected_content.replace(token, original_text) + + # Reconstruct the response with unmasked content + mock_response = MockResponseIterator( + model_response=ModelResponse( + id=last_chunk.id, + object=last_chunk.object, + created=last_chunk.created, + model=last_chunk.model, + choices=[ + Choices( + message=Message( + role="assistant", + content=collected_content, + ), + index=0, + finish_reason="stop", + ) + ], + ), + json_mode=False, + ) + + # Return the reconstructed stream + async for chunk in mock_response: + yield chunk + + except Exception as e: + verbose_proxy_logger.error(f"Error in PII streaming processing: {str(e)}") + # Fallback to original stream on error + async for chunk in response: + yield chunk + def get_presidio_settings_from_request_data( self, data: dict ) -> Optional[PresidioPerRequestConfig]: diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 8f7114837d..214a550dd8 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -211,8 +211,13 @@ class PiiEntityCategoryMap(TypedDict): entities: List[PiiEntityType] -class PresidioConfigModel(BaseModel): - """Configuration parameters for the Presidio PII masking guardrail""" +class GuardrailParamUITypes(str, Enum): + BOOL = "bool" + STR = "str" + + +class PresidioPresidioConfigModelUserInterface(BaseModel): + """Configuration parameters for the Presidio PII masking guardrail on LiteLLM UI""" presidio_analyzer_api_base: Optional[str] = Field( default=None, @@ -222,12 +227,20 @@ class PresidioConfigModel(BaseModel): default=None, description="Base URL for the Presidio anonymizer API", ) + output_parse_pii: Optional[bool] = Field( + default=None, + description="When True, LiteLLM will replace the masked text with the original text in the response", + # extra param to let the ui know this is a boolean + json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL}, + ) + + +class PresidioConfigModel(PresidioPresidioConfigModelUserInterface): + """Configuration parameters for the Presidio PII masking guardrail""" + pii_entities_config: Optional[Dict[PiiEntityType, PiiAction]] = Field( default=None, description="Configuration for PII entity types and actions" ) - output_parse_pii: Optional[bool] = Field( - default=None, description="Whether to parse PII in model outputs" - ) presidio_ad_hoc_recognizers: Optional[str] = Field( default=None, description="Path to a JSON file containing ad-hoc recognizers for Presidio", diff --git a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx index 27ce2039e0..93d97747c8 100644 --- a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx @@ -11,6 +11,14 @@ const { Title, Text, Link } = Typography; const { Option } = Select; const { Step } = Steps; +// Define human-friendly descriptions for each mode +const modeDescriptions = { + pre_call: "Before LLM Call - Runs before the LLM call and checks the input (Recommended)", + during_call: "During LLM Call - Runs in parallel with the LLM call, with response held until check completes", + post_call: "After LLM Call - Runs after the LLM call and checks only the output", + logging_only: "Logging Only - Only runs on logging callbacks without affecting the LLM call" +}; + interface AddGuardrailFormProps { visible: boolean; onClose: () => void; @@ -340,26 +348,52 @@ const AddGuardrailForm: React.FC = ({ - {/* Use the GuardrailProviderFields component to render provider-specific fields */} - - - {guardrailSettings?.supported_modes?.map(mode => ( - + )) || ( <> - - + + + + )} @@ -368,11 +402,20 @@ const AddGuardrailForm: React.FC = ({ - + + + {/* Use the GuardrailProviderFields component to render provider-specific fields */} + ); }; @@ -454,7 +497,7 @@ const AddGuardrailForm: React.FC = ({ form={form} layout="vertical" initialValues={{ - mode: guardrailSettings?.supported_modes?.[0] || "pre_call", + mode: "pre_call", default_on: false }} > diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_provider_fields.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_provider_fields.tsx index 6d62e81f6a..a7758244f4 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_provider_fields.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_provider_fields.tsx @@ -113,6 +113,14 @@ const GuardrailProviderFields: React.FC = ({ ))} + ) : field.type === "bool" ? ( + ) : field.param.includes("password") || field.param.includes("secret") ? (