mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix: address greptile comments
This commit is contained in:
parent
d4172f76d4
commit
31fb3af67a
2 changed files with 63 additions and 4 deletions
|
|
@ -47,12 +47,10 @@ 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).
|
||||
Applied whenever \\n or \\r appear, including in mixed content with real newlines.
|
||||
"""
|
||||
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
|
||||
|
|
@ -353,7 +351,8 @@ class BlockCodeExecutionGuardrail(CustomGuardrail):
|
|||
delta_content += content
|
||||
accumulated += delta_content
|
||||
# Check after every chunk so we block before yielding the chunk that completes a blocked block
|
||||
blocks = self._find_blocks(accumulated)
|
||||
normalized = _normalize_escaped_newlines(accumulated)
|
||||
blocks = self._find_blocks(normalized)
|
||||
for _start, _end, _tag, _body, confidence, action_taken in blocks:
|
||||
if (
|
||||
action_taken == "block"
|
||||
|
|
|
|||
|
|
@ -367,3 +367,63 @@ print(factorial(5)) # Output: 120
|
|||
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
|
||||
mixed = "line1\n```py\\nprint(1)\\n```"
|
||||
normalized = _normalize_escaped_newlines(mixed)
|
||||
assert "```py\n" in normalized
|
||||
assert "print(1)\n" in normalized
|
||||
guardrail = BlockCodeExecutionGuardrail(
|
||||
guardrail_name="test",
|
||||
blocked_languages=["python", "py"],
|
||||
confidence_threshold=0.5,
|
||||
)
|
||||
blocks = guardrail._find_blocks(normalized)
|
||||
assert len(blocks) == 1
|
||||
assert blocks[0][2] == "py"
|
||||
assert blocks[0][5] == "block"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue