scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix: greptile fixes

Browse source
This commit is contained in:
Krrish Dholakia 2026-02-25 10:58:14 -08:00
parent ed5bdabcfd
commit 709815580c
2 changed files with 30 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -13,7 +13,8 @@ from typing import (TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Literal,
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import (CustomGuardrail,
ModifyResponseException)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.base import \
GuardrailConfigModel
@ -46,9 +47,14 @@ 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.
Only applies when text has no real newline but contains literal \\n or \\r (double-encoded).
"""
if not text:
return text
if "\n" in text:
return text
if "\\n" not in text and "\\r" not in 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")
@ -72,6 +78,7 @@ def _is_blocked_language(
if block_all:
# Block all: only allow through if it's explicitly non-executable (we still block but with lower confidence)
return True
# When block_all is False, caller guarantees blocked_languages is non-empty.
if not blocked_languages:
return True
normalized_list = [_normalize_language(t) for t in blocked_languages]
@ -158,12 +165,14 @@ class BlockCodeExecutionGuardrail(CustomGuardrail):
def _find_blocks(
self, text: str
) -> List[Tuple[str, str, float, CodeBlockActionTaken]]:
) -> List[Tuple[int, int, str, str, float, CodeBlockActionTaken]]:
"""
Find all fenced code blocks in text. Returns list of
(language_tag, block_content, confidence, action_taken).
(start, end, language_tag, block_content, confidence, action_taken).
"""
results: List[Tuple[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)
@ -180,7 +189,9 @@ class BlockCodeExecutionGuardrail(CustomGuardrail):
action_taken = "block"
else:
action_taken = "log_only"
results.append((tag or "(none)", body, confidence, action_taken))
results.append(
(m.start(), m.end(), tag or "(none)", body, confidence, action_taken)
)
return results
def _scan_text(
@ -202,29 +213,14 @@ class BlockCodeExecutionGuardrail(CustomGuardrail):
should_raise = False
last_end = 0
parts: List[str] = []
for m in FENCED_BLOCK_RE.finditer(text):
tag = (m.group(1) or "").strip()
tag_in_list = not self.block_all and _normalize_language(tag) in [
_normalize_language(t) for t in (self.blocked_languages or [])
]
is_blocked = _is_blocked_language(
tag, self.blocked_languages, self.block_all
)
confidence = _confidence_for_block(tag, self.block_all, tag_in_list)
if not is_blocked:
action_taken: CodeBlockActionTaken = "allow"
elif confidence >= self.confidence_threshold:
action_taken = "block"
else:
action_taken = "log_only"
for start, end, tag, _body, confidence, action_taken in blocks:
if detections is not None:
detections.append(
cast(
CodeBlockDetection,
{
"type": "code_block",
"language": tag or "(none)",
"language": tag,
"confidence": round(confidence, 2),
"action_taken": action_taken,
},
@ -233,12 +229,12 @@ class BlockCodeExecutionGuardrail(CustomGuardrail):
if action_taken == "block" and self.action == "block":
should_raise = True
parts.append(text[last_end : m.start()])
parts.append(text[last_end:start])
if action_taken == "block":
parts.append(self.MASK_PLACEHOLDER)
else:
parts.append(text[m.start() : m.end()])
last_end = m.end()
parts.append(text[start:end])
last_end = end
parts.append(text[last_end:])
new_text = "".join(parts)
@ -358,7 +354,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail):
accumulated += delta_content
# Check after every chunk so we block before yielding the chunk that completes a blocked block
blocks = self._find_blocks(accumulated)
for _tag, _body, confidence, action_taken in blocks:
for _start, _end, _tag, _body, confidence, action_taken in blocks:
if (
action_taken == "block"
and confidence >= self.confidence_threshold

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

@ -23,7 +23,7 @@ class TestBlockCodeExecutionGuardrail:
)
blocks = guardrail._find_blocks("Here is code:\n```python\nprint(1)\n```\nDone.")
assert len(blocks) == 1
tag, _body, confidence, action_taken = blocks[0]
_start, _end, tag, _body, confidence, action_taken = blocks[0]
assert tag == "python"
assert confidence == 1.0
assert action_taken == "block"
@ -37,7 +37,7 @@ class TestBlockCodeExecutionGuardrail:
)
blocks = guardrail._find_blocks("```\nfoo\n```")
assert len(blocks) == 1
_tag, _body, confidence, action_taken = blocks[0]
_start, _end, _tag, _body, confidence, action_taken = blocks[0]
assert action_taken == "block"
assert confidence in (0.5, 1.0)
@ -50,7 +50,7 @@ class TestBlockCodeExecutionGuardrail:
)
blocks = guardrail._find_blocks("```text\nplain output\n```")
assert len(blocks) == 1
_tag, _body, confidence, action_taken = blocks[0]
_start, _end, _tag, _body, confidence, action_taken = blocks[0]
assert action_taken == "allow"
assert confidence == 0.0
@ -64,7 +64,7 @@ class TestBlockCodeExecutionGuardrail:
# Block with no tag or plaintext tag gets confidence 0.5
blocks = guardrail._find_blocks("```text\nx\n```")
assert len(blocks) == 1
_tag, _body, confidence, action_taken = blocks[0]
_start, _end, _tag, _body, confidence, action_taken = blocks[0]
assert confidence == 0.5
assert action_taken == "log_only"
@ -272,10 +272,10 @@ print(factorial(5)) # Output: 120
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"
assert blocks[0][2] == "python"
assert blocks[0][5] == "block"
assert blocks[1][2] == "python"
assert blocks[1][5] == "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."""