From 5d5409e77aa899887a4fc9c010100de4a61f4b56 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Thu, 26 Feb 2026 20:07:46 +0530 Subject: [PATCH 1/5] fix presidio memory leak --- .../guardrails/guardrail_hooks/presidio.py | 289 ++++++++++++------ .../guardrails/guardrail_initializers.py | 1 + 2 files changed, 199 insertions(+), 91 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 34fbf47253..d4ce2c371e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -427,6 +427,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): analyze_results: Any, output_parse_pii: bool, masked_entity_count: Dict[str, int], + request_data: Optional[Dict] = None, ) -> str: """ Send analysis results to the Presidio anonymizer endpoint to get redacted text @@ -482,13 +483,24 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if item["operator"] == "replace" and output_parse_pii is True: # 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 + str(uuid.uuid4()) + pii_tokens = self.pii_tokens + if request_data is not None: + if "pii_tokens" not in request_data: + request_data["pii_tokens"] = {} + pii_tokens = request_data["pii_tokens"] - self.pii_tokens[replacement] = new_text[ + # Always append a UUID to ensure the replacement token is unique to this request and session. + # This prevents collisions where the LLM might hallucinate a generic token like [PHONE_NUMBER]. + replacement = f"{replacement}_{str(uuid.uuid4())[:12]}" + + pii_tokens[replacement] = new_text[ start:end ] # get text it'll replace + verbose_proxy_logger.info( + f"\033[92mPII Masking\033[0m: Created token {replacement} for original text: '{new_text[start:end]}'" + ) + new_text = new_text[:start] + replacement + new_text[end:] entity_type = item.get("entity_type", None) if entity_type is not None: @@ -525,10 +537,27 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return analyze_results filtered_results: List[PresidioAnalyzeResponseItem] = [] + print( + "DEBUG filter input:", + analyze_results, + " deny_list:", + self.presidio_entities_deny_list, + ) for item in analyze_results: entity_type = item.get("entity_type") - if entity_type and entity_type in self.presidio_entities_deny_list: + deny_list_strings = [ + x.value if hasattr(x, "value") else str(x) + for x in self.presidio_entities_deny_list + ] + str_entity_type = str( + entity_type.value if hasattr(entity_type, "value") else entity_type + ) + print( + f"DEBUG entity_type: {entity_type}, str_entity_type: '{str_entity_type}', deny_strings: {deny_list_strings}" + ) + if entity_type and str_entity_type in deny_list_strings: + print(f"DEBUG Skipping {entity_type} due to deny list") continue if self.presidio_score_thresholds: @@ -621,6 +650,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): analyze_results=analyze_results, output_parse_pii=output_parse_pii, masked_entity_count=masked_entity_count, + request_data=request_data, ) return anonymized_text return redacted_text["text"] @@ -866,14 +896,130 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if isinstance(response, ModelResponse) and not isinstance( response.choices[0], StreamingChoices ): # /chat/completions requests - if isinstance(response.choices[0].message.content, str): - verbose_proxy_logger.debug( - f"self.pii_tokens: {self.pii_tokens}; initial response: {response.choices[0].message.content}" - ) - for key, value in self.pii_tokens.items(): - response.choices[0].message.content = response.choices[ - 0 - ].message.content.replace(key, value) + await self._process_response_for_pii( + response=response, + request_data=data, + mode="unmask", + ) + return response + + async def _process_response_for_pii( + self, + response: ModelResponse, + request_data: dict, + mode: Literal["mask", "unmask"], + ) -> ModelResponse: + """ + Helper to recursively process a ModelResponse for PII. + Handles all choices and tool calls. + """ + pii_tokens = ( + request_data.get("pii_tokens", self.pii_tokens) + if request_data + else self.pii_tokens + ) + presidio_config = self.get_presidio_settings_from_request_data( + request_data or {} + ) + + for choice in response.choices: + message = getattr(choice, "message", None) + if message is None: + continue + + # 1. Process content + content = getattr(message, "content", None) + if isinstance(content, str): + if mode == "unmask": + for token, original_text in pii_tokens.items(): + if token in content: + verbose_proxy_logger.info( + f"\033[94mPII Unmasking\033[0m: Found token {token} in response. Replacing with original text." + ) + content = content.replace(token, original_text) + # FALLBACK: Handle truncated tokens (token cut off by max_tokens) + elif any( + token.startswith(content[i:]) + for i in range( + max(0, len(content) - len(token)), len(content) + ) + if len(content[i:]) > 15 + ): + # If the end of content matches the start of a token, it's likely truncated + for i in range( + max(0, len(content) - len(token)), len(content) + ): + sub = content[i:] + if token.startswith(sub) and len(sub) > 15: + verbose_proxy_logger.info( + f"\033[93mPII Unmasking\033[0m: Found truncated token {sub}... in response. Replacing with original text." + ) + content = content[:i] + original_text + break + message.content = content + elif mode == "mask": + message.content = await self.check_pii( + text=content, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=request_data, + ) + elif isinstance(content, list): + for item in content: + if not isinstance(item, dict): + continue + text_value = item.get("text") + if text_value is None: + continue + if mode == "unmask": + for token, original_text in pii_tokens.items(): + text_value = text_value.replace(token, original_text) + item["text"] = text_value + elif mode == "mask": + item["text"] = await self.check_pii( + text=text_value, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=request_data, + ) + + # 2. Process tool calls + tool_calls = getattr(message, "tool_calls", None) + if tool_calls: + for tool_call in tool_calls: + function = getattr(tool_call, "function", None) + if function and hasattr(function, "arguments"): + args = function.arguments + if isinstance(args, str): + if mode == "unmask": + for token, original_text in pii_tokens.items(): + args = args.replace(token, original_text) + function.arguments = args + elif mode == "mask": + function.arguments = await self.check_pii( + text=args, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=request_data, + ) + + # 3. Process legacy function calls + function_call = getattr(message, "function_call", None) + if function_call and hasattr(function_call, "arguments"): + args = function_call.arguments + if isinstance(args, str): + if mode == "unmask": + for token, original_text in pii_tokens.items(): + args = args.replace(token, original_text) + function_call.arguments = args + elif mode == "mask": + function_call.arguments = await self.check_pii( + text=args, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=request_data, + ) + return response async def _mask_output_response( @@ -891,37 +1037,12 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if response.choices and isinstance(response.choices[0], StreamingChoices): return response - presidio_config = self.get_presidio_settings_from_request_data( - request_data or {} + await self._process_response_for_pii( + response=response, + request_data=request_data, + mode="mask", ) - - for choice in response.choices: - # Type narrowing: StreamingChoices doesn't have .message attribute - if not hasattr(choice, "message"): - continue - content = getattr(choice.message, "content", None) # type: ignore - if content is None: - continue - if isinstance(content, str): - choice.message.content = await self.check_pii( # type: ignore - text=content, - output_parse_pii=False, - presidio_config=presidio_config, - request_data=request_data, - ) - elif isinstance(content, list): - for item in content: - if not isinstance(item, dict): - continue - text_value = item.get("text") - if text_value is None: - continue - item["text"] = await self.check_pii( - text=text_value, - output_parse_pii=False, - presidio_config=presidio_config, - request_data=request_data, - ) + return response return response @@ -934,7 +1055,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """ Process streaming response chunks to unmask PII tokens when needed. """ - from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + from litellm.llms.base_llm.base_model_iterator import ( + convert_model_response_to_streaming, + ) from litellm.main import stream_chunk_builder from litellm.types.utils import ModelResponse @@ -959,45 +1082,16 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return # Apply Presidio masking on the assembled response - presidio_config = self.get_presidio_settings_from_request_data( - request_data or {} - ) - - content_to_mask = "" - if ( - hasattr(assembled_model_response, "choices") - and len(assembled_model_response.choices) > 0 - ): - if hasattr( - assembled_model_response.choices[0], "message" - ) and hasattr( - assembled_model_response.choices[0].message, "content" - ): - content_to_mask = ( - assembled_model_response.choices[0].message.content or "" - ) - - masked_content = await self.check_pii( - text=content_to_mask, - output_parse_pii=False, - presidio_config=presidio_config, + await self._process_response_for_pii( + response=assembled_model_response, request_data=request_data, + mode="mask", ) - if ( - hasattr(assembled_model_response, "choices") - and len(assembled_model_response.choices) > 0 - ): - if hasattr(assembled_model_response.choices[0], "message"): - assembled_model_response.choices[ - 0 - ].message.content = masked_content - - mock_response = MockResponseIterator( - model_response=assembled_model_response + mock_response_stream = convert_model_response_to_streaming( + assembled_model_response ) - async for chunk in mock_response: - yield chunk + yield mock_response_stream return except Exception as e: @@ -1011,7 +1105,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return # --- PII unmasking path (output_parse_pii=True) --- - if not (self.output_parse_pii and self.pii_tokens): + # --- PII unmasking path (output_parse_pii=True) --- + pii_tokens = ( + request_data.get("pii_tokens", self.pii_tokens) + if request_data + else self.pii_tokens + ) + if not (self.output_parse_pii and pii_tokens): async for chunk in response: yield chunk return @@ -1034,20 +1134,27 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk return - # Apply PII unmasking to assembled content - for choice in assembled_model_response.choices: - if hasattr(choice, "message") and hasattr(choice.message, "content"): - content = choice.message.content - if isinstance(content, str): - for token, original_text in self.pii_tokens.items(): - content = content.replace(token, original_text) - choice.message.content = content + # --- PRESERVE USAGE METADATA --- + # stream_chunk_builder might miss usage if it's only in the last chunk + if ( + not hasattr(assembled_model_response, "usage") + or not assembled_model_response.usage + ) and all_chunks: + last_chunk = all_chunks[-1] + if hasattr(last_chunk, "usage") and last_chunk.usage: + assembled_model_response.usage = last_chunk.usage - mock_response = MockResponseIterator( - model_response=assembled_model_response + # Apply PII unmasking to assembled content (unmasking tokens back to original text) + await self._process_response_for_pii( + response=assembled_model_response, + request_data=request_data, + mode="unmask", ) - async for chunk in mock_response: - yield chunk + + mock_response_stream = convert_model_response_to_streaming( + assembled_model_response + ) + yield mock_response_stream except Exception as e: verbose_proxy_logger.error(f"Error in PII streaming processing: {str(e)}") diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 639aebf45c..109f223716 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -93,6 +93,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): presidio_analyzer_api_base=litellm_params.presidio_analyzer_api_base, presidio_anonymizer_api_base=litellm_params.presidio_anonymizer_api_base, presidio_language=litellm_params.presidio_language, + presidio_entities_deny_list=litellm_params.presidio_entities_deny_list, apply_to_output=False, ) params.update(overrides) From 1b632ed473852c0a3f028ec7fe62654bfcacd675 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Thu, 26 Feb 2026 20:24:19 +0530 Subject: [PATCH 2/5] chore: remove debug logs from presidio guardrail --- .../guardrails/guardrail_hooks/presidio.py | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index d4ce2c371e..d8bf2bc306 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -497,10 +497,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): start:end ] # get text it'll replace - verbose_proxy_logger.info( - f"\033[92mPII Masking\033[0m: Created token {replacement} for original text: '{new_text[start:end]}'" - ) - new_text = new_text[:start] + replacement + new_text[end:] entity_type = item.get("entity_type", None) if entity_type is not None: @@ -537,12 +533,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return analyze_results filtered_results: List[PresidioAnalyzeResponseItem] = [] - print( - "DEBUG filter input:", - analyze_results, - " deny_list:", - self.presidio_entities_deny_list, - ) for item in analyze_results: entity_type = item.get("entity_type") @@ -553,11 +543,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): str_entity_type = str( entity_type.value if hasattr(entity_type, "value") else entity_type ) - print( - f"DEBUG entity_type: {entity_type}, str_entity_type: '{str_entity_type}', deny_strings: {deny_list_strings}" - ) if entity_type and str_entity_type in deny_list_strings: - print(f"DEBUG Skipping {entity_type} due to deny list") continue if self.presidio_score_thresholds: @@ -933,9 +919,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if mode == "unmask": for token, original_text in pii_tokens.items(): if token in content: - verbose_proxy_logger.info( - f"\033[94mPII Unmasking\033[0m: Found token {token} in response. Replacing with original text." - ) content = content.replace(token, original_text) # FALLBACK: Handle truncated tokens (token cut off by max_tokens) elif any( @@ -951,9 +934,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ): sub = content[i:] if token.startswith(sub) and len(sub) > 15: - verbose_proxy_logger.info( - f"\033[93mPII Unmasking\033[0m: Found truncated token {sub}... in response. Replacing with original text." - ) content = content[:i] + original_text break message.content = content From 3210da1e2a0bfcfabb6384766d5835875670a7a2 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Fri, 27 Feb 2026 13:43:51 +0530 Subject: [PATCH 3/5] fix: req changes --- .../guardrails/guardrail_hooks/presidio.py | 23 ++++++------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index d8bf2bc306..53d1c5f268 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -533,13 +533,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return analyze_results filtered_results: List[PresidioAnalyzeResponseItem] = [] + deny_list_strings = [ + x.value if hasattr(x, "value") else str(x) + for x in self.presidio_entities_deny_list + ] for item in analyze_results: entity_type = item.get("entity_type") - deny_list_strings = [ - x.value if hasattr(x, "value") else str(x) - for x in self.presidio_entities_deny_list - ] str_entity_type = str( entity_type.value if hasattr(entity_type, "value") else entity_type ) @@ -921,13 +921,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if token in content: content = content.replace(token, original_text) # FALLBACK: Handle truncated tokens (token cut off by max_tokens) - elif any( - token.startswith(content[i:]) - for i in range( - max(0, len(content) - len(token)), len(content) - ) - if len(content[i:]) > 15 - ): + else: # If the end of content matches the start of a token, it's likely truncated for i in range( max(0, len(content) - len(token)), len(content) @@ -1024,8 +1018,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) return response - return response - async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -1084,7 +1076,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk return - # --- PII unmasking path (output_parse_pii=True) --- # --- PII unmasking path (output_parse_pii=True) --- pii_tokens = ( request_data.get("pii_tokens", self.pii_tokens) @@ -1119,8 +1110,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if ( not hasattr(assembled_model_response, "usage") or not assembled_model_response.usage - ) and all_chunks: - last_chunk = all_chunks[-1] + ) and remaining_chunks: + last_chunk = remaining_chunks[-1] if hasattr(last_chunk, "usage") and last_chunk.usage: assembled_model_response.usage = last_chunk.usage From 1badececa31d3a7a3312e6417a74053249ee0e91 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sat, 28 Feb 2026 18:53:37 +0530 Subject: [PATCH 4/5] fix: presidio req change --- .../guardrails/guardrail_hooks/presidio.py | 92 +++++++++++-------- 1 file changed, 54 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 53d1c5f268..849b1ec9ca 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -483,11 +483,16 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if item["operator"] == "replace" and output_parse_pii is True: # 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 - pii_tokens = self.pii_tokens - if request_data is not None: - if "pii_tokens" not in request_data: - request_data["pii_tokens"] = {} - pii_tokens = request_data["pii_tokens"] + if request_data is None: + verbose_proxy_logger.warning( + "Presidio anonymize_text called without request_data — " + "PII tokens cannot be stored per-request. " + "This may indicate a missing caller update." + ) + request_data = {} + if "pii_tokens" not in request_data: + request_data["pii_tokens"] = {} + pii_tokens = request_data["pii_tokens"] # Always append a UUID to ensure the replacement token is unique to this request and session. # This prevents collisions where the LLM might hallucinate a generic token like [PHONE_NUMBER]. @@ -889,6 +894,32 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) return response + @staticmethod + def _unmask_pii_text(text: str, pii_tokens: Dict[str, str]) -> str: + """ + Replace PII tokens in *text* with their original values. + + Includes a fallback for tokens that were truncated by ``max_tokens``: + if the *end* of ``text`` matches the *beginning* of a token and the + overlap is long enough, the truncated suffix is replaced with the + original value. The minimum overlap length is + ``min(20, len(token) // 2)`` to reduce the risk of false positives + when multiple tokens share a common prefix. + """ + for token, original_text in pii_tokens.items(): + if token in text: + text = text.replace(token, original_text) + else: + # FALLBACK: Handle truncated tokens (token cut off by max_tokens) + # Only check at the very end of the text. + min_overlap = min(20, len(token) // 2) + for i in range(max(0, len(text) - len(token)), len(text)): + sub = text[i:] + if token.startswith(sub) and len(sub) >= min_overlap: + text = text[:i] + original_text + break + return text + async def _process_response_for_pii( self, response: ModelResponse, @@ -899,11 +930,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Helper to recursively process a ModelResponse for PII. Handles all choices and tool calls. """ - pii_tokens = ( - request_data.get("pii_tokens", self.pii_tokens) - if request_data - else self.pii_tokens - ) + pii_tokens = request_data.get("pii_tokens", {}) if request_data else {} + if not pii_tokens and mode == "unmask": + verbose_proxy_logger.debug( + "No pii_tokens found in request_data — nothing to unmask" + ) presidio_config = self.get_presidio_settings_from_request_data( request_data or {} ) @@ -917,20 +948,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): content = getattr(message, "content", None) if isinstance(content, str): if mode == "unmask": - for token, original_text in pii_tokens.items(): - if token in content: - content = content.replace(token, original_text) - # FALLBACK: Handle truncated tokens (token cut off by max_tokens) - else: - # If the end of content matches the start of a token, it's likely truncated - for i in range( - max(0, len(content) - len(token)), len(content) - ): - sub = content[i:] - if token.startswith(sub) and len(sub) > 15: - content = content[:i] + original_text - break - message.content = content + message.content = self._unmask_pii_text(content, pii_tokens) elif mode == "mask": message.content = await self.check_pii( text=content, @@ -946,9 +964,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if text_value is None: continue if mode == "unmask": - for token, original_text in pii_tokens.items(): - text_value = text_value.replace(token, original_text) - item["text"] = text_value + item["text"] = self._unmask_pii_text(text_value, pii_tokens) elif mode == "mask": item["text"] = await self.check_pii( text=text_value, @@ -966,9 +982,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): args = function.arguments if isinstance(args, str): if mode == "unmask": - for token, original_text in pii_tokens.items(): - args = args.replace(token, original_text) - function.arguments = args + function.arguments = self._unmask_pii_text( + args, pii_tokens + ) elif mode == "mask": function.arguments = await self.check_pii( text=args, @@ -983,9 +999,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): args = function_call.arguments if isinstance(args, str): if mode == "unmask": - for token, original_text in pii_tokens.items(): - args = args.replace(token, original_text) - function_call.arguments = args + function_call.arguments = self._unmask_pii_text( + args, pii_tokens + ) elif mode == "mask": function_call.arguments = await self.check_pii( text=args, @@ -1077,11 +1093,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return # --- PII unmasking path (output_parse_pii=True) --- - pii_tokens = ( - request_data.get("pii_tokens", self.pii_tokens) - if request_data - else self.pii_tokens - ) + pii_tokens = request_data.get("pii_tokens", {}) if request_data else {} + if not pii_tokens and request_data: + verbose_proxy_logger.debug( + "No pii_tokens in request_data for streaming unmask path" + ) if not (self.output_parse_pii and pii_tokens): async for chunk in response: yield chunk From 1073ba6d13507234bb54f33518039d2a4795cc85 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sat, 28 Feb 2026 19:06:31 +0530 Subject: [PATCH 5/5] fix req changes --- litellm/proxy/guardrails/guardrail_hooks/presidio.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 849b1ec9ca..adfe69315c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -508,7 +508,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): masked_entity_count[entity_type] = ( masked_entity_count.get(entity_type, 0) + 1 ) - return redacted_text["text"] + # When output_parse_pii is True, new_text contains UUID-suffixed + # tokens that match the keys in pii_tokens. Returning + # redacted_text["text"] (Presidio's original output) would send + # un-suffixed tokens to the LLM, making unmasking impossible. + # When output_parse_pii is False, new_text == redacted_text["text"] + # because no UUID suffix is appended. + return new_text else: raise Exception("Invalid anonymizer response: received None") except Exception as e: