mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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.
This commit is contained in:
AUTHENSOR
John Kearney
parent
343e453c2a
commit
c208335a3d
2 changed files with 215 additions and 13 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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", "<US_SSN>")
|
||||
|
||||
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 "<US_SSN>" 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 "<US_SSN>" 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", "<US_SSN>")
|
||||
|
||||
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 "<US_SSN>" 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"] == "<US_SSN>"
|
||||
assert tool_block["input"]["note"] == "customer", "non-PII values unchanged"
|
||||
Loading…
Add table
Reference in a new issue