tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-12 13:05:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(guardrails): cover multi-choice output variants

Browse Source
This commit is contained in:
user
2026-05-01 19:38:04 -07:00
parent f909fa79b5
commit f4e7dde2d8
8 changed files with 296 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py
+23 -9
View File
@@ -254,15 +254,16 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
) -> Any:
from litellm.types.utils import Choices, ModelResponse
if (
isinstance(response, ModelResponse)
and response.choices
and isinstance(response.choices[0], Choices)
):
content = response.choices[0].message.content or ""
await self.async_make_request(
text=content,
)
if isinstance(response, ModelResponse) and response.choices:
for choice in response.choices:
if not isinstance(choice, Choices):
continue
content = _message_content_to_text(choice.message.content)
if not content:
continue
await self.async_make_request(
text=content,
)
return response
async def async_post_call_streaming_hook(
@@ -279,3 +280,16 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
error_returned = json.dumps({"error": e.detail})
return f"data: {error_returned}\n\n"
def _message_content_to_text(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
text_parts = [
item.get("text")
for item in content
if isinstance(item, dict) and isinstance(item.get("text"), str)
]
return "\n".join(part for part in text_parts if part)
return ""
litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
+81 -64
View File
@@ -1873,8 +1873,9 @@ class ContentFilterGuardrail(CustomGuardrail):
and the UI Request Lifecycle panel. Mirrors apply_guardrail's finally-block
contract.
"""
accumulated_full_text = ""
yielded_masked_text_len = 0
accumulated_text_by_choice: Dict[int, str] = {}
yielded_masked_text_len_by_choice: Dict[int, int] = {}
latest_detections_by_choice: Dict[int, List[ContentFilterDetection]] = {}
buffer_size = 50 # Increased buffer to catch patterns split across many chunks
start_time = datetime.now()
@@ -1890,79 +1891,90 @@ class ContentFilterGuardrail(CustomGuardrail):
try:
async for item in response:
if isinstance(item, ModelResponseStream) and item.choices:
delta_content = ""
is_final = False
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
if getattr(choice, "finish_reason", None):
is_final = True
if not (hasattr(choice, "delta") and choice.delta):
continue
accumulated_full_text += delta_content
choice_index = getattr(choice, "index", 0)
if not isinstance(choice_index, int):
choice_index = 0
# Check for blocking or apply masking
# Add a space at the end if it's the final chunk to trigger word boundaries (\b)
text_to_check = accumulated_full_text
if is_final:
text_to_check += " "
content = getattr(choice.delta, "content", None)
is_final = bool(getattr(choice, "finish_reason", None))
if isinstance(content, str) and content:
accumulated_text_by_choice[choice_index] = (
accumulated_text_by_choice.get(choice_index, "")
+ content
)
elif not is_final:
continue
try:
# Reset before each scan: _filter_single_text scans the
# whole accumulated buffer every chunk, so previous-chunk
# matches are guaranteed to be re-found. Keeping only the
# latest scan's detections avoids N× duplication in the
# final log row. BLOCK still records correctly because
# handlers append to detections before raising.
detections.clear()
masked_text = self._filter_single_text(
text_to_check, detections=detections
text_to_check = accumulated_text_by_choice.get(choice_index, "")
if not text_to_check:
continue
# Add a space at the end if it's the final chunk to trigger word boundaries (\b)
text_to_scan = text_to_check + (" " if is_final else "")
choice_detections: List[ContentFilterDetection] = []
try:
# _filter_single_text scans the whole accumulated
# choice buffer every chunk, so previous-chunk
# matches are guaranteed to be re-found. Keeping
# only each choice's latest scan avoids duplicate
# detections in the final log row.
masked_text = self._filter_single_text(
text_to_scan, detections=choice_detections
)
if is_final and masked_text.endswith(" "):
masked_text = masked_text[:-1]
latest_detections_by_choice[choice_index] = (
choice_detections
)
except HTTPException:
latest_detections_by_choice[choice_index] = (
choice_detections
)
raise
except Exception as e:
verbose_proxy_logger.error(
f"ContentFilterGuardrail: Error in masking: {e}"
)
masked_text = text_to_scan # Fallback to current text
# Determine how much can be safely yielded
if is_final:
safe_to_yield_len = len(masked_text)
else:
safe_to_yield_len = max(0, len(masked_text) - buffer_size)
yielded_masked_text_len = yielded_masked_text_len_by_choice.get(
choice_index, 0
)
if is_final and masked_text.endswith(" "):
masked_text = masked_text[:-1]
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.error(
f"ContentFilterGuardrail: Error in masking: {e}"
)
masked_text = text_to_check # Fallback to current text
if safe_to_yield_len > yielded_masked_text_len:
new_masked_content = masked_text[
yielded_masked_text_len:safe_to_yield_len
]
choice.delta.content = new_masked_content
yielded_masked_text_len_by_choice[choice_index] = (
safe_to_yield_len
)
else:
# Hold content by yielding empty content on this choice
# while preserving chunk metadata and other choices.
choice.delta.content = ""
# Determine how much can be safely yielded
if is_final:
safe_to_yield_len = len(masked_text)
else:
safe_to_yield_len = max(0, len(masked_text) - buffer_size)
if safe_to_yield_len > yielded_masked_text_len:
new_masked_content = masked_text[
yielded_masked_text_len:safe_to_yield_len
]
# Modify the chunk to contain only the new masked content
if (
item.choices
and hasattr(item.choices[0], "delta")
and item.choices[0].delta
):
item.choices[0].delta.content = new_masked_content
yielded_masked_text_len = safe_to_yield_len
yield item
else:
# Hold content by yielding empty content chunk (keeps metadata/structure)
if (
item.choices
and hasattr(item.choices[0], "delta")
and item.choices[0].delta
):
item.choices[0].delta.content = ""
yield item
yield item
else:
# Not a ModelResponseStream or no choices - yield as is
yield item
# Any remaining content (should have been handled by is_final, but just in case)
if yielded_masked_text_len < len(accumulated_full_text):
if any(
yielded_masked_text_len_by_choice.get(choice_index, 0)
< len(accumulated_text)
for choice_index, accumulated_text in accumulated_text_by_choice.items()
):
# We already reached the end of the generator
pass
except HTTPException:
@@ -1973,6 +1985,11 @@ class ContentFilterGuardrail(CustomGuardrail):
exception_str = str(e)
raise e
finally:
detections = [
detection
for choice_detections in latest_detections_by_choice.values()
for detection in choice_detections
]
self._count_masked_entities(detections, masked_entity_count)
self._log_guardrail_information(
request_data=request_data,
litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py
+21 -4
View File
@@ -187,11 +187,28 @@ def _extract_user_text(messages: List) -> str:
def _extract_response_text(response: Any) -> str:
"""Extract text from LLM response object."""
"""Extract text from every LLM response choice."""
if hasattr(response, "choices") and response.choices:
choice = response.choices[0]
if hasattr(choice, "message") and choice.message:
return choice.message.content or ""
text_parts: List[str] = []
for choice in response.choices:
if hasattr(choice, "message") and choice.message:
text = _content_to_text(choice.message.content)
if text:
text_parts.append(text)
return "\n".join(text_parts)
return ""
def _content_to_text(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
text_parts = [
block.get("text")
for block in content
if isinstance(block, dict) and isinstance(block.get("text"), str)
]
return " ".join(part for part in text_parts if part)
return ""
litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py
+21 -10
View File
@@ -480,21 +480,32 @@ class XecGuardGuardrail(CustomGuardrail):
choices = response.get("choices")
if not choices:
return None
first = choices[0]
if hasattr(first, "message"):
message = first.message
elif isinstance(first, dict):
message = first.get("message")
text_parts: List[str] = []
for choice in choices:
content = XecGuardGuardrail._extract_choice_content(choice)
text = XecGuardGuardrail._content_to_text(content)
if text:
text_parts.append(text)
return "\n".join(text_parts) or None
@staticmethod
def _extract_choice_content(choice: Any) -> Any:
if hasattr(choice, "message"):
message = choice.message
elif isinstance(choice, dict):
message = choice.get("message")
else:
return None
if message is None:
return None
if hasattr(message, "content"):
content = message.content
elif isinstance(message, dict):
content = message.get("content")
else:
return None
return message.content
if isinstance(message, dict):
return message.get("content")
return None
@staticmethod
def _content_to_text(content: Any) -> Optional[str]:
if isinstance(content, str) and content:
return content
if isinstance(content, list):
tests/guardrails_tests/test_semantic_guard.py
+19
View File
@@ -171,6 +171,25 @@ class TestHelperFunctions:
mock_response.choices[0].message.content = "Hello from LLM"
assert _extract_response_text(mock_response) == "Hello from LLM"
def test_extract_response_text_combines_all_choices(self):
from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.semantic_guard import (
_extract_response_text,
)
first_choice = MagicMock()
first_choice.message.content = "first response"
second_choice = MagicMock()
second_choice.message.content = [
{"type": "text", "text": "second"},
{"type": "text", "text": "response"},
]
mock_response = MagicMock()
mock_response.choices = [first_choice, second_choice]
assert (
_extract_response_text(mock_response) == "first response\nsecond response"
)
def test_extract_response_text_empty(self):
from litellm.proxy.guardrails.guardrail_hooks.semantic_guard.semantic_guard import (
_extract_response_text,
tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py
+46
View File
@@ -232,6 +232,52 @@ async def test_azure_text_moderation_guardrail_post_call_success_hook():
assert mock_async_make_request.call_args.kwargs["text"] == "Hello world"
@pytest.mark.asyncio
async def test_azure_text_moderation_guardrail_post_call_checks_all_choices():
azure_text_moderation_guardrail = AzureContentSafetyTextModerationGuardrail(
guardrail_name="azure_text_moderation",
api_key="azure_text_moderation_api_key",
api_base="azure_text_moderation_api_base",
)
with patch.object(
azure_text_moderation_guardrail, "async_make_request"
) as mock_async_make_request:
mock_async_make_request.side_effect = [
{
"blocklistsMatch": [],
"categoriesAnalysis": [{"category": "Hate", "severity": 0}],
},
HTTPException(
status_code=400,
detail={"error": "blocked second choice"},
),
]
with pytest.raises(HTTPException):
await azure_text_moderation_guardrail.async_post_call_success_hook(
data={},
user_api_key_dict=UserAPIKeyAuth(
api_key="azure_text_moderation_api_key"
),
response=ModelResponse(
choices=[
Choices(
index=0,
message=Message(content="safe response"),
),
Choices(
index=1,
message=Message(content="unsafe response"),
),
]
),
)
assert [
call.kwargs["text"] for call in mock_async_make_request.call_args_list
] == ["safe response", "unsafe response"]
@pytest.mark.asyncio
async def test_azure_text_moderation_guardrail_post_call_streaming_hook():
tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py
+65
View File
@@ -453,6 +453,71 @@ class TestContentFilterGuardrail:
assert "[EMAIL_REDACTED]" in full_content
assert "Contact me at [EMAIL_REDACTED] for info" in full_content
@pytest.mark.asyncio
async def test_streaming_hook_mask_checks_all_choices(self):
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
patterns = [
ContentFilterPattern(
pattern_type="prebuilt",
pattern_name="email",
action=ContentFilterAction.MASK,
),
]
guardrail = ContentFilterGuardrail(
guardrail_name="test-streaming-mask-all-choices",
patterns=patterns,
event_hook=GuardrailEventHooks.during_call,
)
async def mock_stream():
yield ModelResponseStream(
id="chunk1",
choices=[
StreamingChoices(
delta=Delta(content="Contact first@ex"),
index=0,
),
StreamingChoices(
delta=Delta(content="Email second@ex"),
index=1,
),
],
model="gpt-4",
)
yield ModelResponseStream(
id="chunk2",
choices=[
StreamingChoices(
delta=Delta(content="ample.com for help"),
index=0,
finish_reason="stop",
),
StreamingChoices(
delta=Delta(content="ample.com for support"),
index=1,
finish_reason="stop",
),
],
model="gpt-4",
)
content_by_choice = {0: "", 1: ""}
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=MagicMock(),
response=mock_stream(),
request_data={},
):
for choice in chunk.choices:
if choice.delta.content:
content_by_choice[choice.index] += choice.delta.content
assert "first@example.com" not in content_by_choice[0]
assert "second@example.com" not in content_by_choice[1]
assert content_by_choice[0] == "Contact [EMAIL_REDACTED] for help"
assert content_by_choice[1] == "Email [EMAIL_REDACTED] for support"
@pytest.mark.asyncio
async def test_streaming_hook_block(self):
"""
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py
+20 -1
View File
@@ -6,7 +6,6 @@ branch coverage. Network calls are always mocked; the companion live
suite lives in ``test_xecguard_live.py``.
"""
import asyncio
import os
from unittest.mock import MagicMock, patch
@@ -1196,6 +1195,26 @@ class TestXecGuardMessageAssembly:
is None
)
def test_extract_assistant_text_combines_all_choices(self, xecguard_guardrail):
assert (
xecguard_guardrail._extract_assistant_text_from_response(
{
"choices": [
{"message": {"content": "first response"}},
{
"message": {
"content": [
{"type": "text", "text": "second"},
{"type": "text", "text": "response"},
]
}
},
]
}
)
== "first response\nsecond\nresponse"
)
def test_synthesize_user_inputs_not_dict(self, xecguard_guardrail):
assert xecguard_guardrail._synthesize_user_from_inputs("not-dict") is None