From a401155fe0970d774a85c47b177c97d780e83368 Mon Sep 17 00:00:00 2001 From: thomaslwang <24355116@qq.com> Date: Wed, 18 Mar 2026 21:50:53 -0700 Subject: [PATCH 1/5] feat(guardrails): add OpenGuardrails integration Add native guardrail integration for OpenGuardrails, an open-source AI safety platform with 19 risk categories, data leakage prevention, and automatic private model switching. Supported features: - Input/output content detection (prompt injection, risk categories S1-S19) - Sensitive data anonymization with post-call restoration - Private model switching via LiteLLM model routing - Ban policy enforcement - Knowledge base responses Co-Authored-By: Claude Opus 4.6 (1M context) --- .../openguardrails/__init__.py | 47 ++ .../openguardrails/openguardrails.py | 460 ++++++++++++++++++ litellm/types/guardrails.py | 1 + 3 files changed, 508 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/openguardrails/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/openguardrails/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/__init__.py new file mode 100644 index 00000000000..de5d9ee4968 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/__init__.py @@ -0,0 +1,47 @@ +""" +OpenGuardrails Native Integration for LiteLLM + +Full-featured guardrail integration supporting: +- Input/Output detection (19 risk categories + prompt injection) +- Sensitive data anonymization with restoration +- Private model switching (automatic routing to data-safe models) +- Tool call anomaly detection +- Ban policy enforcement +""" + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .openguardrails import OpenGuardrailsGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", + guardrail: "Guardrail", +): + import litellm + + callback = OpenGuardrailsGuardrail( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + guardrail_name=guardrail.get("guardrail_name", "openguardrails"), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + private_model_name=getattr(litellm_params, "private_model_name", None), + ) + + litellm.logging_callback_manager.add_litellm_callback(callback) + return callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.OPENGUARDRAILS.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.OPENGUARDRAILS.value: OpenGuardrailsGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py new file mode 100644 index 00000000000..8dc228a9c6f --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py @@ -0,0 +1,460 @@ +""" +OpenGuardrails Native Guardrail Integration for LiteLLM + +Full-featured integration that supports all OpenGuardrails capabilities: +- Input/Output content detection (19 risk categories + prompt injection) +- Sensitive data anonymization with restoration (including streaming) +- Private model switching via LiteLLM model routing +- Tool call anomaly detection +- Ban policy enforcement +- Knowledge base responses + +This file is intended to be placed at: + litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py + +For the LiteLLM PR submission. +""" + +import json +import os +import re +from typing import Any, Dict, List, Literal, Optional, Union + +import httpx + +from litellm._logging import verbose_proxy_logger +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, + log_guardrail_information, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks + +GUARDRAIL_NAME = "openguardrails" + +# Default private model name in LiteLLM config +DEFAULT_PRIVATE_MODEL_NAME = "og-private-model" + +# Metadata keys for cross-hook state +_META_SESSION_ID = "og_session_id" +_META_RESTORE_MAPPING = "og_restore_mapping" +_META_SKIP_OUTPUT = "og_skip_output_detection" +_META_INPUT_MESSAGES = "og_input_messages" + + +class OpenGuardrailsGuardrail(CustomGuardrail): + """ + OpenGuardrails integration for LiteLLM. + + Provides enterprise AI safety through OpenGuardrails' detection pipeline: + - 19 risk categories (S1-S19) covering security, compliance, and content safety + - Data Leakage Prevention with format-aware detection + - Automatic private model switching for sensitive data + - Real-time streaming placeholder restoration + + Configuration in LiteLLM config.yaml: + guardrails: + - guardrail_name: "openguardrails" + litellm_params: + guardrail: openguardrails + mode: [pre_call, post_call] + api_base: http://og-server:5001 + api_key: sk-xxai-your-key + default_on: true + private_model_name: og-private-model # optional, default: og-private-model + + The private model must also be defined in model_list: + model_list: + - model_name: og-private-model + litellm_params: + model: openai/your-private-model + api_base: https://your-private-endpoint.com + api_key: sk-your-private-key + """ + + def __init__( + self, + api_base: Optional[str] = None, + api_key: Optional[str] = None, + private_model_name: Optional[str] = None, + **kwargs, + ): + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) + + base_url = api_base or os.environ.get("OPENGUARDRAILS_API_BASE") + if not base_url: + raise ValueError( + "api_base is required for OpenGuardrails. " + "Set OPENGUARDRAILS_API_BASE env var or pass in litellm_params." + ) + self.api_base = base_url.rstrip("/") + + resolved_key = api_key or os.environ.get("OPENGUARDRAILS_API_KEY") + if not resolved_key: + raise ValueError( + "api_key is required for OpenGuardrails. " + "Set OPENGUARDRAILS_API_KEY env var or pass in litellm_params." + ) + self.api_key = resolved_key + + self.private_model_name = ( + private_model_name + or os.environ.get("OPENGUARDRAILS_PRIVATE_MODEL", DEFAULT_PRIVATE_MODEL_NAME) + ) + + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + super().__init__(**kwargs) + + verbose_proxy_logger.debug( + "OpenGuardrails initialized: api_base=%s, private_model=%s", + self.api_base, + self.private_model_name, + ) + + # ------------------------------------------------------------------ # + # HTTP helpers # + # ------------------------------------------------------------------ # + + def _build_headers(self) -> dict: + return { + "Content-Type": "application/json", + "Authorization": f"Bearer {self.api_key}", + } + + async def _call_og_api( + self, path: str, payload: dict + ) -> Optional[dict]: + """Call an OpenGuardrails gateway API endpoint.""" + url = f"{self.api_base}{path}" + try: + response = await self.async_handler.post( + url=url, + json=payload, + headers=self._build_headers(), + ) + response.raise_for_status() + return response.json() + except httpx.HTTPStatusError as exc: + verbose_proxy_logger.error( + "OpenGuardrails API error: %s status=%s body=%s", + path, + exc.response.status_code, + exc.response.text[:500], + ) + return None + except Exception as exc: + verbose_proxy_logger.error( + "OpenGuardrails API request failed: %s error=%s", path, str(exc) + ) + return None + + # ------------------------------------------------------------------ # + # Metadata helpers # + # ------------------------------------------------------------------ # + + @staticmethod + def _get_metadata(data: dict) -> dict: + """Get or create the metadata dict inside data.""" + if "metadata" not in data: + data["metadata"] = {} + return data["metadata"] + + @staticmethod + def _get_litellm_metadata(data: dict) -> dict: + """Get metadata from litellm_metadata (used in post-call hooks).""" + return data.get("litellm_metadata", data.get("metadata", {})) + + # ------------------------------------------------------------------ # + # Pre-call hook: input detection # + # ------------------------------------------------------------------ # + + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict, + cache, + data: dict, + call_type: str, + ) -> Optional[dict]: + """ + Run OpenGuardrails input detection before the LLM call. + + Actions: + - block/replace: raise exception (request rejected) + - anonymize: modify data["messages"] with masked content, + store restore_mapping in metadata for post-call restoration + - switch_private_model: change data["model"] to private model + - pass: no modification + """ + messages = data.get("messages") + if not messages: + return data + + # Build request + payload = { + "messages": messages, + "stream": data.get("stream", False), + } + + # Add user_id if available + user_id = None + if user_api_key_dict: + user_id = getattr(user_api_key_dict, "user_id", None) + if user_id: + payload["user_id"] = str(user_id) + + result = await self._call_og_api("/v1/gateway/process-input", payload) + if result is None: + # API unreachable - fail open + verbose_proxy_logger.warning( + "OpenGuardrails unreachable, proceeding without guardrail (fail-open)" + ) + return data + + action = result.get("action", "pass") + detection = result.get("detection_result", {}) + metadata = self._get_metadata(data) + + verbose_proxy_logger.info( + "OpenGuardrails input: action=%s risk=%s", + action, + detection.get("overall_risk_level", "unknown"), + ) + + # Store input messages for output detection context + metadata[_META_INPUT_MESSAGES] = messages + + if action == "block": + reason = self._extract_response_content( + result.get("block_response", {}), + fallback=f"Request blocked by OpenGuardrails (risk={detection.get('overall_risk_level')})", + ) + raise GuardrailRaisedException( + guardrail_name=GUARDRAIL_NAME, + message=reason, + should_wrap_with_default_message=False, + ) + + elif action == "replace": + # Return replacement content as a 200 response (not an error) + content = self._extract_response_content( + result.get("replace_response", {}), + fallback="Content filtered by OpenGuardrails.", + ) + raise ModifyResponseException( + message=content, + model=data.get("model", "unknown"), + request_data=data, + guardrail_name=GUARDRAIL_NAME, + ) + + elif action == "anonymize": + anonymized = result.get("anonymized_messages") + if anonymized: + data["messages"] = anonymized + + # Store restore mapping for post-call restoration + restore_mapping = result.get("restore_mapping") + if restore_mapping: + metadata[_META_RESTORE_MAPPING] = restore_mapping + + session_id = result.get("session_id") + if session_id: + metadata[_META_SESSION_ID] = session_id + + return data + + elif action in ("switch_private_model", "proxy_response"): + if action == "proxy_response": + # OG already called the private model - return its response + proxy_resp = result.get("proxy_response", {}) + content = self._extract_response_content( + proxy_resp, + fallback="Response from private model.", + ) + raise ModifyResponseException( + message=content, + model=data.get("model", "unknown"), + request_data=data, + guardrail_name=GUARDRAIL_NAME, + ) + + # switch_private_model: route to private model via LiteLLM + verbose_proxy_logger.info( + "OpenGuardrails: switching to private model '%s'", + self.private_model_name, + ) + data["model"] = self.private_model_name + + # Skip output detection for private model (data is safe) + metadata[_META_SKIP_OUTPUT] = True + + # If OG also anonymized the messages, apply that + anonymized = result.get("anonymized_messages") + if anonymized: + data["messages"] = anonymized + restore_mapping = result.get("restore_mapping") + if restore_mapping: + metadata[_META_RESTORE_MAPPING] = restore_mapping + session_id = result.get("session_id") + if session_id: + metadata[_META_SESSION_ID] = session_id + + return data + + # action == "pass" + return data + + # ------------------------------------------------------------------ # + # Post-call hook: output detection + restoration # + # ------------------------------------------------------------------ # + + @log_guardrail_information + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict, + response, + ): + """ + Run OpenGuardrails output detection and placeholder restoration + on the LLM response. + + - If restore_mapping exists: restore anonymized placeholders + - If output detection enabled: check response for risks + """ + metadata = self._get_litellm_metadata(data) + + # Skip output detection for private model responses + if metadata.get(_META_SKIP_OUTPUT): + # Still do restoration if needed + restore_mapping = metadata.get(_META_RESTORE_MAPPING) + if restore_mapping: + self._restore_response_content(response, restore_mapping) + return response + + # Extract response content + content = self._extract_model_response_content(response) + if not content: + return response + + # Build output request + payload: Dict[str, Any] = {"content": content} + + session_id = metadata.get(_META_SESSION_ID) + if session_id: + payload["session_id"] = session_id + + restore_mapping = metadata.get(_META_RESTORE_MAPPING) + if restore_mapping: + payload["restore_mapping"] = restore_mapping + + input_messages = metadata.get(_META_INPUT_MESSAGES) + if input_messages: + payload["messages"] = input_messages + + result = await self._call_og_api("/v1/gateway/process-output", payload) + if result is None: + # API unreachable - fail open, but still do local restoration + if restore_mapping: + self._restore_response_content(response, restore_mapping) + return response + + action = result.get("action", "pass") + + verbose_proxy_logger.info("OpenGuardrails output: action=%s", action) + + if action == "block": + reason = self._extract_response_content( + result.get("block_response", {}), + fallback="Response blocked by OpenGuardrails.", + ) + raise GuardrailRaisedException( + guardrail_name=GUARDRAIL_NAME, + message=reason, + should_wrap_with_default_message=False, + ) + + elif action == "restore": + restored = result.get("restored_content", "") + if restored: + self._set_model_response_content(response, restored) + + elif action == "anonymize": + anonymized = result.get("anonymized_content", "") + if anonymized: + self._set_model_response_content(response, anonymized) + + # action == "pass" - no modification + return response + + # ------------------------------------------------------------------ # + # Content extraction / replacement helpers # + # ------------------------------------------------------------------ # + + @staticmethod + def _extract_response_content(resp_obj: dict, fallback: str = "") -> str: + """Extract message content from an OG response body (JSON string).""" + body = resp_obj.get("body", "") + if not body: + return fallback + try: + body_json = json.loads(body) if isinstance(body, str) else body + content = ( + body_json.get("choices", [{}])[0] + .get("message", {}) + .get("content", "") + ) + return content if content else fallback + except (json.JSONDecodeError, IndexError, KeyError): + return fallback + + @staticmethod + def _extract_model_response_content(response) -> str: + """Extract text content from a LiteLLM ModelResponse.""" + try: + choices = getattr(response, "choices", None) + if not choices: + return "" + message = getattr(choices[0], "message", None) + if not message: + return "" + return getattr(message, "content", "") or "" + except (IndexError, AttributeError): + return "" + + @staticmethod + def _set_model_response_content(response, content: str): + """Set text content on a LiteLLM ModelResponse.""" + try: + if response.choices and response.choices[0].message: + response.choices[0].message.content = content + except (IndexError, AttributeError): + pass + + @staticmethod + def _restore_response_content( + response, restore_mapping: Dict[str, str] + ): + """Replace placeholders in response content using restore_mapping.""" + try: + content = response.choices[0].message.content + if not content: + return + for placeholder, original in restore_mapping.items(): + content = content.replace(placeholder, original) + response.choices[0].message.content = content + except (IndexError, AttributeError): + pass diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index f798f05380d..2ccb94a258e 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -80,6 +80,7 @@ class SupportedGuardrailIntegrations(Enum): MCP_END_USER_PERMISSION = "mcp_end_user_permission" BLOCK_CODE_EXECUTION = "block_code_execution" MCP_JWT_SIGNER = "mcp_jwt_signer" + OPENGUARDRAILS = "openguardrails" class Role(Enum): From c71d41ca7360359dd6b6c43f240b31c6b0f73508 Mon Sep 17 00:00:00 2001 From: thomaslwang <24355116@qq.com> Date: Sat, 21 Mar 2026 02:21:40 -0700 Subject: [PATCH 2/5] fix(guardrails): address Greptile review feedback for OpenGuardrails integration - Fix metadata key mismatch: use litellm_metadata consistently (not metadata) so cross-hook state (restore_mapping, skip flags, session_id) is preserved in proxy deployments - Make api_key optional for keyless self-hosted deployments - Add fail_on_error option to block requests when API is unreachable - Make skip_output_for_private_model configurable (default: true) - Remove unused imports (re, List, Literal, Union) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../openguardrails/openguardrails.py | 56 ++++++++++--------- 1 file changed, 29 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py index 8dc228a9c6f..c9ffa56f98b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py @@ -17,8 +17,7 @@ For the LiteLLM PR submission. import json import os -import re -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Dict, Optional import httpx @@ -64,9 +63,11 @@ class OpenGuardrailsGuardrail(CustomGuardrail): guardrail: openguardrails mode: [pre_call, post_call] api_base: http://og-server:5001 - api_key: sk-xxai-your-key + api_key: sk-xxai-your-key # optional for keyless self-hosted deployments default_on: true private_model_name: og-private-model # optional, default: og-private-model + fail_on_error: false # optional, set true to block when API is unreachable + skip_output_for_private_model: true # optional, set false to scan private model output The private model must also be defined in model_list: model_list: @@ -82,6 +83,8 @@ class OpenGuardrailsGuardrail(CustomGuardrail): api_base: Optional[str] = None, api_key: Optional[str] = None, private_model_name: Optional[str] = None, + fail_on_error: bool = False, + skip_output_for_private_model: bool = True, **kwargs, ): self.async_handler = get_async_httpx_client( @@ -96,13 +99,10 @@ class OpenGuardrailsGuardrail(CustomGuardrail): ) self.api_base = base_url.rstrip("/") - resolved_key = api_key or os.environ.get("OPENGUARDRAILS_API_KEY") - if not resolved_key: - raise ValueError( - "api_key is required for OpenGuardrails. " - "Set OPENGUARDRAILS_API_KEY env var or pass in litellm_params." - ) - self.api_key = resolved_key + self.api_key = api_key or os.environ.get("OPENGUARDRAILS_API_KEY") + + self.fail_on_error = fail_on_error + self.skip_output_for_private_model = skip_output_for_private_model self.private_model_name = ( private_model_name @@ -128,10 +128,10 @@ class OpenGuardrailsGuardrail(CustomGuardrail): # ------------------------------------------------------------------ # def _build_headers(self) -> dict: - return { - "Content-Type": "application/json", - "Authorization": f"Bearer {self.api_key}", - } + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + return headers async def _call_og_api( self, path: str, payload: dict @@ -166,15 +166,8 @@ class OpenGuardrailsGuardrail(CustomGuardrail): @staticmethod def _get_metadata(data: dict) -> dict: - """Get or create the metadata dict inside data.""" - if "metadata" not in data: - data["metadata"] = {} - return data["metadata"] - - @staticmethod - def _get_litellm_metadata(data: dict) -> dict: - """Get metadata from litellm_metadata (used in post-call hooks).""" - return data.get("litellm_metadata", data.get("metadata", {})) + """Get or create the litellm_metadata dict inside data.""" + return data.setdefault("litellm_metadata", {}) # ------------------------------------------------------------------ # # Pre-call hook: input detection # @@ -217,7 +210,11 @@ class OpenGuardrailsGuardrail(CustomGuardrail): result = await self._call_og_api("/v1/gateway/process-input", payload) if result is None: - # API unreachable - fail open + if self.fail_on_error: + raise GuardrailRaisedException( + guardrail_name=GUARDRAIL_NAME, + message="OpenGuardrails API unreachable; blocking request (fail_on_error=True)", + ) verbose_proxy_logger.warning( "OpenGuardrails unreachable, proceeding without guardrail (fail-open)" ) @@ -298,8 +295,8 @@ class OpenGuardrailsGuardrail(CustomGuardrail): ) data["model"] = self.private_model_name - # Skip output detection for private model (data is safe) - metadata[_META_SKIP_OUTPUT] = True + if self.skip_output_for_private_model: + metadata[_META_SKIP_OUTPUT] = True # If OG also anonymized the messages, apply that anonymized = result.get("anonymized_messages") @@ -335,7 +332,7 @@ class OpenGuardrailsGuardrail(CustomGuardrail): - If restore_mapping exists: restore anonymized placeholders - If output detection enabled: check response for risks """ - metadata = self._get_litellm_metadata(data) + metadata = self._get_metadata(data) # Skip output detection for private model responses if metadata.get(_META_SKIP_OUTPUT): @@ -367,6 +364,11 @@ class OpenGuardrailsGuardrail(CustomGuardrail): result = await self._call_og_api("/v1/gateway/process-output", payload) if result is None: + if self.fail_on_error: + raise GuardrailRaisedException( + guardrail_name=GUARDRAIL_NAME, + message="OpenGuardrails API unreachable; blocking response (fail_on_error=True)", + ) # API unreachable - fail open, but still do local restoration if restore_mapping: self._restore_response_content(response, restore_mapping) From 42c3fee507b240b3658988176fe6e6238a89b3e7 Mon Sep 17 00:00:00 2001 From: thomaslwang <24355116@qq.com> Date: Sat, 21 Mar 2026 03:13:40 -0700 Subject: [PATCH 3/5] style: apply black formatting to openguardrails.py Co-Authored-By: Claude Opus 4.6 (1M context) --- .../openguardrails/openguardrails.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py index c9ffa56f98b..c42d24bd73d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py @@ -104,9 +104,8 @@ class OpenGuardrailsGuardrail(CustomGuardrail): self.fail_on_error = fail_on_error self.skip_output_for_private_model = skip_output_for_private_model - self.private_model_name = ( - private_model_name - or os.environ.get("OPENGUARDRAILS_PRIVATE_MODEL", DEFAULT_PRIVATE_MODEL_NAME) + self.private_model_name = private_model_name or os.environ.get( + "OPENGUARDRAILS_PRIVATE_MODEL", DEFAULT_PRIVATE_MODEL_NAME ) if "supported_event_hooks" not in kwargs: @@ -133,9 +132,7 @@ class OpenGuardrailsGuardrail(CustomGuardrail): headers["Authorization"] = f"Bearer {self.api_key}" return headers - async def _call_og_api( - self, path: str, payload: dict - ) -> Optional[dict]: + async def _call_og_api(self, path: str, payload: dict) -> Optional[dict]: """Call an OpenGuardrails gateway API endpoint.""" url = f"{self.api_base}{path}" try: @@ -415,9 +412,7 @@ class OpenGuardrailsGuardrail(CustomGuardrail): try: body_json = json.loads(body) if isinstance(body, str) else body content = ( - body_json.get("choices", [{}])[0] - .get("message", {}) - .get("content", "") + body_json.get("choices", [{}])[0].get("message", {}).get("content", "") ) return content if content else fallback except (json.JSONDecodeError, IndexError, KeyError): @@ -447,9 +442,7 @@ class OpenGuardrailsGuardrail(CustomGuardrail): pass @staticmethod - def _restore_response_content( - response, restore_mapping: Dict[str, str] - ): + def _restore_response_content(response, restore_mapping: Dict[str, str]): """Replace placeholders in response content using restore_mapping.""" try: content = response.choices[0].message.content From afc61ba7ad9193120584b5ebc83ba4a352dde19d Mon Sep 17 00:00:00 2001 From: thomaslwang <24355116@qq.com> Date: Sat, 21 Mar 2026 04:17:36 -0700 Subject: [PATCH 4/5] refactor: extract _handle_input_action to fix PLR0915 too-many-statements Co-Authored-By: Claude Opus 4.6 (1M context) --- .../openguardrails/openguardrails.py | 75 ++++++++----------- 1 file changed, 33 insertions(+), 42 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py index c42d24bd73d..927001948f7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openguardrails/openguardrails.py @@ -230,6 +230,12 @@ class OpenGuardrailsGuardrail(CustomGuardrail): # Store input messages for output detection context metadata[_META_INPUT_MESSAGES] = messages + return self._handle_input_action(action, result, detection, data, metadata) + + def _handle_input_action( + self, action: str, result: dict, detection: dict, data: dict, metadata: dict + ) -> dict: + """Process the action returned by OpenGuardrails input detection.""" if action == "block": reason = self._extract_response_content( result.get("block_response", {}), @@ -242,7 +248,6 @@ class OpenGuardrailsGuardrail(CustomGuardrail): ) elif action == "replace": - # Return replacement content as a 200 response (not an error) content = self._extract_response_content( result.get("replace_response", {}), fallback="Content filtered by OpenGuardrails.", @@ -255,62 +260,48 @@ class OpenGuardrailsGuardrail(CustomGuardrail): ) elif action == "anonymize": - anonymized = result.get("anonymized_messages") - if anonymized: - data["messages"] = anonymized - - # Store restore mapping for post-call restoration - restore_mapping = result.get("restore_mapping") - if restore_mapping: - metadata[_META_RESTORE_MAPPING] = restore_mapping - - session_id = result.get("session_id") - if session_id: - metadata[_META_SESSION_ID] = session_id - + self._apply_anonymization(result, data, metadata) return data - elif action in ("switch_private_model", "proxy_response"): - if action == "proxy_response": - # OG already called the private model - return its response - proxy_resp = result.get("proxy_response", {}) - content = self._extract_response_content( - proxy_resp, - fallback="Response from private model.", - ) - raise ModifyResponseException( - message=content, - model=data.get("model", "unknown"), - request_data=data, - guardrail_name=GUARDRAIL_NAME, - ) + elif action == "proxy_response": + proxy_resp = result.get("proxy_response", {}) + content = self._extract_response_content( + proxy_resp, + fallback="Response from private model.", + ) + raise ModifyResponseException( + message=content, + model=data.get("model", "unknown"), + request_data=data, + guardrail_name=GUARDRAIL_NAME, + ) - # switch_private_model: route to private model via LiteLLM + elif action == "switch_private_model": verbose_proxy_logger.info( "OpenGuardrails: switching to private model '%s'", self.private_model_name, ) data["model"] = self.private_model_name - if self.skip_output_for_private_model: metadata[_META_SKIP_OUTPUT] = True - - # If OG also anonymized the messages, apply that - anonymized = result.get("anonymized_messages") - if anonymized: - data["messages"] = anonymized - restore_mapping = result.get("restore_mapping") - if restore_mapping: - metadata[_META_RESTORE_MAPPING] = restore_mapping - session_id = result.get("session_id") - if session_id: - metadata[_META_SESSION_ID] = session_id - + self._apply_anonymization(result, data, metadata) return data # action == "pass" return data + def _apply_anonymization(self, result: dict, data: dict, metadata: dict): + """Apply anonymized messages and store restore mapping in metadata.""" + anonymized = result.get("anonymized_messages") + if anonymized: + data["messages"] = anonymized + restore_mapping = result.get("restore_mapping") + if restore_mapping: + metadata[_META_RESTORE_MAPPING] = restore_mapping + session_id = result.get("session_id") + if session_id: + metadata[_META_SESSION_ID] = session_id + # ------------------------------------------------------------------ # # Post-call hook: output detection + restoration # # ------------------------------------------------------------------ # From 34d0e1e5b358231956cda2df238933748262e1b7 Mon Sep 17 00:00:00 2001 From: thomaslwang <24355116@qq.com> Date: Sat, 21 Mar 2026 06:00:18 -0700 Subject: [PATCH 5/5] ci: retrigger checks