From 9fa4dfbdd36c7f050e38ecbfccb3c98f1ad30f0a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 19 Feb 2024 16:30:44 -0800 Subject: [PATCH] test(test_presidio_pii_masking.py): add more unit tests --- litellm/proxy/_types.py | 9 + litellm/proxy/hooks/presidio_pii_masking.py | 192 +++++++++++--------- litellm/tests/test_presidio_masking.py | 100 ++++++++++ 3 files changed, 215 insertions(+), 86 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 827a25a2b2..bbe20d0479 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -159,6 +159,9 @@ class GenerateKeyRequest(GenerateRequestBase): {} ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + class Config: + protected_namespaces = () + class GenerateKeyResponse(GenerateKeyRequest): key: str @@ -402,6 +405,9 @@ class LiteLLM_VerificationToken(LiteLLMBase): model_spend: Dict = {} model_max_budget: Dict = {} + class Config: + protected_namespaces = () + class UserAPIKeyAuth( LiteLLM_VerificationToken @@ -441,6 +447,9 @@ class LiteLLM_UserTable(LiteLLMBase): values.update({"models": []}) return values + class Config: + protected_namespaces = () + class LiteLLM_SpendLogs(LiteLLMBase): request_id: str diff --git a/litellm/proxy/hooks/presidio_pii_masking.py b/litellm/proxy/hooks/presidio_pii_masking.py index 031f4e2d48..a11ea7f55d 100644 --- a/litellm/proxy/hooks/presidio_pii_masking.py +++ b/litellm/proxy/hooks/presidio_pii_masking.py @@ -29,10 +29,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): user_api_key_cache = None # Class variables or attributes - def __init__(self, mock_testing: bool = False): + def __init__( + self, mock_testing: bool = False, mock_redacted_text: Optional[dict] = None + ): self.pii_tokens: dict = ( {} ) # mapping of PII token to original text - only used with Presidio `replace` operation + + self.mock_redacted_text = mock_redacted_text if mock_testing == True: # for testing purposes only return @@ -67,26 +71,31 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): """ try: async with aiohttp.ClientSession() as session: - # Make the first request to /analyze - analyze_url = f"{self.presidio_analyzer_api_base}/analyze" - verbose_proxy_logger.debug(f"Making request to: {analyze_url}") - analyze_payload = {"text": text, "language": "en"} - redacted_text = None - async with session.post(analyze_url, json=analyze_payload) as response: - analyze_results = await response.json() + if self.mock_redacted_text is not None: + redacted_text = self.mock_redacted_text + else: + # Make the first request to /analyze + analyze_url = f"{self.presidio_analyzer_api_base}analyze" + verbose_proxy_logger.debug(f"Making request to: {analyze_url}") + analyze_payload = {"text": text, "language": "en"} + redacted_text = None + async with session.post( + analyze_url, json=analyze_payload + ) as response: + analyze_results = await response.json() - # Make the second request to /anonymize - anonymize_url = f"{self.presidio_anonymizer_api_base}/anonymize" - verbose_proxy_logger.debug(f"Making request to: {anonymize_url}") - anonymize_payload = { - "text": "hello world, my name is Jane Doe. My number is: 034453334", - "analyzer_results": analyze_results, - } + # Make the second request to /anonymize + anonymize_url = f"{self.presidio_anonymizer_api_base}anonymize" + verbose_proxy_logger.debug(f"Making request to: {anonymize_url}") + anonymize_payload = { + "text": text, + "analyzer_results": analyze_results, + } - async with session.post( - anonymize_url, json=anonymize_payload - ) as response: - redacted_text = await response.json() + async with session.post( + anonymize_url, json=anonymize_payload + ) as response: + redacted_text = await response.json() new_text = text if redacted_text is not None: @@ -99,14 +108,16 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): # check if token in dict # if exists, add a uuid to the replacement token for swapping back to the original text in llm response output parsing if replacement in self.pii_tokens: - replacement = replacement + uuid.uuid4() + replacement = replacement + str(uuid.uuid4()) self.pii_tokens[replacement] = new_text[ start:end ] # get text it'll replace new_text = new_text[:start] + replacement + new_text[end:] - return redacted_text["text"] + return redacted_text["text"] + else: + raise Exception(f"Invalid anonymizer response: {redacted_text}") except Exception as e: traceback.print_exc() raise e @@ -128,79 +139,88 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): For multiple messages in /chat/completions, we'll need to call them in parallel. """ - permissions = user_api_key_dict.permissions - output_parse_pii = permissions.get( - "output_parse_pii", litellm.output_parse_pii - ) # allow key to turn on/off output parsing for pii - no_pii = permissions.get( - "no-pii", None - ) # allow key to turn on/off pii masking (if user is allowed to set pii controls, then they can override the key defaults) + try: + permissions = user_api_key_dict.permissions + output_parse_pii = permissions.get( + "output_parse_pii", litellm.output_parse_pii + ) # allow key to turn on/off output parsing for pii + no_pii = permissions.get( + "no-pii", None + ) # allow key to turn on/off pii masking (if user is allowed to set pii controls, then they can override the key defaults) - if no_pii is None: - # check older way of turning on/off pii - no_pii = not permissions.get("pii", True) + if no_pii is None: + # check older way of turning on/off pii + no_pii = not permissions.get("pii", True) - content_safety = data.get("content_safety", None) - verbose_proxy_logger.debug(f"content_safety: {content_safety}") - ## Request-level turn on/off PII controls ## - if content_safety is not None and isinstance(content_safety, dict): - # pii masking ## - if ( - content_safety.get("no-pii", None) is not None - and content_safety.get("no-pii") == True - ): - # check if user allowed to turn this off - if permissions.get("allow_pii_controls", False) == False: - raise HTTPException( - status_code=400, - detail={"error": "Not allowed to set PII controls per request"}, - ) - else: # user allowed to turn off pii masking - no_pii = content_safety.get("no-pii") - if not isinstance(no_pii, bool): - raise HTTPException( - status_code=400, - detail={"error": "no_pii needs to be a boolean value"}, - ) - ## pii output parsing ## - if content_safety.get("output_parse_pii", None) is not None: - # check if user allowed to turn this off - if permissions.get("allow_pii_controls", False) == False: - raise HTTPException( - status_code=400, - detail={"error": "Not allowed to set PII controls per request"}, - ) - else: # user allowed to turn on/off pii output parsing - output_parse_pii = content_safety.get("output_parse_pii") - if not isinstance(output_parse_pii, bool): + content_safety = data.get("content_safety", None) + verbose_proxy_logger.debug(f"content_safety: {content_safety}") + ## Request-level turn on/off PII controls ## + if content_safety is not None and isinstance(content_safety, dict): + # pii masking ## + if ( + content_safety.get("no-pii", None) is not None + and content_safety.get("no-pii") == True + ): + # check if user allowed to turn this off + if permissions.get("allow_pii_controls", False) == False: raise HTTPException( status_code=400, detail={ - "error": "output_parse_pii needs to be a boolean value" + "error": "Not allowed to set PII controls per request" }, ) - - if no_pii == False: # turn off pii masking - return data - - if call_type == "completion": # /chat/completions requests - messages = data["messages"] - tasks = [] - - for m in messages: - if isinstance(m["content"], str): - tasks.append( - self.check_pii( - text=m["content"], output_parse_pii=output_parse_pii + else: # user allowed to turn off pii masking + no_pii = content_safety.get("no-pii") + if not isinstance(no_pii, bool): + raise HTTPException( + status_code=400, + detail={"error": "no_pii needs to be a boolean value"}, + ) + ## pii output parsing ## + if content_safety.get("output_parse_pii", None) is not None: + # check if user allowed to turn this off + if permissions.get("allow_pii_controls", False) == False: + raise HTTPException( + status_code=400, + detail={ + "error": "Not allowed to set PII controls per request" + }, ) - ) - responses = await asyncio.gather(*tasks) - for index, r in enumerate(responses): - if isinstance(messages[index]["content"], str): - messages[index][ - "content" - ] = r # replace content with redacted string - return data + else: # user allowed to turn on/off pii output parsing + output_parse_pii = content_safety.get("output_parse_pii") + if not isinstance(output_parse_pii, bool): + raise HTTPException( + status_code=400, + detail={ + "error": "output_parse_pii needs to be a boolean value" + }, + ) + + if no_pii == True: # turn off pii masking + return data + + if call_type == "completion": # /chat/completions requests + messages = data["messages"] + tasks = [] + + for m in messages: + if isinstance(m["content"], str): + tasks.append( + self.check_pii( + text=m["content"], output_parse_pii=output_parse_pii + ) + ) + responses = await asyncio.gather(*tasks) + for index, r in enumerate(responses): + if isinstance(messages[index]["content"], str): + messages[index][ + "content" + ] = r # replace content with redacted string + verbose_proxy_logger.debug(f"Redacted pii message: {data['messages']}") + return data + except Exception as e: + verbose_proxy_logger.info(f"An error occurred - {str(e)}") + raise e async def async_post_call_success_hook( self, diff --git a/litellm/tests/test_presidio_masking.py b/litellm/tests/test_presidio_masking.py index f7e1ac12d3..f8ac9820fd 100644 --- a/litellm/tests/test_presidio_masking.py +++ b/litellm/tests/test_presidio_masking.py @@ -63,3 +63,103 @@ async def test_output_parsing(): # asyncio.run(test_output_parsing()) + + +### UNIT TESTS FOR PRESIDIO PII MASKING ### + +input_a_anonymizer_results = { + "text": "hello world, my name is . My number is: ", + "items": [ + { + "start": 48, + "end": 62, + "entity_type": "PHONE_NUMBER", + "text": "", + "operator": "replace", + }, + { + "start": 24, + "end": 32, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + ], +} + +input_b_anonymizer_results = { + "text": "My name is , who are you? Say my name in your response", + "items": [ + { + "start": 11, + "end": 19, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + } + ], +} + + +# Test if PII masking works with input A +@pytest.mark.asyncio +async def test_presidio_pii_masking_input_a(): + """ + Tests to see if correct parts of sentence anonymized + """ + pii_masking = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, mock_redacted_text=input_a_anonymizer_results + ) + + _api_key = "sk-12345" + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) + local_cache = DualCache() + + new_data = await pii_masking.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "messages": [ + { + "role": "user", + "content": "hello world, my name is Jane Doe. My number is: 23r323r23r2wwkl", + } + ] + }, + call_type="completion", + ) + + assert "" in new_data["messages"][0]["content"] + assert "" in new_data["messages"][0]["content"] + + +# Test if PII masking works with input B (also test if the response != A's response) +@pytest.mark.asyncio +async def test_presidio_pii_masking_input_b(): + """ + Tests to see if correct parts of sentence anonymized + """ + pii_masking = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, mock_redacted_text=input_b_anonymizer_results + ) + + _api_key = "sk-12345" + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) + local_cache = DualCache() + + new_data = await pii_masking.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "messages": [ + { + "role": "user", + "content": "My name is Jane Doe, who are you? Say my name in your response", + } + ] + }, + call_type="completion", + ) + + assert "" in new_data["messages"][0]["content"] + assert "" not in new_data["messages"][0]["content"]