diff --git a/docs/my-website/docs/proxy/guardrails/grayswan.md b/docs/my-website/docs/proxy/guardrails/grayswan.md index 7cc75b9f3b..d6efaf1550 100644 --- a/docs/my-website/docs/proxy/guardrails/grayswan.md +++ b/docs/my-website/docs/proxy/guardrails/grayswan.md @@ -73,6 +73,17 @@ Gray Swan can run during `pre_call`, `during_call`, and `post_call` stages. Comb | `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 | + +When using `during_call` with `on_flagged_action: block` or `on_flagged_action: passthrough`: + +- **The LLM call runs in parallel** with the guardrail check using `asyncio.gather` +- **LLM tokens are still consumed** even if the guardrail detects a violation +- The guardrail exception prevents the response from reaching the user, but **does not cancel the running LLM task** +- This means you pay full LLM costs while returning an error/passthrough message to the user + +**Recommendation:** For cost-sensitive applications, use `pre_call` and `post_call` instead of `during_call` for blocking or passthrough modes. Reserve `during_call` for `monitor` mode where you want low-latency logging without impacting the user experience. + + @@ -131,6 +142,24 @@ guardrails: Provides the strongest enforcement by inspecting both prompts and responses. + + + +```yaml +guardrails: + - guardrail_name: "cygnal-passthrough" + litellm_params: + guardrail: grayswan + mode: [pre_call, post_call] + api_key: os.environ/GRAYSWAN_API_KEY + optional_params: + on_flagged_action: passthrough + violation_threshold: 0.5 + default_on: true +``` + +Allows requests to proceed without raising a 400 error when content is flagged. Instead of blocking, the model response content is replaced with a detailed violation message including violation score, violated rules, and detection flags (mutation, IPI). **Supported Response Formats:** OpenAI chat/text completions, Anthropic Messages API. Other response types (embeddings, images, etc.) will log a warning and return unchanged. + @@ -142,7 +171,7 @@ Provides the strongest enforcement by inspecting both prompts and responses. |---------------------------------------|-----------------|-------------| | `api_key` | string | Gray Swan 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), `block` (raise `HTTPException`), or `passthrough` (include detection info in response without blocking). | +| `optional_params.on_flagged_action` | string | `monitor` (log only), `block` (raise `HTTPException`), or `passthrough` (replace response content with violation message, no 400 error). | | `.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. | diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index bb5883dd47..51f7933422 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -35,6 +35,45 @@ if TYPE_CHECKING: dc = DualCache() +class ModifyResponseException(Exception): + """ + Exception raised when a guardrail wants to modify the response. + + This exception carries the synthetic response that should be returned + to the user instead of calling the LLM or instead of the LLM's response. + It should be caught by the proxy and returned with a 200 status code. + + This is a base exception that all guardrails can use to replace responses, + allowing violation messages to be returned as successful responses + rather than errors. + """ + + def __init__( + self, + message: str, + model: str, + request_data: Dict[str, Any], + guardrail_name: Optional[str] = None, + detection_info: Optional[Dict[str, Any]] = None, + ): + """ + Initialize the modify response exception. + + Args: + message: The violation message to return to the user + model: The model that was being called + request_data: The original request data + guardrail_name: Name of the guardrail that raised this exception + detection_info: Additional detection metadata (scores, rules, etc.) + """ + self.message = message + self.model = model + self.request_data = request_data + self.guardrail_name = guardrail_name + self.detection_info = detection_info or {} + super().__init__(message) + + class CustomGuardrail(CustomLogger): def __init__( self, @@ -96,6 +135,50 @@ class CustomGuardrail(CustomLogger): ) return default + def raise_passthrough_exception( + self, + violation_message: str, + request_data: Dict[str, Any], + detection_info: Optional[Dict[str, Any]] = None, + ) -> None: + """ + Raise a passthrough exception for guardrail violations. + + This helper method should be used by guardrails when they detect a violation + in passthrough mode. + + The exception will be caught by the proxy endpoints and converted to a 200 response + with the violation message, preventing the LLM call from being made (pre_call/during_call) + or replacing the LLM response (post_call). + + Args: + violation_message: The formatted violation message to return to the user + request_data: The original request data dictionary + detection_info: Optional dictionary with detection metadata (scores, rules, etc.) + + Raises: + ModifyResponseException: Always raises this exception to short-circuit + the LLM call and return the violation message + + Example: + if violation_detected and self.on_flagged_action == "passthrough": + message = self._format_violation_message(detection_info) + self.raise_passthrough_exception( + violation_message=message, + request_data=data, + detection_info=detection_info + ) + """ + model = request_data.get("model", "unknown") + + raise ModifyResponseException( + message=violation_message, + model=model, + request_data=request_data, + guardrail_name=self.guardrail_name, + detection_info=detection_info, + ) + @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: """ diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 3317560904..53d9bf756f 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -8,6 +8,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.types.utils import TokenCountResponse @@ -65,6 +66,50 @@ async def anthropic_response( # noqa: PLR0915 version=version, ) return result + except ModifyResponseException as e: + # Guardrail flagged content in passthrough mode - return 200 with violation message + _data = e.request_data + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=_data, + ) + + # Create Anthropic-formatted response with violation message + import uuid + from litellm.types.utils import AnthropicMessagesResponse + + _anthropic_response = AnthropicMessagesResponse( + id=f"msg_{str(uuid.uuid4())}", + type="message", + role="assistant", + content=[{"type": "text", "text": e.message}], + model=e.model, + stop_reason="end_turn", + usage={"input_tokens": 0, "output_tokens": 0}, + ) + + if data.get("stream", None) is not None and data["stream"] is True: + # For streaming, use the standard SSE data generator + async def _passthrough_stream_generator(): + yield _anthropic_response + + selected_data_generator = ( + ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=_passthrough_stream_generator(), + user_api_key_dict=user_api_key_dict, + request_data=_data, + proxy_logging_obj=proxy_logging_obj, + ) + ) + + return await create_streaming_response( + generator=selected_data_generator, + media_type="text/event-stream", + headers={}, + ) + + return _anthropic_response except Exception as e: await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 7f9f900a5a..a1cab09209 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -147,7 +147,7 @@ class GraySwanGuardrail(CustomGuardrail): ) return data - await self.run_grayswan_guardrail(payload, data) + await self.run_grayswan_guardrail(payload, data, GuardrailEventHooks.pre_call) add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=self.guardrail_name ) @@ -193,7 +193,9 @@ class GraySwanGuardrail(CustomGuardrail): ) return data - await self.run_grayswan_guardrail(payload, data) + await self.run_grayswan_guardrail( + payload, data, GuardrailEventHooks.during_call + ) add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=self.guardrail_name ) @@ -240,23 +242,57 @@ class GraySwanGuardrail(CustomGuardrail): ) return response - await self.run_grayswan_guardrail(payload, data) + await self.run_grayswan_guardrail(payload, data, GuardrailEventHooks.post_call) - # If passthrough mode and detection info exists, add it to response + # If passthrough mode and detection info exists, replace response content with violation message if self.on_flagged_action == "passthrough" and "metadata" in data: guardrail_detections = data.get("metadata", {}).get( "guardrail_detections", [] ) if guardrail_detections: - # Add guardrail detections to response hidden params for client visibility - hidden_params = getattr(response, "_hidden_params", None) - if hidden_params is not None: - if not hidden_params: - hidden_params = {} - setattr(response, "_hidden_params", hidden_params) + # Replace the model response content with guardrail violation message + violation_message = self._format_violation_message( + guardrail_detections, is_output=True + ) - hidden_params["guardrail_detections"] = guardrail_detections - setattr(response, "_hidden_params", hidden_params) + # Handle ModelResponse (OpenAI-style chat/text completions) + if hasattr(response, "choices") and response.choices: + verbose_proxy_logger.debug( + "Gray Swan Guardrail: Replacing response content in ModelResponse format" + ) + for choice in response.choices: + # Handle chat completion format (message.content) + if hasattr(choice, "message") and hasattr( + choice.message, "content" + ): + choice.message.content = violation_message + # Handle text completion format (text) + elif hasattr(choice, "text"): + choice.text = violation_message + + # Update finish_reason to indicate content filtering + if hasattr(choice, "finish_reason"): + choice.finish_reason = "content_filter" + + # Handle AnthropicMessagesResponse format + elif hasattr(response, "content") and isinstance(response.content, list): # type: ignore + verbose_proxy_logger.debug( + "Gray Swan Guardrail: Replacing response content in Anthropic Messages format" + ) + # Replace content blocks with text block containing violation message + response.content = [ # type: ignore + {"type": "text", "text": violation_message} + ] + # Update stop_reason if present + if hasattr(response, "stop_reason"): + response.stop_reason = "end_turn" # type: ignore + + else: + verbose_proxy_logger.warning( + "Gray Swan Guardrail: Passthrough mode enabled but response format not recognized. " + "Cannot replace content. Response type: %s", + type(response).__name__, + ) add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=self.guardrail_name @@ -267,7 +303,12 @@ class GraySwanGuardrail(CustomGuardrail): # Core GraySwan interaction # ------------------------------------------------------------------ - async def run_grayswan_guardrail(self, payload: dict, data: Optional[dict] = None): + async def run_grayswan_guardrail( + self, + payload: dict, + data: Optional[dict] = None, + hook_type: Optional[GuardrailEventHooks] = None, + ): headers = self._prepare_headers() try: @@ -290,7 +331,7 @@ class GraySwanGuardrail(CustomGuardrail): ) raise GraySwanGuardrailAPIError(str(exc)) from exc - self._process_grayswan_response(result, data) + self._process_grayswan_response(result, data, hook_type) # ------------------------------------------------------------------ # Helpers @@ -324,7 +365,10 @@ class GraySwanGuardrail(CustomGuardrail): return payload def _process_grayswan_response( - self, response_json: Dict[str, Any], data: Optional[dict] = None + self, + response_json: Dict[str, Any], + data: Optional[dict] = None, + hook_type: Optional[GuardrailEventHooks] = None, ) -> None: violation_score = float(response_json.get("violation", 0.0) or 0.0) violated_rules = response_json.get("violated_rules", []) @@ -347,10 +391,17 @@ class GraySwanGuardrail(CustomGuardrail): ) if self.on_flagged_action == "block": + # Determine if violation was in input or output + violation_location = ( + "output" + if hook_type == GuardrailEventHooks.post_call + else "input" + ) raise HTTPException( status_code=400, detail={ "error": "Blocked by Gray Swan Guardrail", + "violation_location": violation_location, "violation": violation_score, "violated_rules": violated_rules, "mutation": mutation_detected, @@ -362,26 +413,90 @@ class GraySwanGuardrail(CustomGuardrail): "Gray Swan Guardrail: Monitoring mode - allowing flagged content to proceed" ) elif self.on_flagged_action == "passthrough": + # Store detection info + detection_info = { + "guardrail": "grayswan", + "flagged": True, + "violation_score": violation_score, + "violated_rules": violated_rules, + "mutation": mutation_detected, + "ipi": ipi_detected, + } + + # For pre_call and during_call, raise exception to short-circuit LLM call + if hook_type in ( + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + ): + verbose_proxy_logger.info( + "Gray Swan Guardrail: Passthrough mode - raising exception to short-circuit LLM call" + ) + violation_message = self._format_violation_message( + [detection_info], is_output=False + ) + self.raise_passthrough_exception( + violation_message=violation_message, + request_data=data or {}, + detection_info=detection_info, + ) + + # For post_call, store in metadata to replace response later verbose_proxy_logger.info( "Gray Swan Guardrail: Passthrough mode - storing detection info in metadata" ) if data is not None: - # Store guardrail detection info in metadata to be included in response if "metadata" not in data: data["metadata"] = {} if "guardrail_detections" not in data["metadata"]: data["metadata"]["guardrail_detections"] = [] - - detection_info = { - "guardrail": "grayswan", - "flagged": True, - "violation_score": violation_score, - "violated_rules": violated_rules, - "mutation": mutation_detected, - "ipi": ipi_detected, - } data["metadata"]["guardrail_detections"].append(detection_info) + def _format_violation_message( + self, guardrail_detections: list, is_output: bool = False + ) -> str: + """ + Format guardrail detections into a user-friendly violation message. + + Args: + guardrail_detections: List of detection info dictionaries + is_output: True if violation is in model output (post_call), False if in input (pre_call/during_call) + + Returns: + Formatted violation message string + """ + if not guardrail_detections: + return "Content was flagged by guardrail" + + # Get the most recent detection (should be from this guardrail) + detection = guardrail_detections[-1] + + violation_score = detection.get("violation_score", 0.0) + violated_rules = detection.get("violated_rules", []) + mutation = detection.get("mutation", False) + ipi = detection.get("ipi", False) + + # Indicate whether violation was in input or output + violation_location = "the model response" if is_output else "input query" + + message_parts = [ + f"Sorry I can't help with that. According to the Gray Swan Cygnal Guardrail, the {violation_location} has a violation score of {violation_score:.2f}.", + ] + + if violated_rules: + message_parts.append( + f"It was violating the rule(s): {', '.join(map(str, violated_rules))}." + ) + + if mutation: + message_parts.append( + "Mutation effort to make the harmful intention disguised was DETECTED." + ) + + if ipi: + message_parts.append("Indirect Prompt Injection was DETECTED.") + + return "\n".join(message_parts) + def _resolve_threshold(self, threshold: Optional[float]) -> float: if threshold is not None: return min(max(threshold, 0.0), 1.0) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6aef4ab493..19cdbf225b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -171,6 +171,7 @@ from litellm.constants import ( ) from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, @@ -4945,6 +4946,43 @@ async def chat_completion( # noqa: PLR0915 return model_dump_with_preserved_fields(result, exclude_unset=True) else: return result + except ModifyResponseException as e: + # Guardrail flagged content in passthrough mode - return 200 with violation message + _data = e.request_data + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=_data, + ) + _chat_response = litellm.ModelResponse() + _chat_response.model = e.model # type: ignore + _chat_response.choices[0].message.content = e.message # type: ignore + _chat_response.choices[0].finish_reason = "content_filter" # type: ignore + + if data.get("stream", None) is not None and data["stream"] is True: + _iterator = litellm.utils.ModelResponseIterator( + model_response=_chat_response, convert_to_delta=True + ) + _streaming_response = litellm.CustomStreamWrapper( + completion_stream=_iterator, + model=e.model, + custom_llm_provider="cached_response", + logging_obj=data.get("litellm_logging_obj", None), + ) + selected_data_generator = select_data_generator( + response=_streaming_response, + user_api_key_dict=user_api_key_dict, + request_data=_data, + ) + + return StreamingResponse( + selected_data_generator, + media_type="text/event-stream", + status_code=200, # Return 200 for passthrough mode + ) + _usage = litellm.Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + _chat_response.usage = _usage # type: ignore + return _chat_response except RejectedRequestError as e: _data = e.request_data await proxy_logging_obj.post_call_failure_hook( @@ -5054,6 +5092,55 @@ async def completion( # noqa: PLR0915 user_api_base=user_api_base, version=version, ) + except ModifyResponseException as e: + # Guardrail flagged content in passthrough mode - return 200 with violation message + _data = e.request_data + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=_data, + ) + + if _data.get("stream", None) is not None and _data["stream"] is True: + _text_response = litellm.ModelResponse() + _text_response.choices[0].text = e.message + _text_response.model = e.model # type: ignore + _usage = litellm.Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + ) + _text_response.usage = _usage # type: ignore + _iterator = litellm.utils.ModelResponseIterator( + model_response=_text_response, convert_to_delta=True + ) + _streaming_response = litellm.TextCompletionStreamWrapper( + completion_stream=_iterator, + model=e.model, + ) + + selected_data_generator = select_data_generator( + response=_streaming_response, + user_api_key_dict=user_api_key_dict, + request_data=_data, + ) + + return StreamingResponse( + selected_data_generator, + media_type="text/event-stream", + status_code=200, # Return 200 for passthrough mode + ) + else: + _response = litellm.TextCompletionResponse() + _response.choices[0].text = e.message + _response.model = e.model # type: ignore + _usage = litellm.Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + ) + _response.usage = _usage # type: ignore + return _response except RejectedRequestError as e: _data = e.request_data await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py b/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py index 6a6e5a1e48..79239e5262 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/grayswan.py @@ -11,8 +11,8 @@ class GraySwanGuardrailConfigModelOptionalParams(BaseModel): """Optional parameters for the Gray Swan guardrail.""" on_flagged_action: Optional[str] = Field( - default="monitor", - description="Action when a violation is detected: 'block' rejects the call, 'monitor' logs only, 'passthrough' includes detection info in response without blocking.", + default="passthrough", + description="Action when a violation is detected: 'block' rejects the call (400 error), 'monitor' logs only, 'passthrough' replaces response content with violation message (200 status).", ) violation_threshold: Optional[float] = Field( default=0.5, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py index 9ee31cf6cb..6dc658827b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -3,6 +3,7 @@ from typing import Optional import pytest from fastapi import HTTPException +from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import ( GraySwanGuardrail, GraySwanGuardrailAPIError, @@ -71,11 +72,27 @@ def test_process_response_blocks_when_threshold_exceeded() -> None: event_hook=GuardrailEventHooks.pre_call, ) + # Test block mode with input violation (pre_call) with pytest.raises(HTTPException) as exc: - guardrail._process_grayswan_response({"violation": 0.5, "violated_rules": [1]}) + guardrail._process_grayswan_response( + {"violation": 0.5, "violated_rules": [1]}, + hook_type=GuardrailEventHooks.pre_call, + ) assert exc.value.status_code == 400 assert exc.value.detail["violation"] == 0.5 + assert exc.value.detail["violation_location"] == "input" + + # Test block mode with output violation (post_call) + with pytest.raises(HTTPException) as exc: + guardrail._process_grayswan_response( + {"violation": 0.5, "violated_rules": [1]}, + hook_type=GuardrailEventHooks.post_call, + ) + + assert exc.value.status_code == 400 + assert exc.value.detail["violation"] == 0.5 + assert exc.value.detail["violation_location"] == "output" class _DummyResponse: @@ -110,7 +127,11 @@ async def test_run_guardrail_posts_payload( captured = {} - def fake_process(response_json: dict, data: Optional[dict] = None) -> None: + def fake_process( + response_json: dict, + data: Optional[dict] = None, + hook_type: Optional[GuardrailEventHooks] = None, + ) -> None: captured["response"] = response_json monkeypatch.setattr(grayswan_guardrail, "_process_grayswan_response", fake_process) @@ -139,8 +160,8 @@ async def test_run_guardrail_raises_api_error( await grayswan_guardrail.run_grayswan_guardrail(payload) -def test_process_response_passthrough_stores_detection_info() -> None: - """Test that passthrough mode stores detection info in metadata without blocking.""" +def test_process_response_passthrough_raises_exception_in_pre_call() -> None: + """Test that passthrough mode raises ModifyResponseException in pre_call hook.""" guardrail = GraySwanGuardrail( guardrail_name="grayswan-passthrough", api_key="test-key", @@ -149,6 +170,64 @@ def test_process_response_passthrough_stores_detection_info() -> None: event_hook=GuardrailEventHooks.pre_call, ) + data = {"messages": [{"role": "user", "content": "test"}], "model": "gpt-4"} + response_json = { + "violation": 0.8, + "violated_rules": [1, 2], + "mutation": True, + "ipi": False, + } + + # Should raise ModifyResponseException + with pytest.raises(ModifyResponseException) as exc: + guardrail._process_grayswan_response( + response_json, data, GuardrailEventHooks.pre_call + ) + + assert "Gray Swan Cygnal Guardrail" in exc.value.message + assert exc.value.model == "gpt-4" + assert exc.value.detection_info["violation_score"] == 0.8 + assert exc.value.detection_info["violated_rules"] == [1, 2] + + +def test_process_response_passthrough_raises_exception_in_during_call() -> None: + """Test that passthrough mode raises ModifyResponseException in during_call hook.""" + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-passthrough", + api_key="test-key", + on_flagged_action="passthrough", + violation_threshold=0.2, + event_hook=GuardrailEventHooks.during_call, + ) + + data = {"messages": [{"role": "user", "content": "test"}], "model": "gpt-4"} + response_json = { + "violation": 0.8, + "violated_rules": [1, 2], + "mutation": True, + "ipi": False, + } + + # Should raise ModifyResponseException + with pytest.raises(ModifyResponseException) as exc: + guardrail._process_grayswan_response( + response_json, data, GuardrailEventHooks.during_call + ) + + assert "Gray Swan Cygnal Guardrail" in exc.value.message + assert exc.value.model == "gpt-4" + + +def test_process_response_passthrough_stores_detection_info_in_post_call() -> None: + """Test that passthrough mode stores detection info in post_call hook (not exception).""" + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-passthrough", + api_key="test-key", + on_flagged_action="passthrough", + violation_threshold=0.2, + event_hook=GuardrailEventHooks.post_call, + ) + data = {"messages": [{"role": "user", "content": "test"}]} response_json = { "violation": 0.8, @@ -157,8 +236,10 @@ def test_process_response_passthrough_stores_detection_info() -> None: "ipi": False, } - # Should not raise an exception - guardrail._process_grayswan_response(response_json, data) + # Should NOT raise an exception in post_call + guardrail._process_grayswan_response( + response_json, data, GuardrailEventHooks.post_call + ) # Verify detection info was stored in metadata assert "metadata" in data @@ -174,8 +255,8 @@ def test_process_response_passthrough_stores_detection_info() -> None: assert detection["ipi"] is False -def test_process_response_passthrough_does_not_store_if_under_threshold() -> None: - """Test that passthrough mode doesn't store anything if violation is under threshold.""" +def test_process_response_passthrough_does_not_raise_if_under_threshold() -> None: + """Test that passthrough mode doesn't raise exception if violation is under threshold.""" guardrail = GraySwanGuardrail( guardrail_name="grayswan-passthrough", api_key="test-key", @@ -184,14 +265,58 @@ def test_process_response_passthrough_does_not_store_if_under_threshold() -> Non event_hook=GuardrailEventHooks.pre_call, ) - data = {"messages": [{"role": "user", "content": "test"}]} + data = {"messages": [{"role": "user", "content": "test"}], "model": "gpt-4"} response_json = { "violation": 0.3, "violated_rules": [], } - # Should not raise an exception - guardrail._process_grayswan_response(response_json, data) + # Should not raise an exception since under threshold + guardrail._process_grayswan_response( + response_json, data, GuardrailEventHooks.pre_call + ) # Should not have any detection info since it didn't exceed threshold assert "guardrail_detections" not in data.get("metadata", {}) + + +def test_format_violation_message() -> None: + """Test that violation message is formatted correctly for input violations.""" + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-passthrough", + api_key="test-key", + on_flagged_action="passthrough", + violation_threshold=0.5, + event_hook=GuardrailEventHooks.pre_call, + ) + + detections = [ + { + "guardrail": "grayswan", + "flagged": True, + "violation_score": 0.85, + "violated_rules": [1, 3, 5], + "mutation": True, + "ipi": False, + } + ] + + # Test input violation message (pre_call/during_call) + message = guardrail._format_violation_message(detections, is_output=False) + + assert "Sorry I can't help with that" in message + assert "Gray Swan Cygnal Guardrail" in message + assert "the input query has a violation score of 0.85" in message + assert "violating the rule(s): 1, 3, 5" in message + assert "Mutation effort to make the harmful intention disguised was DETECTED" in message + # IPI should not be in message since it's False + assert "Indirect Prompt Injection was DETECTED" not in message + + # Test output violation message (post_call) + message = guardrail._format_violation_message(detections, is_output=True) + + assert "Sorry I can't help with that" in message + assert "Gray Swan Cygnal Guardrail" in message + assert "the model response has a violation score of 0.85" in message + assert "violating the rule(s): 1, 3, 5" in message + assert "Mutation effort to make the harmful intention disguised was DETECTED" in message