feat: add GraySwan Guardrails support (#15756)

This commit is contained in:
YutaSaito
2025-10-21 11:13:50 -07:00
committed by GitHub
parent 46d55bd92a
commit d79bdd491f
6 changed files with 768 additions and 0 deletions
@@ -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 LiteLLMs 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 Cygnals 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
+6
View File
@@ -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)