143 lines
4.9 KiB
Python
143 lines
4.9 KiB
Python
"""Collapse an Anthropic event stream into one ``Message``.
|
|
|
|
Mirrors the Anthropic SDK's own accumulator: folds content-block deltas
|
|
into finalized blocks, indexed by their ``content_block`` index.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import TYPE_CHECKING, Any, Literal, cast
|
|
|
|
from anthropic.types import (
|
|
InputJSONDelta,
|
|
Message,
|
|
RawContentBlockDeltaEvent,
|
|
RawContentBlockStartEvent,
|
|
RawContentBlockStopEvent,
|
|
RawMessageDeltaEvent,
|
|
RawMessageStartEvent,
|
|
SignatureDelta,
|
|
TextBlock,
|
|
TextDelta,
|
|
ThinkingBlock,
|
|
ThinkingDelta,
|
|
ToolUseBlock,
|
|
Usage,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
|
|
from beaver_gateway.events.stream import MessageStreamEvent, StopReason
|
|
|
|
|
|
__all__ = ["StreamAccumulator", "accumulate"]
|
|
|
|
|
|
class StreamAccumulator:
|
|
"""Folds a stream of events into one ``Message``, incrementally.
|
|
|
|
Call :meth:`feed` for each event, then :meth:`finalize` once.
|
|
"""
|
|
|
|
__slots__ = (
|
|
"_blocks",
|
|
"_json_buffers",
|
|
"_message_id",
|
|
"_role",
|
|
"_stop_reason",
|
|
"_stop_sequence",
|
|
"_usage",
|
|
)
|
|
|
|
def __init__(self) -> None:
|
|
self._message_id: str = ""
|
|
self._role: str = "assistant"
|
|
self._usage: Usage = Usage(input_tokens=0, output_tokens=0)
|
|
self._blocks: dict[int, dict[str, Any]] = {}
|
|
self._json_buffers: dict[int, str] = {}
|
|
self._stop_reason: str | None = None
|
|
self._stop_sequence: str | None = None
|
|
|
|
def feed(self, ev: MessageStreamEvent) -> None:
|
|
"""Fold one stream event into the in-progress message state.
|
|
|
|
Uses ``isinstance`` rather than ``ev.type == ...`` so ty narrows
|
|
the event type from the class, not a string comparison.
|
|
"""
|
|
if isinstance(ev, RawMessageStartEvent):
|
|
self._message_id = ev.message.id
|
|
self._role = ev.message.role
|
|
self._usage = ev.message.usage
|
|
elif isinstance(ev, RawContentBlockStartEvent):
|
|
self._blocks[ev.index] = ev.content_block.model_dump()
|
|
if self._blocks[ev.index].get("type") == "tool_use":
|
|
self._json_buffers[ev.index] = ""
|
|
elif isinstance(ev, RawContentBlockDeltaEvent):
|
|
blk = self._blocks.setdefault(ev.index, {})
|
|
delta = ev.delta
|
|
if isinstance(delta, TextDelta):
|
|
blk["text"] = blk.get("text", "") + delta.text
|
|
elif isinstance(delta, InputJSONDelta):
|
|
self._json_buffers[ev.index] = (
|
|
self._json_buffers.get(ev.index, "") + delta.partial_json
|
|
)
|
|
elif isinstance(delta, ThinkingDelta):
|
|
blk["thinking"] = blk.get("thinking", "") + delta.thinking
|
|
elif isinstance(delta, SignatureDelta):
|
|
blk["signature"] = delta.signature
|
|
elif isinstance(ev, RawContentBlockStopEvent):
|
|
blk = self._blocks.get(ev.index, {})
|
|
if blk.get("type") == "tool_use":
|
|
raw = self._json_buffers.pop(ev.index, "")
|
|
blk["input"] = json.loads(raw) if raw.strip() else {}
|
|
elif isinstance(ev, RawMessageDeltaEvent):
|
|
self._stop_reason = ev.delta.stop_reason
|
|
self._stop_sequence = ev.delta.stop_sequence
|
|
if ev.usage.output_tokens:
|
|
self._usage = Usage.model_validate(
|
|
{
|
|
**self._usage.model_dump(),
|
|
"output_tokens": ev.usage.output_tokens,
|
|
}
|
|
)
|
|
|
|
def finalize(self, *, model: str) -> Message:
|
|
"""Build the finalized ``Message`` from accumulated block state.
|
|
|
|
``role`` is always ``"assistant"`` at the wire level; the cast
|
|
avoids a runtime check ty would otherwise require.
|
|
"""
|
|
content: list[Any] = []
|
|
for idx in sorted(self._blocks):
|
|
bd = self._blocks[idx]
|
|
btype = bd.get("type")
|
|
if btype == "text":
|
|
content.append(TextBlock.model_validate(bd))
|
|
elif btype == "tool_use":
|
|
content.append(ToolUseBlock.model_validate(bd))
|
|
elif btype == "thinking":
|
|
content.append(ThinkingBlock.model_validate(bd))
|
|
|
|
return Message(
|
|
id=self._message_id or "msg_unknown",
|
|
type="message",
|
|
role=cast("Literal['assistant']", self._role),
|
|
model=model,
|
|
content=content,
|
|
stop_reason=cast("StopReason | None", self._stop_reason),
|
|
stop_sequence=self._stop_sequence,
|
|
usage=self._usage,
|
|
)
|
|
|
|
|
|
async def accumulate(
|
|
events: AsyncIterator[MessageStreamEvent], *, model: str
|
|
) -> Message:
|
|
"""Drain ``events`` into one ``Message`` (used by ``stream=false`` paths)."""
|
|
acc = StreamAccumulator()
|
|
async for ev in events:
|
|
acc.feed(ev)
|
|
return acc.finalize(model=model)
|