1"""Compat bridge: convert `AIMessageChunk` streams to protocol events.23The bridge trusts `AIMessageChunk.content_blocks` as the single4protocol view of any chunk.  That property runs the three-tier lookup5(`output_version == "v1"` short-circuit, registered translator, or6best-effort parsing) and returns a `list[ContentBlock]` for every7well-formed message — whether the provider is a registered partner, an8unregistered community model, or not tagged at all.910Per-chunk `content_blocks` output is a **delta slice**, not accumulated11state: providers in this ecosystem emit SSE-style chunks that each carry12their own increment.  The bridge therefore forwards each slice straight13through as a `content-block-delta` event, and accumulates per-index14state only so the final `content-block-finish` event can report a15finalized block (e.g. `tool_call_chunk` args parsed to a dict).1617Lifecycle::1819    message-start20      -> content-block-start   (first time each index is observed)21      -> content-block-delta*  (per chunk, carrying the slice)22      -> content-block-finish  (finalized block)23    -> message-finish2425Public API:2627- `chunks_to_events` / `achunks_to_events` — for live streams where28  chunks arrive over time.29- `message_to_events` / `amessage_to_events` — for replaying a finalized30  `AIMessage` (cache hit, checkpoint restore, graph-node return value)31  as a synthetic event lifecycle.32"""3334from __future__ import annotations3536import json37from typing import TYPE_CHECKING, Any, cast3839from langchain_protocol.protocol import (40    ContentBlock,41    ContentBlockDeltaData,42    ContentBlockFinishData,43    ContentBlockStartData,44    FinalizedContentBlock,45    InvalidToolCall,46    MessageFinishData,47    MessageMetadata,48    MessagesData,49    MessageStartData,50    ReasoningContentBlock,51    ServerToolCall,52    ServerToolCallChunk,53    TextContentBlock,54    ToolCall,55    ToolCallChunk,56    UsageInfo,57)5859from langchain_core.messages import AIMessageChunk, BaseMessage60from langchain_core.utils._merge import merge_dicts6162if TYPE_CHECKING:63    from collections.abc import AsyncIterator, Iterator6465    from langchain_protocol.protocol import (66        BlockDelta,67        BlockDeltaFields,68        ContentBlockDelta,69        DataDelta,70        ReasoningDelta,71        TextDelta,72    )7374    from langchain_core.messages.ai import UsageMetadata75    from langchain_core.outputs import ChatGenerationChunk767778CompatBlock = dict[str, Any]79"""Internal working type for a content block.8081The bridge works with plain dicts internally because two separate but82structurally similar `ContentBlock` Unions exist — one in83`langchain_core.messages.content` (returned by `msg.content_blocks`),84one in `langchain_protocol.protocol` (the wire/event shape).  They are85not mypy-compatible despite being near-isomorphic.  Passing through86`dict[str, Any]` launders between them.  See `_to_protocol_block` for87the single seam where the laundering cast lives.88"""899091# ---------------------------------------------------------------------------92# Type laundering between core and protocol `ContentBlock` unions93# ---------------------------------------------------------------------------949596def _to_protocol_block(block: CompatBlock) -> ContentBlock:97    """Narrow an internal working dict to a protocol `ContentBlock`.9899    Single seam between the two `ContentBlock` type systems:100    `langchain_core.messages.content` (what `msg.content_blocks`101    returns) and `langchain_protocol.protocol` (what event payloads102    require).  The two Unions overlap structurally but are nominally103    distinct to mypy, so we launder through `dict[str, Any]`.  When the104    Unions are unified, this helper and its finalized counterpart can be105    deleted.106    """107    return cast("ContentBlock", block)108109110def _to_finalized_block(block: CompatBlock) -> FinalizedContentBlock:111    """Counterpart of `_to_protocol_block` for finalized blocks."""112    return cast("FinalizedContentBlock", block)113114115def _to_block_delta_fields(block: CompatBlock) -> BlockDeltaFields:116    """Narrow an internal working dict to protocol block-delta fields."""117    return cast("BlockDeltaFields", block)118119120def _to_content_delta(block: CompatBlock) -> ContentBlockDelta:121    """Convert a content-block slice/snapshot to an explicit protocol delta."""122    btype = block.get("type")123    if btype == "text":124        return cast("TextDelta", {"type": "text-delta", "text": block.get("text", "")})125    if btype == "reasoning":126        return cast(127            "ReasoningDelta",128            {129                "type": "reasoning-delta",130                "reasoning": block.get("reasoning", ""),131            },132        )133    if "data" in block:134        delta = cast("DataDelta", {"type": "data-delta", "data": block.get("data", "")})135        if block.get("encoding") == "base64":136            delta["encoding"] = "base64"137        return delta138    return cast(139        "BlockDelta",140        {141            "type": "block-delta",142            "fields": _to_block_delta_fields(block),143        },144    )145146147# ---------------------------------------------------------------------------148# Block iteration149# ---------------------------------------------------------------------------150151152def _iter_protocol_blocks(msg: BaseMessage) -> list[tuple[Any, CompatBlock]]:153    """Read per-chunk protocol blocks from `msg.content_blocks`.154155    Returns `(key, block)` pairs.  The key is the block's stable identifier156    across the stream: the block's `index` field when present (can be an157    int or a string — some providers use string identifiers like158    `"lc_rs_305f30"`), or the positional index within the message as a159    fallback.  Callers are responsible for allocating wire-level `uint`160    indices; this helper only surfaces the source-side identity.161162    For finalized `AIMessage`, also surfaces `invalid_tool_calls`163    — which `AIMessage.content_blocks` currently omits from its return164    value even though they are a defined protocol block type.165166    The positional fallback is a known fragility: when a provider emits167    blocks without an `index` field (e.g. Anthropic's `_stream` with168    `coerce_content_to_string=True`, where text chunks lose their169    source-side index), every such chunk gets positional key 0 and170    successive chunks merge into one block. This works correctly for171    single-type streams (pure-text responses merge cleanly) because all172    chunks share the same key and the open-block logic collapses them.173    It would miscategorise a stream that mixed indexed structured174    blocks with non-indexed coerced-text blocks, since an indexed175    block with `index == 0` would collide with the anonymous text176    block's positional-0 key.  In the anthropic integration this177    cannot currently occur: coerce-to-string mode is only selected178    when no tools, thinking, or documents are present, and any of179    those flips the stream to structured mode where every block180    carries an integer index.  A native `_stream_chat_model_events`181    hook per provider (or a bridge-level "continue the open block when182    the source has no identity" rule) would close the gap if another183    integration ever emits mixed content.184    """185    try:186        raw = msg.content_blocks187    except Exception:188        return []189190    result: list[tuple[Any, CompatBlock]] = []191    for i, block in enumerate(raw):192        if not isinstance(block, dict):193            continue  # type: ignore[unreachable]194        explicit_idx = block.get("index")195        if explicit_idx is None:196            # No source-side identity. Bucket by (sentinel, block type,197            # positional `i`) so two blocks of different types at the198            # same position across chunks (e.g. Gemini emitting a199            # reasoning block in one chunk and a `tool_call` in the200            # next, both at positional 0 because each chunk carries one201            # block) get distinct wire blocks. Without this, the second202            # type's incoming block hits `_accumulate`'s self-contained203            # `else` branch and clobbers the first. Same-type chunks204            # still share the bucket and merge cleanly, which is what205            # streaming text / reasoning relies on.206            key: Any = ("__lc_no_index__", block.get("type"), i)207        else:208            key = explicit_idx209        result.append((key, dict(block)))210211    if not isinstance(msg, AIMessageChunk):212        # Finalized AIMessage: pull invalid_tool_calls from the dedicated213        # field — AIMessage.content_blocks does not currently include them.214        for itc in getattr(msg, "invalid_tool_calls", None) or []:215            itc_block: CompatBlock = {"type": "invalid_tool_call"}216            for key_name in ("id", "name", "args", "error"):217                if itc.get(key_name) is not None:218                    itc_block[key_name] = itc[key_name]219            result.append((len(result), itc_block))220221    return result222223224# ---------------------------------------------------------------------------225# Per-block helpers226# ---------------------------------------------------------------------------227228229# Fields that can carry large payloads (inline base64 media, parsed args,230# arbitrary dicts).  Stripped from `content-block-start` for self-contained231# block types so the payload rides on `content-block-finish` alone instead232# of being serialized twice on the wire.233_HEAVY_FIELDS = frozenset({"args", "data", "output", "transcript", "value"})234235236def _start_skeleton(block: CompatBlock) -> ContentBlock:237    """Empty-content placeholder for the `content-block-start` event.238239    Deltaable block types (text, reasoning, the `_chunk` tool variants)240    get an empty payload so the lifecycle's "start" signal is distinct241    from the first incremental delta.  Self-contained types (image,242    audio, video, file, non_standard, finalized tool calls) drop their243    heavy payload fields; those are carried by `content-block-finish`.244    Correlation fields (id, name, toolCallId) and small metadata245    (mime_type, url, status, …) are preserved on the start event.246    """247    btype = block.get("type", "text")248    if btype == "text":249        return TextContentBlock(type="text", text="")250    if btype == "reasoning":251        return ReasoningContentBlock(type="reasoning", reasoning="")252    if btype == "tool_call_chunk":253        return ToolCallChunk(254            type="tool_call_chunk",255            id=block.get("id"),256            name=block.get("name"),257            args="",258        )259    if btype == "server_tool_call_chunk":260        s_skel = ServerToolCallChunk(261            type="server_tool_call_chunk",262            args="",263        )264        if block.get("id") is not None:265            s_skel["id"] = block["id"]266        if block.get("name") is not None:267            s_skel["name"] = block["name"]268        return s_skel269270    stripped: CompatBlock = {k: v for k, v in block.items() if k not in _HEAVY_FIELDS}271    # Restore required-but-heavy fields with minimal placeholders so the272    # start event still validates against the CDDL shape of the block type.273    if btype in {"tool_call", "server_tool_call"}:274        stripped["args"] = {}275    elif btype == "non_standard":276        stripped["value"] = {}277    return _to_protocol_block(stripped)278279280def _should_emit_delta(block: CompatBlock) -> bool:281    """Whether a per-chunk block carries content worth a delta event.282283    Deltaable types emit only when they have fresh content.  Self-contained284    / already-finalized types skip the delta entirely — the `finish`285    event carries them.286    """287    btype = block.get("type")288    if btype == "text":289        return bool(block.get("text"))290    if btype == "reasoning":291        return bool(block.get("reasoning"))292    if btype in {"tool_call_chunk", "server_tool_call_chunk"}:293        return bool(294            block.get("args") or block.get("id") or block.get("name"),295        )296    if "data" in block:297        return bool(block.get("data"))298    return False299300301def _accumulate(state: CompatBlock | None, delta: CompatBlock) -> CompatBlock:302    """Merge a per-chunk delta slice into accumulated per-index state.303304    Used only for the finalization pass — live delta events are emitted305    directly from the per-chunk block, without round-tripping through306    accumulated state.307    """308    if state is None:309        return dict(delta)310    btype = state.get("type")311    dtype = delta.get("type")312    if btype == "text" and dtype == "text":313        state["text"] = state.get("text", "") + delta.get("text", "")314        # Providers may send non-text fields (like `id`, or annotations)315        # on later deltas. Merging (not replacing) keeps earlier keys316        # intact while picking up these late-arriving fields.317        for key, value in delta.items():318            if key in {"type", "text"} or value is None:319                continue320            if key == "extras" and isinstance(value, dict):321                state["extras"] = {**(state.get("extras") or {}), **value}322            else:323                state[key] = value324    elif btype == "reasoning" and dtype == "reasoning":325        state["reasoning"] = state.get("reasoning", "") + delta.get("reasoning", "")326        # Providers may ship non-text fields on later deltas. Claude's327        # `signature_delta` arrives after the reasoning text, surfaced328        # as `extras.signature`; merging (not replacing) keeps earlier329        # keys intact.330        for key, value in delta.items():331            if key in {"type", "reasoning"} or value is None:332                continue333            if key == "extras" and isinstance(value, dict):334                state["extras"] = {**(state.get("extras") or {}), **value}335            else:336                state[key] = value337    elif btype in {"tool_call_chunk", "server_tool_call_chunk"} and dtype == btype:338        state["args"] = (state.get("args", "") or "") + (delta.get("args") or "")339        if delta.get("id") is not None:340            state["id"] = delta["id"]341        if delta.get("name") is not None:342            state["name"] = delta["name"]343    elif btype == dtype and "data" in delta:344        state["data"] = (state.get("data", "") or "") + (delta.get("data") or "")345        for key, value in delta.items():346            if key in {"type", "data"} or value is None:347                continue348            if key == "extras" and isinstance(value, dict):349                state["extras"] = {**(state.get("extras") or {}), **value}350            else:351                state[key] = value352    else:353        # Self-contained or already-finalized types: replace wholesale.354        state.clear()355        state.update(delta)356    return state357358359def finalize_tool_call_chunk(360    *,361    raw_args: str | None,362    id_: str | None,363    name: str | None,364    extras: dict[str, Any],365    finalized_type: str,366) -> FinalizedContentBlock:367    """Parse accumulated tool-chunk args into a finalized block.368369    Shared between the compat bridge's `_finalize_block` and the370    `ChatModelStream` end-of-stream sweep. Parses `raw_args` as JSON:371    on success builds the requested finalized type (`tool_call` or372    `server_tool_call`) with provider-specific fields (`extras`)373    preserved; on failure falls back to `invalid_tool_call` carrying374    the raw string so downstream consumers can still introspect the375    malformed payload.376377    Args:378        raw_args: Accumulated partial-JSON string; `None` or empty379            treated as `{}`.380        id_: Tool-call id collected across chunks.381        name: Tool name collected across chunks.382        extras: Provider-specific fields to carry onto the finalized383            block. Callers are responsible for having already dropped384            keys they don't want propagated (notably `type`, `id`,385            `name`, `args`, and `index` on client-side `tool_call`).386        finalized_type: `"tool_call"` or `"server_tool_call"`.387388    Returns:389        A `ToolCall`, `ServerToolCall`, or `InvalidToolCall` — the390        latter when `raw_args` is non-empty but not valid JSON.391    """392    raw = raw_args or "{}"393    try:394        parsed = json.loads(raw) if raw else {}395    except (json.JSONDecodeError, TypeError):396        invalid = InvalidToolCall(397            type="invalid_tool_call",398            id=id_,399            name=name,400            args=raw,401            error="Failed to parse tool call arguments as JSON",402        )403        invalid.update(extras)  # type: ignore[typeddict-item]404        return invalid405    if finalized_type == "tool_call":406        finalized_tc = ToolCall(407            type="tool_call",408            id=id_ or "",409            name=name or "",410            args=parsed,411        )412        finalized_tc.update(extras)  # type: ignore[typeddict-item]413        return finalized_tc414    finalized_stc = ServerToolCall(415        type="server_tool_call",416        id=id_ or "",417        name=name or "",418        args=parsed,419    )420    finalized_stc.update(extras)  # type: ignore[typeddict-item]421    return finalized_stc422423424def _finalize_block(block: CompatBlock) -> FinalizedContentBlock:425    """Promote chunk variants to their finalized form.426427    `tool_call_chunk` becomes `tool_call` — or `invalid_tool_call`428    if the accumulated `args` don't parse as JSON.429    `server_tool_call_chunk` becomes `server_tool_call` under the same430    rule.  Everything else passes through: text/reasoning blocks carry431    their accumulated snapshot, and self-contained types are already in432    their terminal shape.433    """434    btype = block.get("type")435    if btype in {"tool_call_chunk", "server_tool_call_chunk"}:436        # Carry provider-specific fields from the accumulated chunk onto437        # the finalized block. Drop the chunk-only keys we rewrite438        # explicitly. `index` is stripped on client-side439        # `tool_call` / `invalid_tool_call` finalizations to match v1440        # (`AIMessage.init_tool_calls` rebuilds tool_call blocks without441        # `index`), preventing `merge_lists` from re-merging further442        # chunks into an already-parsed args dict. `server_tool_call`443        # retains `index` because v1's `init_server_tool_calls`444        # finalizes in-place and preserves it.445        client_tool_call = btype == "tool_call_chunk"446        extras_drop = {"type", "id", "name", "args"}447        if client_tool_call:448            extras_drop |= {"index"}449        extras = {450            k: v for k, v in block.items() if k not in extras_drop and v is not None451        }452        return finalize_tool_call_chunk(453            raw_args=block.get("args"),454            id_=block.get("id"),455            name=block.get("name"),456            extras=extras,457            finalized_type="tool_call" if client_tool_call else "server_tool_call",458        )459    return _to_finalized_block(block)460461462# ---------------------------------------------------------------------------463# Metadata, usage, finish-reason464# ---------------------------------------------------------------------------465466467def _extract_start_metadata(response_metadata: dict[str, Any]) -> MessageMetadata:468    """Pull provider/model hints for the `message-start` event."""469    metadata: MessageMetadata = {}470    if "model_provider" in response_metadata:471        metadata["provider"] = response_metadata["model_provider"]472    if "model_name" in response_metadata:473        metadata["model"] = response_metadata["model_name"]474    return metadata475476477def _accumulate_usage(current: UsageInfo | None, delta: UsageMetadata) -> UsageInfo:478    """Sum usage counts and merge detail dicts across chunks.479480    `delta` is a chunk's `usage_metadata`; `current` is the running total.481    Both sides are read and written by literal key so the typed shape is482    preserved end to end — no `dict[str, Any]` detour.483    """484    new: UsageInfo = current if current is not None else {}485    if "input_tokens" in delta:486        new["input_tokens"] = new.get("input_tokens", 0) + delta["input_tokens"]487    if "output_tokens" in delta:488        new["output_tokens"] = new.get("output_tokens", 0) + delta["output_tokens"]489    if "total_tokens" in delta:490        new["total_tokens"] = new.get("total_tokens", 0) + delta["total_tokens"]491    input_details = delta.get("input_token_details")492    if input_details:493        merged_input = new.get("input_token_details", {})494        merged_input.update(input_details)495        new["input_token_details"] = merged_input496    output_details = delta.get("output_token_details")497    if output_details:498        merged_output = new.get("output_token_details", {})499        merged_output.update(output_details)500        new["output_token_details"] = merged_output501    return new502503504def _isolate_usage(usage: UsageInfo | None) -> UsageInfo | None:505    """Copy usage for the event so consumers can't mutate the source message.506507    The replay path (`message_to_events`) feeds the live `msg.usage_metadata`,508    so the emitted event must not share its dicts: copy the top level plus the509    nested `input_token_details` / `output_token_details` to de-alias it. The510    streaming accumulator already owns the dicts it builds, so the copy is a511    harmless no-op on that path.512    """513    if not usage:514        return None515    result: UsageInfo = usage.copy()516    input_details = result.get("input_token_details")517    if input_details is not None:518        result["input_token_details"] = input_details.copy()519    output_details = result.get("output_token_details")520    if output_details is not None:521        result["output_token_details"] = output_details.copy()522    return result523524525# ---------------------------------------------------------------------------526# Event builders527# ---------------------------------------------------------------------------528529530def _build_message_start(531    msg: BaseMessage,532    message_id: str | None,533) -> MessageStartData:534    start_data = MessageStartData(event="message-start", role="ai", id="")535    resolved_id = message_id if message_id is not None else getattr(msg, "id", None)536    if resolved_id:537        start_data["id"] = resolved_id538    start_metadata = _extract_start_metadata(msg.response_metadata or {})539    if start_metadata:540        start_data["metadata"] = start_metadata541    return start_data542543544def _build_message_finish(545    *,546    usage: UsageInfo | None,547    response_metadata: dict[str, Any] | None,548    additional_kwargs: dict[str, Any] | None = None,549) -> MessageFinishData:550    # Protocol 0.0.9 removed the top-level `reason` field from551    # `MessageFinishData`; the provider's raw `finish_reason` /552    # `stop_reason` now rides inside `metadata` alongside other553    # response metadata. Pass it through unchanged.554    finish_data: dict[str, Any] = {"event": "message-finish"}555    usage_info = _isolate_usage(usage)556    if usage_info is not None:557        finish_data["usage"] = usage_info558    if response_metadata:559        finish_data["metadata"] = dict(response_metadata)560    # `additional_kwargs` is an off-spec extension on the message-finish561    # event (parallel to `metadata`, which `MessageFinishData` also doesn't562    # formally declare but the consumer reads). It carries provider-side563    # kwargs that don't map onto a typed protocol field — notably Gemini's564    # `__gemini_function_call_thought_signatures__`, which the model565    # requires on follow-up turns to replay prior thinking. Without this,566    # streaming-assembled messages would silently drop data that567    # `ainvoke` preserves, breaking multi-turn streaming flows.568    if additional_kwargs:569        finish_data["additional_kwargs"] = dict(additional_kwargs)570    return cast("MessageFinishData", finish_data)571572573def _finalize_and_build_finish(574    wire_idx: int,575    block: CompatBlock,576) -> MessagesData:577    """Finalize a block and wrap it in a `content-block-finish` event."""578    return ContentBlockFinishData(579        event="content-block-finish",580        index=wire_idx,581        content=_finalize_block(block),582    )583584585# ---------------------------------------------------------------------------586# Main generators587# ---------------------------------------------------------------------------588589590def chunks_to_events(591    chunks: Iterator[ChatGenerationChunk],592    *,593    message_id: str | None = None,594) -> Iterator[MessagesData]:595    """Convert a stream of `ChatGenerationChunk` to protocol events.596597    Blocks are tracked independently by source-side identifier. Providers598    such as Anthropic can interleave parallel tool-call chunks by index, so599    each first-seen block gets a `content-block-start`, deltas keep their600    stable wire index, and all open blocks are finalized at message end.601    Source-side identifiers (from the block's `index` field, which may be602    int or string) are translated to sequential `uint` wire indices.603604    Args:605        chunks: Iterator of `ChatGenerationChunk` from `_stream()`.606        message_id: Optional stable message ID.607608    Yields:609        `MessagesData` lifecycle events.610    """611    started = False612    blocks: dict[Any, tuple[int, CompatBlock]] = {}613    next_wire_idx = 0614    usage: UsageInfo | None = None615    response_metadata: dict[str, Any] = {}616    additional_kwargs: dict[str, Any] = {}617618    for chunk in chunks:619        msg = chunk.message620        if not isinstance(msg, AIMessageChunk):621            continue622623        # The v1 `stream()` wrapper merges `generation_info` into624        # `response_metadata` before yielding (`chat_models.py` via625        # `_gen_info_and_msg_metadata`). We bypass that wrapper by reading626        # `_stream` directly, so reproduce the merge here with the same627        # priority: `generation_info` first, then `message.response_metadata`628        # overlays. This is how provider fields like `model_name`,629        # `system_fingerprint`, and `finish_reason` reach the bridge when630        # a provider emits them via `generation_info` instead of the631        # message's `response_metadata`.632        merged_rm: dict[str, Any] = {633            **(chunk.generation_info or {}),634            **(msg.response_metadata or {}),635        }636        if merged_rm:637            response_metadata.update(merged_rm)638639        # Carry chunks' `additional_kwargs` through to the assembled640        # message. Provider-side fields that don't map onto a typed641        # protocol block (e.g. Gemini's per-tool-call thought signatures)642        # live here on non-streaming `ainvoke` results; dropping them on643        # the streaming path silently diverges multi-turn behavior. Use644        # `merge_dicts` because the same key can arrive in pieces across645        # chunks (e.g. an accumulating `function_call`), matching how646        # `AIMessageChunk` merges itself.647        if msg.additional_kwargs:648            additional_kwargs = merge_dicts(additional_kwargs, msg.additional_kwargs)649650        if not started:651            started = True652            yield _build_message_start(msg, message_id)653654        for key, block in _iter_protocol_blocks(msg):655            if key not in blocks:656                wire_idx = next_wire_idx657                next_wire_idx += 1658                blocks[key] = (wire_idx, dict(block))659                yield ContentBlockStartData(660                    event="content-block-start",661                    index=wire_idx,662                    content=_start_skeleton(block),663                )664            else:665                wire_idx, existing = blocks[key]666                blocks[key] = (wire_idx, _accumulate(existing, block))667            if _should_emit_delta(block):668                wire_idx, current = blocks[key]669                is_block_delta = block.get("type") in {670                    "tool_call_chunk",671                    "server_tool_call_chunk",672                }673                delta_source = current if is_block_delta else block674                yield ContentBlockDeltaData(675                    event="content-block-delta",676                    index=wire_idx,677                    delta=_to_content_delta(delta_source or block),678                )679680        if msg.usage_metadata:681            usage = _accumulate_usage(usage, msg.usage_metadata)682683    if not started:684        return685686    for wire_idx, block in blocks.values():687        yield _finalize_and_build_finish(wire_idx, block)688689    yield _build_message_finish(690        usage=usage,691        response_metadata=response_metadata,692        additional_kwargs=additional_kwargs,693    )694695696async def achunks_to_events(697    chunks: AsyncIterator[ChatGenerationChunk],698    *,699    message_id: str | None = None,700) -> AsyncIterator[MessagesData]:701    """Async variant of `chunks_to_events`."""702    started = False703    blocks: dict[Any, tuple[int, CompatBlock]] = {}704    next_wire_idx = 0705    usage: UsageInfo | None = None706    response_metadata: dict[str, Any] = {}707    additional_kwargs: dict[str, Any] = {}708709    async for chunk in chunks:710        msg = chunk.message711        if not isinstance(msg, AIMessageChunk):712            continue713714        # See sync twin for rationale: merge `generation_info` into the715        # accumulated `response_metadata` with the same priority as the716        # v1 `stream()` wrapper.717        merged_rm: dict[str, Any] = {718            **(chunk.generation_info or {}),719            **(msg.response_metadata or {}),720        }721        if merged_rm:722            response_metadata.update(merged_rm)723724        # See sync twin: carry chunk `additional_kwargs` through so725        # provider-specific data (e.g. Gemini thought signatures) reaches726        # the assembled message instead of being dropped.727        if msg.additional_kwargs:728            additional_kwargs = merge_dicts(additional_kwargs, msg.additional_kwargs)729730        if not started:731            started = True732            yield _build_message_start(msg, message_id)733734        for key, block in _iter_protocol_blocks(msg):735            if key not in blocks:736                wire_idx = next_wire_idx737                next_wire_idx += 1738                blocks[key] = (wire_idx, dict(block))739                yield ContentBlockStartData(740                    event="content-block-start",741                    index=wire_idx,742                    content=_start_skeleton(block),743                )744            else:745                wire_idx, existing = blocks[key]746                blocks[key] = (wire_idx, _accumulate(existing, block))747            if _should_emit_delta(block):748                wire_idx, current = blocks[key]749                is_block_delta = block.get("type") in {750                    "tool_call_chunk",751                    "server_tool_call_chunk",752                }753                delta_source = current if is_block_delta else block754                yield ContentBlockDeltaData(755                    event="content-block-delta",756                    index=wire_idx,757                    delta=_to_content_delta(delta_source or block),758                )759760        if msg.usage_metadata:761            usage = _accumulate_usage(usage, msg.usage_metadata)762763    if not started:764        return765766    for wire_idx, block in blocks.values():767        yield _finalize_and_build_finish(wire_idx, block)768769    yield _build_message_finish(770        usage=usage,771        response_metadata=response_metadata,772        additional_kwargs=additional_kwargs,773    )774775776def message_to_events(777    msg: BaseMessage,778    *,779    message_id: str | None = None,780) -> Iterator[MessagesData]:781    """Replay a finalized message as a synthetic event lifecycle.782783    For a message returned whole (from a graph node, checkpoint, or784    cache), produce the same `message-start` / per-block /785    `message-finish` event stream a live call would produce.  Consumers786    downstream see a uniform event shape regardless of source.787788    Text and reasoning blocks emit a single `content-block-delta` with789    the full accumulated content.  Already-finalized blocks (tool_call,790    server_tool_call, image, etc.) skip the delta and rely on the791    `content-block-finish` event alone.792793    Args:794        msg: The finalized message — typically an `AIMessage`.795        message_id: Optional stable message ID; falls back to `msg.id`.796797    Yields:798        `MessagesData` lifecycle events.799    """800    response_metadata = msg.response_metadata or {}801    yield _build_message_start(msg, message_id)802803    for wire_idx, (_key, block) in enumerate(_iter_protocol_blocks(msg)):804        yield ContentBlockStartData(805            event="content-block-start",806            index=wire_idx,807            content=_start_skeleton(block),808        )809        if _should_emit_delta(block):810            yield ContentBlockDeltaData(811                event="content-block-delta",812                index=wire_idx,813                delta=_to_content_delta(block),814            )815        yield ContentBlockFinishData(816            event="content-block-finish",817            index=wire_idx,818            content=_finalize_block(block),819        )820821    yield _build_message_finish(822        usage=getattr(msg, "usage_metadata", None),823        response_metadata=response_metadata,824    )825826827async def amessage_to_events(828    msg: BaseMessage,829    *,830    message_id: str | None = None,831) -> AsyncIterator[MessagesData]:832    """Async variant of `message_to_events`."""833    for event in message_to_events(msg, message_id=message_id):834        yield event835836837__all__ = [838    "CompatBlock",839    "achunks_to_events",840    "amessage_to_events",841    "chunks_to_events",842    "finalize_tool_call_chunk",843    "message_to_events",844]