scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

fix(guardrails): keep the presidio output masker from unmasking after an in-memory update

Browse source
This commit is contained in:
mateo-berri 2026-09-02 13:32:21 -07:00
parent 53a3d74b92
commit 6fae4b3c39
3 changed files with 25 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -1633,6 +1633,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
Update the guardrails litellm params in memory
"""
super().update_in_memory_litellm_params(litellm_params)
if self.apply_to_output:
self.output_parse_pii = False
if litellm_params.pii_entities_config:
self.pii_entities_config = litellm_params.pii_entities_config
if litellm_params.presidio_score_thresholds:

12
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
View file

@ -3129,6 +3129,18 @@ def test_update_in_memory_applies_analyze_chunk_size():
assert guardrail.presidio_analyze_chunk_size_bytes == 99_000
def test_update_in_memory_keeps_output_masker_from_unmasking():
masker = _OPTIONAL_PresidioPIIMasking(mock_testing=True, apply_to_output=True, output_parse_pii=False)
unmasker = _OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True)
params = LitellmParams(guardrail="presidio", mode="pre_call", output_parse_pii=True)
masker.update_in_memory_litellm_params(params)
unmasker.update_in_memory_litellm_params(params)
assert (masker.apply_to_output, masker.output_parse_pii) == (True, False)
assert (unmasker.apply_to_output, unmasker.output_parse_pii) == (False, True)
def test_merge_drops_truncated_same_type_fragment_from_overlap():
"""A boundary entity seen truncated by chunk 1 and whole by chunk 2 must
merge to the single full span; keeping both overlapping spans corrupts the

13
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

@ -566,7 +566,14 @@ def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_st
try:
handler.initialize_guardrail(_presidio_db_guardrail({"EMAIL_ADDRESS": "MASK", "IP_ADDRESS": "MASK"}))
tracked = _presidio_callbacks_in(litellm.callbacks)
roles_before = [(callback.apply_to_output, callback.event_hook) for callback in tracked]
roles_before = [
(callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in tracked
]
assert roles_before == [
(False, True, [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call]),
(False, True, GuardrailEventHooks.post_call),
(True, False, GuardrailEventHooks.post_call),
]
updated = Guardrail(
guardrail_id=PRESIDIO_SIBLINGS_GID,
@ -585,7 +592,9 @@ def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_st
handler.update_in_memory_guardrail(guardrail_id=PRESIDIO_SIBLINGS_GID, guardrail=updated)
assert [callback.pii_entities_config for callback in tracked] == [{"EMAIL_ADDRESS": "MASK"}] * 3
assert [(callback.apply_to_output, callback.event_hook) for callback in tracked] == roles_before
assert [
(callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in tracked
] == roles_before
assert _presidio_callbacks_in(litellm.callbacks) == tracked
finally:
for cb_list, snapshot in zip(lists, snapshots):