fix(guardrails): flush privaite held tails onto a bare finish chunk and drain them at end of stream

litellm/proxy/guardrails/guardrail_hooks/privaite/privaite.py

@@ -34,7 +34,8 @@ See https://github.com/crp4222/PrivAiTe for the package.
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
+import copy
+from collections.abc import Callable, Iterator
 from typing import Any, cast
 
 from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -116,6 +117,144 @@ def _obj_set(obj: Any, key: str, value: Any) -> None:
         setattr(obj, key, value)
 
 
+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
+    or audio object)."""
+    holder = _obj_get(delta, carrier)
+    if holder is None:
+        _obj_set(delta, carrier, {field: remaining})
+    else:
+        _obj_set(holder, field, (_obj_get(holder, field) or "") + remaining)
+
+
+def _append_tool_remainder(delta: object, tc_index: int, remaining: str) -> None:
+    """Append a held argument fragment under its tool_call index, creating the
+    slot when the delta has none: the client reassembles arguments per index."""
+    calls = _obj_get(delta, "tool_calls")
+    if calls is None:
+        calls = []
+        _obj_set(delta, "tool_calls", calls)
+    for call in calls:
+        if (_obj_get(call, "index") or 0) == tc_index:
+            _append_nested(call, "function", "arguments", remaining)
+            return
+    calls.append({"index": tc_index, "function": {"arguments": remaining}})
+
+
+def _append_remainder(delta: object, key: tuple, remaining: str) -> None:
+    """Append a restore buffer's held tail onto a streamed delta. The key is
+    (kind, choice index[, tool_call index]) as built by the streaming hook."""
+    kind = key[0]
+    if kind == "tool":
+        _append_tool_remainder(delta, key[2], remaining)
+    elif kind == "fc":
+        _append_nested(delta, "function_call", "arguments", remaining)
+    elif kind == "audio":
+        _append_nested(delta, "audio", "transcript", remaining)
+    else:
+        # content and the text fields (reasoning_content, reasoning, refusal)
+        _obj_set(delta, kind, (_obj_get(delta, kind) or "") + remaining)
+
+
+_DELTA_CHANNELS = (
+    "content",
+    "reasoning_content",
+    "reasoning",
+    "refusal",
+    "audio",
+    "tool_calls",
+    "function_call",
+)
+
+
+def _drain_chunk(template: object, key: tuple, remaining: str) -> object:
+    """A trailing chunk carrying one held tail after the provider closed the
+    stream without a finish_reason: a clone of the last chunk of that choice
+    (same id, model and object as the rest of the stream) reduced to that one
+    choice, its delta cleared, no finish_reason and no usage."""
+    chunk = copy.deepcopy(template)
+    index = key[1]
+    choice = next(c for c in _obj_get(chunk, "choices") or [] if (_obj_get(c, "index") or 0) == index)
+    _obj_set(chunk, "choices", [choice])
+    if _obj_get(chunk, "usage") is not None:
+        _obj_set(chunk, "usage", None)
+    _obj_set(choice, "finish_reason", None)
+    delta = _obj_get(choice, "delta")
+    for field in _DELTA_CHANNELS:
+        if _obj_get(delta, field) is not None:
+            _obj_set(delta, field, None)
+    _append_remainder(delta, key, remaining)
+    return chunk
+
+
+class _StreamRestorer:
+    """Restore state for one streamed response.
+
+    One de-anonymizer buffer per streamed segment, keyed by (kind, choice
+    index, ...): with n>1 the provider interleaves chunks for different
+    choices, and tool-call arguments stream as fragments per tool_call index;
+    each segment keeps its own boundary buffer so a placeholder split across
+    chunks reassembles without mixing segments. The buffer holds back partial
+    placeholders that span chunk boundaries; whatever remains is flushed onto
+    the chunk that carries finish_reason (by the channel's own restore when
+    the finish delta carries it, by the sweep in restore_chunk when it does
+    not: the common finish shape is a bare `delta: {}`), or, when the provider
+    closes the stream without a finish_reason, drained as trailing chunks
+    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:
+        self._mapping = mapping
+        self._make_buffer = make_buffer
+        self._buffers: dict = {}
+        self._last_chunk: dict = {}
+
+    def restore(self, key: tuple, text: str, finished: bool) -> str:
+        if key not in self._buffers:
+            self._buffers[key] = self._make_buffer(self._mapping)
+        deanon = self._buffers[key]
+        out = deanon.feed(text) if text else ""
+        if finished:
+            # The channel is closed: drop its buffer so a sweep never flushes
+            # it a second time.
+            out += self._buffers.pop(key).flush()
+        return out
+
+    def close_open(self, index: int | None) -> tuple[tuple[tuple, str], ...]:
+        """Flush and drop every buffer still open for this choice (every choice
+        when index is None); returns the non-empty tails."""
+        tails = []
+        for key in [k for k in self._buffers if index is None or k[1] == index]:
+            remaining = self._buffers.pop(key).flush()
+            if remaining:
+                tails.append((key, remaining))
+        return tuple(tails)
+
+    def restore_chunk(self, chunk: object, restore_delta: Callable[..., None]) -> None:
+        """Restore every choice of one chunk in place; on a finished choice,
+        append the tails of the channels its finish delta does not carry."""
+        for choice in getattr(chunk, "choices", None) or []:
+            delta = getattr(choice, "delta", None)
+            if delta is None:
+                continue
+            index = getattr(choice, "index", 0) or 0
+            self._last_chunk[index] = chunk
+            finished = getattr(choice, "finish_reason", None) is not None
+            restore_delta(delta, index, finished, self.restore)
+            if finished:
+                for key, remaining in self.close_open(index):
+                    _append_remainder(delta, key, remaining)
+
+    def drain(self) -> Iterator[object]:
+        """Trailing chunks for the tails still held after the provider closed
+        the stream without a finish_reason for their choice."""
+        for key, remaining in self.close_open(None):
+            template = self._last_chunk.get(key[1])
+            if template is not None:
+                yield _drain_chunk(template, key, remaining)
+
+
 class PrivaiteGuardrail(CustomGuardrail):
     def __init__(self, **kwargs: Any) -> None:
         self.preset = kwargs.pop("preset", None) or "onnx"
