diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3fb09cde9314..5f205df487db 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -178,6 +178,7 @@ ProxyConfig = Any from litellm.proxy.litellm_pre_call_utils import ( add_litellm_data_to_request, + refresh_proxy_server_request_body_snapshot, reject_url_valued_destination, ) from litellm.types.utils import ( @@ -1862,6 +1863,12 @@ async def common_processing_pre_call_logic( call_type=route_type, ) + # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may + # have mutated `self.data` in place, and the audit-trail snapshot taken in + # add_litellm_data_to_request predates that mutation. + refresh_proxy_server_request_body_snapshot(self.data) + verbose_proxy_logger.debug("receiving data: %s", self.data) + if "messages" in self.data and self.data["messages"]: logging_obj.update_messages(self.data["messages"]) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index c3b7498d9ec9..bcee45355e34 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -788,7 +788,7 @@ def run_in_new_loop(): async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: """ - Masks the input before logging to langfuse, datadog, etc. + Masks the input and output before logging to langfuse, datadog, etc. """ if call_type == "completion" or call_type == "acompletion": # /chat/completions requests messages: Final[list | None] = kwargs.get("messages", None) @@ -847,6 +847,19 @@ async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> verbose_proxy_logger.debug("Presidio PII Masking: Redacted pii message: %s", messages) kwargs["messages"] = messages + if ( + isinstance(result, ModelResponse) + and result.choices + and not isinstance(result.choices[0], StreamingChoices) + ): + await self._process_response_for_pii(response=result, request_data=kwargs, mode="mask") + elif self._is_anthropic_message_response(result): + await self._process_anthropic_response_for_pii( + response=cast(dict, result), # cast-ok: _is_anthropic_message_response narrows via isinstance + request_data=kwargs, + mode="mask", + ) + return kwargs, result async def async_post_call_success_hook( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index c10990818671..4525adb82f35 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1622,6 +1622,32 @@ def apply_client_tag_policy_pre_auth( ) +def refresh_proxy_server_request_body_snapshot( + data: dict, # mutable-ok: mutates proxy_server_request.body in place on the shared request dict +) -> None: + """ + Re-snapshot ``data["proxy_server_request"]["body"]`` from the current state of ``data``. + + ``add_litellm_data_to_request`` takes the initial snapshot before guardrails + (pre_call_hook) run. A guardrail that masks PII/PCI in place (e.g. Presidio) + mutates ``data`` afterward, so callers that persist ``proxy_server_request.body`` + for audit/spend-tracking purposes must call this again post-guardrail, or the + persisted body silently bypasses whatever masking the guardrail applied. + + By the time a caller refreshes post-guardrail, ``litellm.utils.function_setup`` + has already stamped ``data["litellm_logging_obj"]`` with a live (non-serializable) + ``Logging`` instance, so it must be excluded here the same way ``secret_fields`` + and ``proxy_server_request`` are. + """ + proxy_server_request = data.get("proxy_server_request") + if not isinstance(proxy_server_request, dict): + return + _body_snapshot_exclude = ( + frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj"}) | _TRANSPORT_ONLY_CREDENTIAL_KEYS + ) + proxy_server_request["body"] = {k: v for k, v in data.items() if k not in _body_snapshot_exclude} + + async def add_litellm_data_to_request( data: dict, request: Request, @@ -1802,8 +1828,6 @@ async def add_litellm_data_to_request( cache_dict: Final = parse_cache_control(cache_control_header) data["ttl"] = cache_dict.get("s-maxage") - verbose_proxy_logger.debug("receiving data: %s", data) - # requester_metadata is snapshotted AFTER the strip below so # downstream consumers (e.g. PANW guardrail reading user_ip / # profile_id) don't see attacker-injected admin slots preserved in @@ -1863,9 +1887,7 @@ async def add_litellm_data_to_request( # self-reference — body.proxy_server_request.body would be the same # dict as body, producing an infinite traversal loop for any consumer # that walks the structure. - _body_snapshot_exclude = frozenset({"secret_fields", "proxy_server_request"}) | _TRANSPORT_ONLY_CREDENTIAL_KEYS - _body_snapshot: Final = {k: v for k, v in data.items() if k not in _body_snapshot_exclude} - data["proxy_server_request"]["body"] = _body_snapshot + refresh_proxy_server_request_body_snapshot(data) # Snapshot the requester-supplied metadata for downstream consumers. # Taking the deepcopy after the user_api_key_* / _pipeline_managed_guardrails diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 60be3be5e8b5..acb43bc5b744 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -582,6 +582,52 @@ async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): print("✓ Logging hook multiple content items test passed") +@pytest.mark.asyncio +async def test_logging_hook_masks_the_response_too(presidio_guardrail): + """ + Regression: async_logging_hook only masked kwargs["messages"] (the request) and + left `result` (the model's response) completely untouched, so in `logging_only` + mode any PII in the assistant's reply was logged to langfuse/datadog/etc. in the + clear. The hook's own docstring promises masking "before logging" for both input + and output. + """ + + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): + return text.replace("4111-1111-1111-1111", "[CREDIT_CARD]") + + presidio_guardrail.check_pii = mock_check_pii + + test_kwargs = { + "messages": [{"role": "user", "content": "hello"}], + "model": "gpt-4", + } + response = ModelResponse( + id="1", + object="chat.completion", + created=0, + model="gpt-test", + choices=[ + Choices( + message=Message( + role="assistant", + content="Sure, your card is 4111-1111-1111-1111", + ), + index=0, + finish_reason="stop", + ) + ], + ) + + _, result_response = await presidio_guardrail.async_logging_hook( + kwargs=test_kwargs, + result=response, + call_type="completion", + ) + + assert "[CREDIT_CARD]" in result_response.choices[0].message.content + assert "4111-1111-1111-1111" not in result_response.choices[0].message.content + + @pytest.mark.asyncio async def test_logging_only_does_not_mask_pre_call_request( mock_user_api_key, mock_cache diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 3c738aa164c9..f9ba91a246ec 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -323,6 +323,62 @@ async def mock_common_processing_pre_call_logic(user_api_key_dict, data, call_ty pytest.fail("litellm_call_id is not a valid UUID") assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"] + @pytest.mark.asyncio + async def test_common_processing_pre_call_logic_refreshes_proxy_server_request_body_after_guardrails( + self, monkeypatch + ): + """ + A guardrail (e.g. Presidio PII masking) mutates data["messages"] in place inside + pre_call_hook. The proxy_server_request.body snapshot is taken before that hook + runs, so it must be refreshed afterward or SpendLogs (when store_prompts_in_spend_logs + is enabled) persists the raw pre-guardrail body, bypassing the masking entirely. + """ + processing_obj = ProxyBaseLLMRequestProcessing(data={}) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + raw_messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return { + "messages": raw_messages, + "proxy_server_request": { + "url": "http://testserver/chat/completions", + "method": "POST", + "body": {"messages": raw_messages}, + }, + } + + async def mock_pre_call_hook(user_api_key_dict, data, call_type): + data["messages"] = [{"role": "user", "content": "my ssn is "}] + return data + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=MagicMock(spec=ProxyConfig), + route_type="acompletion", + ) + + persisted_body = returned_data["proxy_server_request"]["body"] + assert persisted_body["messages"] == returned_data["messages"] + assert "123-45-6789" not in json.dumps(persisted_body["messages"]) + # litellm_logging_obj is stamped onto `data` by function_setup between the + # initial snapshot and pre_call_hook; it must never leak into the persisted + # audit body, which needs to stay plain-JSON-serializable end to end. + assert "litellm_logging_obj" not in persisted_body + json.dumps(persisted_body) + def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch): mock_set_active_span_tag = MagicMock(return_value=True) import litellm.proxy.dd_span_tagger diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 111f11f85bad..501b03eae0fa 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -710,6 +710,54 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r ) +def test_refresh_proxy_server_request_body_snapshot_picks_up_guardrail_masking(): + """ + Regression: proxy_server_request['body'] is snapshotted by + add_litellm_data_to_request BEFORE guardrails (e.g. Presidio PII masking) run + in pre_call_hook. Without a refresh after pre_call_hook, the persisted body + silently bypasses whatever masking the guardrail applied, so raw PII/PCI + lands in SpendLogs when store_prompts_in_spend_logs is enabled. + """ + from litellm.proxy.litellm_pre_call_utils import ( + refresh_proxy_server_request_body_snapshot, + ) + + class _FakeLoggingObj: + """Stands in for the live, non-JSON-serializable Logging instance that + litellm.utils.function_setup stamps onto `data` between the initial + snapshot and pre_call_hook.""" + + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}], + "secret_fields": {"raw_headers": {"authorization": "Bearer sk-secret"}}, + "litellm_logging_obj": _FakeLoggingObj(), + "proxy_server_request": { + "url": "http://localhost/v1/chat/completions", + "method": "POST", + "body": { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}], + }, + }, + } + + # Simulate a PII-masking guardrail mutating `messages` in place, like Presidio's + # async_pre_call_hook does, after the initial snapshot was already taken. + data["messages"] = [{"role": "user", "content": "my ssn is "}] + + refresh_proxy_server_request_body_snapshot(data) + + refreshed_body = data["proxy_server_request"]["body"] + assert refreshed_body["messages"] == data["messages"] + # Still excludes secrets, self-reference, and the live logging object, same as + # the initial snapshot -- and proves the persisted body stays JSON-serializable. + assert "secret_fields" not in refreshed_body + assert "proxy_server_request" not in refreshed_body + assert "litellm_logging_obj" not in refreshed_body + assert "123-45-6789" not in json.dumps(refreshed_body) + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_strips_string_encoded_admin_injection(): """Regression: metadata arriving as a JSON string (multipart/form-data or