cloakapi 0.5.1__py3-none-any.whl

cloakapi/_stream_detok.py

@@ -0,0 +1,788 @@
+"""Boundary-safe streaming detokenisation — ONE implementation, four call sites.
+
+WHY THIS MODULE EXISTS
+----------------------
+Detokenising each SSE delta INDEPENDENTLY cannot see a replacement that
+straddles a chunk boundary::
+
+    delta 1: "...mail emily.rivera@w"
+    delta 2: "estmarch.co.uk, ok?"
+
+Neither delta contains the whole key, so neither is rewritten and the STAND-IN
+reaches the caller. Under ``max`` that is at least visible (a stray
+``<<CLOAK:...>>`` looks obviously wrong); under the default ``balanced`` the
+stand-in is a plausible value, so the caller silently reads a fake name, address
+or email as though the model had produced it.
+
+The SDK has FOUR streaming paths — Anthropic and OpenAI, each sync and async.
+The defect was present in all four because each had its own copy of the
+per-delta logic. Putting the sequencing here means a future fix cannot land in
+two of the four and be believed to have landed in all of them; the only thing
+the clients supply is how to read text out of an event and how to put it back.
+
+THE END-OF-STREAM CONTRACT
+--------------------------
+:class:`~cloakapi.pii.tokenizer.StreamDetokenizer` withholds a short tail, so it
+MUST be flushed or the last characters of the reply are silently truncated —
+trading a corruption bug for a data-loss bug. The tail is released BEFORE the
+first terminal event rather than after the loop, so a caller that stops reading
+at ``message_stop`` / ``finish_reason`` cannot miss it.
+"""
+from __future__ import annotations
+
+import contextlib
+from collections.abc import AsyncIterator, Callable, Iterator
+from typing import Any
+
+# A "shape" is the small amount a caller must tell us about its event type:
+#   read(event)             -> str | None   (the text carried, or None)
+#   rewrite(event, text)    -> event        (same event, new text)
+#   is_terminal(event)      -> bool         (stream is ending with this event)
+#   make_tail(text, template) -> event      (a synthetic event carrying the tail)
+Shape = tuple[
+    Callable[[Any], "str | dict[int, str] | None"],
+    Callable[[Any, Any], Any],
+    Callable[[Any], bool],
+    Callable[[Any, Any], Any],
+]
+
+
+def _choice_detokenizer(template: Any, streams: dict[int, Any]) -> Any:
+    """Use the supplied instance for the first choice and clone its map for others."""
+    if not streams:
+        return template
+    try:
+        return type(template)(getattr(template, "_map", {}))
+    except TypeError:
+        return type(template)()
+
+
+def detok_stream(events: Iterator[Any], sd: Any, shape: Shape) -> Iterator[Any]:
+    """Synchronous boundary-safe pass over ``events``."""
+    read, rewrite, is_terminal, make_tail = shape
+    if getattr(read, "_cloak_openai_shape", False):
+        yield from _detok_openai_stream(events, sd)
+        return
+    if getattr(read, "_cloak_anthropic_shape", False):
+        yield from _detok_anthropic_stream(events, sd)
+        return
+    flushed = False
+    last: Any = None
+    tail: Any
+    per_choice: dict[int, Any] = {}
+    per_choice_flushed: set[int] = set()
+    multi_choice = False
+    try:
+        for event in events:
+            text = read(event)
+            if isinstance(text, dict):
+                multi_choice = True
+                last = event
+                output: dict[int, str] = {}
+                terminal_indices = {
+                    choice.index
+                    for choice in getattr(event, "choices", ())
+                    if getattr(choice, "finish_reason", None) is not None
+                }
+                if is_terminal(event) and not terminal_indices:
+                    tail = {
+                        index: stream.flush()
+                        for index, stream in per_choice.items()
+                        if index not in per_choice_flushed
+                    }
+                    per_choice_flushed.update(tail)
+                    tail = {index: value for index, value in tail.items() if value}
+                    if tail:
+                        yield make_tail(tail, event)
+                terminal_tails: dict[int, str] = {}
+                for index, part in text.items():
+                    if index not in per_choice:
+                        per_choice[index] = _choice_detokenizer(sd, per_choice)
+                    stream = per_choice[index]
+                    emitted = stream.feed(part) or ""
+                    if index in terminal_indices and index not in per_choice_flushed:
+                        tail = stream.flush() or ""
+                        if not part and not emitted and tail:
+                            terminal_tails[index] = tail
+                        else:
+                            emitted += tail
+                        per_choice_flushed.add(index)
+                    output[index] = emitted
+                for index in terminal_indices:
+                    if index in per_choice and index not in per_choice_flushed:
+                        output[index] = output.get(index, "") + (per_choice[index].flush() or "")
+                        per_choice_flushed.add(index)
+                if terminal_tails:
+                    yield make_tail(terminal_tails, event)
+                yield rewrite(event, output)
+                continue
+            # A finish chunk can carry the last text. Flushing before feed()
+            # leaves that suffix in the buffer after the one-shot flush.
+            if is_terminal(event) and not flushed and text:
+                emitted = sd.feed(text) or ""
+                tail = sd.flush() or ""
+                flushed = True
+                last = event
+                yield rewrite(event, emitted + tail)
+                continue
+            if is_terminal(event) and not flushed:
+                tail = sd.flush()
+                flushed = True
+                if tail:
+                    yield make_tail(tail, event)
+            last = event
+            yield rewrite(event, sd.feed(text)) if text is not None else event
+    except Exception as exc:
+        if getattr(exc, "code", None) == "stream_truncated" and not flushed:
+            if multi_choice:
+                tails = {
+                    index: stream.flush()
+                    for index, stream in per_choice.items()
+                    if index not in per_choice_flushed
+                }
+                per_choice_flushed.update(tails)
+                tails = {index: value for index, value in tails.items() if value}
+                flushed = True
+                if tails and last is not None:
+                    yield make_tail(tails, last)
+            else:
+                tail = sd.flush()
+                flushed = True
+                if tail and last is not None:
+                    yield make_tail(tail, last)
+        raise
+    finally:
+        close = getattr(events, "close", None)
+        if callable(close):
+            with contextlib.suppress(Exception):
+                close()
+    if multi_choice:
+        tails = {
+            index: stream.flush()
+            for index, stream in per_choice.items()
+            if index not in per_choice_flushed
+        }
+        tails = {index: value for index, value in tails.items() if value}
+        if tails and last is not None:
+            yield make_tail(tails, last)
+    elif not flushed:
+        tail = sd.flush()
+        if tail and last is not None:
+            yield make_tail(tail, last)
+
+
+async def adetok_stream(events: AsyncIterator[Any], sd: Any, shape: Shape) -> AsyncIterator[Any]:
+    """Asynchronous twin of :func:`detok_stream`. Same sequencing, same rules."""
+    read, rewrite, is_terminal, make_tail = shape
+    if getattr(read, "_cloak_openai_shape", False):
+        openai_events = _adetok_openai_stream(events, sd)
+        try:
+            async for chunk in openai_events:
+                yield chunk
+        finally:
+            close = getattr(openai_events, "aclose", None)
+            if callable(close):
+                with contextlib.suppress(Exception):
+                    await close()
+        return
+    if getattr(read, "_cloak_anthropic_shape", False):
+        anthropic_events = _adetok_anthropic_stream(events, sd)
+        try:
+            async for event in anthropic_events:
+                yield event
+        finally:
+            close = getattr(anthropic_events, "aclose", None)
+            if callable(close):
+                with contextlib.suppress(Exception):
+                    await close()
+        return
+    flushed = False
+    last: Any = None
+    tail: Any
+    per_choice: dict[int, Any] = {}
+    per_choice_flushed: set[int] = set()
+    multi_choice = False
+    try:
+        async for event in events:
+            text = read(event)
+            if isinstance(text, dict):
+                multi_choice = True
+                last = event
+                output: dict[int, str] = {}
+                terminal_indices = {
+                    choice.index
+                    for choice in getattr(event, "choices", ())
+                    if getattr(choice, "finish_reason", None) is not None
+                }
+                if is_terminal(event) and not terminal_indices:
+                    tail = {
+                        index: stream.flush()
+                        for index, stream in per_choice.items()
+                        if index not in per_choice_flushed
+                    }
+                    per_choice_flushed.update(tail)
+                    tail = {index: value for index, value in tail.items() if value}
+                    if tail:
+                        yield make_tail(tail, event)
+                terminal_tails: dict[int, str] = {}
+                for index, part in text.items():
+                    if index not in per_choice:
+                        per_choice[index] = _choice_detokenizer(sd, per_choice)
+                    stream = per_choice[index]
+                    emitted = stream.feed(part) or ""
+                    if index in terminal_indices and index not in per_choice_flushed:
+                        tail = stream.flush() or ""
+                        if not part and not emitted and tail:
+                            terminal_tails[index] = tail
+                        else:
+                            emitted += tail
+                        per_choice_flushed.add(index)
+                    output[index] = emitted
+                for index in terminal_indices:
+                    if index in per_choice and index not in per_choice_flushed:
+                        output[index] = output.get(index, "") + (per_choice[index].flush() or "")
+                        per_choice_flushed.add(index)
+                if terminal_tails:
+                    yield make_tail(terminal_tails, event)
+                yield rewrite(event, output)
+                continue
+            # Same ordering as detok_stream: feed a non-empty terminal chunk
+            # before the one-shot flush, or its held suffix never leaves.
+            if is_terminal(event) and not flushed and text:
+                emitted = sd.feed(text) or ""
+                tail = sd.flush() or ""
+                flushed = True
+                last = event
+                yield rewrite(event, emitted + tail)
+                continue
+            if is_terminal(event) and not flushed:
+                tail = sd.flush()
+                flushed = True
+                if tail:
+                    yield make_tail(tail, event)
+            last = event
+            yield rewrite(event, sd.feed(text)) if text is not None else event
+    except Exception as exc:
+        if getattr(exc, "code", None) == "stream_truncated" and not flushed:
+            if multi_choice:
+                tails = {
+                    index: stream.flush()
+                    for index, stream in per_choice.items()
+                    if index not in per_choice_flushed
+                }
+                per_choice_flushed.update(tails)
+                tails = {index: value for index, value in tails.items() if value}
+                flushed = True
+                if tails and last is not None:
+                    yield make_tail(tails, last)
+            else:
+                tail = sd.flush()
+                flushed = True
+                if tail and last is not None:
+                    yield make_tail(tail, last)
+        raise
+    finally:
+        close = getattr(events, "aclose", None)
+        if callable(close):
+            with contextlib.suppress(Exception):
+                await close()
+    if multi_choice:
+        tails = {
+            index: stream.flush()
+            for index, stream in per_choice.items()
+            if index not in per_choice_flushed
+        }
+        tails = {index: value for index, value in tails.items() if value}
+        if tails and last is not None:
+            yield make_tail(tails, last)
+    elif not flushed:
+        tail = sd.flush()
+        if tail and last is not None:
+            yield make_tail(tail, last)
+
+
+def anthropic_shape() -> Shape:
+    """Event shape for the Anthropic ``content_block_delta`` stream."""
+    from .types import ContentBlockDelta
+
+    def read(event: Any) -> str | None:
+        if isinstance(event, ContentBlockDelta):
+            d = event.delta
+            if isinstance(d, dict):
+                text = d.get("text")
+                if isinstance(text, str):
+                    return text
+        return None
+
+    def rewrite(event: Any, text: str) -> Any:
+        return ContentBlockDelta(
+            type=event.type, index=event.index, delta={**event.delta, "text": text}
+        )
+
+    def is_terminal(event: Any) -> bool:
+        # Only events a caller stops on. message_start and content_block_start
+        # also have no text; treating them as terminal used the single flush
+        # before any delta, so the held tail never left.
+        from .types import ErrorEvent, MessageStop
+
+        return isinstance(event, (MessageStop, ErrorEvent))
+
+    def make_tail(text: str, template: Any) -> Any:
+        return ContentBlockDelta(
+            type="content_block_delta",
+            index=getattr(template, "index", 0) or 0,
+            delta={"type": "text_delta", "text": text},
+        )
+
+    # Text and input_json_delta are separate streams per content block. The
+    # generic single-buffer path would glue a held text tail onto tool JSON.
+    read._cloak_anthropic_shape = True  # type: ignore[attr-defined]
+    return read, rewrite, is_terminal, make_tail
+
+
+def openai_shape() -> Shape:
+    """Event shape for the OpenAI ``chat.completion.chunk`` stream."""
+    from .types import ChatCompletionChunk, ChunkDelta, ErrorEvent, StreamChoice
+
+    def read(chunk: Any) -> dict[int, str]:
+        if isinstance(chunk, ErrorEvent):
+            return {}
+        texts: dict[int, str] = {}
+        for choice in chunk.choices:
+            if isinstance(choice.delta.content, str):
+                texts[choice.index] = choice.delta.content
+        return texts
+
+    def rewrite(chunk: Any, text: Any) -> Any:
+        if isinstance(chunk, ErrorEvent):
+            return chunk
+        new_choices = []
+        for choice in chunk.choices:
+            if choice.index in text:
+                choice = choice.model_copy(
+                    update={
+                        "delta": choice.delta.model_copy(
+                            update={"content": text[choice.index]}
+                        )
+                    }
+                )
+            new_choices.append(choice)
+        return chunk.model_copy(update={"choices": new_choices})
+
+    def is_terminal(chunk: Any) -> bool:
+        return isinstance(chunk, ErrorEvent) or any(c.finish_reason for c in chunk.choices)
+
+    def make_tail(text: Any, template: Any) -> Any:
+        if isinstance(template, ErrorEvent):
+            template = ChatCompletionChunk(
+                id="", object="chat.completion.chunk", created=0, model="", choices=[]
+            )
+        if isinstance(text, dict):
+            return ChatCompletionChunk(
+                id=template.id,
+                object=template.object,
+                created=template.created,
+                model=template.model,
+                choices=[
+                    StreamChoice(
+                        index=index,
+                        delta=ChunkDelta(role=None, content=tail),
+                        finish_reason=None,
+                    )
+                    for index, tail in sorted(text.items())
+                ],
+                usage=None,
+            )
+        return ChatCompletionChunk(
+            id=template.id,
+            object=template.object,
+            created=template.created,
+            model=template.model,
+            choices=[
+                StreamChoice(
+                    index=0, delta=ChunkDelta(role=None, content=text), finish_reason=None
+                )
+            ],
+            usage=None,
+        )
+
+    # OpenAI chunks can carry several choices, and each choice is an independent
+    # output stream. Route the generic entry points through per-choice state
+    # instead of feeding choice 0 into one shared detokenizer.
+    read._cloak_openai_shape = True  # type: ignore[attr-defined]
+    return read, rewrite, is_terminal, make_tail
+
+
+def _openai_tail(text: str, template: Any, index: int) -> Any:
+    from .types import ChatCompletionChunk, ChunkDelta, StreamChoice
+
+    return ChatCompletionChunk(
+        id=template.id,
+        object=template.object,
+        created=template.created,
+        model=template.model,
+        choices=[
+            StreamChoice(
+                index=index, delta=ChunkDelta(role=None, content=text), finish_reason=None
+            )
+        ],
+        usage=None,
+    )
+
+
+def _delta_field(delta: Any) -> tuple[str, str] | None:
+    """``(field, text)`` for an Anthropic delta that carries restorable text."""
+    if not isinstance(delta, dict):
+        return None
+    if isinstance(delta.get("text"), str):
+        return "text", delta["text"]
+    if delta.get("type") == "input_json_delta" and isinstance(delta.get("partial_json"), str):
+        return "partial_json", delta["partial_json"]
+    return None
+
+
+def _anthropic_delta(event: Any, field: str, text: str) -> Any:
+    from .types import ContentBlockDelta
+
+    delta = dict(event.delta) if isinstance(event.delta, dict) else {}
+    delta[field] = text
+    if field == "partial_json":
+        delta["type"] = "input_json_delta"
+    elif "type" not in delta:
+        delta["type"] = "text_delta"
+    return ContentBlockDelta(type="content_block_delta", index=event.index, delta=delta)
+
+
+def _anthropic_tail(index: int, field: str, text: str) -> Any:
+    from .types import ContentBlockDelta
+
+    if field == "partial_json":
+        delta = {"type": "input_json_delta", "partial_json": text}
+    else:
+        delta = {"type": "text_delta", "text": text}
+    return ContentBlockDelta(type="content_block_delta", index=index, delta=delta)
+
+
+def _detok_anthropic_stream(events: Iterator[Any], sd: Any) -> Iterator[Any]:
+    """Restore text and ``input_json_delta`` per content block, across chunks."""
+    from .pii.tokenizer import StreamDetokenizer
+    from .types import ContentBlockDelta, ContentBlockStop, MessageStop
+
+    token_map = getattr(sd, "_map", {})
+    states: dict[int, Any] = {}
+    fields: dict[int, str] = {}
+    flushed: set[int] = set()
+
+    def stream_for(index: int) -> Any:
+        state = states.get(index)
+        if state is None:
+            state = sd if not states else StreamDetokenizer(token_map)
+            states[index] = state
+        return state
+
+    def flush_index(index: int) -> Any | None:
+        if index in flushed or index not in states:
+            flushed.add(index)
+            return None
+        tail = states[index].flush()
+        flushed.add(index)
+        if tail:
+            return _anthropic_tail(index, fields.get(index, "text"), tail)
+        return None
+
+    def flush_all() -> Iterator[Any]:
+        for index in list(states):
+            tail = flush_index(index)
+            if tail is not None:
+                yield tail
+
+    try:
+        for event in events:
+            if isinstance(event, ContentBlockStop):
+                tail = flush_index(event.index)
+                if tail is not None:
+                    yield tail
+                yield event
+                continue
+            if isinstance(event, MessageStop):
+                yield from flush_all()
+                yield event
+                continue
+            if isinstance(event, ContentBlockDelta):
+                found = _delta_field(event.delta)
+                if found is not None:
+                    field, text = found
+                    fields[event.index] = field
+                    emitted = stream_for(event.index).feed(text)
+                    yield _anthropic_delta(event, field, emitted)
+                    continue
+            yield event
+    except Exception as exc:
+        if getattr(exc, "code", None) == "stream_truncated":
+            yield from flush_all()
+        raise
+    finally:
+        close = getattr(events, "close", None)
+        if callable(close):
+            with contextlib.suppress(Exception):
+                close()
+    yield from flush_all()
+
+
+async def _adetok_anthropic_stream(events: AsyncIterator[Any], sd: Any) -> AsyncIterator[Any]:
+    """Async twin of :func:`_detok_anthropic_stream`."""
+    from .pii.tokenizer import StreamDetokenizer
+    from .types import ContentBlockDelta, ContentBlockStop, MessageStop
+
+    token_map = getattr(sd, "_map", {})
+    states: dict[int, Any] = {}
+    fields: dict[int, str] = {}
+    flushed: set[int] = set()
+
+    def stream_for(index: int) -> Any:
+        state = states.get(index)
+        if state is None:
+            state = sd if not states else StreamDetokenizer(token_map)
+            states[index] = state
+        return state
+
+    def take_tail(index: int) -> Any | None:
+        if index in flushed or index not in states:
+            flushed.add(index)
+            return None
+        tail = states[index].flush()
+        flushed.add(index)
+        if tail:
+            return _anthropic_tail(index, fields.get(index, "text"), tail)
+        return None
+
+    try:
+        async for event in events:
+            if isinstance(event, ContentBlockStop):
+                tail = take_tail(event.index)
+                if tail is not None:
+                    yield tail
+                yield event
+                continue
+            if isinstance(event, MessageStop):
+                for index in list(states):
+                    tail = take_tail(index)
+                    if tail is not None:
+                        yield tail
+                yield event
+                continue
+            if isinstance(event, ContentBlockDelta):
+                found = _delta_field(event.delta)
+                if found is not None:
+                    field, text = found
+                    fields[event.index] = field
+                    emitted = stream_for(event.index).feed(text)
+                    yield _anthropic_delta(event, field, emitted)
+                    continue
+            yield event
+    except Exception as exc:
+        if getattr(exc, "code", None) == "stream_truncated":
+            for index in list(states):
+                tail = take_tail(index)
+                if tail is not None:
+                    yield tail
+        raise
+    finally:
+        close = getattr(events, "aclose", None)
+        if callable(close):
+            with contextlib.suppress(Exception):
+                await close()
+    for index in list(states):
+        tail = take_tail(index)
+        if tail is not None:
+            yield tail
+
+
+def _restore_openai_tool_args(
+    delta: Any,
+    choice_index: int,
+    states: dict[tuple[int, int], Any],
+    token_map: dict[str, str],
+    *,
+    finishing: bool,
+    flushed: set[tuple[int, int]],
+) -> Any:
+    """Restore ``delta.tool_calls[].function.arguments``, including a held tail."""
+    from .pii.tokenizer import StreamDetokenizer
+
+    calls = getattr(delta, "tool_calls", None)
+    new_calls: list[Any] = []
+    changed = False
+    seen: set[tuple[int, int]] = set()
+    if isinstance(calls, list):
+        for call in calls:
+            if not isinstance(call, dict):
+                new_calls.append(call)
+                continue
+            tool_index = call.get("index", 0)
+            if not isinstance(tool_index, int):
+                tool_index = 0
+            key = (choice_index, tool_index)
+            seen.add(key)
+            fn = call.get("function")
+            arguments = fn.get("arguments") if isinstance(fn, dict) else None
+            if not isinstance(arguments, str):
+                new_calls.append(call)
+                continue
+            state = states.get(key)
+            if state is None:
+                state = StreamDetokenizer(token_map)
+                states[key] = state
+            emitted = state.feed(arguments) or ""
+            if finishing and key not in flushed:
+                emitted += state.flush() or ""
+                flushed.add(key)
+            new_calls.append({**call, "function": {**fn, "arguments": emitted}})
+            changed = True
+    if finishing:
+        for key, state in list(states.items()):
+            if key[0] != choice_index or key in flushed or key in seen:
+                continue
+            tail = state.flush() or ""
+            flushed.add(key)
+            if tail:
+                new_calls.append({"index": key[1], "function": {"arguments": tail}})
+                changed = True
+    if not changed:
+        return delta
+    return delta.model_copy(update={"tool_calls": new_calls})
+
+
+def _detok_openai_stream(events: Iterator[Any], sd: Any) -> Iterator[Any]:
+    """Detokenize each OpenAI choice with its own boundary buffer."""
+    from .pii.tokenizer import StreamDetokenizer
+    from .types import ErrorEvent
+
+    token_map = getattr(sd, "_map", {})
+    states: dict[int, StreamDetokenizer] = {}
+    arg_states: dict[tuple[int, int], Any] = {}
+    arg_flushed: set[tuple[int, int]] = set()
+    last: dict[int, Any] = {}
+    flushed: set[int] = set()
+
+    def flush_pending(template: Any) -> Iterator[Any]:
+        for index, state in list(states.items()):
+            if index in flushed:
+                continue
+            tail = state.flush()
+            flushed.add(index)
+            if tail and index in last:
+                yield _openai_tail(tail, last.get(index, template), index)
+
+    try:
+        for chunk in events:
+            if isinstance(chunk, ErrorEvent):
+                yield from flush_pending(chunk)
+                yield chunk
+                continue
+            new_choices = []
+            for choice in chunk.choices:
+                index = choice.index
+                last[index] = chunk
+                state = states.get(index)
+                if state is None:
+                    state = sd if not states else StreamDetokenizer(token_map)
+                    states[index] = state
+                content = choice.delta.content
+                if isinstance(content, str):
+                    content = state.feed(content)
+                if choice.finish_reason and index not in flushed:
+                    tail = state.flush()
+                    flushed.add(index)
+                    if tail and not content:
+                        yield _openai_tail(tail, chunk, index)
+                    else:
+                        content = (content or "") + tail
+                delta = choice.delta
+                if isinstance(choice.delta.content, str) or content:
+                    delta = choice.delta.model_copy(update={"content": content})
+                delta = _restore_openai_tool_args(
+                    delta, index, arg_states, token_map,
+                    finishing=bool(choice.finish_reason), flushed=arg_flushed,
+                )
+                new_choices.append(choice.model_copy(update={"delta": delta}))
+            yield chunk.model_copy(update={"choices": new_choices})
+    except Exception as exc:
+        if getattr(exc, "code", None) == "stream_truncated":
+            yield from flush_pending(next(iter(last.values()), None))
+        raise
+    finally:
+        close = getattr(events, "close", None)
+        if callable(close):
+            with contextlib.suppress(Exception):
+                close()
+    yield from flush_pending(next(iter(last.values()), None))
+
+
+async def _adetok_openai_stream(events: AsyncIterator[Any], sd: Any) -> AsyncIterator[Any]:
+    """Async per-choice twin of :func:`_detok_openai_stream`."""
+    from .pii.tokenizer import StreamDetokenizer
+    from .types import ErrorEvent
+
+    token_map = getattr(sd, "_map", {})
+    states: dict[int, StreamDetokenizer] = {}
+    arg_states: dict[tuple[int, int], Any] = {}
+    arg_flushed: set[tuple[int, int]] = set()
+    last: dict[int, Any] = {}
+    flushed: set[int] = set()
+
+    async def flush_pending(template: Any) -> AsyncIterator[Any]:
+        for index, state in list(states.items()):
+            if index in flushed:
+                continue
+            tail = state.flush()
+            flushed.add(index)
+            if tail and index in last:
+                yield _openai_tail(tail, last.get(index, template), index)
+
+    try:
+        async for chunk in events:
+            if isinstance(chunk, ErrorEvent):
+                async for tail in flush_pending(chunk):
+                    yield tail
+                yield chunk
+                continue
+            new_choices = []
+            for choice in chunk.choices:
+                index = choice.index
+                last[index] = chunk
+                state = states.get(index)
+                if state is None:
+                    state = sd if not states else StreamDetokenizer(token_map)
+                    states[index] = state
+                content = choice.delta.content
+                if isinstance(content, str):
+                    content = state.feed(content)
+                if choice.finish_reason and index not in flushed:
+                    tail = state.flush()
+                    flushed.add(index)
+                    if tail and not content:
+                        yield _openai_tail(tail, chunk, index)
+                    else:
+                        content = (content or "") + tail
+                delta = choice.delta
+                if isinstance(choice.delta.content, str) or content:
+                    delta = choice.delta.model_copy(update={"content": content})
+                delta = _restore_openai_tool_args(
+                    delta, index, arg_states, token_map,
+                    finishing=bool(choice.finish_reason), flushed=arg_flushed,
+                )
+                new_choices.append(choice.model_copy(update={"delta": delta}))
+            yield chunk.model_copy(update={"choices": new_choices})
+    except Exception as exc:
+        if getattr(exc, "code", None) == "stream_truncated":
+            async for tail in flush_pending(next(iter(last.values()), None)):
+                yield tail
+        raise
+    finally:
+        close = getattr(events, "aclose", None)
+        if callable(close):
+            with contextlib.suppress(Exception):
+                await close()
+    async for tail in flush_pending(next(iter(last.values()), None)):
+        yield tail