diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index f382bd84925..4ec7c3f0748 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -21,6 +21,7 @@ without it. See https://github.com/crp4222/PrivAiTe. from __future__ import annotations import asyncio +import copy import itertools from collections.abc import AsyncGenerator, AsyncIterable, Iterable, Iterator, Sequence from typing import TYPE_CHECKING, Final, TypeAlias @@ -228,6 +229,13 @@ class PrivaiteGuardrail(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, response: object, ) -> object: + """Restores the original values in a copy of the response. + + LiteLLM keeps the object it hands this hook for its response cache and its + success logs, and a cache write may serialize it after this hook ran. Restoring it + in place would put this caller's values in the cache under a key built from the + anonymized request, where another caller's identical-looking request hits them. + """ if not self.deanonymize: return response pairs: Final = _pop_restore_map(data) @@ -235,12 +243,13 @@ class PrivaiteGuardrail(CustomGuardrail): return response mapping: Final = restore_mapping(pairs) engine: Final = await self._engine_for(self._languages()) - for choice in read_list(response, "choices"): + restored: Final = copy.deepcopy(response) + for choice in read_list(restored, "choices"): message = read_field(choice, "message") if message is not None: await restore_message(message, engine, mapping) - await restore_responses_output(response, engine, mapping) - return response + await restore_responses_output(restored, engine, mapping) + return restored async def async_post_call_streaming_iterator_hook( self, @@ -248,6 +257,12 @@ class PrivaiteGuardrail(CustomGuardrail): response: AsyncIterable[ModelResponseStream], request_data: JsonObject, ) -> AsyncGenerator[ModelResponseStream, None]: + """Restores the original values in copies of the chunks. + + The stream wrapper keeps every chunk it yields and rebuilds the whole response from + them for the cache and the success logs once the stream ends, so the chunks + themselves keep their placeholders, as in the non-streaming hook. + """ pairs: Final = _pop_restore_map(request_data) if not self.deanonymize or not pairs: async for chunk in response: @@ -260,8 +275,9 @@ class PrivaiteGuardrail(CustomGuardrail): restore_buffer_factory(), ) async for chunk in response: - restorer.restore_chunk(chunk) - yield chunk + restored = copy.deepcopy(chunk) + restorer.restore_chunk(restored) + yield restored for drained in restorer.drain(): yield drained diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index 6fb70e02171..5a19ce91fe2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -1424,3 +1424,37 @@ async def test_engine_for_a_privaite_without_pii_engine_recommends_the_install(m with pytest.raises(RuntimeError, match=r"pip install 'privaite>=0\.7\.2'"): await _make_guardrail()._engine_for(["en"]) + + +@pytest.mark.asyncio +async def test_post_call_restores_a_copy_and_leaves_the_response_litellm_caches_and_logs(): + gr = _make_guardrail() + message = types.SimpleNamespace(content="Sending to ", tool_calls=None, function_call=None) + response = types.SimpleNamespace(choices=[types.SimpleNamespace(message=message)]) + + out = await gr.async_post_call_success_hook({"metadata": {"privaite_map": _FAKES}}, None, response) + + assert out.choices[0].message.content == "Sending to Marie Dupont" + assert message.content == "Sending to ", "the object LiteLLM caches and logs must keep the placeholder" + + +@pytest.mark.asyncio +async def test_streaming_restores_copies_and_leaves_the_chunks_litellm_reassembles(): + gr = _make_guardrail() + source = [ + _choice_chunk(_bare_delta(content="Hi "], ( + "the chunks LiteLLM keeps to rebuild the logged and cached response must keep the placeholder" + )