From f4729027fd08504a5cba30043ef4be5febe85cfe Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:22:07 +0200 Subject: [PATCH] fix(guardrails): restore privaite responses on a copy The post-call hooks wrote the restored originals into the objects LiteLLM handed them. LiteLLM keeps those objects, the response for its cache write and success logs and every stream chunk in the wrapper's chunk list, which it reassembles into the cached and logged response at the end of the stream. A cached reply could then carry one caller's values under a key built from the anonymized request, where another caller's identical-looking request would hit them Both hooks now restore a deep copy and return or yield that, so what LiteLLM keeps holds placeholders only. The proxy adopts a guardrail's returned response, so callers still get the restored values --- .../guardrail_hooks/privaite/privaite.py | 26 +++++++++++--- .../guardrail_hooks/test_privaite.py | 34 +++++++++++++++++++ 2 files changed, 55 insertions(+), 5 deletions(-) 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" + )