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

refactor: extract _mask_tool_call_arguments to stay under PLR0915 limit

Browse source 
Moves the tool-call argument masking loop out of async_pre_call_hook
into a private helper to keep the method at ≤50 statements.
This commit is contained in:
AUTHENSOR 2026-06-18 21:44:52 -05:00
parent a23b13ae87
commit 83eb1f9f73
1 changed files with 49 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -724,6 +724,54 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
masked_entity_count=masked_entity_count,
)
async def _mask_tool_call_arguments(
self,
messages: List[Any],
presidio_config: Optional[dict],
data: dict,
) -> None:
"""Mask PII in tool_call/function_call arguments across all messages."""
tool_tasks: List[Any] = []
tool_targets: List[Any] = []
for m in messages:
if not isinstance(m, dict):
continue
for tool_call in m.get("tool_calls") or []:
function = (
tool_call.get("function")
if isinstance(tool_call, dict)
else None
)
if isinstance(function, dict) and isinstance(
function.get("arguments"), str
):
tool_tasks.append(
self.check_pii(
text=function["arguments"],
output_parse_pii=self.output_parse_pii,
presidio_config=presidio_config,
request_data=data,
)
)
tool_targets.append((function, "arguments"))
function_call = m.get("function_call")
if isinstance(function_call, dict) and isinstance(
function_call.get("arguments"), str
):
tool_tasks.append(
self.check_pii(
text=function_call["arguments"],
output_parse_pii=self.output_parse_pii,
presidio_config=presidio_config,
request_data=data,
)
)
tool_targets.append((function_call, "arguments"))
if tool_tasks:
tool_results = await asyncio.gather(*tool_tasks)
for (container, key), result in zip(tool_targets, tool_results):
container[key] = result
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -802,48 +850,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
elif isinstance(content, list) and content_idx_optional is not None:
messages[msg_idx]["content"][content_idx_optional]["text"] = r
# Also mask PII in tool_call / function_call arguments. Gather
# coroutines to match the parallel dispatch used for content above.
tool_tasks: List[Any] = []
tool_targets: List[Any] = []
for m in messages:
if not isinstance(m, dict):
continue
for tool_call in m.get("tool_calls") or []:
function = (
tool_call.get("function")
if isinstance(tool_call, dict)
else None
)
if isinstance(function, dict) and isinstance(
function.get("arguments"), str
):
tool_tasks.append(
self.check_pii(
text=function["arguments"],
output_parse_pii=self.output_parse_pii,
presidio_config=presidio_config,
request_data=data,
)
)
tool_targets.append((function, "arguments"))
function_call = m.get("function_call")
if isinstance(function_call, dict) and isinstance(
function_call.get("arguments"), str
):
tool_tasks.append(
self.check_pii(
text=function_call["arguments"],
output_parse_pii=self.output_parse_pii,
presidio_config=presidio_config,
request_data=data,
)
)
tool_targets.append((function_call, "arguments"))
if tool_tasks:
tool_results = await asyncio.gather(*tool_tasks)
for (container, key), result in zip(tool_targets, tool_results):
container[key] = result
await self._mask_tool_call_arguments(messages, presidio_config, data)
verbose_proxy_logger.debug(
f"Presidio PII Masking: Redacted pii message: {data['messages']}"