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 15fd6c12570..aa1aab56d7d 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 @@ -31,6 +31,8 @@ if TYPE_CHECKING: LANGUAGE_ALIASES: Dict[str, str] = { "js": "javascript", "py": "python", + "sh": "bash", + "ts": "typescript", } # Tags that indicate non-executable / plain text (lower confidence when block-all) @@ -415,9 +417,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): Find all fenced code blocks in text. Returns list of (start, end, language_tag, block_content, confidence, action_taken). """ - results: List[ - Tuple[int, int, str, str, float, CodeBlockActionTaken] - ] = [] + results: List[Tuple[int, int, str, str, float, CodeBlockActionTaken]] = [] for m in FENCED_BLOCK_RE.finditer(text): tag = (m.group(1) or "").strip() body = m.group(2) @@ -494,7 +494,9 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): "type": "code_block", "language": tag, "confidence": round(confidence, 2), - "action_taken": "block" if effective_block else action_taken, + "action_taken": ( + "block" if effective_block else action_taken + ), }, ) ) @@ -609,32 +611,3 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): event_type=event_type, tracing_detail=GuardrailTracingDetail(**tracing_kw), # type: ignore[typeddict-item] ) - - async def async_post_call_streaming_iterator_hook( - self, - user_api_key_dict: Any, - response: Any, - request_data: dict, - ) -> AsyncGenerator[ModelResponseStream, None]: - """Accumulate streamed content and block as soon as a complete fenced code block is detected (before yielding that chunk).""" - accumulated = "" - async for item in response: - if isinstance(item, ModelResponseStream) and item.choices: - delta_content = "" - for choice in item.choices: - if hasattr(choice, "delta") and choice.delta: - content = getattr(choice.delta, "content", None) - if content and isinstance(content, str): - delta_content += content - accumulated += delta_content - # Check after every chunk so we block before yielding the chunk that completes a blocked block - normalized = _normalize_escaped_newlines(accumulated) - blocks = self._find_blocks(normalized) - for _start, _end, _tag, _body, confidence, action_taken in blocks: - 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/proxy/guardrails/guardrail_hooks/block_code_execution.py b/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py index 56cdacc99e5..68a2e4a7de6 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py @@ -8,13 +8,13 @@ from .base import GuardrailConfigModel CodeBlockActionTaken = Literal["block", "allow", "log_only"] -# Supported language tags for the blocked_languages multiselect dropdown +# Supported language tags for the blocked_languages multiselect dropdown. +# Only canonical names are listed; LANGUAGE_ALIASES in the guardrail normalizes +# aliases (e.g. js→javascript, sh→bash) when matching. BLOCKED_LANGUAGES_OPTIONS = [ "python", "javascript", - "js", "bash", - "sh", "ruby", "go", "java", 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 6a4a7040884..e849c30ba7a 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 @@ -330,92 +330,6 @@ print(factorial(5)) # Output: 120 exc_info.value ).lower() - @pytest.mark.asyncio - async def test_streaming_hook_blocks_before_yielding_chunk_that_completes_block( - self, - ): - """Streaming hook runs block check after every chunk and raises before yielding the chunk that completes a blocked fenced block.""" - from litellm.types.utils import (Delta, ModelResponseStream, - StreamingChoices) - - guardrail = BlockCodeExecutionGuardrail( - guardrail_name="test", - blocked_languages=["python"], - action="block", - confidence_threshold=0.5, - ) - - # Chunks that form "```python\nprint(1)\n```" when concatenated - async def mock_stream(): - yield ModelResponseStream( - choices=[StreamingChoices(delta=Delta(content="```python\n"))], - ) - yield ModelResponseStream( - choices=[StreamingChoices(delta=Delta(content="print(1)\n"))], - ) - yield ModelResponseStream( - choices=[StreamingChoices(delta=Delta(content="```"))], - ) - - request_data = {"model": "gpt-4", "metadata": {}} - yielded_chunks = [] - - with pytest.raises(HTTPException) as exc_info: - async for chunk in guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=None, - response=mock_stream(), - request_data=request_data, - ): - yielded_chunks.append(chunk) - - assert exc_info.value.status_code == 400 - assert "code block" in (exc_info.value.detail or {}).get("error", "") - # The chunk that completes the block (third chunk) must not have been yielded - assert len(yielded_chunks) == 2 - - @pytest.mark.asyncio - async def test_streaming_hook_blocks_when_accumulated_has_literal_backslash_n( - self, - ): - """Streaming hook normalizes escaped newlines before detection; blocks when literal \\n forms a complete code block.""" - from litellm.types.utils import (Delta, ModelResponseStream, - StreamingChoices) - - guardrail = BlockCodeExecutionGuardrail( - guardrail_name="test", - blocked_languages=["python"], - action="block", - confidence_threshold=0.5, - ) - - # Chunks that when concatenated form "```python\\nprint(1)\\n```" (literal backslash-n) - async def mock_stream(): - yield ModelResponseStream( - choices=[StreamingChoices(delta=Delta(content="```python\\n"))], - ) - yield ModelResponseStream( - choices=[StreamingChoices(delta=Delta(content="print(1)\\n"))], - ) - yield ModelResponseStream( - choices=[StreamingChoices(delta=Delta(content="```"))], - ) - - request_data = {"model": "gpt-4", "metadata": {}} - yielded_chunks = [] - - with pytest.raises(HTTPException) as exc_info: - async for chunk in guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=None, - response=mock_stream(), - request_data=request_data, - ): - yielded_chunks.append(chunk) - - assert exc_info.value.status_code == 400 - assert "code block" in (exc_info.value.detail or {}).get("error", "") - # Block is detected after normalization; chunk that completes the block not yielded - assert len(yielded_chunks) == 2 - def test_normalize_escaped_newlines_mixed_content_detects_block(self): """Mixed content (real newlines and literal \\n) is normalized so code block is detected.""" # Text with real newline then a fence using literal \n after language tag