@@ -639,33 +778,12 @@ class PrivaiteGuardrail(CustomGuardrail):
         for fake, original in fakes.items():
             mapping.add(original, fake, "PII")
 
-        # One de-anonymizer buffer per streamed segment, keyed by (kind, choice
-        # index, ...). With n>1 the provider interleaves chunks for different
-        # choices, and tool-call arguments stream as fragments per tool_call
-        # index; each segment keeps its own boundary buffer so a placeholder split
-        # across chunks reassembles without mixing segments.
-        buffers: dict = {}
-
-        def _restore(key: tuple, text: str, finished: bool) -> str:
-            deanon = buffers.get(key)
-            if deanon is None:
-                deanon = buffers[key] = StreamingDeAnonymizer(mapping)
-            out = deanon.feed(text) if text else ""
-            if finished:
-                out += deanon.flush()
-            return out
-
-        # The buffer holds back partial placeholders that span chunk boundaries.
-        # Whatever remains is flushed onto the chunk that carries finish_reason.
+        restorer = _StreamRestorer(mapping, StreamingDeAnonymizer)
         async for chunk in response:
-            for choice in getattr(chunk, "choices", None) or []:
-                delta = getattr(choice, "delta", None)
-                if delta is None:
-                    continue
-                index = getattr(choice, "index", 0) or 0
-                finished = getattr(choice, "finish_reason", None) is not None
-                self._restore_delta(delta, index, finished, _restore)
+            restorer.restore_chunk(chunk, self._restore_delta)
             yield chunk
