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:
commit
3bf32e8e5c
@ -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 LiteLLM’s `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 Pillar’s 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)
|
||||
|
||||
@ -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,
|
||||
}
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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
|
||||
# ============================================================================
|
||||
Loading…
Reference in New Issue
Block a user