beaver-gateway/src/beaver_gateway/events/stream.py

"""Streaming-event protocol shared by every backend adapter.

Aliases over :mod:`anthropic.types` plus ``build_*`` helpers that keep
verbose constructors out of adapters.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

from anthropic.types import (
    InputJSONDelta,
    Message,
    MessageDeltaUsage,
    RawContentBlockDeltaEvent,
    RawContentBlockStartEvent,
    RawContentBlockStopEvent,
    RawMessageDeltaEvent,
    RawMessageStartEvent,
    RawMessageStopEvent,
    RawMessageStreamEvent,
    SignatureDelta,
    TextBlock,
    TextDelta,
    ThinkingBlock,
    ThinkingDelta,
    ToolUseBlock,
    Usage,
)
from anthropic.types.raw_message_delta_event import Delta as MessageDelta

if TYPE_CHECKING:
    from collections.abc import Mapping

StopReason = Literal[
    "end_turn", "max_tokens", "stop_sequence", "tool_use", "pause_turn", "refusal"
]
"""Anthropic-canonical stop reasons. Backends map their native vocabulary into this."""


MessageStreamEvent = RawMessageStreamEvent
"""Single public alias for the event family yielded by ``Backend.complete``."""


__all__ = [
    "MessageDelta",
    "MessageDeltaUsage",
    "MessageStreamEvent",
    "RawContentBlockDeltaEvent",
    "RawContentBlockStartEvent",
    "RawContentBlockStopEvent",
    "RawMessageDeltaEvent",
    "RawMessageStartEvent",
    "RawMessageStopEvent",
    "StopReason",
    "Usage",
    "build_content_block_stop",
    "build_input_json_delta",
    "build_message_delta",
    "build_message_start",
    "build_message_stop",
    "build_signature_delta",
    "build_text_block_start",
    "build_text_delta",
    "build_thinking_block_start",
    "build_thinking_delta",
    "build_tool_use_block_start",
]


def build_message_start(
    *, message_id: str, model: str, input_tokens: int = 0, output_tokens: int = 0
) -> RawMessageStartEvent:
    """Open a turn. Token counts get finalized later via ``message_delta``."""
    return RawMessageStartEvent(
        type="message_start",
        message=Message(
            id=message_id,
            type="message",
            role="assistant",
            model=model,
            content=[],
            stop_reason=None,
            stop_sequence=None,
            usage=Usage(input_tokens=input_tokens, output_tokens=output_tokens),
        ),
    )


def build_text_block_start(index: int) -> RawContentBlockStartEvent:
    return RawContentBlockStartEvent(
        type="content_block_start",
        index=index,
        content_block=TextBlock(type="text", text="", citations=None),
    )


def build_text_delta(index: int, text: str) -> RawContentBlockDeltaEvent:
    return RawContentBlockDeltaEvent(
        type="content_block_delta",
        index=index,
        delta=TextDelta(type="text_delta", text=text),
    )


def build_tool_use_block_start(
    index: int, *, tool_use_id: str, name: str
) -> RawContentBlockStartEvent:
    return RawContentBlockStartEvent(
        type="content_block_start",
        index=index,
        content_block=ToolUseBlock(
            type="tool_use", id=tool_use_id, name=name, input={}
        ),
    )


def build_input_json_delta(index: int, partial_json: str) -> RawContentBlockDeltaEvent:
    return RawContentBlockDeltaEvent(
        type="content_block_delta",
        index=index,
        delta=InputJSONDelta(type="input_json_delta", partial_json=partial_json),
    )


def build_thinking_block_start(index: int) -> RawContentBlockStartEvent:
    return RawContentBlockStartEvent(
        type="content_block_start",
        index=index,
        content_block=ThinkingBlock(type="thinking", thinking="", signature=""),
    )


def build_thinking_delta(index: int, thinking: str) -> RawContentBlockDeltaEvent:
    return RawContentBlockDeltaEvent(
        type="content_block_delta",
        index=index,
        delta=ThinkingDelta(type="thinking_delta", thinking=thinking),
    )


def build_signature_delta(index: int, signature: str) -> RawContentBlockDeltaEvent:
    return RawContentBlockDeltaEvent(
        type="content_block_delta",
        index=index,
        delta=SignatureDelta(type="signature_delta", signature=signature),
    )


def build_content_block_stop(index: int) -> RawContentBlockStopEvent:
    return RawContentBlockStopEvent(type="content_block_stop", index=index)


def build_message_delta(
    *,
    stop_reason: StopReason | None,
    usage: Mapping[str, int] | None = None,
    stop_sequence: str | None = None,
) -> RawMessageDeltaEvent:
    """Close a turn. ``usage`` is a partial dict (input/output tokens, cache_*)."""
    return RawMessageDeltaEvent(
        type="message_delta",
        delta=MessageDelta(stop_reason=stop_reason, stop_sequence=stop_sequence),
        usage=MessageDeltaUsage.model_validate(
            dict(usage) if usage else {"output_tokens": 0}
        ),
    )


def build_message_stop() -> RawMessageStopEvent:
    return RawMessageStopEvent(type="message_stop")