[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 2615f04afd.

.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

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

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]:

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",

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<AddGuardrailFormProps> = ({
           </Select>
         </Form.Item>
 
-        {/* Use the GuardrailProviderFields component to render provider-specific fields */}
-        <GuardrailProviderFields 
-          selectedProvider={selectedProvider} 
-          accessToken={accessToken} 
-          providerParams={providerParams}
-        />
-
         <Form.Item
           name="mode"
           label="Mode"
           tooltip="How the guardrail should be applied"
           rules={[{ required: true, message: 'Please select a mode' }]}
         >
-          <Select>
+          <Select
+            optionLabelProp="label"
+            mode="multiple"
+          >
             {guardrailSettings?.supported_modes?.map(mode => (
-              <Option key={mode} value={mode}>{mode}</Option>
+              <Option key={mode} value={mode} label={mode}>
+                <div>
+                  <div>
+                    <strong>{mode}</strong>
+                    {mode === 'pre_call' && <Tag color="green" style={{ marginLeft: '8px' }}>Recommended</Tag>}
+                  </div>
+                  <div style={{ fontSize: '12px', color: '#888' }}>{modeDescriptions[mode as keyof typeof modeDescriptions]}</div>
+                </div>
+              </Option>
             )) || (
               <>
-                <Option value="pre_call">pre_call</Option>
-                <Option value="post_call">post_call</Option>
+                <Option value="pre_call" label="pre_call">
+                  <div>
+                    <div><strong>pre_call</strong> <Tag color="green">Recommended</Tag></div>
+                    <div style={{ fontSize: '12px', color: '#888' }}>{modeDescriptions.pre_call}</div>
+                  </div>
+                </Option>
+                <Option value="during_call" label="during_call">
+                  <div>
+                    <div><strong>during_call</strong></div>
+                    <div style={{ fontSize: '12px', color: '#888' }}>{modeDescriptions.during_call}</div>
+                  </div>
+                </Option>
+                <Option value="post_call" label="post_call">
+                  <div>
+                    <div><strong>post_call</strong></div>
+                    <div style={{ fontSize: '12px', color: '#888' }}>{modeDescriptions.post_call}</div>
+                  </div>
+                </Option>
+                <Option value="logging_only" label="logging_only">
+                  <div>
+                    <div><strong>logging_only</strong></div>
+                    <div style={{ fontSize: '12px', color: '#888' }}>{modeDescriptions.logging_only}</div>
+                  </div>
+                </Option>
               </>
             )}
           </Select>
@@ -368,11 +402,20 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({
         <Form.Item
           name="default_on"
           label="Always On"
-          tooltip="If enabled, this guardrail will be applied to all requests by default"
-          valuePropName="checked"
+          tooltip="If enabled, this guardrail will be applied to all requests by default."
         >
-          <Switch />
+          <Select>
+            <Select.Option value={true}>Yes</Select.Option>
+            <Select.Option value={false}>No</Select.Option>
+          </Select>
         </Form.Item>
+
+        {/* Use the GuardrailProviderFields component to render provider-specific fields */}
+        <GuardrailProviderFields 
+          selectedProvider={selectedProvider} 
+          accessToken={accessToken} 
+          providerParams={providerParams}
+        />
       </>
     );
   };
@@ -454,7 +497,7 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({
         form={form}
         layout="vertical"
         initialValues={{
-          mode: guardrailSettings?.supported_modes?.[0] || "pre_call",
+          mode: "pre_call",
           default_on: false
         }}
       >

ui/litellm-dashboard/src/components/guardrails/guardrail_provider_fields.tsx

@@ -113,6 +113,14 @@ const GuardrailProviderFields: React.FC<GuardrailProviderFieldsProps> = ({
                 </Select.Option>
               ))}
             </Select>
+          ) : field.type === "bool" ? (
+            <Select
+              placeholder={field.description}
+              defaultValue={field.default_value}
+            >
+              <Select.Option value="True">True</Select.Option>
+              <Select.Option value="False">False</Select.Option>
+            </Select>
           ) : field.param.includes("password") || field.param.includes("secret") ? (
             <TextInput
               placeholder={field.description}