+        for drained in restorer.drain():
+            yield drained
 
     async def async_post_call_failure_hook(
         self,

tests/unit/proxy/guardrails/guardrail_hooks/test_privaite.py

@@ -10,6 +10,7 @@ wiring end to end (registry, engine caching, pre/post/streaming hooks) without
 re-testing PrivAiTe's detection internals.
 """
 
+import json
 import sys
 import types
 
@@ -1103,3 +1104,178 @@ async def test_streaming_restores_audio_transcript():
     chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data))
     transcript = "".join(c.choices[0].delta.audio.transcript or "" for c in chunks if c.choices[0].delta.audio)
     assert transcript == "Hi Marie Dupont"
+
+
+def _bare_delta(**fields):
+    """A streamed delta carrying only the given fields (the others None), the
+    shape of the finish chunk most providers send: delta {} + finish_reason."""
+    base = dict(content=None, tool_calls=None, function_call=None)
+    base.update(fields)
+    return types.SimpleNamespace(**base)
+
+
+def _choice_chunk(delta, finish=None, index=0):
+    return types.SimpleNamespace(choices=[types.SimpleNamespace(index=index, delta=delta, finish_reason=finish)])
+
+
+def _get(obj, *path):
+    """Walk dict-or-object attributes; the guardrail may create a carrier as a
+    plain dict on a finish delta that had none."""
+    for key in path:
+        if obj is None:
+            return None
+        obj = obj.get(key) if isinstance(obj, dict) else getattr(obj, key, None)
+    return obj
+
+
+# The restorer holds back a tail that could begin a placeholder ("<PER" here),
+# so a channel ending on one loses it unless the end of the stream flushes it.
+
+
+@pytest.mark.parametrize("field", ["reasoning_content", "reasoning", "refusal"])
+async def test_streaming_flushes_text_field_tail_onto_a_bare_finish_chunk(field):
+    gr = _make_guardrail()
+    request_data = {"metadata": {"privaite_map": _FAKES}}
+
+    async def _source():
+        yield _choice_chunk(_bare_delta(**{field: "Considering <PERSON_1>, then <PER"}))
+        yield _choice_chunk(_bare_delta(**{field: None}), finish="stop")
+
+    chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data))
+    out = "".join(getattr(choice.delta, field, None) or "" for chunk in chunks for choice in chunk.choices)
+    assert out == "Considering Marie Dupont, then <PER"
+
+
+async def test_streaming_flushes_tool_tail_onto_a_bare_finish_chunk():
+    # A provider that stops mid-argument (finish_reason "length") sends the
+    # finish on a chunk with no tool_calls slot. The held fragment must still
+    # go out, on that chunk, under the same tool_call index: the client
+    # reassembles arguments per index and must see every byte it was sent.
+    gr = _make_guardrail()
+    request_data = {"metadata": {"privaite_map": _FAKES}}
+
+    async def _source():
+        call = types.SimpleNamespace(index=2, function=types.SimpleNamespace(arguments='{"to": "<EMAIL'))
+        yield _choice_chunk(_bare_delta(tool_calls=[call]))
+        yield _choice_chunk(_bare_delta(), finish="length")
+
+    chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data))
+    args = {}
+    for chunk in chunks:
+        for choice in chunk.choices:
+            for tc in choice.delta.tool_calls or []:
+                slot = _get(tc, "index")
+                args[slot] = args.get(slot, "") + (_get(tc, "function", "arguments") or "")
+    assert args == {2: '{"to": "<EMAIL'}
+    assert chunks[-1].choices[0].finish_reason == "length"
+    assert chunks[-1].choices[0].delta.tool_calls
+
+
+async def test_streaming_flushes_function_call_tail_onto_a_bare_finish_chunk():
+    gr = _make_guardrail()
+    request_data = {"metadata": {"privaite_map": _FAKES}}
+
+    async def _source():
+        yield _choice_chunk(_bare_delta(function_call=types.SimpleNamespace(arguments='{"n": "<PER')))
+        yield _choice_chunk(_bare_delta(), finish="length")
+
+    chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data))
+    out = "".join(
+        _get(choice.delta, "function_call", "arguments") or "" for chunk in chunks for choice in chunk.choices
+    )
+    assert out == '{"n": "<PER'
+
+
+async def test_streaming_drains_held_tails_when_the_stream_ends_without_finish_reason():
+    # Some providers close the stream without ever sending finish_reason. What
+    # the buffers still hold goes out as trailing chunks instead of vanishing
+    # with the generator: one per held channel, cloned from the last chunk of
+    # its choice, without that chunk's usage and without a finish_reason.
+    gr = _make_guardrail()
+    request_data = {"metadata": {"privaite_map": _FAKES}}
+
+    async def _source():
+        first = _choice_chunk(_bare_delta(content="Ask <PERSON_1> or <PER", audio={"transcript": "Call <EMAIL"}))
+        first.usage = types.SimpleNamespace(total_tokens=3)
+        yield first
+        call = types.SimpleNamespace(index=0, function=types.SimpleNamespace(arguments='{"to": "<EMAIL'))
+        yield _choice_chunk(_bare_delta(tool_calls=[call]), index=1)
+
+    chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data))
+    content = audio = args = ""
+    for chunk in chunks:
+        for choice in chunk.choices:
+            assert choice.finish_reason is None
+            if choice.index == 0:
+                content += choice.delta.content or ""
+                audio += _get(choice.delta, "audio", "transcript") or ""
+            else:
+                for tc in choice.delta.tool_calls or []:
+                    args += _get(tc, "function", "arguments") or ""
+    assert content == "Ask Marie Dupont or <PER"
+    assert audio == "Call <EMAIL"
+    assert args == '{"to": "<EMAIL'
+    assert chunks[0].usage.total_tokens == 3
+    assert all(getattr(chunk, "usage", None) is None for chunk in chunks[2:])
+    assert len(chunks) == 5  # 2 from the provider + content, audio and tool tails
+
+
+async def test_streaming_finish_chunk_with_empty_slots_gets_the_tails_under_them():
+    # A finish chunk may carry the tool_call slot and the function_call holder
+    # with empty arguments. Their restore skips an empty fragment, so the sweep
+    # appends the held tail under the existing holder: no duplicate slot.
+    gr = _make_guardrail()
+    request_data = {"metadata": {"privaite_map": _FAKES}}
+
+    def _slots(args, fc_args):
+        return _bare_delta(
+            tool_calls=[types.SimpleNamespace(index=2, function=types.SimpleNamespace(arguments=args))],
+            function_call=types.SimpleNamespace(arguments=fc_args),
+        )
+
+    async def _source():
+        yield _choice_chunk(_slots('{"to": "<EMAIL', '{"n": "<PER'))
+        yield _choice_chunk(_slots("", ""), finish="length")
+
+    chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data))
+    finish = chunks[-1].choices[0].delta
+    assert len(finish.tool_calls) == 1
+    assert finish.tool_calls[0].index == 2
+    assert finish.tool_calls[0].function.arguments == "<EMAIL"
+    assert finish.function_call.arguments == "<PER"
+
+
+async def test_streaming_tails_on_real_litellm_chunks_stay_serializable():
+    # The same two shapes on litellm's own chunk types: a carrier created on a
+    # bare finish Delta and a drained clone must both still serialize the way
+    # the proxy does it (model_dump_json).
+    from litellm.types.utils import ChatCompletionDeltaToolCall, Delta, Function, ModelResponseStream, StreamingChoices
+
+    gr = _make_guardrail()
+
+    def _chunk(delta, finish=None):
+        choice = StreamingChoices(index=0, delta=delta, finish_reason=finish)
+        return ModelResponseStream(id="chunk", model="m", choices=[choice])
+
+    async def _source():
+        call = ChatCompletionDeltaToolCall(index=0, type="function", function=Function(arguments='{"to": "<EMAIL'))
+        yield _chunk(Delta(content=None, tool_calls=[call]))
+        yield _chunk(Delta(content=None), finish="length")
+
+    async def _no_finish():
+        yield _chunk(Delta(content=None, reasoning_content="then <PER"))
+
+    request_data = {"metadata": {"privaite_map": _FAKES}}
+    chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _source(), request_data))
+    payloads = [json.loads(chunk.model_dump_json(exclude_none=True)) for chunk in chunks]
+    finish = payloads[-1]["choices"][0]
+    assert finish["finish_reason"] == "length"
+    assert finish["delta"]["tool_calls"][0]["index"] == 0
+    assert finish["delta"]["tool_calls"][0]["function"]["arguments"] == "<EMAIL"
+
+    request_data = {"metadata": {"privaite_map": _FAKES}}
+    chunks = await _collect(gr.async_post_call_streaming_iterator_hook(None, _no_finish(), request_data))
+    payloads = [json.loads(chunk.model_dump_json(exclude_none=True)) for chunk in chunks]
+    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"