From 84bd8a37358fce5f0d36df8887b8d5e3d0377dbf Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sat, 7 Mar 2026 13:16:01 +0530 Subject: [PATCH] fix: address req changes --- tests/guardrails_tests/test_presidio_pii.py | 226 ++++++++++---------- 1 file changed, 109 insertions(+), 117 deletions(-) diff --git a/tests/guardrails_tests/test_presidio_pii.py b/tests/guardrails_tests/test_presidio_pii.py index 0d730288e63..eda0c7bb5b5 100644 --- a/tests/guardrails_tests/test_presidio_pii.py +++ b/tests/guardrails_tests/test_presidio_pii.py @@ -1,20 +1,19 @@ import sys import os -import io, asyncio import pytest -import time from litellm import mock_completion -from unittest.mock import MagicMock, AsyncMock, patch +from unittest.mock import patch + sys.path.insert(0, os.path.abspath("../..")) import litellm -from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking, PresidioPerRequestConfig +from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, + PresidioPerRequestConfig, +) from litellm.types.guardrails import PiiEntityType, PiiAction from litellm.proxy._types import UserAPIKeyAuth from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError -from litellm.types.utils import CallTypes as LitellmCallTypes - - @pytest.mark.asyncio @@ -26,42 +25,37 @@ async def test_presidio_with_entities_config(): PiiEntityType.CREDIT_CARD: PiiAction.MASK, PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Test text with different PII types test_text = "My credit card number is 4111-1111-1111-1111, my email is test@example.com, and my phone is 555-123-4567" - + # Test the analyze request configuration analyze_request = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify entities were passed correctly assert "entities" in analyze_request assert set(analyze_request["entities"]) == set(pii_entities_config.keys()) - + # Test the check_pii method - this will call the actual Presidio API redacted_text = await presidio_guardrail.check_pii( - text=test_text, - output_parse_pii=True, - presidio_config=None, - request_data={} + text=test_text, output_parse_pii=True, presidio_config=None, request_data={} ) - + # Verify PII has been masked/replaced/redacted in the result assert "4111-1111-1111-1111" not in redacted_text assert "test@example.com" not in redacted_text # Since this entity is not in the config, it should not be masked assert "555-123-4567" in redacted_text - + # The specific replacements will vary based on Presidio's implementation print(f"Redacted text: {redacted_text}") @@ -73,10 +67,12 @@ async def test_presidio_apply_guardrail(): presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config={}, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - test_text = "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" + test_text = ( + "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" + ) response = await presidio_guardrail.apply_guardrail( inputs={"texts": [test_text]}, request_data={}, @@ -91,6 +87,7 @@ async def test_presidio_apply_guardrail(): assert "4111-1111-1111-1111" not in modified_text assert "test@example.com" not in modified_text + @pytest.mark.asyncio async def test_presidio_with_blocked_entities(): """Test for Presidio guardrail with blocked entities - requires actual Presidio API""" @@ -100,36 +97,33 @@ async def test_presidio_with_blocked_entities(): PiiEntityType.CREDIT_CARD: PiiAction.BLOCK, # This entity should cause a block PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, # This entity should be masked } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Test text with blocked PII type - test_text = "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" - + test_text = ( + "My credit card number is 4111-1111-1111-1111 and my email is test@example.com" + ) + # Verify the analyze request configuration analyze_request = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify entities were passed correctly assert "entities" in analyze_request assert set(analyze_request["entities"]) == set(pii_entities_config.keys()) - + # Test that BlockedPiiEntityError is raised when check_pii is called with pytest.raises(BlockedPiiEntityError) as excinfo: await presidio_guardrail.check_pii( - text=test_text, - output_parse_pii=True, - presidio_config=None, - request_data={} + text=test_text, output_parse_pii=True, presidio_config=None, request_data={} ) - + # Verify the error contains the correct entity type assert excinfo.value.entity_type == PiiEntityType.CREDIT_CARD assert excinfo.value.guardrail_name == presidio_guardrail.guardrail_name @@ -143,37 +137,40 @@ async def test_presidio_pre_call_hook_with_blocked_entities(): PiiEntityType.CREDIT_CARD: PiiAction.BLOCK, # This entity should cause a block PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, # This entity should be masked } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Create a sample chat completion request with PII data data = { "messages": [ {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com."} + { + "role": "user", + "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com.", + }, ], - "model": "gpt-3.5-turbo" + "model": "gpt-3.5-turbo", } - + # Mock objects needed for the pre-call hook user_api_key_dict = UserAPIKeyAuth(api_key="test_key") cache = DualCache() - + # Call the pre-call hook and expect BlockedPiiEntityError with pytest.raises(BlockedPiiEntityError) as excinfo: await presidio_guardrail.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, data=data, - call_type="completion" + call_type="completion", ) - + print(f"got error: {excinfo}") - + # Verify the error contains the correct entity type assert excinfo.value.entity_type == PiiEntityType.CREDIT_CARD assert excinfo.value.guardrail_name == presidio_guardrail.guardrail_name @@ -188,44 +185,46 @@ async def test_presidio_pre_call_hook_with_different_call_types(call_type): PiiEntityType.CREDIT_CARD: PiiAction.MASK, PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, } - + presidio_guardrail = _OPTIONAL_PresidioPIIMasking( pii_entities_config=pii_entities_config, presidio_analyzer_api_base=os.environ.get("PRESIDIO_ANALYZER_API_BASE"), - presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE") + presidio_anonymizer_api_base=os.environ.get("PRESIDIO_ANONYMIZER_API_BASE"), ) - + # Create a sample request with PII data data = { "messages": [ {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567"} + { + "role": "user", + "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567", + }, ], - "model": "gpt-3.5-turbo" + "model": "gpt-3.5-turbo", } - + # Mock objects needed for the pre-call hook user_api_key_dict = UserAPIKeyAuth(api_key="test_key") cache = DualCache() - + # Call the pre-call hook with the specified call type modified_data = await presidio_guardrail.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=cache, - data=data, - call_type=call_type + user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type=call_type ) - + # Verify the messages have been modified to mask PII - assert modified_data["messages"][0]["content"] == "You are a helpful assistant." # System prompt should be unchanged - + assert ( + modified_data["messages"][0]["content"] == "You are a helpful assistant." + ) # System prompt should be unchanged + user_message = modified_data["messages"][1]["content"] assert "4111-1111-1111-1111" not in user_message assert "test@example.com" not in user_message # Since this entity is not in the config, it should not be masked assert "555-123-4567" in user_message - + print(f"Modified user message for call_type={call_type}: {user_message}") @@ -243,7 +242,7 @@ def test_validate_environment_missing_http(base_url): # Use patch.dict to temporarily modify environment variables only for this test env_vars = { "PRESIDIO_ANALYZER_API_BASE": f"{base_url}/analyze", - "PRESIDIO_ANONYMIZER_API_BASE": f"{base_url}/anonymize" + "PRESIDIO_ANONYMIZER_API_BASE": f"{base_url}/anonymize", } with patch.dict(os.environ, env_vars): pii_masking.validate_environment() @@ -294,8 +293,12 @@ async def test_output_parsing(): new_response = await pii_masking.async_post_call_success_hook( user_api_key_dict=UserAPIKeyAuth(), data={ - "messages": [{"role": "system", "content": "You are an helpfull assistant"}], - "pii_tokens": {"": "Jane Doe", "": "034453334"}, + "messages": [ + {"role": "system", "content": "You are an helpfull assistant"} + ], + "metadata": { + "pii_tokens": {"": "Jane Doe", "": "034453334"} + }, }, response=response, ) @@ -440,24 +443,26 @@ async def test_presidio_pii_masking_logging_output_only_no_pre_api_hook(): @pytest.mark.asyncio -@patch.dict(os.environ, { - "PRESIDIO_ANALYZER_API_BASE": "http://localhost:5002", - "PRESIDIO_ANONYMIZER_API_BASE": "http://localhost:5001" -}) +@patch.dict( + os.environ, + { + "PRESIDIO_ANALYZER_API_BASE": "http://localhost:5002", + "PRESIDIO_ANONYMIZER_API_BASE": "http://localhost:5001", + }, +) async def test_presidio_pii_masking_logging_output_only_logged_response_guardrails_config(): from typing import Dict, List, Optional import litellm from litellm.proxy.guardrails.init_guardrails import initialize_guardrails from litellm.types.guardrails import ( - GuardrailItem, GuardrailItemSpec, GuardrailEventHooks, ) litellm.set_verbose = True # Environment variables are now patched via the decorator instead of setting them directly - + guardrails_config: List[Dict[str, GuardrailItemSpec]] = [ { "pii_masking": { @@ -499,60 +504,53 @@ async def test_presidio_pii_masking_logging_output_only_logged_response_guardrai async def test_presidio_language_configuration(): """Test that presidio_language parameter is properly set and used in analyze requests""" litellm._turn_on_debug() - + # Test with German language using mock testing to avoid API calls presidio_guardrail_de = _OPTIONAL_PresidioPIIMasking( pii_entities_config={}, presidio_language="de", - mock_testing=True # This bypasses the API validation + mock_testing=True, # This bypasses the API validation ) - + test_text = "Meine Telefonnummer ist +49 30 12345678" - + # Test the analyze request configuration analyze_request = presidio_guardrail_de._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify the language is set to German assert analyze_request["language"] == "de" assert analyze_request["text"] == test_text - + # Test with Spanish language presidio_guardrail_es = _OPTIONAL_PresidioPIIMasking( - pii_entities_config={}, - presidio_language="es", - mock_testing=True + pii_entities_config={}, presidio_language="es", mock_testing=True ) - + test_text_es = "Mi número de teléfono es +34 912 345 678" - + analyze_request_es = presidio_guardrail_es._get_presidio_analyze_request_payload( - text=test_text_es, - presidio_config=None, - request_data={} + text=test_text_es, presidio_config=None, request_data={} ) - + # Verify the language is set to Spanish assert analyze_request_es["language"] == "es" assert analyze_request_es["text"] == test_text_es - + # Test default language (English) when not specified presidio_guardrail_default = _OPTIONAL_PresidioPIIMasking( - pii_entities_config={}, - mock_testing=True + pii_entities_config={}, mock_testing=True ) - + test_text_en = "My phone number is +1 555-123-4567" - - analyze_request_default = presidio_guardrail_default._get_presidio_analyze_request_payload( - text=test_text_en, - presidio_config=None, - request_data={} + + analyze_request_default = ( + presidio_guardrail_default._get_presidio_analyze_request_payload( + text=test_text_en, presidio_config=None, request_data={} + ) ) - + # Verify the language defaults to English assert analyze_request_default["language"] == "en" assert analyze_request_default["text"] == test_text_en @@ -562,36 +560,30 @@ async def test_presidio_language_configuration(): async def test_presidio_language_configuration_with_per_request_override(): """Test that per-request language configuration overrides the default configured language""" litellm._turn_on_debug() - + # Set up guardrail with German as default language presidio_guardrail = _OPTIONAL_PresidioPIIMasking( - pii_entities_config={}, - presidio_language="de", - mock_testing=True + pii_entities_config={}, presidio_language="de", mock_testing=True ) - + test_text = "Test text with PII" - + # Test with per-request config overriding the default language presidio_config = PresidioPerRequestConfig(language="fr") - + analyze_request = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=presidio_config, - request_data={} + text=test_text, presidio_config=presidio_config, request_data={} ) - + # Verify the per-request language (French) overrides the default (German) assert analyze_request["language"] == "fr" assert analyze_request["text"] == test_text - + # Test without per-request config - should use default language analyze_request_default = presidio_guardrail._get_presidio_analyze_request_payload( - text=test_text, - presidio_config=None, - request_data={} + text=test_text, presidio_config=None, request_data={} ) - + # Verify the default language (German) is used assert analyze_request_default["language"] == "de" assert analyze_request_default["text"] == test_text