diff --git a/docs/my-website/docs/proxy/guardrails/grayswan.md b/docs/my-website/docs/proxy/guardrails/grayswan.md new file mode 100644 index 0000000000..325df16915 --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/grayswan.md @@ -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 | + + + + +```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. + + + + +```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. + + + + +```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. + + + + +--- + +## 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. | diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py new file mode 100644 index 0000000000..ffb38b1ac3 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py @@ -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", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py new file mode 100644 index 0000000000..247b144a14 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -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 diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 115cce9fb9..5e718f7801 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py b/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py new file mode 100644 index 0000000000..93836f2308 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py @@ -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" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py new file mode 100644 index 0000000000..d6bd0a251b --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -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)