From c208335a3d484c0d1aefc9e56a2dcea8b243e363 Mon Sep 17 00:00:00 2001 From: AUTHENSOR Date: Wed, 17 Jun 2026 23:05:48 -0500 Subject: [PATCH] fix(guardrails): mask PII in tool-call arguments for Presidio pre-call and Anthropic response Previously, Presidio's async_pre_call_hook only scanned message content strings, leaving tool_calls[*].function.arguments and function_call.arguments unmasked. Similarly, _process_anthropic_response_for_pii skipped tool_use block input dicts. Add check_pii sweeps over both tool-call argument channels so PII in structured tool I/O is masked the same way as free-text content. Includes regression tests covering both channels. --- .../guardrails/guardrail_hooks/presidio.py | 79 ++++++++-- .../test_presidio_tool_call_pii_masking.py | 149 ++++++++++++++++++ 2 files changed, 215 insertions(+), 13 deletions(-) create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio_tool_call_pii_masking.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index fc414ab7b54..846667c10b5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -802,6 +802,39 @@ 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. The content loop + # above only inspects message content, so PII embedded in tool-call + # arguments would otherwise reach the provider unmasked -- this mirrors the + # response-side handling in _process_response_for_pii. + 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 + ): + function["arguments"] = await self.check_pii( + text=function["arguments"], + output_parse_pii=self.output_parse_pii, + presidio_config=presidio_config, + request_data=data, + ) + function_call = m.get("function_call") + if isinstance(function_call, dict) and isinstance( + function_call.get("arguments"), str + ): + function_call["arguments"] = await self.check_pii( + text=function_call["arguments"], + output_parse_pii=self.output_parse_pii, + presidio_config=presidio_config, + request_data=data, + ) + verbose_proxy_logger.debug( f"Presidio PII Masking: Redacted pii message: {data['messages']}" ) @@ -1016,20 +1049,40 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return response for block in content: - if not isinstance(block, dict) or block.get("type") != "text": + if not isinstance(block, dict): continue - text_value = block.get("text") - if text_value is None: - continue - if mode == "unmask": - block["text"] = self._unmask_pii_text(text_value, pii_tokens) - elif mode == "mask": - block["text"] = await self.check_pii( - text=text_value, - output_parse_pii=False, - presidio_config=presidio_config, - request_data=request_data, - ) + block_type = block.get("type") + if block_type == "text": + text_value = block.get("text") + if text_value is None: + continue + if mode == "unmask": + block["text"] = self._unmask_pii_text(text_value, pii_tokens) + elif mode == "mask": + block["text"] = await self.check_pii( + text=text_value, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=request_data, + ) + elif block_type == "tool_use": + # Mirror the OpenAI-format path (_process_response_for_pii), which masks + # tool_call arguments: a tool_use block's `input` dict can carry PII the + # model produced, and must be masked too -- not just sibling text blocks. + tool_input = block.get("input") + if isinstance(tool_input, dict): + for key, value in list(tool_input.items()): + if not isinstance(value, str): + continue + if mode == "unmask": + tool_input[key] = self._unmask_pii_text(value, pii_tokens) + elif mode == "mask": + tool_input[key] = await self.check_pii( + text=value, + output_parse_pii=False, + presidio_config=presidio_config, + request_data=request_data, + ) return response diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio_tool_call_pii_masking.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio_tool_call_pii_masking.py new file mode 100644 index 00000000000..396e5408b90 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio_tool_call_pii_masking.py @@ -0,0 +1,149 @@ +""" +Unit tests for Presidio PII masking of tool-call arguments. + +Covers two channels that were previously unmasked: +1. Pre-call: tool_calls[*].function.arguments and function_call.arguments +2. Anthropic response: tool_use block input values +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from unittest.mock import MagicMock + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.presidio import ( + _OPTIONAL_PresidioPIIMasking, +) +from litellm.types.guardrails import PiiAction, PiiEntityType + + +@pytest.fixture +def presidio_guardrail(): + return _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=False, + pii_entities_config={ + PiiEntityType.US_SSN: PiiAction.MASK, + PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK, + }, + ) + + +@pytest.fixture +def mock_user_api_key(): + return UserAPIKeyAuth(api_key="test_key", user_id="test_user") + + +@pytest.fixture +def mock_cache(): + return MagicMock() + + +@pytest.mark.asyncio +async def test_precall_masks_pii_in_tool_call_arguments( + presidio_guardrail, mock_user_api_key, mock_cache +): + """ + async_pre_call_hook must mask PII in tool_calls[*].function.arguments + and function_call.arguments, not just message content. + """ + test_data = { + "messages": [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "save_record", + "arguments": '{"ssn": "123-45-6789", "name": "Alice"}', + }, + } + ], + }, + { + "role": "assistant", + "content": None, + "function_call": { + "name": "save_record", + "arguments": '{"ssn": "123-45-6789"}', + }, + }, + ], + "model": "gpt-4", + } + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + return text.replace("123-45-6789", "") + + presidio_guardrail.check_pii = mock_check_pii + + result = await presidio_guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key, + cache=mock_cache, + data=test_data, + call_type="completion", + ) + + tc_args = result["messages"][0]["tool_calls"][0]["function"]["arguments"] + assert "123-45-6789" not in tc_args, "tool_calls arguments should be masked" + assert "" in tc_args + + fc_args = result["messages"][1]["function_call"]["arguments"] + assert "123-45-6789" not in fc_args, "function_call arguments should be masked" + assert "" in fc_args + + +@pytest.mark.asyncio +async def test_anthropic_response_masks_pii_in_tool_use_input(presidio_guardrail): + """ + _process_anthropic_response_for_pii must mask PII in tool_use block + input values, not just text blocks. The method takes a raw Anthropic + message dict (type=="message"), not a ModelResponse. + """ + response = { + "id": "msg_01", + "type": "message", + "role": "assistant", + "content": [ + {"type": "text", "text": "Saving SSN 078-05-1120 now."}, + { + "type": "tool_use", + "id": "tu_1", + "name": "save_record", + "input": {"ssn": "078-05-1120", "note": "customer"}, + }, + ], + "model": "claude-sonnet-4-20250514", + "stop_reason": "end_turn", + } + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + return text.replace("078-05-1120", "") + + presidio_guardrail.check_pii = mock_check_pii + + result = await presidio_guardrail._process_anthropic_response_for_pii( + response=response, + request_data={}, + mode="mask", + ) + + text_block = result["content"][0] + tool_block = result["content"][1] + + assert "078-05-1120" not in text_block["text"], "text block should be masked" + assert "" in text_block["text"] + + assert ( + "078-05-1120" not in tool_block["input"]["ssn"] + ), "tool_use input should be masked" + assert tool_block["input"]["ssn"] == "" + assert tool_block["input"]["note"] == "customer", "non-PII values unchanged"