mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
fix(presidio): restore masked PII in streamed /v1/messages responses
With output_parse_pii, streamed Anthropic /v1/messages responses came back with the <ENTITY_N> tokens still in them, for two reasons The restore callback was routed through unified_guardrail, whose block_only streaming mode drops text rewrites. It now sets mask_response_content so Presidio's own streaming hook runs The SSE unmask also worked one transport chunk at a time, so a token split across deltas (<EMAIL_ADDRESS_ then 2>) leaked. AnthropicSSEUnmasker now reassembles events across chunks, holds back a partial trailing token per content block, and flushes it before content_block_stop. It also restores input_json_delta tool arguments and leaves signed thinking_delta untouched
This commit is contained in:
loitragg
parent
254c2f6b3b
commit
f10ed0f651
4 changed files with 394 additions and 2 deletions
|
|
@ -48,6 +48,7 @@ from litellm.proxy.guardrails.anthropic_sse import (
|
|||
is_anthropic_sse_stream,
|
||||
model_response_text,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.presidio_anthropic_sse import AnthropicSSEUnmasker
|
||||
from litellm.types.guardrails import (
|
||||
GuardrailEventHooks,
|
||||
LitellmParams,
|
||||
|
|
@ -1588,6 +1589,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
|||
|
||||
remaining_chunks: list[ModelResponseStream] = []
|
||||
saw_non_chat_chunk = False
|
||||
sse_unmasker: Final = AnthropicSSEUnmasker(pii_tokens) if pii_tokens else None
|
||||
try:
|
||||
async for chunk in response:
|
||||
if isinstance(chunk, ModelResponseStream):
|
||||
|
|
@ -1596,8 +1598,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
|||
else:
|
||||
remaining_chunks.append(chunk)
|
||||
elif isinstance(chunk, bytes):
|
||||
if pii_tokens:
|
||||
yield self._unmask_sse_bytes_chunk(chunk, pii_tokens)
|
||||
if sse_unmasker is not None:
|
||||
if unmasked_bytes := sse_unmasker.feed(chunk):
|
||||
yield unmasked_bytes
|
||||
else:
|
||||
yield chunk
|
||||
continue
|
||||
|
|
@ -1615,6 +1618,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
|||
saw_non_chat_chunk = True
|
||||
yield chunk
|
||||
|
||||
if sse_unmasker is not None and (tail := sse_unmasker.flush()):
|
||||
yield tail
|
||||
|
||||
if saw_non_chat_chunk:
|
||||
return
|
||||
|
||||
|
|
@ -1645,6 +1651,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
|||
verbose_proxy_logger.error("Error in PII streaming processing: %s", e)
|
||||
for chunk in remaining_chunks:
|
||||
yield chunk
|
||||
if sse_unmasker is not None and (tail := sse_unmasker.flush()):
|
||||
yield tail
|
||||
|
||||
async def async_post_call_streaming_iterator_hook(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,158 @@
|
|||
"""
|
||||
Stateful Presidio PII unmasking for Anthropic ``/v1/messages`` SSE streams.
|
||||
|
||||
Claude splits tokens like ``<EMAIL_ADDRESS_2>`` across deltas, and byte chunks
|
||||
need not end on an event boundary, so unmasking each chunk on its own misses
|
||||
them. This reassembles events, holds back a trailing partial token per content
|
||||
block until it completes, and flushes held text before ``content_block_stop``.
|
||||
|
||||
``thinking_delta`` is left alone: thinking blocks are signed, and rewriting
|
||||
them would invalidate the signature when the client sends them back.
|
||||
"""
|
||||
|
||||
import codecs
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from functools import reduce
|
||||
from typing import Final
|
||||
|
||||
_EVENT_SPLIT: Final = re.compile(r"(\r?\n\r?\n)")
|
||||
|
||||
# delta.type -> (field holding the text, whether values land inside a JSON string)
|
||||
_UNMASKABLE_DELTAS: Final[Mapping[str, tuple[str, bool]]] = {
|
||||
"text_delta": ("text", False),
|
||||
"input_json_delta": ("partial_json", True),
|
||||
}
|
||||
|
||||
_FLUSH_ALL_EVENTS: Final = frozenset({"message_delta", "message_stop", "error"})
|
||||
|
||||
_MAX_BUFFERED_CHARS: Final = 1_000_000
|
||||
|
||||
|
||||
class AnthropicSSEUnmasker:
|
||||
"""Feed raw SSE byte chunks in order, then call ``flush`` once at end of stream."""
|
||||
|
||||
def __init__(self, pii_tokens: Mapping[str, str]) -> None:
|
||||
self._tokens: Final = dict(pii_tokens)
|
||||
self._max_token_len: Final = max((len(t) for t in self._tokens), default=0)
|
||||
self._decoder: Final = codecs.getincrementaldecoder("utf-8")(errors="surrogateescape")
|
||||
self._buffer: str = ""
|
||||
self._held: dict[int, tuple[str, str]] = {} # mutable-ok: per-block stream state, index -> (delta type, text)
|
||||
|
||||
def feed(self, chunk: bytes) -> bytes:
|
||||
pieces: Final = _EVENT_SPLIT.split(self._buffer + self._decoder.decode(chunk))
|
||||
rest: Final = pieces[-1]
|
||||
processed: Final = "".join(
|
||||
tuple(self._process_event(event, terminator) for event, terminator in zip(pieces[:-1:2], pieces[1::2]))
|
||||
)
|
||||
overflow: Final = len(rest) >= _MAX_BUFFERED_CHARS
|
||||
self._buffer = "" if overflow else rest
|
||||
return _encode(processed + (rest if overflow else ""))
|
||||
|
||||
def flush(self) -> bytes:
|
||||
held: Final = self._flush_all_held()
|
||||
tail: Final = self._buffer + self._decoder.decode(b"", final=True)
|
||||
self._buffer = ""
|
||||
return _encode(held + tail)
|
||||
|
||||
def _process_event(self, event_text: str, terminator: str) -> str:
|
||||
original: Final = event_text + terminator
|
||||
lines: Final = tuple(event_text.split("\n"))
|
||||
data_indexes: Final = tuple(i for i, line in enumerate(lines) if line.startswith("data:"))
|
||||
if len(data_indexes) != 1:
|
||||
return original
|
||||
data_index: Final = data_indexes[0]
|
||||
line_ending: Final = "\r" if lines[data_index].endswith("\r") else ""
|
||||
try:
|
||||
event: Final[object] = json.loads(lines[data_index][5:].rstrip("\r").removeprefix(" "))
|
||||
except ValueError:
|
||||
return original
|
||||
if not isinstance(event, dict):
|
||||
return original
|
||||
|
||||
match event.get("type"):
|
||||
case "content_block_stop" if isinstance(event.get("index"), int):
|
||||
return self._flush_held(event["index"]) + original
|
||||
case event_type if event_type in _FLUSH_ALL_EVENTS:
|
||||
return self._flush_all_held() + original
|
||||
case "content_block_delta":
|
||||
rewritten: Final = self._unmask_delta(event)
|
||||
if rewritten is None:
|
||||
return original
|
||||
data_line: Final = "data: " + json.dumps(rewritten, ensure_ascii=False) + line_ending
|
||||
return "\n".join((*lines[:data_index], data_line, *lines[data_index + 1 :])) + terminator
|
||||
case _:
|
||||
return original
|
||||
|
||||
def _unmask_delta(self, event: dict[str, object]) -> dict[str, object] | None:
|
||||
"""Return the rewritten event, or None to pass the original through untouched."""
|
||||
index: Final = event.get("index")
|
||||
delta: Final = event.get("delta")
|
||||
if not isinstance(index, int) or not isinstance(delta, dict):
|
||||
return None
|
||||
delta_type: Final = delta.get("type")
|
||||
spec: Final = _UNMASKABLE_DELTAS.get(delta_type) if isinstance(delta_type, str) else None
|
||||
if spec is None or not isinstance(delta.get(spec[0]), str):
|
||||
return None
|
||||
field, escape = spec
|
||||
text: Final[str] = delta[field]
|
||||
|
||||
_, previously_held = self._held.pop(index, ("", ""))
|
||||
emit, hold = self._split_partial_token(previously_held + text)
|
||||
if hold:
|
||||
self._held[index] = (delta_type, hold)
|
||||
unmasked: Final = self._replace_tokens(emit, escape)
|
||||
if unmasked == text:
|
||||
return None
|
||||
return {**event, "delta": {**delta, field: unmasked}}
|
||||
|
||||
def _split_partial_token(self, text: str) -> tuple[str, str]:
|
||||
"""Split off a trailing prefix of a known token, e.g. ``"<EMAIL_ADD"``."""
|
||||
start: Final = text.rfind("<")
|
||||
suffix: Final = text[start:] if start != -1 else ""
|
||||
is_partial_token: Final = (
|
||||
bool(suffix)
|
||||
and ">" not in suffix
|
||||
and len(suffix) < self._max_token_len
|
||||
and any(token.startswith(suffix) for token in self._tokens)
|
||||
)
|
||||
return (text[:start], suffix) if is_partial_token else (text, "")
|
||||
|
||||
def _replace_tokens(self, text: str, escape: bool) -> str:
|
||||
return reduce(
|
||||
lambda acc, item: acc.replace(item[0], _render(item[1], escape)),
|
||||
self._tokens.items(),
|
||||
text,
|
||||
)
|
||||
|
||||
def _flush_all_held(self) -> str:
|
||||
return "".join(tuple(self._flush_held(index) for index in sorted(self._held)))
|
||||
|
||||
def _flush_held(self, index: int) -> str:
|
||||
held: Final = self._held.pop(index, None)
|
||||
if held is None:
|
||||
return ""
|
||||
delta_type, text = held
|
||||
field, escape = _UNMASKABLE_DELTAS[delta_type]
|
||||
# Block ended mid-token (e.g. max_tokens): same truncation rule as
|
||||
# _OPTIONAL_PresidioPIIMasking._unmask_pii_text.
|
||||
truncated_value: Final = next(
|
||||
(
|
||||
value
|
||||
for token, value in self._tokens.items()
|
||||
if token.startswith(text) and len(text) >= min(20, len(token) // 2)
|
||||
),
|
||||
None,
|
||||
)
|
||||
final_text: Final = text if truncated_value is None else _render(truncated_value, escape)
|
||||
event: Final = {"type": "content_block_delta", "index": index, "delta": {"type": delta_type, field: final_text}}
|
||||
return "event: content_block_delta\ndata: " + json.dumps(event, ensure_ascii=False) + "\n\n"
|
||||
|
||||
|
||||
def _render(value: str, escape: bool) -> str:
|
||||
return json.dumps(value, ensure_ascii=False)[1:-1] if escape else value
|
||||
|
||||
|
||||
def _encode(text: str) -> bytes:
|
||||
return text.encode("utf-8", errors="surrogateescape")
|
||||
|
|
@ -176,6 +176,8 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail) ->
|
|||
output_parse_pii=True,
|
||||
event_hook=_presidio_output_mode(litellm_params.mode, include_mcp=True),
|
||||
_callback_role="restore",
|
||||
# Forces its own streaming hook: unified_guardrail's block_only mode drops the unmask on /v1/messages
|
||||
mask_response_content=True,
|
||||
)
|
||||
if run_input and litellm_params.output_parse_pii
|
||||
else None
|
||||
|
|
|
|||
|
|
@ -0,0 +1,224 @@
|
|||
"""Tests for stateful Presidio unmasking of Anthropic /v1/messages SSE streams."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking
|
||||
from litellm.proxy.guardrails.guardrail_hooks.presidio_anthropic_sse import AnthropicSSEUnmasker
|
||||
|
||||
TOKENS = {"<PERSON_1>": "Maria Gonzalez", "<EMAIL_ADDRESS_2>": "maria@example.com"}
|
||||
|
||||
|
||||
def _event(payload: dict) -> bytes:
|
||||
return f"event: {payload['type']}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
def _text_delta(text: str, index: int = 0) -> bytes:
|
||||
return _event({"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}})
|
||||
|
||||
|
||||
def _json_delta(partial: str, index: int = 1) -> bytes:
|
||||
return _event(
|
||||
{"type": "content_block_delta", "index": index, "delta": {"type": "input_json_delta", "partial_json": partial}}
|
||||
)
|
||||
|
||||
|
||||
def _stop(index: int = 0) -> bytes:
|
||||
return _event({"type": "content_block_stop", "index": index})
|
||||
|
||||
|
||||
def _run(chunks: list[bytes], tokens: dict[str, str] = TOKENS) -> bytes:
|
||||
unmasker = AnthropicSSEUnmasker(tokens)
|
||||
return b"".join(unmasker.feed(c) for c in chunks) + unmasker.flush()
|
||||
|
||||
|
||||
def _deltas(raw: bytes, field: str = "text", index: int | None = None) -> str:
|
||||
"""Concatenate the given delta field across every event in an SSE stream."""
|
||||
out = []
|
||||
for line in raw.decode().splitlines():
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
event = json.loads(line[6:])
|
||||
if event.get("type") == "content_block_delta" and (index is None or event["index"] == index):
|
||||
out.append(event["delta"].get(field, ""))
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def test_whole_token_in_one_delta():
|
||||
raw = _run([_text_delta("Hello <PERSON_1>!"), _stop()])
|
||||
assert _deltas(raw) == "Hello Maria Gonzalez!"
|
||||
|
||||
|
||||
def test_token_split_across_deltas():
|
||||
# Observed live from claude-sonnet-4-6: " and `<EMAIL_ADDRESS_" + "2>` but"
|
||||
raw = _run([_text_delta("Hi <PER"), _text_delta("SON_1>, mail <EMAIL_ADDRESS_"), _text_delta("2> now"), _stop()])
|
||||
assert _deltas(raw) == "Hi Maria Gonzalez, mail maria@example.com now"
|
||||
assert "<" not in _deltas(raw)
|
||||
|
||||
|
||||
def test_token_split_one_char_per_delta():
|
||||
token = "<EMAIL_ADDRESS_2>"
|
||||
raw = _run([_text_delta("x ")] + [_text_delta(c) for c in token] + [_text_delta(" y"), _stop()])
|
||||
assert _deltas(raw) == "x maria@example.com y"
|
||||
|
||||
|
||||
def test_held_text_is_never_emitted_early():
|
||||
unmasker = AnthropicSSEUnmasker(TOKENS)
|
||||
first = unmasker.feed(_text_delta("Dear <PERSON_"))
|
||||
assert _deltas(first) == "Dear "
|
||||
assert "PERSON" not in first.decode()
|
||||
|
||||
|
||||
def test_lone_angle_bracket_that_is_not_a_token_is_released():
|
||||
raw = _run([_text_delta("if a <"), _text_delta(" b then"), _stop()])
|
||||
assert _deltas(raw) == "if a < b then"
|
||||
|
||||
|
||||
def test_unknown_tag_is_not_held():
|
||||
unmasker = AnthropicSSEUnmasker(TOKENS)
|
||||
assert _deltas(unmasker.feed(_text_delta("use <div"))) == "use <div"
|
||||
|
||||
|
||||
def test_held_text_flushed_before_block_stop():
|
||||
raw = _run([_text_delta("cut <PE"), _stop()])
|
||||
lines = [line for line in raw.decode().splitlines() if line.startswith("data: ")]
|
||||
types = [json.loads(line[6:])["type"] for line in lines]
|
||||
assert types == ["content_block_delta", "content_block_delta", "content_block_stop"]
|
||||
# Too short to be unambiguous, so the prefix is released as-is.
|
||||
assert _deltas(raw) == "cut <PE"
|
||||
|
||||
|
||||
def test_truncated_token_at_block_end_restores_value():
|
||||
# max_tokens cut the token off; long enough prefix -> restore, matching
|
||||
# _OPTIONAL_PresidioPIIMasking._unmask_pii_text's truncation fallback.
|
||||
raw = _run([_text_delta("mail <EMAIL_ADDRESS"), _stop()])
|
||||
assert _deltas(raw) == "mail maria@example.com"
|
||||
|
||||
|
||||
def test_held_text_flushed_at_end_of_stream_without_stop():
|
||||
raw = _run([_text_delta("x <PE")])
|
||||
assert _deltas(raw) == "x <PE"
|
||||
|
||||
|
||||
def test_message_delta_flushes_all_blocks():
|
||||
unmasker = AnthropicSSEUnmasker(TOKENS)
|
||||
unmasker.feed(_text_delta("a <PER", index=0))
|
||||
out = unmasker.feed(_event({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}))
|
||||
assert _deltas(out) == "<PER"
|
||||
|
||||
|
||||
def test_blocks_are_tracked_independently():
|
||||
raw = _run(
|
||||
[
|
||||
_text_delta("A <PER", index=0),
|
||||
_text_delta("B <EMAIL_", index=2),
|
||||
_text_delta("SON_1>", index=0),
|
||||
_text_delta("ADDRESS_2>", index=2),
|
||||
_stop(0),
|
||||
_stop(2),
|
||||
]
|
||||
)
|
||||
assert _deltas(raw, index=0) == "A Maria Gonzalez"
|
||||
assert _deltas(raw, index=2) == "B maria@example.com"
|
||||
|
||||
|
||||
def test_input_json_delta_split_token_is_json_escaped():
|
||||
tokens = {"<PERSON_1>": 'Bob "B" \\ Smith'}
|
||||
raw = _run([_json_delta('{"to": "<PER'), _json_delta('SON_1>"}'), _stop(1)], tokens)
|
||||
assert json.loads(_deltas(raw, field="partial_json")) == {"to": 'Bob "B" \\ Smith'}
|
||||
|
||||
|
||||
def test_thinking_delta_untouched():
|
||||
chunk = _event(
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "<PERSON_1>"}}
|
||||
)
|
||||
assert _run([chunk]) == chunk
|
||||
|
||||
|
||||
def test_unchanged_events_pass_through_byte_for_byte():
|
||||
chunks = [
|
||||
_event({"type": "message_start", "message": {"id": "msg_1"}}),
|
||||
b'event: ping\ndata: {"type": "ping"}\n\n',
|
||||
_text_delta("no pii here"),
|
||||
_stop(),
|
||||
_event({"type": "message_stop"}),
|
||||
]
|
||||
assert _run(chunks) == b"".join(chunks)
|
||||
|
||||
|
||||
def test_event_split_across_byte_chunks():
|
||||
whole = _text_delta("Hello <PERSON_1>") + _stop()
|
||||
pieces = [whole[i : i + 7] for i in range(0, len(whole), 7)]
|
||||
assert _deltas(_run(pieces)) == "Hello Maria Gonzalez"
|
||||
|
||||
|
||||
def test_multibyte_character_split_across_byte_chunks():
|
||||
tokens = {"<PERSON_1>": "José"}
|
||||
event = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "¡Hola <PERSON_1>! ñ"}}
|
||||
whole = f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + _stop()
|
||||
cut = whole.index("ñ".encode()) + 1 # split inside the 2-byte ñ
|
||||
raw = _run([whole[:cut], whole[cut:]], tokens)
|
||||
assert _deltas(raw) == "¡Hola José! ñ"
|
||||
assert "Jos\\u" not in raw.decode()
|
||||
|
||||
|
||||
def test_crlf_framing():
|
||||
event = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi <PERSON_1>"}}
|
||||
chunk = ("event: content_block_delta\r\ndata: " + json.dumps(event) + "\r\n\r\n").encode()
|
||||
raw = _run([chunk])
|
||||
assert raw.endswith(b"\r\n\r\n")
|
||||
assert json.loads(raw.decode().split("data: ", 1)[1].strip())["delta"]["text"] == "Hi Maria Gonzalez"
|
||||
|
||||
|
||||
def test_malformed_and_non_sse_bytes_pass_through():
|
||||
chunks = [b"data: {not json}\n\n", b"\xff\xfe raw tail"]
|
||||
assert _run(chunks) == b"".join(chunks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_pii_unmasking_handles_split_tokens():
|
||||
guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True)
|
||||
request_data = {"metadata": {"pii_tokens": dict(TOKENS)}}
|
||||
|
||||
async def stream():
|
||||
yield _text_delta("Dear <PERSON")
|
||||
yield _text_delta("_1>, see <EMAIL_ADDRESS_2")
|
||||
yield _text_delta(">.")
|
||||
yield _stop()
|
||||
|
||||
out = b"".join([c async for c in guardrail._stream_pii_unmasking(stream(), request_data)])
|
||||
assert _deltas(out) == "Dear Maria Gonzalez, see maria@example.com."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_pii_unmasking_flushes_when_stream_ends_mid_token():
|
||||
guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True)
|
||||
request_data = {"metadata": {"pii_tokens": dict(TOKENS)}}
|
||||
|
||||
async def stream():
|
||||
yield _text_delta("x <PE")
|
||||
|
||||
out = b"".join([c async for c in guardrail._stream_pii_unmasking(stream(), request_data)])
|
||||
assert _deltas(out) == "x <PE"
|
||||
|
||||
|
||||
def test_unmask_output_callback_runs_native_streaming_hook(monkeypatch):
|
||||
"""ProxyLogging only runs a guardrail's own streaming hook on /v1/messages when
|
||||
mask_response_content is set; otherwise unified_guardrail drops the unmask."""
|
||||
import litellm
|
||||
from litellm.proxy.guardrails.guardrail_initializers import initialize_presidio
|
||||
from litellm.types.guardrails import LitellmParams
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
params = LitellmParams(
|
||||
guardrail="presidio",
|
||||
mode="pre_call",
|
||||
presidio_filter_scope="input",
|
||||
output_parse_pii=True,
|
||||
presidio_analyzer_api_base="http://analyzer",
|
||||
presidio_anonymizer_api_base="http://anonymizer",
|
||||
)
|
||||
input_cb, unmask_cb = initialize_presidio(params, {"guardrail_name": "g"})
|
||||
assert unmask_cb.mask_response_content is True
|
||||
assert input_cb.mask_response_content is False
|
||||
Loading…
Add table
Reference in a new issue