From f4e7dde2d8ff150c969fea80c7fd663482458a89 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 19:38:04 -0700 Subject: [PATCH] fix(guardrails): cover multi-choice output variants --- .../guardrail_hooks/azure/text_moderation.py | 32 ++-- .../litellm_content_filter/content_filter.py | 145 ++++++++++-------- .../semantic_guard/semantic_guard.py | 25 ++- .../guardrail_hooks/xecguard/xecguard.py | 31 ++-- tests/guardrails_tests/test_semantic_guard.py | 19 +++ .../azure/test_azure_text_moderation.py | 46 ++++++ .../content_filter/test_content_filter.py | 65 ++++++++ .../guardrail_hooks/test_xecguard.py | 21 ++- 8 files changed, 296 insertions(+), 88 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index a465c5428d..55b446eb52 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -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 "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 3ef99e4ad6..d6d2e01494 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py index be48991500..21c01da029 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py @@ -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 "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py index 5c374540e2..5acfc5403c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/xecguard/xecguard.py @@ -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): diff --git a/tests/guardrails_tests/test_semantic_guard.py b/tests/guardrails_tests/test_semantic_guard.py index 44d8e417f5..a7e6230d02 100644 --- a/tests/guardrails_tests/test_semantic_guard.py +++ b/tests/guardrails_tests/test_semantic_guard.py @@ -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, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py index 5e3cb76f44..94c7cefaeb 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py @@ -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(): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index f7377926ef..fb952d4b18 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -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): """ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py index b663544238..60c595c64e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_xecguard.py @@ -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