diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py index 7fa8f4a5a42..c3691d80db5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py @@ -1,6 +1,6 @@ """Block Code Execution guardrail: blocks or masks fenced code blocks by language.""" -from typing import TYPE_CHECKING, Literal, cast +from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union, cast from litellm.types.guardrails import (GuardrailEventHooks, SupportedGuardrailIntegrations) @@ -17,6 +17,22 @@ DEFAULT_EVENT_HOOKS = [ ] +def _get_param( + litellm_params: "LitellmParams", + guardrail: "Guardrail", + key: str, + default: Any = None, +) -> Any: + """Get a param from litellm_params, with fallback to raw guardrail litellm_params (for extra fields not on LitellmParams).""" + value = getattr(litellm_params, key, default) + if value is not None: + return value + raw = guardrail.get("litellm_params") + if isinstance(raw, dict) and key in raw: + return raw[key] + return default + + def initialize_guardrail( litellm_params: "LitellmParams", guardrail: "Guardrail", @@ -30,15 +46,25 @@ def initialize_guardrail( "Block Code Execution guardrail requires a guardrail_name" ) - blocked_languages = getattr(litellm_params, "blocked_languages", None) + blocked_languages: Optional[List[str]] = cast( + Optional[List[str]], + _get_param(litellm_params, guardrail, "blocked_languages"), + ) action = cast( - Literal["block", "mask"], getattr(litellm_params, "action", "block") + Literal["block", "mask"], + _get_param(litellm_params, guardrail, "action", "block"), ) confidence_threshold = float( - getattr(litellm_params, "confidence_threshold", 0.7) + cast( + Union[int, float, str], + _get_param(litellm_params, guardrail, "confidence_threshold", 0.5), + ) + ) + mode = _get_param(litellm_params, guardrail, "mode") + event_hook = cast( + Optional[Union[Literal["pre_call", "post_call", "during_call"], List[str]]], + mode if mode is not None else DEFAULT_EVENT_HOOKS, ) - mode = getattr(litellm_params, "mode", None) - event_hook = mode if mode is not None else DEFAULT_EVENT_HOOKS instance = BlockCodeExecutionGuardrail( guardrail_name=guardrail_name, @@ -46,7 +72,7 @@ def initialize_guardrail( action=action, confidence_threshold=confidence_threshold, event_hook=event_hook, - default_on=getattr(litellm_params, "default_on", False), + default_on=bool(_get_param(litellm_params, guardrail, "default_on", False)), ) litellm.logging_callback_manager.add_litellm_callback(instance) return instance diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py index 33edf5735f9..3af8a435a39 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py @@ -61,6 +61,20 @@ NON_EXECUTABLE_TAGS: frozenset = frozenset( FENCED_BLOCK_RE = re.compile(r"```(\w*)\n(.*?)```", re.DOTALL) +def _normalize_escaped_newlines(text: str) -> str: + """ + Replace literal escaped newlines (backslash + n or backslash + r) with real newlines. + API/JSON payloads sometimes deliver newlines as the two-character sequence \\n. + """ + if not text: + return text + # Order matters: replace \r\n first so we don't produce extra \n from \r then \n + text = text.replace("\\r\\n", "\n") + text = text.replace("\\n", "\n") + text = text.replace("\\r", "\n") + return text + + def _normalize_language(tag: str) -> str: """Normalize language tag (lowercase, resolve aliases).""" tag = (tag or "").strip().lower() @@ -115,7 +129,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): guardrail_name: Optional[str] = None, blocked_languages: Optional[List[str]] = None, action: Literal["block", "mask"] = "block", - confidence_threshold: float = 0.7, + confidence_threshold: float = 0.5, event_hook: Optional[ Union[Literal["pre_call", "post_call", "during_call"], List[str]] ] = None, @@ -123,9 +137,9 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): **kwargs: Any, ) -> None: # Normalize to type expected by CustomGuardrail - _event_hook: Optional[ - Union[GuardrailEventHooks, List[GuardrailEventHooks]] - ] = None + _event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]] = ( + None + ) if event_hook is not None: if isinstance(event_hook, list): _event_hook = [ @@ -158,9 +172,12 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): def get_config_model() -> Optional[type[GuardrailConfigModel]]: from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import \ BlockCodeExecutionGuardrailConfigModel + return BlockCodeExecutionGuardrailConfigModel - def _find_blocks(self, text: str) -> List[Tuple[str, str, float, CodeBlockActionTaken]]: + def _find_blocks( + self, text: str + ) -> List[Tuple[str, str, float, CodeBlockActionTaken]]: """ Find all fenced code blocks in text. Returns list of (language_tag, block_content, confidence, action_taken). @@ -196,6 +213,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): """ if not text: return text, False + text = _normalize_escaped_newlines(text) blocks = self._find_blocks(text) if not blocks: return text, False @@ -245,10 +263,10 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): new_text = "".join(parts) return new_text, should_raise - def _raise_block_error(self, language: str, is_output: bool, request_data: dict) -> None: - msg = ( - f"Content blocked: executable code block detected (language: {language})" - ) + def _raise_block_error( + self, language: str, is_output: bool, request_data: dict + ) -> None: + msg = f"Content blocked: executable code block detected (language: {language})" if is_output: raise HTTPException( status_code=400, @@ -356,7 +374,10 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): # Run detection on full accumulated text (streaming: block only, no mask) blocks = self._find_blocks(accumulated) for _tag, _body, confidence, action_taken in blocks: - if action_taken == "block" and confidence >= self.confidence_threshold: + if ( + action_taken == "block" + and confidence >= self.confidence_threshold + ): lang = _tag or "unknown" self._raise_block_error(lang, True, request_data) yield item diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 47b54a8a0fe..29ed07520d4 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -5,24 +5,20 @@ from typing import Any, Dict, List, Literal, Optional, Union from pydantic import BaseModel, ConfigDict, Field, field_validator 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, -) -from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( - IBMGuardrailsBaseConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( - ContentFilterCategoryConfig, -) -from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( - QualifireGuardrailConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( - ToolPermissionGuardrailConfigModel, -) +from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import \ + BlockCodeExecutionGuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import \ + EnkryptAIGuardrailConfigs +from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import \ + GraySwanGuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.ibm import \ + IBMGuardrailsBaseConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import \ + ContentFilterCategoryConfig +from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import \ + QualifireGuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import \ + ToolPermissionGuardrailConfigModel """ Pydantic object defining how to set guardrails on litellm proxy @@ -700,6 +696,7 @@ class LitellmParams( EnkryptAIGuardrailConfigs, IBMGuardrailsBaseConfigModel, QualifireGuardrailConfigModel, + BlockCodeExecutionGuardrailConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") mode: Union[str, List[str], Mode] = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py b/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py index 641892cb674..89a796db921 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py @@ -54,7 +54,7 @@ class BlockCodeExecutionGuardrailConfigModel(GuardrailConfigModel): description="'block' raises an error; 'mask' replaces the code block with a placeholder.", ) confidence_threshold: float = Field( - default=0.7, + default=0.5, ge=0.0, le=1.0, description="Only block or mask when detection confidence >= this value; below threshold, allow or log_only.", diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py index f17a8a224c2..0a54f12b535 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py @@ -3,8 +3,15 @@ import pytest from fastapi import HTTPException +from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.proxy.guardrails.guardrail_hooks.block_code_execution import ( - DEFAULT_EVENT_HOOKS, BlockCodeExecutionGuardrail, initialize_guardrail) + DEFAULT_EVENT_HOOKS, + BlockCodeExecutionGuardrail, + initialize_guardrail, +) +from litellm.proxy.guardrails.guardrail_hooks.block_code_execution.block_code_execution import ( + _normalize_escaped_newlines, +) from litellm.types.guardrails import GuardrailEventHooks @@ -112,6 +119,42 @@ class TestBlockCodeExecutionGuardrail: assert "[CODE_BLOCK_REDACTED]" in result["texts"][0] assert "x=1" not in result["texts"][0] + @pytest.mark.asyncio + async def test_execute_python_factorial_string_blocked(self): + """Guardrail blocks the exact 'execute \"```python...' string with two python blocks (real newlines).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.5, + ) + # Exact user payload; newlines are real so regex ```(\w*)\n(.*?)``` matches + text = ( + 'execute "```python\n' + "def factorial(n: int) -> int:\n" + ' """Return the factorial of n."""\n' + ' if n < 0:\n' + ' raise ValueError("n must be non-negative")\n' + " if n in (0, 1):\n" + " return 1\n" + " return n * factorial(n - 1)\n" + '```\n\n' + "Example usage:\n" + "```python\n" + "print(factorial(5)) # Output: 120\n" + '```"' + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = {"texts": [text]} + # pre_call (request) raises ModifyResponseException; post_call (response) raises HTTPException + with pytest.raises((HTTPException, ModifyResponseException)) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + assert "python" in str(exc_info.value).lower() + @pytest.mark.asyncio async def test_factorial_scenario_blocked(self): """Exact user scenario: Python factorial snippet in markdown is blocked when python in list.""" @@ -198,3 +241,90 @@ print(factorial(5)) # Output: 120 assert instance.event_hook == DEFAULT_EVENT_HOOKS assert GuardrailEventHooks.pre_call.value in instance.event_hook assert GuardrailEventHooks.post_call.value in instance.event_hook + + def test_normalize_escaped_newlines_converts_backslash_n_to_newline(self): + """Literal \\n in text is converted to real newline so regex can match code blocks.""" + raw = 'execute this "```python\\ndef factorial(n):\\n return 1\\n```"' + normalized = _normalize_escaped_newlines(raw) + assert "\\n" not in normalized + assert "\n" in normalized + assert "```python\n" in normalized + + def test_find_blocks_detects_python_block_with_escaped_newlines(self): + """_find_blocks finds a block when text uses literal \\n instead of real newlines.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + confidence_threshold=0.7, + ) + # Text as received from API with escaped newlines (e.g. JSON-decoded string) + text_with_escaped = ( + 'execute this "```python\\n' + 'def factorial(n: int) -> int:\\n' + ' """Return the factorial of n."""\\n' + ' if n < 0:\\n' + ' raise ValueError("n must be non-negative")\\n' + " if n in (0, 1):\\n" + " return 1\\n" + " return n * factorial(n - 1)\\n" + '```\\n\\n' + 'Example usage:\\n' + '```python\\n' + 'print(factorial(5)) # Output: 120\\n' + '```"' + ) + normalized = _normalize_escaped_newlines(text_with_escaped) + blocks = guardrail._find_blocks(normalized) + assert len(blocks) == 2 + assert blocks[0][0] == "python" + assert blocks[0][3] == "block" + assert blocks[1][0] == "python" + assert blocks[1][3] == "block" + + def test_scan_text_blocks_and_masks_when_text_has_escaped_newlines(self): + """_scan_text detects blocks and applies block/mask when newlines are literal \\n.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="mask", + confidence_threshold=0.5, + ) + text_with_escaped = 'execute "```python\\nprint(1)\\n```"' + new_text, should_raise = guardrail._scan_text(text_with_escaped) + assert "[CODE_BLOCK_REDACTED]" in new_text + assert "print(1)" not in new_text + assert should_raise is False # action is mask + + @pytest.mark.asyncio + async def test_apply_guardrail_blocks_when_text_has_escaped_newlines(self): + """apply_guardrail blocks request/response when code block uses literal \\n (e.g. from API).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.5, + ) + text_with_escaped = ( + 'execute this "```python\\n' + 'def factorial(n: int) -> int:\\n' + ' """Return the factorial of n."""\\n' + " if n in (0, 1):\\n" + " return 1\\n" + " return n * factorial(n - 1)\\n" + '```\\n\\n' + 'Example usage:\\n' + '```python\\n' + 'print(factorial(5)) # Output: 120\\n' + '```"' + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = {"texts": [text_with_escaped]} + with pytest.raises((HTTPException, ModifyResponseException)) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + assert "python" in str(exc_info.value).lower() or "code" in str( + exc_info.value + ).lower()