From 4af91b7b4f2ef78bdab734eacd3eba7c4bc61447 Mon Sep 17 00:00:00 2001 From: crp4222 <121295348+crp4222@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:38:32 +0200 Subject: [PATCH] fix(guardrails): keep privaite-restored tool-call arguments valid JSON, in responses and in streams --- .../guardrail_hooks/privaite/privaite.py | 67 ++++++++++++--- .../guardrail_hooks/test_privaite.py | 83 ++++++++++++++++++- 2 files changed, 138 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py index 25fb5f718af..d2bcb3ce476 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py +++ b/litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py @@ -35,6 +35,7 @@ from __future__ import annotations import asyncio import copy +import json from collections.abc import Callable, Iterator from typing import Any, cast @@ -117,6 +118,43 @@ def _obj_set(obj: Any, key: str, value: Any) -> None: setattr(obj, key, value) +def _json_escape(value: str) -> str: + """The JSON string-literal encoding of value, without the surrounding + quotes: what a restored original must look like when spliced into streamed + argument fragments. Mirrors privaite.streaming.buffer.json_escape.""" + return json.dumps(value, ensure_ascii=False)[1:-1] + + +async def _restore_json_tree(engine: Any, value: Any, mapping: Any) -> Any: + """Restore every string leaf of a parsed JSON document (keys are not + rewritten, mirroring the core engine's walker).""" + if isinstance(value, str): + return await engine.process_response(value, mapping) + if isinstance(value, dict): + return {key: await _restore_json_tree(engine, item, mapping) for key, item in value.items()} + if isinstance(value, list): + return [await _restore_json_tree(engine, item, mapping) for item in value] + return value + + +async def _restore_arguments(engine: Any, arguments: str, mapping: Any) -> str: + """Restore a tool/function call's `arguments`: a JSON document inside a + string. Plain substitution would splice a raw quote, backslash or newline + from the original into a JSON string literal and the client's json.loads of + the arguments would fail; restore on the parsed tree and re-encode instead + (core parity: PIIEngine._deanonymize_arguments). When nothing changes, or + the string is not JSON, the previous behaviour is kept byte for byte.""" + plain = await engine.process_response(arguments, mapping) + if plain == arguments: + return arguments + try: + parsed = json.loads(arguments) + except ValueError: + # Not JSON: the plain string restore is all there is. + return plain + return json.dumps(await _restore_json_tree(engine, parsed, mapping), ensure_ascii=False) + + def _append_nested(delta: object, carrier: str, field: str, remaining: str) -> None: """Append text to delta.., creating the carrier as a plain dict when the delta has none (a bare finish delta carries no function_call @@ -204,15 +242,20 @@ class _StreamRestorer: cloned from the last chunk seen for that choice. A held tail is never dropped.""" - def __init__(self, mapping: object, make_buffer: Callable[[object], object]) -> None: + def __init__(self, mapping: object, json_mapping: object, make_buffer: Callable[..., object]) -> None: self._mapping = mapping + # For the argument channels: their fragments are JSON source text, so the + # original has to arrive JSON-escaped or the client's parse of the + # reassembled arguments fails on a quote, a backslash or a newline. + # (Placeholders need no escaping, so the fakes are the same on both sides.) + self._json_mapping = json_mapping self._make_buffer = make_buffer self._buffers: dict = {} self._last_chunk: dict = {} - def restore(self, key: tuple, text: str, finished: bool) -> str: + def restore(self, key: tuple, text: str, finished: bool, json_fragment: bool = False) -> str: if key not in self._buffers: - self._buffers[key] = self._make_buffer(self._mapping) + self._buffers[key] = self._make_buffer(self._json_mapping if json_fragment else self._mapping) deanon = self._buffers[key] out = deanon.feed(text) if text else "" if finished: @@ -668,12 +711,12 @@ class PrivaiteGuardrail(CustomGuardrail): continue args = getattr(fn, "arguments", None) if args: - fn.arguments = await engine.process_response(args, mapping) + fn.arguments = await _restore_arguments(engine, args, mapping) function_call = getattr(message, "function_call", None) if function_call is not None: fc_args = getattr(function_call, "arguments", None) if fc_args: - function_call.arguments = await engine.process_response(fc_args, mapping) + function_call.arguments = await _restore_arguments(engine, fc_args, mapping) async def _restore_responses_output(self, response: Any, engine: Any, mapping: Any) -> None: """Restore originals in a Responses API result: output_text content @@ -685,7 +728,7 @@ class PrivaiteGuardrail(CustomGuardrail): _obj_set(block, "text", await engine.process_response(text, mapping)) args = _obj_get(item, "arguments") if isinstance(args, str) and args: - _obj_set(item, "arguments", await engine.process_response(args, mapping)) + _obj_set(item, "arguments", await _restore_arguments(engine, args, mapping)) async def async_post_call_success_hook(self, data: dict, user_api_key_dict: Any, response: Any) -> Any: if not self.deanonymize: @@ -753,13 +796,15 @@ class PrivaiteGuardrail(CustomGuardrail): if args: tc_index = getattr(tool_call, "index", 0) or 0 # Pass `finished` so a fragment on the same chunk as finish_reason - # flushes its held-back tail instead of dropping it. - fn.arguments = restore(("tool", index, tc_index), args, finished) + # flushes its held-back tail instead of dropping it. json_fragment: + # the fragment is JSON source text, so the restored original must + # be spliced in JSON-escaped. + fn.arguments = restore(("tool", index, tc_index), args, finished, True) function_call = getattr(delta, "function_call", None) if function_call is not None: fc_args = getattr(function_call, "arguments", None) if fc_args: - function_call.arguments = restore(("fc", index), fc_args, finished) + function_call.arguments = restore(("fc", index), fc_args, finished, True) async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: Any, response: Any, request_data: dict @@ -775,10 +820,12 @@ class PrivaiteGuardrail(CustomGuardrail): from privaite.streaming.buffer import StreamingDeAnonymizer mapping = PIIMapping() + json_mapping = PIIMapping() for fake, original in fakes.items(): mapping.add(original, fake, "PII") + json_mapping.add(_json_escape(original), fake, "PII") - restorer = _StreamRestorer(mapping, StreamingDeAnonymizer) + restorer = _StreamRestorer(mapping, json_mapping, StreamingDeAnonymizer) async for chunk in response: restorer.restore_chunk(chunk, self._restore_delta) yield chunk diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py index acc9877b555..c37cec6bbe1 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py @@ -117,12 +117,13 @@ class _FakeStreamingDeAnonymizer: buffer isolation observable in tests. """ - def __init__(self, _mapping): + def __init__(self, mapping): + self._fakes = mapping.get_all_fakes() self._buf = "" def feed(self, text): self._buf += text - for fake, original in _FAKES.items(): + for fake, original in self._fakes.items(): self._buf = self._buf.replace(fake, original) lt = self._buf.rfind("<") if lt != -1 and ">" not in self._buf[lt:]: @@ -1279,3 +1280,81 @@ async def test_streaming_tails_on_real_litellm_chunks_stay_serializable(): assert [p["choices"][0]["delta"].get("reasoning_content") for p in payloads] == ["then ", "", _SPECIALS_ORIGINAL) + gr = _make_guardrail() + data = {"metadata": {"privaite_map": _FAKES}} + message = types.SimpleNamespace( + content=None, + tool_calls=[ + types.SimpleNamespace(function=types.SimpleNamespace(arguments='{"name": "", "n": 1}')), + # not JSON: the plain string restore is all there is, byte for byte + types.SimpleNamespace(function=types.SimpleNamespace(arguments="name=")), + ], + function_call=types.SimpleNamespace(arguments='["", {"k": ""}]'), + ) + response = types.SimpleNamespace( + choices=[types.SimpleNamespace(message=message)], + output=[{"type": "function_call", "arguments": '{"who": ""}'}], + ) + + out = await gr.async_post_call_success_hook(data, None, response) + + msg = out.choices[0].message + assert json.loads(msg.tool_calls[0].function.arguments) == {"name": _SPECIALS_ORIGINAL, "n": 1} + assert msg.tool_calls[1].function.arguments == f"name={_SPECIALS_ORIGINAL}" + assert json.loads(msg.function_call.arguments) == [_SPECIALS_ORIGINAL, {"k": "Marie Dupont"}] + assert json.loads(out.output[0]["arguments"]) == {"who": _SPECIALS_ORIGINAL} + + +async def test_post_call_untouched_arguments_are_kept_byte_for_byte(): + # No placeholder in the arguments: the string is returned as received, not + # re-encoded (key order, spacing and escapes stay the provider's). + gr = _make_guardrail() + data = {"metadata": {"privaite_map": _FAKES}} + raw = '{"b":1,\n "a": "x\\u00e9"}' + message = types.SimpleNamespace( + content=None, + tool_calls=[types.SimpleNamespace(function=types.SimpleNamespace(arguments=raw))], + function_call=None, + ) + response = types.SimpleNamespace(choices=[types.SimpleNamespace(message=message)]) + out = await gr.async_post_call_success_hook(data, None, response) + assert out.choices[0].message.tool_calls[0].function.arguments == raw + + +async def test_streaming_tool_arguments_stay_valid_json_with_special_chars(monkeypatch): + # Streamed argument fragments are JSON source text: the original spliced in + # must be its JSON-escaped form, so the reassembled arguments still parse. + # Content keeps the raw original. The placeholder is split across chunks. + monkeypatch.setitem(_FAKES, "", _SPECIALS_ORIGINAL) + gr = _make_guardrail() + request_data = {"metadata": {"privaite_map": _FAKES}} + + def _delta(args, content=None): + call = types.SimpleNamespace(index=0, function=types.SimpleNamespace(arguments=args)) + return _bare_delta(content=content, tool_calls=[call], function_call=types.SimpleNamespace(arguments=args)) + + async def _source(): + yield _choice_chunk(_delta('{"name": ""}', content="ON_2>"), finish="stop") + + chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data)) + tool_args = "".join( + tc.function.arguments for chunk in chunks for choice in chunk.choices for tc in (choice.delta.tool_calls or []) + ) + fc_args = "".join(choice.delta.function_call.arguments for chunk in chunks for choice in chunk.choices) + content = "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) + assert json.loads(tool_args) == {"name": _SPECIALS_ORIGINAL} + assert json.loads(fc_args) == {"name": _SPECIALS_ORIGINAL} + assert content == f"Hi {_SPECIALS_ORIGINAL}"