mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 14:23:44 +00:00
feat: add GraySwan Guardrails support (#15756)
This commit is contained in:
@@ -0,0 +1,147 @@
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# GraySwan Cygnal Guardrail
|
||||
|
||||
Use [GraySwan Cygnal](https://docs.grayswan.ai/cygnal/monitor-requests) to continuously monitor conversations for policy violations, indirect prompt injection (IPI), jailbreak attempts, and other safety risks.
|
||||
|
||||
Cygnal returns a `violation` score between `0` and `1` (higher means more likely to violate policy), plus metadata such as violated rule indices, mutation detection, and IPI flags. LiteLLM can automatically block or monitor requests based on this signal.
|
||||
|
||||
---
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Obtain Credentials
|
||||
|
||||
1. Create a GraySwan account and generate a Cygnal API key.
|
||||
2. Configure environment variables for the LiteLLM proxy host:
|
||||
|
||||
```bash
|
||||
export GRAYSWAN_API_KEY="your-grayswan-key"
|
||||
```
|
||||
|
||||
### 2. Configure `config.yaml`
|
||||
|
||||
Add a guardrail entry that references the GraySwan integration. Below is a balanced example that monitors both input and output but only blocks once the violation score reaches the configured threshold.
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: openai/gpt-4.1-mini
|
||||
litellm_params:
|
||||
model: openai/gpt-4.1-mini
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
guardrails:
|
||||
- guardrail_name: "cygnal-monitor"
|
||||
litellm_params:
|
||||
guardrail: grayswan
|
||||
mode: [pre_call, post_call] # monitor both input and output
|
||||
api_key: os.environ/GRAYSWAN_API_KEY
|
||||
optional_params:
|
||||
on_flagged_action: monitor # or "block"
|
||||
violation_threshold: 0.5 # score >= threshold is flagged
|
||||
reasoning_mode: hybrid # off | hybrid | thinking
|
||||
categories:
|
||||
safety: "Detect jailbreaks and policy violations"
|
||||
policy_id: "your-cygnal-policy-id"
|
||||
default_on: true
|
||||
|
||||
general_settings:
|
||||
master_key: "your-litellm-master-key"
|
||||
|
||||
litellm_settings:
|
||||
set_verbose: true
|
||||
```
|
||||
|
||||
### 3. Launch the Proxy
|
||||
|
||||
```bash
|
||||
litellm --config config.yaml --port 4000
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Choosing Guardrail Modes
|
||||
|
||||
GraySwan can run during `pre_call`, `during_call`, and `post_call` stages. Combine modes based on your latency and coverage requirements.
|
||||
|
||||
| Mode | When it Runs | Protects | Typical Use Case |
|
||||
|--------------|-------------------|-----------------------|------------------|
|
||||
| `pre_call` | Before LLM call | User input only | Block prompt injection before it reaches the model |
|
||||
| `during_call`| Parallel to call | User input only | Low-latency monitoring without blocking |
|
||||
| `post_call` | After response | Full conversation | Scan output for policy violations, leaked secrets, or IPI |
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="monitor" label="Monitor Only">
|
||||
|
||||
```yaml
|
||||
guardrails:
|
||||
- guardrail_name: "cygnal-monitor-only"
|
||||
litellm_params:
|
||||
guardrail: grayswan
|
||||
mode: "during_call"
|
||||
api_key: os.environ/GRAYSWAN_API_KEY
|
||||
optional_params:
|
||||
on_flagged_action: monitor
|
||||
violation_threshold: 0.6
|
||||
default_on: true
|
||||
```
|
||||
|
||||
Best for visibility without blocking. Alerts are logged via LiteLLM’s standard logging callbacks.
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="block-input" label="Block Input">
|
||||
|
||||
```yaml
|
||||
guardrails:
|
||||
- guardrail_name: "cygnal-block-input"
|
||||
litellm_params:
|
||||
guardrail: grayswan
|
||||
mode: "pre_call"
|
||||
api_key: os.environ/GRAYSWAN_API_KEY
|
||||
optional_params:
|
||||
on_flagged_action: block
|
||||
violation_threshold: 0.4
|
||||
categories:
|
||||
pii: "Detect sensitive data"
|
||||
default_on: true
|
||||
```
|
||||
|
||||
Stops malicious or sensitive prompts before any tokens are generated.
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="full-coverage" label="Full Coverage">
|
||||
|
||||
```yaml
|
||||
guardrails:
|
||||
- guardrail_name: "cygnal-full-coverage"
|
||||
litellm_params:
|
||||
guardrail: grayswan
|
||||
mode: [pre_call, post_call]
|
||||
api_key: os.environ/GRAYSWAN_API_KEY
|
||||
optional_params:
|
||||
on_flagged_action: block
|
||||
violation_threshold: 0.5
|
||||
reasoning_mode: thinking
|
||||
policy_id: "policy-id-from-grayswan"
|
||||
default_on: true
|
||||
```
|
||||
|
||||
Provides the strongest enforcement by inspecting both prompts and responses.
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
---
|
||||
|
||||
## Configuration Reference
|
||||
|
||||
| Parameter | Type | Description |
|
||||
|---------------------------------------|-----------------|-------------|
|
||||
| `api_key` | string | GraySwan Cygnal API key. Reads from `GRAYSWAN_API_KEY` if omitted. |
|
||||
| `mode` | string or list | Guardrail stages (`pre_call`, `during_call`, `post_call`). |
|
||||
| `optional_params.on_flagged_action` | string | `monitor` (log only) or `block` (raise `HTTPException`). |
|
||||
| `.optional_params.violation_threshold`| number (0-1) | Scores at or above this value are considered violations. |
|
||||
| `optional_params.reasoning_mode` | string | `off`, `hybrid`, or `thinking`. Enables Cygnal’s reasoning capabilities. |
|
||||
| `optional_params.categories` | object | Map of custom category names to descriptions. |
|
||||
| `optional_params.policy_id` | string | GraySwan policy identifier. |
|
||||
@@ -0,0 +1,74 @@
|
||||
"""GraySwan Cygnal guardrail integration for LiteLLM."""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
from .grayswan import (
|
||||
GraySwanGuardrail,
|
||||
GraySwanGuardrailAPIError,
|
||||
GraySwanGuardrailMissingSecrets,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def initialize_guardrail(
|
||||
litellm_params: "LitellmParams", guardrail: "Guardrail"
|
||||
) -> GraySwanGuardrail:
|
||||
import litellm
|
||||
|
||||
guardrail_name = guardrail.get("guardrail_name")
|
||||
if not guardrail_name:
|
||||
raise ValueError("GraySwan guardrail requires a guardrail_name")
|
||||
|
||||
optional_params = getattr(litellm_params, "optional_params", None)
|
||||
|
||||
grayswan_guardrail = GraySwanGuardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
api_key=litellm_params.api_key,
|
||||
api_base=litellm_params.api_base,
|
||||
on_flagged_action=_get_config_value(
|
||||
litellm_params, optional_params, "on_flagged_action"
|
||||
),
|
||||
violation_threshold=_get_config_value(
|
||||
litellm_params, optional_params, "violation_threshold"
|
||||
),
|
||||
reasoning_mode=_get_config_value(
|
||||
litellm_params, optional_params, "reasoning_mode"
|
||||
),
|
||||
categories=_get_config_value(litellm_params, optional_params, "categories"),
|
||||
policy_id=_get_config_value(litellm_params, optional_params, "policy_id"),
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
)
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_callback(grayswan_guardrail)
|
||||
return grayswan_guardrail
|
||||
|
||||
|
||||
def _get_config_value(litellm_params, optional_params, attribute_name):
|
||||
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.GRAYSWAN.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
|
||||
guardrail_class_registry = {
|
||||
SupportedGuardrailIntegrations.GRAYSWAN.value: GraySwanGuardrail,
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GraySwanGuardrail",
|
||||
"GraySwanGuardrailAPIError",
|
||||
"GraySwanGuardrailMissingSecrets",
|
||||
"initialize_guardrail",
|
||||
]
|
||||
@@ -0,0 +1,365 @@
|
||||
"""GraySwan Cygnal guardrail integration."""
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, Literal, Optional, Union
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import LLMResponseTypes
|
||||
|
||||
|
||||
class GraySwanGuardrailMissingSecrets(Exception):
|
||||
"""Raised when the GraySwan API key is missing."""
|
||||
|
||||
|
||||
class GraySwanGuardrailAPIError(Exception):
|
||||
"""Raised when the GraySwan API returns an error."""
|
||||
|
||||
|
||||
class GraySwanGuardrail(CustomGuardrail):
|
||||
"""
|
||||
Guardrail that calls GraySwan's Cygnal monitoring endpoint.
|
||||
|
||||
see: https://docs.grayswan.ai/cygnal/monitor-requests
|
||||
"""
|
||||
|
||||
SUPPORTED_ON_FLAGGED_ACTIONS = {"block", "monitor"}
|
||||
DEFAULT_ON_FLAGGED_ACTION = "monitor"
|
||||
BASE_API_URL = "https://api.grayswan.ai"
|
||||
MONITOR_PATH = "/cygnal/monitor"
|
||||
SUPPORTED_REASONING_MODES = {"off", "hybrid", "thinking"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrail_name: Optional[str] = "grayswan",
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
on_flagged_action: Optional[str] = None,
|
||||
violation_threshold: Optional[float] = None,
|
||||
reasoning_mode: Optional[str] = None,
|
||||
categories: Optional[Dict[str, str]] = None,
|
||||
policy_id: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self.async_handler = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback
|
||||
)
|
||||
|
||||
api_key_value = api_key or os.getenv("GRAYSWAN_API_KEY")
|
||||
if not api_key_value:
|
||||
raise GraySwanGuardrailMissingSecrets(
|
||||
"GraySwan API key missing. Set `GRAYSWAN_API_KEY` or pass `api_key`."
|
||||
)
|
||||
self.api_key: str = api_key_value
|
||||
|
||||
base = api_base or os.getenv("GRAYSWAN_API_BASE") or self.BASE_API_URL
|
||||
self.api_base = base.rstrip("/")
|
||||
self.monitor_url = f"{self.api_base}{self.MONITOR_PATH}"
|
||||
|
||||
action = on_flagged_action
|
||||
if action and action.lower() in self.SUPPORTED_ON_FLAGGED_ACTIONS:
|
||||
self.on_flagged_action = action.lower()
|
||||
else:
|
||||
if action:
|
||||
verbose_proxy_logger.warning(
|
||||
"GraySwan Guardrail: Unsupported on_flagged_action '%s', defaulting to '%s'.",
|
||||
action,
|
||||
self.DEFAULT_ON_FLAGGED_ACTION,
|
||||
)
|
||||
self.on_flagged_action = self.DEFAULT_ON_FLAGGED_ACTION
|
||||
|
||||
self.violation_threshold = self._resolve_threshold(violation_threshold)
|
||||
self.reasoning_mode = self._resolve_reasoning_mode(reasoning_mode)
|
||||
self.categories = categories
|
||||
self.policy_id = policy_id
|
||||
|
||||
supported_event_hooks = [
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.during_call,
|
||||
GuardrailEventHooks.post_call,
|
||||
]
|
||||
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
supported_event_hooks=supported_event_hooks,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Guardrail hook entry points
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@log_guardrail_information
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache,
|
||||
data: dict,
|
||||
call_type: Literal[
|
||||
"completion",
|
||||
"text_completion",
|
||||
"embeddings",
|
||||
"image_generation",
|
||||
"moderation",
|
||||
"audio_transcription",
|
||||
"pass_through_endpoint",
|
||||
"rerank",
|
||||
"mcp_call",
|
||||
"anthropic_messages",
|
||||
],
|
||||
) -> Optional[Union[Exception, str, dict]]:
|
||||
if (
|
||||
self.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.pre_call
|
||||
)
|
||||
is not True
|
||||
):
|
||||
return data
|
||||
|
||||
verbose_proxy_logger.debug("GraySwan Guardrail: pre-call hook triggered")
|
||||
|
||||
messages = data.get("messages")
|
||||
if not messages:
|
||||
verbose_proxy_logger.debug("GraySwan Guardrail: No messages in data")
|
||||
return data
|
||||
|
||||
dynamic_body = self.get_guardrail_dynamic_request_body_params(data) or {}
|
||||
|
||||
payload = self._prepare_payload(messages, dynamic_body)
|
||||
if payload is None:
|
||||
verbose_proxy_logger.debug(
|
||||
"GraySwan Guardrail: no content to scan; skipping request"
|
||||
)
|
||||
return data
|
||||
|
||||
await self.run_grayswan_guardrail(payload)
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
return data
|
||||
|
||||
@log_guardrail_information
|
||||
async def async_moderation_hook(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
call_type: Literal[
|
||||
"completion",
|
||||
"embeddings",
|
||||
"image_generation",
|
||||
"moderation",
|
||||
"audio_transcription",
|
||||
"responses",
|
||||
"mcp_call",
|
||||
"anthropic_messages",
|
||||
],
|
||||
) -> Optional[Union[Exception, str, dict]]:
|
||||
if (
|
||||
self.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.during_call
|
||||
)
|
||||
is not True
|
||||
):
|
||||
return data
|
||||
|
||||
verbose_proxy_logger.debug("GraySwan Guardrail: during-call hook triggered")
|
||||
|
||||
messages = data.get("messages")
|
||||
if not messages:
|
||||
verbose_proxy_logger.debug("GraySwan Guardrail: No messages in data")
|
||||
return data
|
||||
|
||||
dynamic_body = self.get_guardrail_dynamic_request_body_params(data) or {}
|
||||
|
||||
payload = self._prepare_payload(messages, dynamic_body)
|
||||
if payload is None:
|
||||
verbose_proxy_logger.debug(
|
||||
"GraySwan Guardrail: no content to scan; skipping request"
|
||||
)
|
||||
return data
|
||||
|
||||
await self.run_grayswan_guardrail(payload)
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
return data
|
||||
|
||||
@log_guardrail_information
|
||||
async def async_post_call_success_hook(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: LLMResponseTypes,
|
||||
) -> LLMResponseTypes:
|
||||
if (
|
||||
self.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.post_call
|
||||
)
|
||||
is not True
|
||||
):
|
||||
return response
|
||||
|
||||
verbose_proxy_logger.debug("GraySwan Guardrail: post-call hook triggered")
|
||||
|
||||
response_dict = response.model_dump() if hasattr(response, "model_dump") else {}
|
||||
response_messages = [
|
||||
msg if isinstance(msg, dict) else msg.model_dump()
|
||||
for choice in response_dict.get("choices", [])
|
||||
if isinstance(choice, dict)
|
||||
for msg in [choice.get("message")]
|
||||
if msg is not None
|
||||
]
|
||||
|
||||
if not response_messages:
|
||||
verbose_proxy_logger.debug(
|
||||
"GraySwan Guardrail: no response messages detected; skipping post-call scan"
|
||||
)
|
||||
return response
|
||||
|
||||
dynamic_body = self.get_guardrail_dynamic_request_body_params(data) or {}
|
||||
|
||||
payload = self._prepare_payload(response_messages, dynamic_body)
|
||||
if payload is None:
|
||||
verbose_proxy_logger.debug(
|
||||
"GraySwan Guardrail: no content to scan; skipping request"
|
||||
)
|
||||
return response
|
||||
|
||||
await self.run_grayswan_guardrail(payload)
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
return response
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Core GraySwan interaction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def run_grayswan_guardrail(self, payload: dict):
|
||||
headers = self._prepare_headers()
|
||||
|
||||
try:
|
||||
response = await self.async_handler.post(
|
||||
url=self.monitor_url,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=30.0,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
verbose_proxy_logger.debug(
|
||||
"GraySwan Guardrail: monitor response %s", safe_dumps(result)
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as exc: # pragma: no cover - depends on HTTP client behaviour
|
||||
verbose_proxy_logger.exception(
|
||||
"GraySwan Guardrail: API request failed: %s", exc
|
||||
)
|
||||
raise GraySwanGuardrailAPIError(str(exc)) from exc
|
||||
|
||||
self._process_grayswan_response(result)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _prepare_headers(self) -> Dict[str, str]:
|
||||
return {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"grayswan-api-key": self.api_key,
|
||||
}
|
||||
|
||||
def _prepare_payload(
|
||||
self, messages: list[dict], dynamic_body: dict
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
payload: Dict[str, Any] = {}
|
||||
payload["messages"] = messages
|
||||
|
||||
categories = dynamic_body.get("categories") or self.categories
|
||||
if categories:
|
||||
payload["categories"] = categories
|
||||
|
||||
policy_id = dynamic_body.get("policy_id") or self.policy_id
|
||||
if policy_id:
|
||||
payload["policy_id"] = policy_id
|
||||
|
||||
reasoning_mode = dynamic_body.get("reasoning_mode") or self.reasoning_mode
|
||||
if reasoning_mode:
|
||||
payload["reasoning_mode"] = reasoning_mode
|
||||
|
||||
return payload
|
||||
|
||||
def _process_grayswan_response(self, response_json: Dict[str, Any]) -> None:
|
||||
violation_score = float(response_json.get("violation", 0.0) or 0.0)
|
||||
violated_rules = response_json.get("violated_rules", [])
|
||||
mutation_detected = response_json.get("mutation")
|
||||
ipi_detected = response_json.get("ipi")
|
||||
|
||||
flagged = violation_score >= self.violation_threshold
|
||||
if not flagged:
|
||||
verbose_proxy_logger.debug(
|
||||
"GraySwan Guardrail: request passed (score=%s, rules=%s)",
|
||||
violation_score,
|
||||
violated_rules,
|
||||
)
|
||||
return
|
||||
|
||||
verbose_proxy_logger.warning(
|
||||
"GraySwan Guardrail: violation score %.3f exceeds threshold %.3f",
|
||||
violation_score,
|
||||
self.violation_threshold,
|
||||
)
|
||||
|
||||
if self.on_flagged_action == "block":
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "Blocked by GraySwan Guardrail",
|
||||
"violation": violation_score,
|
||||
"violated_rules": violated_rules,
|
||||
"mutation": mutation_detected,
|
||||
"ipi": ipi_detected,
|
||||
},
|
||||
)
|
||||
|
||||
def _resolve_threshold(self, threshold: Optional[float]) -> float:
|
||||
if threshold is not None:
|
||||
return min(max(threshold, 0.0), 1.0)
|
||||
return 0.5
|
||||
|
||||
def _resolve_reasoning_mode(self, candidate: Optional[str]) -> Optional[str]:
|
||||
if candidate is None:
|
||||
return None
|
||||
normalised = candidate.strip().lower()
|
||||
if normalised in self.SUPPORTED_REASONING_MODES:
|
||||
return normalised
|
||||
verbose_proxy_logger.warning(
|
||||
"GraySwan Guardrail: ignoring unsupported reasoning_mode '%s'",
|
||||
candidate,
|
||||
)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_config_model():
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import (
|
||||
GraySwanGuardrailConfigModel,
|
||||
)
|
||||
|
||||
return GraySwanGuardrailConfigModel
|
||||
@@ -8,6 +8,9 @@ from typing_extensions import Required, TypedDict
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import (
|
||||
EnkryptAIGuardrailConfigs,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import (
|
||||
GraySwanGuardrailConfigModel,
|
||||
)
|
||||
|
||||
"""
|
||||
Pydantic object defining how to set guardrails on litellm proxy
|
||||
@@ -35,6 +38,7 @@ class SupportedGuardrailIntegrations(Enum):
|
||||
PANGEA = "pangea"
|
||||
LASSO = "lasso"
|
||||
PILLAR = "pillar"
|
||||
GRAYSWAN = "grayswan"
|
||||
PANW_PRISMA_AIRS = "panw_prisma_airs"
|
||||
AZURE_PROMPT_SHIELD = "azure/prompt_shield"
|
||||
AZURE_TEXT_MODERATIONS = "azure/text_moderations"
|
||||
@@ -45,6 +49,7 @@ class SupportedGuardrailIntegrations(Enum):
|
||||
JAVELIN = "javelin"
|
||||
ENKRYPTAI = "enkryptai"
|
||||
|
||||
|
||||
class Role(Enum):
|
||||
SYSTEM = "system"
|
||||
ASSISTANT = "assistant"
|
||||
@@ -518,6 +523,7 @@ class LitellmParams(
|
||||
LakeraV2GuardrailConfigModel,
|
||||
LassoGuardrailConfigModel,
|
||||
PillarGuardrailConfigModel,
|
||||
GraySwanGuardrailConfigModel,
|
||||
NomaGuardrailConfigModel,
|
||||
ToolPermissionGuardrailConfigModel,
|
||||
JavelinGuardrailConfigModel,
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
"""GraySwan guardrail configuration models."""
|
||||
|
||||
from typing import Dict, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class GraySwanGuardrailConfigModelOptionalParams(BaseModel):
|
||||
"""Optional parameters for the GraySwan guardrail."""
|
||||
|
||||
on_flagged_action: Optional[str] = Field(
|
||||
default="monitor",
|
||||
description="Action when a violation is detected: 'block' rejects the call, 'monitor' logs only.",
|
||||
)
|
||||
violation_threshold: Optional[float] = Field(
|
||||
default=0.5,
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
description="Threshold between 0 and 1 at which GraySwan violations trigger the configured action.",
|
||||
)
|
||||
reasoning_mode: Optional[str] = Field(
|
||||
default=None,
|
||||
description="GraySwan reasoning mode override. Accepted values: 'off', 'hybrid', 'thinking'.",
|
||||
)
|
||||
policy_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="GraySwan policy identifier to apply during monitoring.",
|
||||
)
|
||||
categories: Optional[Dict[str, str]] = Field(
|
||||
default=None,
|
||||
description="Default GraySwan category definitions to send with each request.",
|
||||
)
|
||||
|
||||
|
||||
class GraySwanGuardrailConfigModel(
|
||||
GuardrailConfigModel[GraySwanGuardrailConfigModelOptionalParams]
|
||||
):
|
||||
"""Configuration parameters for the GraySwan guardrail."""
|
||||
|
||||
api_key: Optional[str] = Field(
|
||||
default=None,
|
||||
description="API key for GraySwan. Reads from the `GRAYSWAN_API_KEY` environment variable when omitted.",
|
||||
)
|
||||
api_base: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Override for the GraySwan API base URL. Defaults to https://api.grayswan.ai and can be set via `GRAYSWAN_API_BASE`.",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "GraySwan Guardrail"
|
||||
@@ -0,0 +1,123 @@
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import (
|
||||
GraySwanGuardrail,
|
||||
GraySwanGuardrailAPIError,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def grayswan_guardrail() -> GraySwanGuardrail:
|
||||
return GraySwanGuardrail(
|
||||
guardrail_name="grayswan-test",
|
||||
api_key="test-key",
|
||||
on_flagged_action="monitor",
|
||||
violation_threshold=0.5,
|
||||
categories={"safety": "general policy"},
|
||||
reasoning_mode="hybrid",
|
||||
policy_id="default-policy",
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_payload_uses_dynamic_overrides(grayswan_guardrail: GraySwanGuardrail) -> None:
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
dynamic_body = {
|
||||
"categories": {"custom": "override"},
|
||||
"policy_id": "dynamic-policy",
|
||||
"reasoning_mode": "thinking",
|
||||
}
|
||||
|
||||
payload = grayswan_guardrail._prepare_payload(messages, dynamic_body)
|
||||
|
||||
assert payload["messages"] == messages
|
||||
assert payload["categories"] == {"custom": "override"}
|
||||
assert payload["policy_id"] == "dynamic-policy"
|
||||
assert payload["reasoning_mode"] == "thinking"
|
||||
|
||||
|
||||
def test_prepare_payload_falls_back_to_guardrail_defaults(grayswan_guardrail: GraySwanGuardrail) -> None:
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
|
||||
payload = grayswan_guardrail._prepare_payload(messages, {})
|
||||
|
||||
assert payload["categories"] == {"safety": "general policy"}
|
||||
assert payload["policy_id"] == "default-policy"
|
||||
assert payload["reasoning_mode"] == "hybrid"
|
||||
|
||||
|
||||
def test_process_response_does_not_block_under_threshold(grayswan_guardrail: GraySwanGuardrail) -> None:
|
||||
grayswan_guardrail._process_grayswan_response({"violation": 0.3, "violated_rules": []})
|
||||
|
||||
|
||||
def test_process_response_blocks_when_threshold_exceeded() -> None:
|
||||
guardrail = GraySwanGuardrail(
|
||||
guardrail_name="grayswan-block",
|
||||
api_key="test-key",
|
||||
on_flagged_action="block",
|
||||
violation_threshold=0.2,
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
guardrail._process_grayswan_response({"violation": 0.5, "violated_rules": [1]})
|
||||
|
||||
assert exc.value.status_code == 400
|
||||
assert exc.value.detail["violation"] == 0.5
|
||||
|
||||
|
||||
class _DummyResponse:
|
||||
def __init__(self, payload: dict):
|
||||
self._payload = payload
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
return None
|
||||
|
||||
def json(self) -> dict:
|
||||
return self._payload
|
||||
|
||||
|
||||
class _DummyClient:
|
||||
def __init__(self, payload: dict):
|
||||
self.payload = payload
|
||||
self.calls: list[dict] = []
|
||||
|
||||
async def post(self, *, url: str, headers: dict, json: dict, timeout: float):
|
||||
self.calls.append({"url": url, "headers": headers, "json": json, "timeout": timeout})
|
||||
return _DummyResponse(self.payload)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_guardrail_posts_payload(monkeypatch, grayswan_guardrail: GraySwanGuardrail) -> None:
|
||||
dummy_client = _DummyClient({"violation": 0.1})
|
||||
grayswan_guardrail.async_handler = dummy_client
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_process(response_json: dict) -> None:
|
||||
captured["response"] = response_json
|
||||
|
||||
monkeypatch.setattr(grayswan_guardrail, "_process_grayswan_response", fake_process)
|
||||
|
||||
payload = {"messages": [{"role": "user", "content": "test"}]}
|
||||
|
||||
await grayswan_guardrail.run_grayswan_guardrail(payload)
|
||||
|
||||
assert dummy_client.calls[0]["json"] == payload
|
||||
assert captured["response"] == {"violation": 0.1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_guardrail_raises_api_error(grayswan_guardrail: GraySwanGuardrail) -> None:
|
||||
class _FailingClient:
|
||||
async def post(self, **_kwargs):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
grayswan_guardrail.async_handler = _FailingClient()
|
||||
|
||||
payload = {"messages": [{"role": "user", "content": "test"}]}
|
||||
|
||||
with pytest.raises(GraySwanGuardrailAPIError):
|
||||
await grayswan_guardrail.run_grayswan_guardrail(payload)
|
||||
Reference in New Issue
Block a user