feature: update pillar security integration to support no persistence mode in litellm proxy

feature: update pillar security integration to support no persistence mode in litellm proxy
This commit is contained in:
Krish Dholakia 2025-10-16 12:13:37 -07:00 committed by GitHub
commit 3bf32e8e5c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 360 additions and 15 deletions

View File

@ -38,13 +38,17 @@ model_list:
api_key: os.environ/OPENAI_API_KEY
guardrails:
- guardrail_name: "pillar-minitor-everything" # you can change my name
- guardrail_name: "pillar-monitor-everything" # you can change my name
litellm_params:
guardrail: pillar
mode: [pre_call, post_call] # Monitor both input and output
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
persist_session: true # Keep conversations visible in Pillar dashboard
async_mode: false # Request synchronous verdicts
include_scanners: true # Return scanner category breakdown
include_evidence: true # Include detailed findings for triage
default_on: true # Enable for all requests
general_settings:
@ -104,10 +108,14 @@ 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: "block" # Block malicious requests
persist_session: true # Keep records for investigation
async_mode: false # Require an immediate verdict
include_scanners: true # Understand which rule triggered
include_evidence: true # Capture concrete evidence
default_on: true # Enable for all requests
general_settings:
master_key: "your-master-key-here"
master_key: "YOUR_LITELLM_PROXY_MASTER_KEY"
litellm_settings:
set_verbose: true
@ -136,10 +144,14 @@ 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
persist_session: false # Skip dashboard storage for low latency
async_mode: false # Still receive results inline
include_scanners: false # Minimal payload for performance
include_evidence: false # Omit details to keep responses light
default_on: true # Enable for all requests
general_settings:
master_key: "your-secure-master-key-here"
master_key: "YOUR_LITELLM_PROXY_MASTER_KEY"
litellm_settings:
set_verbose: true # Enable detailed logging
@ -169,10 +181,14 @@ 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: "block" # Block threats on input and output
persist_session: true # Preserve conversations in Pillar dashboard
async_mode: false # Require synchronous approval
include_scanners: true # Inspect which scanners fired
include_evidence: true # Include detailed evidence for auditing
default_on: true # Enable for all requests
general_settings:
master_key: "your-secure-master-key-here"
master_key: "YOUR_LITELLM_PROXY_MASTER_KEY"
litellm_settings:
set_verbose: true # Enable detailed logging
@ -229,19 +245,139 @@ Logs the violation but allows the request to proceed:
on_flagged_action: "monitor"
```
## Advanced Configuration
**Quick takeaways**
- Every request still runs *all* Pillar scanners; these options only change what comes back.
- Choose richer responses when you need audit trails, lighter responses when latency or cost matters.
- Blocking is controlled by LiteLLMs `on_flagged_action` configuration—Pillar headers do not change block/monitor behaviour.
Pillar Security executes the full scanner suite on each call. The settings below tune the Protect response headers LiteLLM sends, letting you balance fidelity, retention, and latency.
### Response Control
#### Data Retention (`persist_session`)
```yaml
persist_session: false # Default: true
```
- **Why**: Controls whether Pillar stores session data for dashboard visibility.
- **Set false for**: Ephemeral testing, privacy-sensitive interactions.
- **Set true for**: Production monitoring, compliance, historical review (default behaviour).
- **Impact**: `false` means the conversation will *not* appear in the Pillar dashboard.
#### Response Detail Level
The following toggles grow the payload size without changing detection behaviour.
```yaml
include_scanners: true # → plr_scanners (default true in LiteLLM)
include_evidence: true # → plr_evidence (default true in LiteLLM)
```
- **Minimal response** (`include_scanners=false`, `include_evidence=false`)
```json
{
"session_id": "abc-123",
"flagged": true
}
```
Use when you only care about whether Pillar detected a threat.
> **📝 Note:** `flagged: true` means Pillars scanners recommend blocking. Pillar only reports this verdict—LiteLLM enforces your policy via the `on_flagged_action` configuration (no Pillar header controls it):
> - `on_flagged_action: "block"` → LiteLLM raises a 400 guardrail error
> - `on_flagged_action: "monitor"` → LiteLLM logs the threat but still returns the LLM response
- **Scanner breakdown** (`include_scanners=true`)
```json
{
"session_id": "abc-123",
"flagged": true,
"scanners": {
"jailbreak": true,
"prompt_injection": false,
"pii": false,
"secret": false,
"toxic_language": false
/* ... more categories ... */
}
}
```
Use when you need to know which categories triggered.
- **Full context** (both toggles true)
```json
{
"session_id": "abc-123",
"flagged": true,
"scanners": { /* ... */ },
"evidence": [
{
"category": "jailbreak",
"type": "prompt_injection",
"evidence": "Ignore previous instructions",
"metadata": { "start_idx": 0, "end_idx": 28 }
}
]
}
```
Ideal for debugging, audit logs, or compliance exports.
### Processing Mode (`async_mode`)
```yaml
async_mode: true # Default: false
```
- **Why**: Queue the request for background processing instead of waiting for a synchronous verdict.
- **Response shape**:
```json
{
"status": "queued",
"session_id": "abc-123",
"position": 1
}
```
- **Set true for**: Large batch jobs, latency-tolerant pipelines.
- **Set false for**: Real-time user flows (default).
- ⚠️ **Note**: Async mode returns only a 202 queue acknowledgment (no flagged verdict). LiteLLM treats that as “no block,” so the pre-call hook always allows the request. Use async mode only for post-call or monitor-only workflows where delayed review is acceptable.
### Complete Examples
```yaml
guardrails:
# Production: full fidelity & dashboard visibility
- guardrail_name: "pillar-production"
litellm_params:
guardrail: pillar
mode: [pre_call, post_call]
persist_session: true
include_scanners: true
include_evidence: true
on_flagged_action: "block"
# Testing: lightweight, no persistence
- guardrail_name: "pillar-testing"
litellm_params:
guardrail: pillar
mode: pre_call
persist_session: false
include_scanners: false
include_evidence: false
on_flagged_action: "monitor"
```
Keep in mind that LiteLLM forwards these values as the documented `plr_*` headers, so any direct HTTP integrations outside the proxy can reuse the same guidance.
## Examples
<Tabs>
<TabItem value="safe" label="Simple Safe Request">
**Safe requset**
**Safe request**
```bash
# Test with safe content
curl -X POST "http://localhost:4000/v1/chat/completions" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-master-key-here" \
-H "Authorization: Bearer YOUR_LITELLM_PROXY_MASTER_KEY" \
-d '{
"model": "gpt-4.1-mini",
"messages": [{"role": "user", "content": "Hello! Can you tell me a joke?"}],
@ -300,7 +436,7 @@ curl -X POST "http://localhost:4000/v1/chat/completions" \
```bash
curl -X POST "http://localhost:4000/v1/chat/completions" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-master-key-here" \
-H "Authorization: Bearer YOUR_LITELLM_PROXY_MASTER_KEY" \
-d '{
"model": "gpt-4.1-mini",
"messages": [
@ -350,7 +486,7 @@ curl -X POST "http://localhost:4000/v1/chat/completions" \
```bash
curl -X POST "http://localhost:4000/v1/chat/completions" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-master-key-here" \
-H "Authorization: Bearer YOUR_LITELLM_PROXY_MASTER_KEY" \
-d '{
"model": "gpt-4.1-mini",
"messages": [
@ -405,4 +541,4 @@ Feel free to contact us at support@pillar.security
- [Pillar Security API Docs](https://docs.pillar.security/docs/api/introduction)
- [Pillar Security Dashboard](https://app.pillar.security)
- [Pillar Security Website](https://pillar.security)
- [LiteLLM Docs](https://docs.litellm.ai)
- [LiteLLM Docs](https://docs.litellm.ai)

View File

@ -23,6 +23,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
if not guardrail_name:
raise ValueError("Pillar guardrail name is required")
optional_params = getattr(litellm_params, "optional_params", None)
_pillar_callback = PillarGuardrail(
guardrail_name=guardrail_name,
api_key=litellm_params.api_key,
@ -30,12 +32,34 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
on_flagged_action=getattr(litellm_params, "on_flagged_action", "monitor"),
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
async_mode=_get_config_value(
litellm_params, optional_params, "async_mode"
),
persist_session=_get_config_value(
litellm_params, optional_params, "persist_session"
),
include_scanners=_get_config_value(
litellm_params, optional_params, "include_scanners"
),
include_evidence=_get_config_value(
litellm_params, optional_params, "include_evidence"
),
)
litellm.logging_callback_manager.add_litellm_callback(_pillar_callback)
return _pillar_callback
def _get_config_value(litellm_params, optional_params, attribute_name):
"""Return guardrail configuration value prioritising optional params when present."""
if optional_params is not None:
value = getattr(optional_params, attribute_name, None)
if value is not None:
return value
return getattr(litellm_params, attribute_name, None)
guardrail_initializer_registry = {
SupportedGuardrailIntegrations.PILLAR.value: initialize_guardrail,
}

View File

@ -69,6 +69,10 @@ class PillarGuardrail(CustomGuardrail):
api_key: Optional[str] = None,
api_base: Optional[str] = None,
on_flagged_action: Optional[str] = None,
async_mode: Optional[bool] = None,
persist_session: Optional[bool] = None,
include_scanners: Optional[bool] = None,
include_evidence: Optional[bool] = None,
**kwargs,
) -> None:
"""
@ -110,6 +114,31 @@ class PillarGuardrail(CustomGuardrail):
f"Pillar Guardrail: Initialized with on_flagged_action: {self.on_flagged_action}"
)
self.async_mode = self._resolve_bool_config(
provided_value=async_mode,
env_var="PILLAR_ASYNC",
default=None,
setting_name="async_mode",
)
self.persist_session = self._resolve_bool_config(
provided_value=persist_session,
env_var="PILLAR_PERSIST",
default=None,
setting_name="persist_session",
)
self.include_scanners = self._resolve_bool_config(
provided_value=include_scanners,
env_var="PILLAR_INCLUDE_SCANNERS",
default=True,
setting_name="include_scanners",
)
self.include_evidence = self._resolve_bool_config(
provided_value=include_evidence,
env_var="PILLAR_INCLUDE_EVIDENCE",
default=True,
setting_name="include_evidence",
)
# Define supported event hooks
supported_event_hooks = [
GuardrailEventHooks.pre_call,
@ -347,12 +376,74 @@ class PillarGuardrail(CustomGuardrail):
"Content-Type": "application/json",
}
# Add Pillar-specific headers for enhanced response data
headers["plr_evidence"] = "true"
headers["plr_scanners"] = "true"
# Add Pillar-specific headers based on configuration
self._set_bool_header(headers, "plr_scanners", self.include_scanners)
self._set_bool_header(headers, "plr_evidence", self.include_evidence)
self._set_bool_header(headers, "plr_async", self.async_mode)
self._set_bool_header(headers, "plr_persist", self.persist_session)
return headers
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:
return
headers[header_name] = "true" if value else "false"
def _resolve_bool_config(
self,
provided_value: Optional[Union[bool, str, int]],
env_var: Optional[str],
default: Optional[bool],
setting_name: str,
) -> Optional[bool]:
"""Resolve configuration precedence: explicit value -> environment -> default."""
if provided_value is not None:
try:
return self._parse_bool_value(provided_value)
except ValueError:
verbose_proxy_logger.warning(
"Pillar Guardrail: Invalid boolean value '%s' for %s, falling back to default.",
provided_value,
setting_name,
)
return default
if env_var:
env_value = os.getenv(env_var)
if env_value is not None:
try:
return self._parse_bool_value(env_value)
except ValueError:
verbose_proxy_logger.warning(
"Pillar Guardrail: Invalid boolean env value '%s' for %s, falling back to default.",
env_value,
env_var,
)
return default
return default
@staticmethod
def _parse_bool_value(value: Union[bool, str, int]) -> bool:
"""Normalise various truthy/falsey inputs to a strict boolean."""
if isinstance(value, bool):
return value
if isinstance(value, int):
return bool(value)
value_str = str(value).strip().lower()
if value_str in {"true", "1", "yes", "y", "on"}:
return True
if value_str in {"false", "0", "no", "n", "off"}:
return False
raise ValueError(f"Unrecognised boolean value: {value}")
def _extract_model_and_provider(self, data: dict) -> Tuple[str, str]:
"""
Extract the model and provider from the request data.

View File

@ -360,6 +360,22 @@ class PillarGuardrailConfigModel(BaseModel):
default="monitor",
description="Action to take when content is flagged: 'block' (raise exception) or 'monitor' (log only)",
)
async_mode: Optional[bool] = Field(
default=None,
description="Set to True to request asynchronous analysis (sets `plr_async` header). Defaults to provider behaviour when omitted.",
)
persist_session: Optional[bool] = Field(
default=None,
description="Controls Pillar session persistence (sets `plr_persist` header). Set to False to disable persistence.",
)
include_scanners: Optional[bool] = Field(
default=True,
description="Include scanner category summaries in responses (sets `plr_scanners` header).",
)
include_evidence: Optional[bool] = Field(
default=True,
description="Include detailed evidence payloads in responses (sets `plr_evidence` header).",
)
class NomaGuardrailConfigModel(BaseModel):

View File

@ -15,6 +15,22 @@ class PillarGuardrailConfigModelOptionalParams(BaseModel):
default="monitor",
description="Action to take when content is flagged: 'block' (raise exception) or 'monitor' (log only). If not provided, the `PILLAR_ON_FLAGGED_ACTION` environment variable is checked, defaults to 'monitor'.",
)
async_mode: Optional[bool] = Field(
default=None,
description="Set to True to request asynchronous analysis (sets `plr_async` header).",
)
persist_session: Optional[bool] = Field(
default=None,
description="Set to False to disable session persistence (sets `plr_persist` header).",
)
include_scanners: Optional[bool] = Field(
default=True,
description="Include scanner summaries in response payloads (sets `plr_scanners` header).",
)
include_evidence: Optional[bool] = Field(
default=True,
description="Include detailed evidence objects in response payloads (sets `plr_evidence` header).",
)
class PillarGuardrailConfigModel(

View File

@ -8,8 +8,12 @@ and following LiteLLM testing patterns and best practices.
# Standard library imports
import os
import sys
from typing import Dict
from unittest.mock import Mock, patch
# Add parent directory to path for imports
sys.path.insert(0, os.path.abspath("../../.."))
# Third-party imports
import pytest
from fastapi.exceptions import HTTPException
@ -26,9 +30,6 @@ from litellm.proxy.guardrails.guardrail_hooks.pillar import (
)
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
# Add parent directory to path for imports
sys.path.insert(0, os.path.abspath("../.."))
# ============================================================================
# FIXTURES
@ -221,6 +222,18 @@ def mock_llm_response():
return mock_response
@pytest.fixture
def pillar_async_response():
"""Fixture providing an asynchronous Pillar API queue response."""
return Response(
json={"status": "queued", "session_id": "async-session", "position": 1},
status_code=202,
request=Request(
method="POST", url="https://api.pillar.security/api/v1/protect"
),
)
@pytest.fixture
def mock_llm_response_with_tools():
"""Fixture providing a mock LLM response with tool calls."""
@ -440,6 +453,55 @@ async def test_post_call_hook_with_tool_calls(
assert result == mock_llm_response_with_tools
# =========================================================================
# HEADER CONFIGURATION TESTS
# =========================================================================
@pytest.mark.asyncio
async def test_pre_call_hook_custom_header_overrides(
sample_request_data,
user_api_key_dict,
dual_cache,
pillar_async_response,
):
"""Ensure configuration values translate into correct Protect headers."""
guardrail = PillarGuardrail(
guardrail_name="pillar-header-test",
api_key="test-pillar-key",
api_base="https://api.pillar.security",
on_flagged_action="monitor",
persist_session=False,
async_mode=True,
include_scanners=False,
include_evidence=False,
)
captured_headers: Dict[str, str] = {}
async def _mock_post(*args, **kwargs):
captured_headers.update(kwargs.get("headers", {}))
return pillar_async_response
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new=_mock_post,
):
result = await guardrail.async_pre_call_hook(
data=sample_request_data,
cache=dual_cache,
user_api_key_dict=user_api_key_dict,
call_type="completion",
)
assert result == sample_request_data
assert captured_headers.get("plr_persist") == "false"
assert captured_headers.get("plr_async") == "true"
assert captured_headers.get("plr_scanners") == "false"
assert captured_headers.get("plr_evidence") == "false"
# ============================================================================
# EDGE CASE TESTS
# ============================================================================