From ca79ebe7d2eeb9ff5662b3cd01d7778e03c31db0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 1 Nov 2025 18:00:54 -0700 Subject: [PATCH] UI - Fix regression where Guardrail Entity Could not be selected and entity was not displayed (#16165) * fix PiiEntityCategoryMap * fix OpenAIChatCompletionsHandler * fix lint --- litellm/integrations/custom_guardrail.py | 4 +- .../chat/guardrail_translation/handler.py | 11 ++- .../proxy/guardrails/guardrail_endpoints.py | 11 ++- .../guardrail_hooks/bedrock_guardrails.py | 14 ++- .../guardrail_hooks/enkryptai/enkryptai.py | 3 +- .../guardrails/guardrail_hooks/noma/noma.py | 8 +- .../guardrails/guardrail_hooks/presidio.py | 3 +- litellm/types/guardrails.py | 10 +- .../guardrail_hooks/test_presidio.py | 92 +++++++++++++++++++ 9 files changed, 138 insertions(+), 18 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 86eed1747be..c3e1a31c3ef 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -285,7 +285,7 @@ class CustomGuardrail(CustomLogger): data, self.event_hook ) if result is not None: - return result + return result return True def _event_hook_is_event_type(self, event_type: GuardrailEventHooks) -> bool: @@ -403,6 +403,7 @@ class CustomGuardrail(CustomLogger): text: str, language: Optional[str] = None, entities: Optional[List[PiiEntityType]] = None, + request_data: Optional[dict] = None, ) -> str: """ Apply your guardrail logic to the given text @@ -411,6 +412,7 @@ class CustomGuardrail(CustomLogger): text: The text to apply the guardrail to language: The language of the text entities: The entities to mask, optional + request_data: The request data dictionary to store guardrail metadata Any of the custom guardrails can override this method to provide custom guardrail logic diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index ec25f491d8e..b01f9f1b980 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -63,6 +63,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tasks=tasks, task_mappings=task_mappings, guardrail_to_apply=guardrail_to_apply, + request_data=data, ) # Step 2: Run all guardrail tasks in parallel @@ -88,6 +89,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tasks: List, task_mappings: List[Tuple[int, Optional[int]]], guardrail_to_apply: "CustomGuardrail", + request_data: Optional[Dict[str, Any]] = None, ) -> None: """ Extract text content from a message and create guardrail tasks. @@ -100,7 +102,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # Simple string content - tasks.append(guardrail_to_apply.apply_guardrail(text=content)) + tasks.append(guardrail_to_apply.apply_guardrail(text=content, request_data=request_data)) task_mappings.append((msg_idx, None)) elif isinstance(content, list): @@ -109,7 +111,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): text_str = content_item.get("text", None) if text_str is None: continue - tasks.append(guardrail_to_apply.apply_guardrail(text=text_str)) + tasks.append(guardrail_to_apply.apply_guardrail(text=text_str, request_data=request_data)) task_mappings.append((msg_idx, int(content_idx))) async def _apply_guardrail_responses_to_input( @@ -217,6 +219,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tasks: List, task_mappings: List[Tuple[int, Optional[int]]], guardrail_to_apply: "CustomGuardrail", + request_data: Optional[Dict[str, Any]] = None, ) -> None: """ Extract text content from a response choice and create guardrail tasks. @@ -233,7 +236,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if choice.message.content and isinstance(choice.message.content, str): # Simple string content tasks.append( - guardrail_to_apply.apply_guardrail(text=choice.message.content) + guardrail_to_apply.apply_guardrail(text=choice.message.content, request_data=request_data) ) task_mappings.append((choice_idx, None)) @@ -242,7 +245,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for content_idx, content_item in enumerate(choice.message.content): content_text = content_item.get("text") if content_text: - tasks.append(guardrail_to_apply.apply_guardrail(text=content_text)) + tasks.append(guardrail_to_apply.apply_guardrail(text=content_text, request_data=request_data)) task_mappings.append((choice_idx, int(content_idx))) async def _apply_guardrail_responses_to_output( diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 7f793424193..a09185c6e3e 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -697,12 +697,15 @@ async def get_guardrail_ui_settings(): # Convert the PII_ENTITY_CATEGORIES_MAP to the format expected by the UI category_maps = [] for category, entities in PII_ENTITY_CATEGORIES_MAP.items(): - category_maps.append({"category": category, "entities": entities}) + category_maps.append({ + "category": category.value, + "entities": [entity.value for entity in entities] + }) return GuardrailUIAddGuardrailSettings( - supported_entities=list(PiiEntityType), - supported_actions=list(PiiAction), - supported_modes=list(GuardrailEventHooks), + supported_entities=[entity.value for entity in PiiEntityType], + supported_actions=[action.value for action in PiiAction], + supported_modes=[mode.value for mode in GuardrailEventHooks], pii_entity_categories=category_maps, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index e163d3dfbf2..9b71c49f2a9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -1156,22 +1156,34 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): text: str, language: Optional[str] = None, entities: Optional[List[PiiEntityType]] = None, + request_data: Optional[dict] = None, ) -> str: """ Apply Bedrock guardrail to the given text for testing purposes. This method allows users to test Bedrock guardrails without making actual LLM calls. It creates a mock request and response to test the guardrail functionality. + + Args: + text: The text to analyze + language: Optional language parameter (not used by Bedrock) + entities: Optional entities parameter (not used by Bedrock) + request_data: Optional request data dictionary for logging metadata """ try: verbose_proxy_logger.debug("Bedrock Guardrail: Applying guardrail") mock_messages: List[AllMessageValues] = [ ChatCompletionUserMessage(role="user", content=text) ] + + # Use provided request_data or create a mock one for testing + if request_data is None: + request_data = {"messages": mock_messages} + bedrock_response = await self.make_bedrock_api_request( source="INPUT", messages=mock_messages, - request_data={"messages": mock_messages}, + request_data=request_data, ) if bedrock_response.get("action") == "BLOCKED": diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py index 9ffb243d229..d41599c1e19 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py @@ -491,10 +491,11 @@ class EnkryptAIGuardrails(CustomGuardrail): text: str, language: Optional[str] = None, entities: Optional[List[PiiEntityType]] = None, + request_data: Optional[dict] = None, ) -> str: result = await self._call_enkryptai_guardrails( prompt=text, - request_data={}, + request_data=request_data or {}, ) # Process the guardrails response processed_result = self._process_enkryptai_guardrails_response(result) diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index e049ca6a13c..df28b0df8a4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -856,6 +856,7 @@ class NomaGuardrail(CustomGuardrail): text: str, language: Optional[str] = None, entities: Optional[List[PiiEntityType]] = None, + request_data: Optional[dict] = None, ) -> str: """ Apply Noma guardrail to the given text for testing purposes. @@ -867,6 +868,7 @@ class NomaGuardrail(CustomGuardrail): text: The text to analyze language: Optional language parameter (not used by Noma) entities: Optional entities parameter (not used by Noma) + request_data: Optional request data dictionary for logging metadata Returns: The original text if allowed, or anonymized text if available @@ -884,11 +886,15 @@ class NomaGuardrail(CustomGuardrail): # Create payload for Noma API payload = {"request": {"text": text}} + # Use provided request_data or create a mock one for testing + if request_data is None: + request_data = {"messages": [{"role": "user", "content": text}]} + # Call Noma API response_json = await self._call_noma_api( payload=payload, llm_request_id=None, - request_data={"messages": [{"role": "user", "content": text}]}, + request_data=request_data, user_auth=mock_user_auth, extra_data={}, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 2b436681a8e..0f21dfc82d5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -702,6 +702,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): text: str, language: Optional[str] = None, entities: Optional[List[PiiEntityType]] = None, + request_data: Optional[dict] = None, ) -> str: """ UI will call this function to check: @@ -712,7 +713,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): text=text, output_parse_pii=self.output_parse_pii, presidio_config=None, - request_data={}, + request_data=request_data or {}, ) return text diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 0fe7a2c4170..f832fdfd7df 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -231,8 +231,8 @@ PII_ENTITY_CATEGORIES_MAP = { class PiiEntityCategoryMap(TypedDict): - category: PiiEntityCategory - entities: List[PiiEntityType] + category: str + entities: List[str] class GuardrailParamUITypes(str, Enum): @@ -611,9 +611,9 @@ class ListGuardrailsResponse(BaseModel): class GuardrailUIAddGuardrailSettings(BaseModel): - supported_entities: List[PiiEntityType] - supported_actions: List[PiiAction] - supported_modes: List[GuardrailEventHooks] + supported_entities: List[str] + supported_actions: List[str] + supported_modes: List[str] pii_entity_categories: List[PiiEntityCategoryMap] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 1acf0080e3f..234f3a951cd 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -533,6 +533,98 @@ async def test_logging_hook_multiple_content_items(presidio_guardrail): print("✓ Logging hook multiple content items test passed") +@pytest.mark.asyncio +async def test_presidio_sets_guardrail_information_in_request_data(): + """ + Test that Presidio populates guardrail information into request_data metadata. + + This validates that add_standard_logging_guardrail_information_to_request_data + correctly sets the guardrail information that will be used for logging. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", + output_parse_pii=True, + ) + + request_data = { + "messages": [{"role": "user", "content": "Test"}], + "model": "gpt-4o", + "metadata": {}, + } + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + assert request_data is not None + + presidio.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider="presidio", + guardrail_json_response=[], + request_data=request_data, + guardrail_status="success", + start_time=1234567890.0, + end_time=1234567891.0, + duration=1.0, + masked_entity_count={"EMAIL_ADDRESS": 1, "PERSON": 1}, + ) + + return text + + with patch.object(presidio, 'check_pii', mock_check_pii): + await presidio.apply_guardrail( + text="Test message", + request_data=request_data, + ) + + assert "metadata" in request_data + assert "standard_logging_guardrail_information" in request_data["metadata"] + + guardrail_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert "masked_entity_count" in guardrail_info + assert guardrail_info["masked_entity_count"]["EMAIL_ADDRESS"] == 1 + assert guardrail_info["masked_entity_count"]["PERSON"] == 1 + + print("✓ Presidio sets guardrail_information in request_data") + + +@pytest.mark.asyncio +async def test_request_data_flows_to_apply_guardrail(): + """ + Test that request_data is correctly passed to apply_guardrail method. + + This validates the fix where guardrail translation handler passes data + as request_data to apply_guardrail so guardrails can store metadata for logging. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + guardrail_name="test_presidio", + output_parse_pii=True, + ) + + request_data = { + "messages": [{"role": "user", "content": "Test message"}], + "model": "gpt-4o", + "metadata": {}, + } + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + assert request_data is not None, "request_data should be passed to check_pii" + assert "metadata" in request_data, "request_data should have metadata" + + request_data.setdefault("metadata", {}) + request_data["metadata"]["test_flag"] = "passed_correctly" + + return text + + with patch.object(presidio, 'check_pii', mock_check_pii): + result = await presidio.apply_guardrail( + text="Test message", + request_data=request_data, + ) + + assert "metadata" in request_data + assert request_data["metadata"].get("test_flag") == "passed_correctly" + + print("✓ request_data correctly passed to apply_guardrail") + + if __name__ == "__main__": # Run tests asyncio.run(