mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(guardrails): keep privaite-restored tool-call arguments valid JSON, in responses and in streams
This commit is contained in:
crp4222
parent
6d216ac068
commit
4af91b7b4f
2 changed files with 138 additions and 12 deletions
|
|
@ -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.<carrier>.<field>, 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
|
||||
|
|
|
|||
|
|
@ -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 ", "<PER"]
|
||||
assert "finish_reason" not in payloads[-1]["choices"][0]
|
||||
assert payloads[-1]["id"] == "chunk"
|
||||
|
||||
|
||||
# An original carrying the characters a JSON string literal must escape.
|
||||
_SPECIALS_ORIGINAL = 'Jean "JJ" O\'Neil\\n\tnext'
|
||||
|
||||
|
||||
async def test_post_call_restores_tool_arguments_as_valid_json(monkeypatch):
|
||||
# A restored original holding a quote, a backslash or a newline must be
|
||||
# spliced JSON-escaped, on every argument carrier: tool_calls, the legacy
|
||||
# function_call and a Responses function_call output item. A plain
|
||||
# substitution would leave the client's json.loads failing on the quote.
|
||||
monkeypatch.setitem(_FAKES, "<PERSON_2>", _SPECIALS_ORIGINAL)
|
||||
gr = _make_guardrail()
|
||||
data = {"metadata": {"privaite_map": _FAKES}}
|
||||
message = types.SimpleNamespace(
|
||||
content=None,
|
||||
tool_calls=[
|
||||
types.SimpleNamespace(function=types.SimpleNamespace(arguments='{"name": "<PERSON_2>", "n": 1}')),
|
||||
# not JSON: the plain string restore is all there is, byte for byte
|
||||
types.SimpleNamespace(function=types.SimpleNamespace(arguments="name=<PERSON_2>")),
|
||||
],
|
||||
function_call=types.SimpleNamespace(arguments='["<PERSON_2>", {"k": "<PERSON_1>"}]'),
|
||||
)
|
||||
response = types.SimpleNamespace(
|
||||
choices=[types.SimpleNamespace(message=message)],
|
||||
output=[{"type": "function_call", "arguments": '{"who": "<PERSON_2>"}'}],
|
||||
)
|
||||
|
||||
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, "<PERSON_2>", _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": "<PERS', content="Hi <PERS"))
|
||||
yield _choice_chunk(_delta('ON_2>"}', 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}"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue