scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

refactor(presidio): rename request-scoped pii token helpers

Browse source 
Rename _get_or_create_request_scoped_pii_tokens to _ensure_request_scoped_pii_tokens
and _get_request_scoped_pii_tokens to _resolve_request_scoped_pii_tokens so names
reflect anchoring vs read-with-metadata sync behavior.

Made-with: Cursor
This commit is contained in:
Milan 2026-04-16 22:58:39 +03:00
parent e3e4956050
commit 013d360b3e
No known key found for this signature in database
1 changed files with 14 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -504,7 +504,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
"This may indicate a missing caller update."
)
request_data = {}
pii_tokens = self._get_or_create_request_scoped_pii_tokens(
pii_tokens = self._ensure_request_scoped_pii_tokens(
request_data
)
@ -931,12 +931,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
)
return response
def _get_or_create_request_scoped_pii_tokens(
def _ensure_request_scoped_pii_tokens(
self, request_data: Dict[str, Any]
) -> Dict[str, str]:
"""
Return request-scoped token map and ensure it is anchored in all
metadata locations used in pre/post guardrail hooks.
Return the request-scoped token map, creating it if needed, and anchor
the same dict under both ``metadata`` and ``litellm_metadata`` for hooks
that only read one of those namespaces.
"""
metadata = request_data.get("metadata")
if not isinstance(metadata, dict):
@ -955,11 +956,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
return cast(Dict[str, str], pii_tokens)
def _get_request_scoped_pii_tokens(
def _resolve_request_scoped_pii_tokens(
self, request_data: Optional[Dict[str, Any]]
) -> Dict[str, str]:
"""
Read request-scoped pii token map with fallback across known metadata paths.
Resolve the request-scoped token map from ``metadata`` or ``litellm_metadata``.
If the map exists only under ``metadata``, assigns the same dict onto
``litellm_metadata`` so post-call paths that read the latter stay in sync.
"""
if not request_data:
return {}
@ -1026,7 +1030,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
Process an Anthropic native message dict for PII masking/unmasking.
Handles content blocks with type == "text".
"""
pii_tokens = self._get_request_scoped_pii_tokens(request_data)
pii_tokens = self._resolve_request_scoped_pii_tokens(request_data)
if not pii_tokens and mode == "unmask":
verbose_proxy_logger.debug(
"No pii_tokens found for Anthropic response unmask"
@ -1067,7 +1071,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
Helper to recursively process a ModelResponse for PII.
Handles all choices and tool calls.
"""
pii_tokens = self._get_request_scoped_pii_tokens(request_data)
pii_tokens = self._resolve_request_scoped_pii_tokens(request_data)
if not pii_tokens and mode == "unmask":
verbose_proxy_logger.debug(
"No pii_tokens found in request_data — nothing to unmask"
@ -1297,7 +1301,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
yield chunk
return
pii_tokens = self._get_request_scoped_pii_tokens(request_data)
pii_tokens = self._resolve_request_scoped_pii_tokens(request_data)
if not pii_tokens and request_data:
verbose_proxy_logger.debug(
"No pii_tokens found in request_data for streaming unmask path"
@ -1360,7 +1364,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
# When input_type is "response" and pii_tokens are available,
# unmask the text instead of masking it.
pii_tokens = self._get_request_scoped_pii_tokens(request_data)
pii_tokens = self._resolve_request_scoped_pii_tokens(request_data)
new_texts = []
if input_type == "response":