[feat]: graceful degradation for pillar service when using litellm (#15857)

* graceful degradation for pillar service when using litellm * remove unnecessary mode * simplify docs * final fixes * lint fixes * fix linting
2026-08-13 06:23:32 +00:00 · 2025-10-27 19:51:29 -07:00
parent e27bab3238
commit 647f2f5d86
5 changed files with 289 additions and 74 deletions
@@ -29,7 +29,7 @@ Use Pillar Security for comprehensive LLM security including:

 Add Pillar Security to your `config.yaml`:

-**🌟 Recommended Configuration (Dual Mode):**
+**🌟 Recommended Configuration:**
 ```yaml
 model_list:
  - model_name: gpt-4.1-mini
@@ -45,6 +45,8 @@ guardrails:
      api_key: os.environ/PILLAR_API_KEY            # Your Pillar API key
      api_base: os.environ/PILLAR_API_BASE          # Pillar API endpoint
      on_flagged_action: "monitor"                  # Log threats but allow requests
+      fallback_on_error: "allow"                    # Gracefully degrade if Pillar is down (default)
+      timeout: 5.0                                  # Timeout for Pillar API calls in seconds (default)
      persist_session: true                         # Keep conversations visible in Pillar dashboard
      async_mode: false                             # Request synchronous verdicts
      include_scanners: true                        # Return scanner category breakdown
@@ -207,6 +209,8 @@ You can configure Pillar Security using environment variables:
 export PILLAR_API_KEY="your_api_key_here"
 export PILLAR_API_BASE="https://api.pillar.security"
 export PILLAR_ON_FLAGGED_ACTION="monitor"
+export PILLAR_FALLBACK_ON_ERROR="allow"
+export PILLAR_TIMEOUT="30.0"
 ```

 ### Session Tracking
@@ -245,6 +249,66 @@ Logs the violation but allows the request to proceed:
 on_flagged_action: "monitor"
 ```

+### Resilience and Error Handling
+
+#### Graceful Degradation (`fallback_on_error`)
+
+Control what happens when the Pillar API is unavailable (network errors, timeouts, service outages):
+
+```yaml
+fallback_on_error: "allow"  # Default - recommended for production resilience
+```
+
+**Available Options:**
+
+- **`allow` (Default - Recommended)**: Proceed without scanning when Pillar is unavailable
+  - **No service interruption** if Pillar is down
+  - **Best for production** where availability is critical
+  - Security scans are skipped during outages (logged as warnings)
+
+  ```yaml
+  guardrails:
+    - guardrail_name: "pillar-resilient"
+      litellm_params:
+        guardrail: pillar
+        fallback_on_error: "allow"  # Graceful degradation
+  ```
+
+- **`block`**: Reject all requests when Pillar is unavailable
+  - **Fail-secure approach** - no request proceeds without scanning
+  - **Service interruption** during Pillar outages
+  - Returns 503 Service Unavailable error
+
+  ```yaml
+  guardrails:
+    - guardrail_name: "pillar-fail-secure"
+      litellm_params:
+        guardrail: pillar
+        fallback_on_error: "block"  # Fail secure
+  ```
+
+#### Timeout Configuration
+
+Configure how long to wait for Pillar API responses:
+
+**Example Configurations:**
+
+```yaml
+# Production: Default - Fast with graceful degradation
+guardrails:
+  - guardrail_name: "pillar-production"
+    litellm_params:
+      guardrail: pillar
+      timeout: 5.0               # Default - fast failure detection
+      fallback_on_error: "allow"  # Graceful degradation (required)
+```
+
+**Environment Variables:**
+```bash
+export PILLAR_FALLBACK_ON_ERROR="allow"
+export PILLAR_TIMEOUT="5.0"
+```
+
 ## Advanced Configuration

 **Quick takeaways**
@@ -44,6 +44,10 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
        include_evidence=_get_config_value(
            litellm_params, optional_params, "include_evidence"
        ),
+        fallback_on_error=_get_config_value(
+            litellm_params, optional_params, "fallback_on_error"
+        ),
+        timeout=_get_config_value(litellm_params, optional_params, "timeout"),
    )
    litellm.logging_callback_manager.add_litellm_callback(_pillar_callback)

@@ -60,7 +60,10 @@ class PillarGuardrail(CustomGuardrail):

    SUPPORTED_ON_FLAGGED_ACTIONS = ["block", "monitor"]
    DEFAULT_ON_FLAGGED_ACTION = "monitor"
+    SUPPORTED_FALLBACK_ACTIONS = ["allow", "block"]
+    DEFAULT_FALLBACK_ACTION = "allow"
    BASE_API_URL = "https://api.pillar.security"
+    DEFAULT_TIMEOUT = 5.0  # 5 seconds - fast failure detection with graceful degradation

    def __init__(
        self,
@@ -72,6 +75,8 @@ class PillarGuardrail(CustomGuardrail):
        persist_session: Optional[bool] = None,
        include_scanners: Optional[bool] = None,
        include_evidence: Optional[bool] = None,
+        fallback_on_error: Optional[str] = None,
+        timeout: Optional[float] = None,
        **kwargs,
    ) -> None:
        """
@@ -82,11 +87,11 @@ class PillarGuardrail(CustomGuardrail):
            api_key: Pillar API key
            api_base: Pillar API base URL
            on_flagged_action: Action to take when content is flagged ('block' or 'monitor')
+            fallback_on_error: Action when API errors occur ('allow' or 'block')
+            timeout: Timeout for API calls in seconds
            **kwargs: Additional arguments passed to parent class
        """
-        self.async_handler = get_async_httpx_client(
-            llm_provider=httpxSpecialProvider.GuardrailCallback
-        )
+        self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
        self.api_key = api_key or os.environ.get("PILLAR_API_KEY")

        if self.api_key is None:
@@ -104,14 +109,10 @@ class PillarGuardrail(CustomGuardrail):
            self.on_flagged_action = action
        else:
            if action:
-                verbose_proxy_logger.warning(
-                    f"Invalid action '{action}', using default"
-                )
+                verbose_proxy_logger.warning(f"Invalid action '{action}', using default")
            self.on_flagged_action = self.DEFAULT_ON_FLAGGED_ACTION

-        verbose_proxy_logger.debug(
-            f"Pillar Guardrail: Initialized with on_flagged_action: {self.on_flagged_action}"
-        )
+        verbose_proxy_logger.debug(f"Pillar Guardrail: Initialized with on_flagged_action: {self.on_flagged_action}")

        self.async_mode = self._resolve_bool_config(
            provided_value=async_mode,
@@ -138,6 +139,32 @@ class PillarGuardrail(CustomGuardrail):
            setting_name="include_evidence",
        )

+        # Validate and set fallback_on_error
+        action = fallback_on_error or os.environ.get("PILLAR_FALLBACK_ON_ERROR")
+        if action and action in self.SUPPORTED_FALLBACK_ACTIONS:
+            self.fallback_on_error = action
+        else:
+            if action:
+                verbose_proxy_logger.warning(
+                    f"Invalid fallback action '{action}', using default '{self.DEFAULT_FALLBACK_ACTION}'"
+                )
+            self.fallback_on_error = self.DEFAULT_FALLBACK_ACTION
+
+        verbose_proxy_logger.debug(f"Pillar Guardrail: Initialized with fallback_on_error: {self.fallback_on_error}")
+
+        # Set timeout with graceful fallback on invalid configuration
+        if timeout is not None:
+            self.timeout = timeout
+        else:
+            try:
+                self.timeout = float(os.environ.get("PILLAR_TIMEOUT", str(self.DEFAULT_TIMEOUT)))
+            except (ValueError, TypeError):
+                verbose_proxy_logger.warning(
+                    f"Pillar Guardrail: Invalid PILLAR_TIMEOUT value '{os.environ.get('PILLAR_TIMEOUT')}', "
+                    f"falling back to default {self.DEFAULT_TIMEOUT}s"
+                )
+                self.timeout = self.DEFAULT_TIMEOUT
+
        # Define supported event hooks
        supported_event_hooks = [
            GuardrailEventHooks.pre_call,
@@ -191,18 +218,14 @@ class PillarGuardrail(CustomGuardrail):
        """
        event_type = GuardrailEventHooks.pre_call
        if self.should_run_guardrail(data=data, event_type=event_type) is not True:
-            verbose_proxy_logger.debug(
-                f"Pillar Guardrail: Pre-call scanning disabled for {self.guardrail_name}"
-            )
+            verbose_proxy_logger.debug(f"Pillar Guardrail: Pre-call scanning disabled for {self.guardrail_name}")
            return data

        verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook")
        result = await self.run_pillar_guardrail(data)

        # Add guardrail name to response headers
-        add_guardrail_to_applied_guardrails_header(
-            request_data=data, guardrail_name=self.guardrail_name
-        )
+        add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)

        return result

@@ -238,18 +261,14 @@ class PillarGuardrail(CustomGuardrail):
        """
        event_type = GuardrailEventHooks.during_call
        if self.should_run_guardrail(data=data, event_type=event_type) is not True:
-            verbose_proxy_logger.debug(
-                f"Pillar Guardrail: During-call scanning disabled for {self.guardrail_name}"
-            )
+            verbose_proxy_logger.debug(f"Pillar Guardrail: During-call scanning disabled for {self.guardrail_name}")
            return data

        verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook")
        result = await self.run_pillar_guardrail(data)

        # Add guardrail name to response headers
-        add_guardrail_to_applied_guardrails_header(
-            request_data=data, guardrail_name=self.guardrail_name
-        )
+        add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)

        return result

@@ -276,9 +295,7 @@ class PillarGuardrail(CustomGuardrail):
        """
        event_type = GuardrailEventHooks.post_call
        if self.should_run_guardrail(data=data, event_type=event_type) is not True:
-            verbose_proxy_logger.debug(
-                f"Pillar Guardrail: Post-call scanning disabled for {self.guardrail_name}"
-            )
+            verbose_proxy_logger.debug(f"Pillar Guardrail: Post-call scanning disabled for {self.guardrail_name}")
            return response

        verbose_proxy_logger.debug("Pillar Guardrail: Post-call hook")
@@ -286,15 +303,11 @@ class PillarGuardrail(CustomGuardrail):
        # Extract response messages in the format Pillar expects
        response_dict = response.model_dump() if hasattr(response, "model_dump") else {}  # type: ignore[union-attr]
        response_messages = [
-            choice.get("message")
-            for choice in response_dict.get("choices", [])
-            if choice.get("message")
+            choice.get("message") for choice in response_dict.get("choices", []) if choice.get("message")
        ]

        if not response_messages:
-            verbose_proxy_logger.debug(
-                "Pillar Guardrail: No response content to scan, skipping post-call analysis"
-            )
+            verbose_proxy_logger.debug("Pillar Guardrail: No response content to scan, skipping post-call analysis")
            return response

        # Create complete conversation: original messages + response messages
@@ -305,9 +318,7 @@ class PillarGuardrail(CustomGuardrail):
        await self.run_pillar_guardrail(post_call_data)

        # Add guardrail name to response headers
-        add_guardrail_to_applied_guardrails_header(
-            request_data=data, guardrail_name=self.guardrail_name
-        )
+        add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)

        return response

@@ -326,14 +337,11 @@ class PillarGuardrail(CustomGuardrail):
            Original data if safe or in monitor mode

        Raises:
-            PillarGuardrailAPIError: If the Pillar API call fails
-            HTTPException: If content is flagged and action is 'block'
+            HTTPException: If content is flagged and action is 'block', or if API fails and fallback_on_error is 'block'
        """
        # Check if messages are present
        if not data.get("messages"):
-            verbose_proxy_logger.debug(
-                "Pillar Guardrail: No messages detected, bypassing security scan"
-            )
+            verbose_proxy_logger.debug("Pillar Guardrail: No messages detected, bypassing security scan")
            return data

        try:
@@ -350,19 +358,51 @@ class PillarGuardrail(CustomGuardrail):
            return data

        except Exception as e:
+            # If it's already an HTTPException from content being flagged, re-raise it
            if isinstance(e, HTTPException):
                raise e
-            verbose_proxy_logger.error(
-                f"Pillar Guardrail: API communication failed - {str(e)}"
-            )
-            raise PillarGuardrailAPIError(
-                f"Pillar Guardrail scan failed - unable to verify request safety: {str(e)}"
-            )
+
+            # Handle API communication errors based on fallback_on_error setting
+            verbose_proxy_logger.error(f"Pillar Guardrail: API communication failed - {str(e)}")
+
+            return self._handle_api_error(e, data)

    # =========================================================================
    # PRIVATE HELPER METHODS (In logical order of usage)
    # =========================================================================

+    def _handle_api_error(self, error: Exception, data: dict) -> dict:
+        """
+        Handle API errors based on fallback_on_error configuration.
+
+        Args:
+            error: The exception that occurred during API communication
+            data: Original request data
+
+        Returns:
+            Original data if fallback_on_error is 'allow'
+
+        Raises:
+            HTTPException: If fallback_on_error is 'block'
+        """
+        if self.fallback_on_error == "allow":
+            verbose_proxy_logger.warning(
+                "Pillar Guardrail: API unavailable, proceeding without scanning (fallback_on_error=allow)"
+            )
+            return data
+        else:  # fallback_on_error == "block"
+            verbose_proxy_logger.warning(
+                "Pillar Guardrail: API unavailable, blocking request (fallback_on_error=block)"
+            )
+            raise HTTPException(
+                status_code=503,
+                detail={
+                    "error": "Pillar Security Guardrail Unavailable",
+                    "message": "Security scanning service is temporarily unavailable and fallback is set to block",
+                    "original_error": str(error),
+                },
+            )
+
    def _prepare_headers(self) -> Dict[str, str]:
        """Prepare headers for the Pillar API request."""
        if not self.api_key:
@@ -385,9 +425,7 @@ class PillarGuardrail(CustomGuardrail):

        return headers

-    def _set_bool_header(
-        self, headers: Dict[str, str], header_name: str, value: Optional[bool]
-    ) -> None:
+    def _set_bool_header(self, headers: Dict[str, str], header_name: str, value: Optional[bool]) -> None:
        """Apply a boolean value as a lowercase string HTTP header when provided."""

        if value is None:
@@ -520,9 +558,7 @@ class PillarGuardrail(CustomGuardrail):
        )
        return payload

-    async def _call_pillar_api(
-        self, headers: Dict[str, str], payload: Dict[str, Any]
-    ) -> Dict[str, Any]:
+    async def _call_pillar_api(self, headers: Dict[str, str], payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Call the Pillar API and return the response.

@@ -540,21 +576,17 @@ class PillarGuardrail(CustomGuardrail):
            url=f"{self.api_base}/api/v1/protect",
            headers=headers,
            json=payload,
-            timeout=30.0,
+            timeout=self.timeout,
        )
        response.raise_for_status()
        res = response.json()

        flagged = res.get("flagged")
        session_id = res.get("session_id")
-        verbose_proxy_logger.debug(
-            f"Pillar Guardrail: Analysis complete - flagged={flagged}, session={session_id}"
-        )
+        verbose_proxy_logger.debug(f"Pillar Guardrail: Analysis complete - flagged={flagged}, session={session_id}")
        return res

-    def _process_pillar_response(
-        self, pillar_response: Dict[str, Any], original_data: dict
-    ) -> None:
+    def _process_pillar_response(self, pillar_response: Dict[str, Any], original_data: dict) -> None:
        """
        Process the Pillar API response and handle detections based on configuration.

@@ -573,9 +605,7 @@ class PillarGuardrail(CustomGuardrail):
        # Store session_id from Pillar response for potential reuse
        pillar_session_id = pillar_response.get("session_id")
        if pillar_session_id:
-            verbose_proxy_logger.debug(
-                f"Pillar Guardrail: Received session_id from server: {pillar_session_id}"
-            )
+            verbose_proxy_logger.debug(f"Pillar Guardrail: Received session_id from server: {pillar_session_id}")
            # Store in request metadata for use in subsequent hooks
            if "metadata" not in original_data:
                original_data["metadata"] = {}
@@ -587,13 +617,9 @@ class PillarGuardrail(CustomGuardrail):
            if self.on_flagged_action == "block":
                self._raise_pillar_detection_exception(pillar_response)
            elif self.on_flagged_action == "monitor":
-                verbose_proxy_logger.info(
-                    "Pillar Guardrail: Monitoring mode - allowing flagged content to proceed"
-                )
+                verbose_proxy_logger.info("Pillar Guardrail: Monitoring mode - allowing flagged content to proceed")

-    def _raise_pillar_detection_exception(
-        self, pillar_response: Dict[str, Any]
-    ) -> None:
+    def _raise_pillar_detection_exception(self, pillar_response: Dict[str, Any]) -> None:
        """
        Raise an HTTPException for Pillar security detections.

@@ -613,9 +639,7 @@ class PillarGuardrail(CustomGuardrail):
            },
        }

-        verbose_proxy_logger.warning(
-            "Pillar Guardrail: Request blocked - Security threats detected"
-        )
+        verbose_proxy_logger.warning("Pillar Guardrail: Request blocked - Security threats detected")

        raise HTTPException(status_code=400, detail=error_detail)

@@ -31,6 +31,14 @@ class PillarGuardrailConfigModelOptionalParams(BaseModel):
        default=True,
        description="Include detailed evidence objects in response payloads (sets `plr_evidence` header).",
    )
+    fallback_on_error: Optional[str] = Field(
+        default=None,
+        description="Action to take when Pillar API is unavailable or errors: 'allow' (proceed without scanning) or 'block' (reject request with 503 error). If not provided, the `PILLAR_FALLBACK_ON_ERROR` environment variable is checked, defaults to 'allow'.",
+    )
+    timeout: Optional[float] = Field(
+        default=None,
+        description="Timeout in seconds for Pillar API calls. If not provided, the `PILLAR_TIMEOUT` environment variable is checked, defaults to 5.0 seconds.",
+    )


 class PillarGuardrailConfigModel(
@@ -526,8 +526,12 @@ async def test_empty_messages(pillar_guardrail_instance, user_api_key_dict, dual
 async def test_api_error_handling(
    pillar_guardrail_instance, sample_request_data, user_api_key_dict, dual_cache
 ):
-    """Test handling of API connection errors."""
-    with pytest.raises(PillarGuardrailAPIError) as excinfo:
+    """Test handling of API connection errors with block fallback."""
+    # Note: pillar_guardrail_instance has fallback_on_error defaulting to "allow"
+    # so this test sets it to "block" to test error handling
+    pillar_guardrail_instance.fallback_on_error = "block"  # Set to block for this test
+
+    with pytest.raises(HTTPException) as excinfo:
        with patch(
            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
            side_effect=Exception("Connection error"),
@@ -539,8 +543,119 @@ async def test_api_error_handling(
                call_type="completion",
            )

-    assert "unable to verify request safety" in str(excinfo.value)
-    assert "Connection error" in str(excinfo.value)
+    assert excinfo.value.status_code == 503
+    assert "Pillar Security Guardrail Unavailable" in str(excinfo.value.detail)
+
+
+@pytest.mark.asyncio
+async def test_api_error_fallback_allow(env_setup):
+    """Test fallback_on_error='allow' allows requests when API is down."""
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-fallback-allow",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+        fallback_on_error="allow",
+    )
+
+    sample_data = {
+        "model": "gpt-4",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        side_effect=Exception("Connection timeout"),
+    ):
+        result = await guardrail.async_pre_call_hook(
+            data=sample_data,
+            cache=DualCache(),
+            user_api_key_dict=UserAPIKeyAuth(),
+            call_type="completion",
+        )
+
+    # Should proceed without scanning
+    assert result == sample_data
+
+
+@pytest.mark.asyncio
+async def test_api_error_fallback_block(env_setup):
+    """Test fallback_on_error='block' blocks requests when API is down."""
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-fallback-block",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+        fallback_on_error="block",
+    )
+
+    sample_data = {
+        "model": "gpt-4",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    with pytest.raises(HTTPException) as excinfo:
+        with patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+            side_effect=Exception("Connection timeout"),
+        ):
+            await guardrail.async_pre_call_hook(
+                data=sample_data,
+                cache=DualCache(),
+                user_api_key_dict=UserAPIKeyAuth(),
+                call_type="completion",
+            )
+
+    # Should block with 503 Service Unavailable
+    assert excinfo.value.status_code == 503
+    assert "Pillar Security Guardrail Unavailable" in str(excinfo.value.detail)
+
+
+@pytest.mark.asyncio
+async def test_custom_timeout_configuration(env_setup):
+    """Test custom timeout configuration."""
+    custom_timeout = 10.0
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-custom-timeout",
+        api_key="test-pillar-key",
+        api_base="https://api.pillar.security",
+        timeout=custom_timeout,
+    )
+
+    assert guardrail.timeout == custom_timeout
+
+
+def test_fallback_on_error_env_variable(monkeypatch):
+    """Test fallback_on_error can be set via environment variable."""
+    monkeypatch.setenv("PILLAR_API_KEY", "test-key")
+    monkeypatch.setenv("PILLAR_FALLBACK_ON_ERROR", "block")
+
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-env-fallback",
+    )
+
+    assert guardrail.fallback_on_error == "block"
+
+
+def test_timeout_env_variable(monkeypatch):
+    """Test timeout can be set via environment variable."""
+    monkeypatch.setenv("PILLAR_API_KEY", "test-key")
+    monkeypatch.setenv("PILLAR_TIMEOUT", "15.0")
+
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-env-timeout",
+    )
+
+    assert guardrail.timeout == 15.0
+
+
+def test_invalid_fallback_action_defaults_to_allow(env_setup):
+    """Test invalid fallback_on_error value defaults to 'allow'."""
+    guardrail = PillarGuardrail(
+        guardrail_name="pillar-invalid-fallback",
+        api_key="test-pillar-key",
+        fallback_on_error="invalid_action",
+    )
+
+    assert guardrail.fallback_on_error == "allow"


@pytest.mark.asyncio