scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

feat: working block code execution guardrail

Browse source
This commit is contained in:
Krrish Dholakia 2026-02-24 19:16:39 -08:00
parent 20fdedf44f
commit adda8874e3
5 changed files with 211 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py
View file

@ -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

41
litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py
View file

@ -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

33
litellm/types/guardrails.py
View file

@ -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(

2
litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py
View file

@ -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.",

132
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py
View file

@ -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()