"""Per-message streaming objects for content-block protocol events. `ChatModelStream` is the synchronous variant returned by `BaseChatModel.stream_events(version="v3")`. `AsyncChatModelStream` is the asynchronous variant returned by `BaseChatModel.astream_events(version="v3")`. Both expose typed projection properties (`.text`, `.reasoning`, `.tool_calls`, `.usage`, `.output`) that accumulate protocol events as they arrive. Projections can be iterated for deltas or drained for the final accumulated value. Raw protocol events are also available via direct iteration on the stream object (replay-buffer semantics — multiple independent consumers supported). """ from __future__ import annotations import asyncio import contextlib from typing import TYPE_CHECKING, Any, cast from langchain_core.language_models._compat_bridge import finalize_tool_call_chunk from langchain_core.messages import AIMessage if TYPE_CHECKING: from collections.abc import Awaitable, Callable, Generator, Iterator, Mapping from langchain_protocol.protocol import ( ContentBlockDeltaData, ContentBlockFinishData, FinalizedContentBlock, InvalidToolCall, MessageFinishData, MessageMetadata, MessagesData, MessageStartData, ReasoningContentBlock, ServerToolCallChunk, TextContentBlock, ToolCall, ToolCallChunk, UsageInfo, ) from typing_extensions import Self # --------------------------------------------------------------------------- # Tool-call chunk helpers (shared by tool_call_chunk and server_tool_call_chunk) # --------------------------------------------------------------------------- def _merge_chunk_into_store( store: dict[int, dict[str, Any]], idx: int, block: dict[str, Any], ) -> None: """Merge a tool-call-chunk delta: sticky id/name, concat args.""" existing = store.get(idx, {}) if block.get("id") and "id" not in existing: existing["id"] = block["id"] if block.get("name") and "name" not in existing: existing["name"] = block["name"] existing["args"] = existing.get("args", "") + (block.get("args") or "") store[idx] = existing def _merge_block_delta_into_store( store: dict[int, dict[str, Any]], idx: int, fields: dict[str, Any], ) -> None: """Shallow-merge a block-delta snapshot into an indexed chunk store.""" existing = store.get(idx, {}) for key, value in fields.items(): if value is not None: existing[key] = value store[idx] = existing def _event_content_block(data: Mapping[str, Any]) -> dict[str, Any] | None: """Return start/finish content, tolerating the pre-delta field name.""" block = data.get("content") or data.get("content_block") return block if isinstance(block, dict) else None def _legacy_block_to_delta(block: Mapping[str, Any]) -> dict[str, Any]: """Convert the old content-block delta shape to an explicit delta.""" btype = block.get("type") if btype == "text": return {"type": "text-delta", "text": block.get("text", "")} if btype == "reasoning": return { "type": "reasoning-delta", "reasoning": block.get("reasoning", ""), } if "data" in block: delta = {"type": "data-delta", "data": block.get("data", "")} if block.get("encoding") == "base64": delta["encoding"] = "base64" return delta return {"type": "legacy-block-delta", "fields": block} def _event_delta(data: Mapping[str, Any]) -> dict[str, Any] | None: """Return an explicit delta, converting legacy content-block deltas.""" delta = data.get("delta") if isinstance(delta, dict): return delta block = data.get("content_block") if isinstance(block, dict): return _legacy_block_to_delta(block) return None def _sweep_chunk_store( store: dict[int, dict[str, Any]], *, finalized_type: str, finalized_blocks: dict[int, FinalizedContentBlock], tool_calls_acc: list[ToolCall] | None, invalid_acc: list[InvalidToolCall], ) -> None: """Parse each unswept chunk's `args`; record as `finalized_type` or invalid. `tool_calls_acc` is only populated when `finalized_type == "tool_call"` (server-side calls don't surface through `.tool_calls`). Deliberately does not backfill `index` onto finalized tool-call blocks: matches v1 (`AIMessage.init_tool_calls` drops `index` when substituting `tool_call_chunk` → `tool_call`) and prevents `merge_lists` from re-merging further chunks into an already-parsed args dict. """ for idx in sorted(store): chunk = store[idx] # Carry over any non-finalize-rewritten fields the chunk collected # (e.g., `extras`). `_merge_chunk_into_store` only populates # `id` / `name` / `args`, so this is empty in practice today; # future provider-specific fields would flow through here. extras = { k: v for k, v in chunk.items() if k not in {"type", "id", "name", "args"} and v is not None } final_block = finalize_tool_call_chunk( raw_args=chunk.get("args"), id_=chunk.get("id"), name=chunk.get("name"), extras=extras, finalized_type=finalized_type, ) if final_block["type"] == "invalid_tool_call": invalid_acc.append(final_block) elif tool_calls_acc is not None and finalized_type == "tool_call": tool_calls_acc.append(cast("ToolCall", final_block)) finalized_blocks[idx] = final_block store.clear() # --------------------------------------------------------------------------- # Projection base — shared producer API # --------------------------------------------------------------------------- class _ProjectionBase: """Shared state and producer API for sync and async projections. The `push` / `complete` / `fail` methods are the producer-side API — called by the stream as events arrive. Subclasses add the consumer protocol (sync iteration or async iteration + await). `done` and `error` are safe read-only views of the terminal state for iterators and other siblings that need to observe lifecycle without reaching into the underlying fields. """ __slots__ = ("_deltas", "_done", "_error", "_final_set", "_final_value") def __init__(self) -> None: """Initialize empty projection state.""" self._deltas: list[Any] = [] self._final_value: Any = None self._final_set: bool = False self._done: bool = False self._error: BaseException | None = None @property def done(self) -> bool: """Whether the projection has finished (successfully or via error).""" return self._done @property def error(self) -> BaseException | None: """The terminal error, if any.""" return self._error def push(self, delta: Any) -> None: """Append a delta value. Producer-side API.""" self._deltas.append(delta) def complete(self, final_value: Any) -> None: """Set the final accumulated value and mark as done. Producer-side API.""" self._final_value = final_value self._final_set = True self._done = True def fail(self, error: BaseException) -> None: """Mark as errored. Producer-side API.""" self._error = error self._done = True # --------------------------------------------------------------------------- # Sync projections # --------------------------------------------------------------------------- class SyncProjection(_ProjectionBase): """Sync iterable of deltas with pull-based backpressure. Follows the same `_request_more` convention as langgraph's `EventLog`: when the cursor catches up to the buffer and the projection is not done, it calls `_request_more()` to pull more events from the producer. Each call to `__iter__` creates a new cursor at position 0. Multiple iterators replay all deltas from the start. """ __slots__ = ("_ensure_started", "_request_more") def __init__(self) -> None: """Initialize with no pull callback.""" super().__init__() self._ensure_started: Callable[[], None] | None = None self._request_more: Callable[[], bool] | None = None def set_start(self, cb: Callable[[], None] | None) -> None: """Install a lazy-start callback invoked on first consumption.""" self._ensure_started = cb def set_request_more(self, cb: Callable[[], bool] | None) -> None: """Install the pull callback the iterator uses to drain the source.""" self._request_more = cb def __iter__(self) -> Iterator[Any]: """Yield deltas, pulling via `_request_more` when caught up.""" if self._ensure_started is not None: self._ensure_started() cursor = 0 while True: if cursor < len(self._deltas): yield self._deltas[cursor] cursor += 1 elif self._error is not None: raise self._error elif self._done: return elif self._request_more is not None: while cursor >= len(self._deltas) and not self._done: if not self._request_more(): break if cursor >= len(self._deltas): if self._error is not None: raise self._error return else: return def get(self) -> Any: """Drain via `_request_more` and return the final value.""" if self._ensure_started is not None: self._ensure_started() if not self._done and self._request_more is not None: while not self._done: if not self._request_more(): break if self._error is not None: raise self._error return self._final_value class SyncTextProjection(SyncProjection): """String-specialized sync projection. Adds `__str__`, `__bool__`, `__repr__` for ergonomic use with `.text` and `.reasoning` projections. """ __slots__ = () def __str__(self) -> str: """Drain and return the full accumulated string.""" val = self.get() return val if val is not None else "" def __bool__(self) -> bool: """Return whether any deltas have been pushed.""" return len(self._deltas) > 0 def __repr__(self) -> str: """Return repr of the accumulated text so far.""" if self._final_set: return repr(self._final_value) return repr("".join(self._deltas)) # --------------------------------------------------------------------------- # Async projection # --------------------------------------------------------------------------- class AsyncProjection(_ProjectionBase): """Async iterable of deltas that is also awaitable for the final value. Uses an `asyncio.Event` to notify consumers of state changes. Each waiter — the awaitable (`__await__`) and each async iterator cursor — shares the event and re-checks its own condition on wake. The event is cleared before a waiter awaits, so stale "something happened" signals don't cause spin loops. This is single-loop only — producers and consumers must share an event loop. If cross-thread wake is ever required, revert to a list-of-futures pattern with `call_soon_threadsafe`. """ __slots__ = ("_arequest_more", "_ensure_started", "_event") def __init__(self) -> None: """Initialize with an un-set event and no pump callback.""" super().__init__() self._event = asyncio.Event() self._arequest_more: Callable[[], Awaitable[bool]] | None = None self._ensure_started: Callable[[], Awaitable[None]] | None = None def set_start(self, cb: Callable[[], Awaitable[None]] | None) -> None: """Install a lazy-start callback invoked on first consumption.""" self._ensure_started = cb def set_arequest_more(self, cb: Callable[[], Awaitable[bool]] | None) -> None: """Wire the async pull callback iterators use to drive the source. Mirrors `SyncProjection.set_request_more`. Under caller-driven streaming, consumers call this callback when their buffer is empty so that the owning graph advances one step. Args: cb: Async no-arg callable returning `True` when a new event was produced, `False` when the source is exhausted. Pass `None` to unwire. """ self._arequest_more = cb def push(self, delta: Any) -> None: """Append a delta and notify waiters.""" super().push(delta) self._event.set() def complete(self, final_value: Any) -> None: """Set the final value, mark done, and notify waiters.""" super().complete(final_value) self._event.set() def fail(self, error: BaseException) -> None: """Mark errored and notify waiters.""" super().fail(error) self._event.set() # -- Async iterable (yields deltas) ------------------------------------ def __aiter__(self) -> _AsyncProjectionIterator: """Return an async iterator over deltas.""" return _AsyncProjectionIterator(self) # -- Awaitable (returns final value) ----------------------------------- def __await__(self) -> Generator[Any, None, Any]: """Await the final accumulated value.""" return self._await_impl().__await__() async def _await_impl(self) -> Any: """Wait until the final value is set and return it. When a caller-driven pump is wired via `set_arequest_more`, drive it instead of blocking on `self._event`; otherwise fall back to the event (used by tests that dispatch manually). """ if self._ensure_started is not None: await self._ensure_started() while not self._final_set: if self._error is not None: raise self._error if self._arequest_more is not None: if not await self._arequest_more() and not self._final_set: # Pump exhausted without completing this projection — # nothing more will arrive. Return current state and # let callers observe the missing final via the # returned None / unset error. break else: self._event.clear() await self._event.wait() if self._error is not None: raise self._error return self._final_value class _AsyncProjectionIterator: """Async iterator over an `AsyncProjection`'s deltas.""" __slots__ = ("_offset", "_proj") def __init__(self, proj: AsyncProjection) -> None: """Initialize cursor at position 0.""" self._proj = proj self._offset = 0 def __aiter__(self) -> _AsyncProjectionIterator: """Return self for the async iteration protocol.""" return self async def __anext__(self) -> Any: """Return the next delta, awaiting if necessary. When the projection has an `_arequest_more` pump wired, drain it in an inner loop (mirrors `SyncProjection.__iter__`) until this cursor advances or the pump reports exhaustion. Without a pump, fall back to waiting on the shared event. """ proj = self._proj if proj._ensure_started is not None: # noqa: SLF001 await proj._ensure_started() # noqa: SLF001 while True: # Direct access to the projection's internal list/event is # intentional — the iterator is the projection's sidekick and # depends on reading the shared buffer by cursor. if self._offset < len(proj._deltas): # noqa: SLF001 item = proj._deltas[self._offset] # noqa: SLF001 self._offset += 1 return item if proj.error is not None: raise proj.error if proj.done: raise StopAsyncIteration if proj._arequest_more is not None: # noqa: SLF001 # Caller-driven: drive the producer. Pump may land new # deltas for a sibling projection — loop until our cursor # advances, the projection terminates, or the pump is # exhausted. while ( self._offset >= len(proj._deltas) # noqa: SLF001 and not proj.done ): if not await proj._arequest_more(): # noqa: SLF001 break if ( self._offset >= len(proj._deltas) # noqa: SLF001 and not proj.done ): if proj.error is not None: raise proj.error raise StopAsyncIteration else: proj._event.clear() # noqa: SLF001 await proj._event.wait() # noqa: SLF001 # --------------------------------------------------------------------------- # Sync stream # --------------------------------------------------------------------------- class _ChatModelStreamBase: """Shared state and event dispatch for chat-model streams. Holds accumulated protocol state (text, reasoning, tool calls, usage, metadata) and the event-dispatch machinery that drives the typed projections. `ChatModelStream` (sync) and `AsyncChatModelStream` (async) inherit from this base and add the projection types and consumer APIs for their flavor. """ # Projection instances — concrete subclasses create them as sync or # async variants in their own __init__ after calling super(). _text_proj: _ProjectionBase _reasoning_proj: _ProjectionBase _tool_calls_proj: _ProjectionBase def __init__( self, *, namespace: list[str] | None = None, node: str | None = None, message_id: str | None = None, ) -> None: self._namespace = namespace or [] self._node = node self._message_id = message_id # Accumulated state self._text_acc: str = "" self._reasoning_acc: str = "" # Per-block text / reasoning storage keyed by wire index. Used to # populate the finalized block payload without cross-contaminating # other blocks of the same type in the same message. Without # per-block storage the message-wide accumulator would bleed # earlier block text into later finalized blocks. self._text_per_block: dict[int, str] = {} self._reasoning_per_block: dict[int, str] = {} self._tool_call_chunks: dict[int, dict[str, Any]] = {} self._tool_calls_acc: list[ToolCall] = [] self._invalid_tool_calls_acc: list[InvalidToolCall] = [] self._server_tool_call_chunks: dict[int, dict[str, Any]] = {} # Ordered snapshot of every finalized block, keyed by event index. # Single source of truth for .output.content. Typed accumulators # (text/reasoning/tool_calls/invalid_tool_calls) continue to serve # the public projections. self._blocks: dict[int, FinalizedContentBlock] = {} self._usage_value: UsageInfo | None = None self._start_metadata: MessageMetadata | None = None self._finish_metadata: dict[str, Any] | None = None self._additional_kwargs: dict[str, Any] | None = None self._done: bool = False self._error: BaseException | None = None self._output_message: AIMessage | None = None # Raw event replay buffer self._events: list[MessagesData] = [] # -- Common properties ------------------------------------------------ @property def namespace(self) -> list[str]: """Graph namespace path for this message.""" return self._namespace @property def node(self) -> str | None: """Graph node that produced this message.""" return self._node @property def message_id(self) -> str | None: """Stable message identifier.""" return self._message_id def set_message_id(self, message_id: str) -> None: """Assign the stable message identifier once the run starts. Called by the stream driver (`stream_events(version="v3")` / `astream_events(version="v3")`) after `on_chat_model_start` produces a run id. Not intended for end-user code. """ self._message_id = message_id @property def done(self) -> bool: """Whether the stream has finished.""" return self._done @property def has_events(self) -> bool: """Whether any protocol events have been recorded.""" return bool(self._events) @property def output_message(self) -> AIMessage | None: """The assembled message if the stream has finished, else `None`. Unlike `ChatModelStream.output` (which blocks until the stream finishes), this never pumps, blocks, or raises. Intended for the stream driver (`stream_events(version="v3")` and its async equivalent) to check whether the stream produced a message before firing `on_llm_end` callbacks. """ return self._output_message # -- Event ingestion (public) ------------------------------------------ def dispatch(self, event: Mapping[str, Any]) -> None: """Route a protocol event to the appropriate internal handler. Public entry point for feeding events into the stream. Called by the stream driver (the `stream_events(version="v3")` pump and its async equivalent) and by any observer or test that needs to inject protocol events. """ self._record_event(event) event_type = event.get("event") if event_type == "message-start": self._push_message_start(cast("MessageStartData", event)) elif event_type == "content-block-delta": self._push_content_block_delta(cast("ContentBlockDeltaData", event)) elif event_type == "content-block-finish": self._push_content_block_finish(cast("ContentBlockFinishData", event)) elif event_type == "message-finish": self._finish(cast("MessageFinishData", event)) elif event_type == "error": self.fail(RuntimeError(event.get("message", "Unknown error"))) # content-block-start is informational — no accumulation needed # -- Internal push API (called by dispatch) ---------------------------- def _record_event(self, event: Mapping[str, Any]) -> None: """Append a raw event to the replay buffer.""" self._events.append(cast("MessagesData", event)) def _push_message_start(self, data: MessageStartData) -> None: """Process a `message-start` event.""" self._start_metadata = data.get("metadata") message_id = data.get("id") if message_id: self._message_id = message_id def _push_content_block_delta(self, data: ContentBlockDeltaData) -> None: """Process a `content-block-delta` event.""" delta = _event_delta(data) if delta is None: return event_idx = data.get("index") dtype = delta.get("type", "") if dtype == "text-delta": delta_text = delta.get("text", "") if delta_text: self._text_acc += delta_text if event_idx is not None: self._text_per_block[event_idx] = ( self._text_per_block.get(event_idx, "") + delta_text ) self._text_proj.push(delta_text) elif dtype == "reasoning-delta": delta_r = delta.get("reasoning", "") if delta_r: self._reasoning_acc += delta_r if event_idx is not None: self._reasoning_per_block[event_idx] = ( self._reasoning_per_block.get(event_idx, "") + delta_r ) self._reasoning_proj.push(delta_r) elif dtype == "block-delta": fields = delta.get("fields") if not isinstance(fields, dict): return btype = fields.get("type", "") if btype == "tool_call_chunk": tcc = cast("ToolCallChunk", fields) idx = data.get("index") if idx is None: idx = tcc.get("index", len(self._tool_call_chunks)) # type: ignore[unreachable] _merge_block_delta_into_store(self._tool_call_chunks, idx, dict(tcc)) chunk_block: ToolCallChunk = { "type": "tool_call_chunk", "id": tcc.get("id"), "name": tcc.get("name"), "args": tcc.get("args"), } if "index" in tcc: chunk_block["index"] = tcc["index"] self._tool_calls_proj.push(chunk_block) elif btype == "server_tool_call_chunk": stcc = cast("ServerToolCallChunk", fields) idx = data.get("index") if idx is None: idx = len(self._server_tool_call_chunks) # type: ignore[unreachable] _merge_block_delta_into_store( self._server_tool_call_chunks, idx, dict(stcc), ) elif dtype == "legacy-block-delta": fields = delta.get("fields") if not isinstance(fields, dict): return btype = fields.get("type", "") if btype == "tool_call_chunk": tcc = cast("ToolCallChunk", fields) idx = data.get("index") if idx is None: idx = tcc.get("index", len(self._tool_call_chunks)) # type: ignore[unreachable] _merge_chunk_into_store(self._tool_call_chunks, idx, dict(tcc)) legacy_chunk_block: ToolCallChunk = { "type": "tool_call_chunk", "id": tcc.get("id"), "name": tcc.get("name"), "args": tcc.get("args"), } if "index" in tcc: legacy_chunk_block["index"] = tcc["index"] self._tool_calls_proj.push(legacy_chunk_block) elif btype == "server_tool_call_chunk": stcc = cast("ServerToolCallChunk", fields) idx = data.get("index") if idx is None: idx = len(self._server_tool_call_chunks) # type: ignore[unreachable] _merge_chunk_into_store( self._server_tool_call_chunks, idx, dict(stcc), ) elif dtype == "data-delta": # Binary/modal payload deltas are reflected in the final # content-block finish event; there is no dedicated projection. return else: # Transitional legacy path for old `content_block` deltas that # should not be reachable after `_event_delta` conversion, kept # here for custom in-tree test fixtures or third-party emitters. block = data.get("content_block") if not isinstance(block, dict): return btype = block.get("type", "") if btype != "tool_call_chunk": return tcc = cast("ToolCallChunk", block) idx = data.get("index") if idx is None: idx = tcc.get("index", len(self._tool_call_chunks)) # type: ignore[unreachable] _merge_chunk_into_store(self._tool_call_chunks, idx, dict(tcc)) fallback_chunk_block: ToolCallChunk = { "type": "tool_call_chunk", "id": tcc.get("id"), "name": tcc.get("name"), "args": tcc.get("args"), } if "index" in tcc: fallback_chunk_block["index"] = tcc["index"] self._tool_calls_proj.push(fallback_chunk_block) def _resolve_block_text(self, idx: int | None, full_text: str) -> str: """Return authoritative text for a single text block at `idx`. Prefers per-block delta accumulation; reconciles with the finish event's `full_text` when the provider emits authoritative text that differs from what the deltas built up. Does not mutate `self._text_acc` (the delta-sum accumulator) — the message-wide projection value is derived from per-block storage at `_finish` time, so reconciliation remains correct regardless of finish ordering across blocks. """ if idx is None: # No wire index — legacy behavior: use the message-wide # accumulator. Preserved for pre-index semantics; not # exercised by the compat bridge or any in-tree provider. if full_text and full_text != self._text_acc: self._text_acc = full_text return self._text_acc existing = self._text_per_block.get(idx, "") if full_text and full_text != existing: if not existing: # No deltas arrived for this block — surface the full # text as a single delta so the stream projection # reflects it. self._text_acc += full_text self._text_proj.push(full_text) elif full_text.startswith(existing): # Authoritative text extends the partial deltas — emit # the tail so delta consumers see the completion. tail = full_text[len(existing) :] self._text_acc += tail self._text_proj.push(tail) # else: authoritative text replaces the partial deltas # entirely. No corrective delta is emitted (semantics # would be ambiguous mid-stream). `_text_acc` is not # spliced — the final value is computed from per-block # storage at `_finish`, so this remains correct even when # other blocks have added to `_text_acc` in between. self._text_per_block[idx] = full_text return self._text_per_block.get(idx, "") def _resolve_block_reasoning(self, idx: int | None, full_r: str) -> str: """Return authoritative reasoning text for a single block at `idx`. Mirrors `_resolve_block_text` for the reasoning projection. """ if idx is None: if full_r and full_r != self._reasoning_acc: self._reasoning_acc = full_r return self._reasoning_acc existing = self._reasoning_per_block.get(idx, "") if full_r and full_r != existing: if not existing: self._reasoning_acc += full_r self._reasoning_proj.push(full_r) elif full_r.startswith(existing): tail = full_r[len(existing) :] self._reasoning_acc += tail self._reasoning_proj.push(tail) self._reasoning_per_block[idx] = full_r return self._reasoning_per_block.get(idx, "") def _push_content_block_finish(self, data: ContentBlockFinishData) -> None: """Process a `content-block-finish` event.""" block = _event_content_block(data) if block is None: return btype = block.get("type", "") idx = data.get("index") finalized: FinalizedContentBlock | None = None if btype == "text": text_block = cast("TextContentBlock", block) full_text = text_block.get("text", "") block_text = self._resolve_block_text(idx, full_text) finalized = cast( "FinalizedContentBlock", { **text_block, "type": "text", "text": block_text, }, ) elif btype == "reasoning": reasoning_block = cast("ReasoningContentBlock", block) full_r = reasoning_block.get("reasoning", "") block_reasoning = self._resolve_block_reasoning(idx, full_r) # Keep provider-specific fields alongside the accumulated # reasoning text. Anthropic's `signature` arrives under # `extras` and is required on follow-up turns. Only overwrite # `reasoning` when we have accumulated content; OpenAI can # emit a reasoning block with no text deltas, and writing an # empty string there makes downstream serializers synthesize # an empty summary entry. finalized_dict: dict[str, Any] = {**reasoning_block, "type": "reasoning"} if block_reasoning: finalized_dict["reasoning"] = block_reasoning finalized = cast("FinalizedContentBlock", finalized_dict) elif btype == "tool_call": tcb = cast("ToolCall", block) # Preserve provider-specific fields (extras, etc.) on the # content block. `_assemble_message` separately projects the # minimal {id, name, args, type} shape onto # `AIMessage.tool_calls`. Strip `index` to match v1 # (`AIMessage.init_tool_calls` rebuilds the block without # `index`); see `_finalize_block` in `_compat_bridge.py`. tc = cast( "ToolCall", { **{k: v for k, v in tcb.items() if k != "index"}, "type": "tool_call", "id": tcb.get("id", ""), "name": tcb.get("name", ""), "args": tcb.get("args", {}), }, ) self._tool_calls_acc.append(tc) if idx is not None and idx in self._tool_call_chunks: del self._tool_call_chunks[idx] finalized = tc elif btype == "invalid_tool_call": itc = cast("InvalidToolCall", block) # Strip `index` on the stored block to stay symmetric with # the `tool_call` path. itc = cast( "InvalidToolCall", {k: v for k, v in itc.items() if k != "index"}, ) self._invalid_tool_calls_acc.append(itc) # Critical: drop the stale chunk so _finish's sweep doesn't revive # it as an empty-args ToolCall. if idx is not None and idx in self._tool_call_chunks: del self._tool_call_chunks[idx] if idx is not None and idx in self._server_tool_call_chunks: del self._server_tool_call_chunks[idx] finalized = itc elif btype in { "server_tool_call", "server_tool_result", "image", "audio", "video", "file", "non_standard", }: if btype == "server_tool_call" and idx is not None: self._server_tool_call_chunks.pop(idx, None) finalized = cast("FinalizedContentBlock", block) if finalized is not None and idx is not None: # Backfill the wire index onto the finalized block when the # source didn't supply one. `langchain_core.utils._merge`'s # block-merger (used by `AIMessageChunk.__add__` / # `add_ai_message_chunks`) keys on `block["index"]` to group # deltas into the same output block — without it, a v2- # assembled `AIMessage` that later re-enters the chunk # aggregation path won't merge cleanly. Client-side # `tool_call` / `invalid_tool_call` blocks are excluded: v1 # finalization drops `index` on them so further deltas # cannot clobber already-parsed args, and v2 mirrors that. if btype not in {"tool_call", "invalid_tool_call"}: finalized.setdefault("index", idx) self._blocks[idx] = finalized def _finish(self, data: MessageFinishData) -> None: """Process a `message-finish` event.""" self._done = True self._usage_value = data.get("usage") self._finish_metadata = cast("dict[str, Any] | None", data.get("metadata")) # Off-spec extension carrying provider-side `additional_kwargs` # that don't map onto a typed protocol field (e.g. Gemini's # `__gemini_function_call_thought_signatures__`). The compat # bridge emits this on `message-finish` so the assembled message # carries the same data `ainvoke` would have preserved. self._additional_kwargs = cast( "dict[str, Any] | None", cast("dict[str, Any]", data).get("additional_kwargs"), ) # Finalize any unswept chunks — both client- and server-side. _sweep_chunk_store( self._tool_call_chunks, finalized_type="tool_call", finalized_blocks=self._blocks, tool_calls_acc=self._tool_calls_acc, invalid_acc=self._invalid_tool_calls_acc, ) _sweep_chunk_store( self._server_tool_call_chunks, finalized_type="server_tool_call", finalized_blocks=self._blocks, tool_calls_acc=None, invalid_acc=self._invalid_tool_calls_acc, ) # Prefer the per-block sum when any indexed text / reasoning # arrived — it stays correct regardless of finish ordering and # of whether finish events carried authoritative text that # differed from the deltas. Fall back to the delta-sum # accumulator only for the legacy no-index path. if self._text_per_block: text_final = "".join( self._text_per_block[i] for i in sorted(self._text_per_block) ) else: text_final = self._text_acc if self._reasoning_per_block: reasoning_final = "".join( self._reasoning_per_block[i] for i in sorted(self._reasoning_per_block) ) else: reasoning_final = self._reasoning_acc self._text_proj.complete(text_final) self._reasoning_proj.complete(reasoning_final) self._tool_calls_proj.complete(self._tool_calls_acc) self._output_message = self._assemble_message() def fail(self, error: BaseException) -> None: """Mark the stream as errored and propagate to all projections. Public API — called by the stream driver (`stream_events(version="v3")` / `astream_events(version="v3")`) when the underlying producer raises, by `dispatch` when an `error` protocol event arrives, and by cancellation paths. """ self._done = True self._error = error self._text_proj.fail(error) self._reasoning_proj.fail(error) self._tool_calls_proj.fail(error) def _assemble_message(self) -> AIMessage: """Build an `AIMessage` from accumulated state. Content is built from `self._blocks`, an index-ordered snapshot of finalized protocol blocks. The bare-string fast path is used when the message has exactly one `text` block (the common chat case); otherwise content is a list of protocol-shape block dicts. """ content: Any if not self._blocks: # No protocol blocks ever arrived. Fall back to the accumulated # text (possibly empty) as bare-string content. content = self._text_acc else: # `ChatModelStream` is the v1 content-block surface: content # is always a list of protocol blocks when any block arrived. # Do not collapse a single text block down to a bare string — # that would drop block-level fields (`id`, `index`, # annotations, extras) that downstream serializers need to # round-trip the message on a follow-up turn. ordered_blocks = [self._blocks[idx] for idx in sorted(self._blocks)] content = [dict(b) for b in ordered_blocks] response_metadata: dict[str, Any] = {} if self._start_metadata: if "provider" in self._start_metadata: response_metadata["model_provider"] = self._start_metadata["provider"] if "model" in self._start_metadata: response_metadata["model_name"] = self._start_metadata["model"] if self._finish_metadata: response_metadata.update(self._finish_metadata) # Pin `output_version` last: `stream_events(version="v3")` always # assembles content as v1 protocol blocks, regardless of the # provider's configured output format. # A provider-supplied `output_version` in finish metadata (e.g. # `"responses/v1"` from `ChatOpenAI(use_responses_api=True, ...)`) would # otherwise cause `AIMessage.content_blocks` to re-run the wrong # translator on already-v1 content. response_metadata["output_version"] = "v1" tool_calls = [ { "id": tc.get("id", ""), "name": tc.get("name", ""), "args": tc.get("args", {}), "type": "tool_call", } for tc in self._tool_calls_acc ] invalid_tool_calls = [ { "type": "invalid_tool_call", "id": itc.get("id") or None, "name": itc.get("name") or None, "args": itc.get("args") or None, "error": itc.get("error"), } for itc in self._invalid_tool_calls_acc ] message_kwargs: dict[str, Any] = { "content": content, "id": self._message_id, "tool_calls": tool_calls, "invalid_tool_calls": invalid_tool_calls, "usage_metadata": self._usage_value, "response_metadata": response_metadata, } if self._additional_kwargs: message_kwargs["additional_kwargs"] = dict(self._additional_kwargs) return AIMessage(**message_kwargs) # --------------------------------------------------------------------------- # Sync stream # --------------------------------------------------------------------------- class ChatModelStream(_ChatModelStreamBase): """Synchronous per-message streaming object for a single LLM response. Returned by `BaseChatModel.stream_events(version="v3")`. Content-block protocol events are fed into this object and accumulated into typed projections. Projections (always return the same cached object): - `.text` — iterable of `str` deltas; `str()` for full text - `.reasoning` — same as `.text` for reasoning content - `.tool_calls` — iterable of `ToolCallChunk` deltas; `.get()` returns `list[ToolCall]` - `.output` — blocking property, returns assembled `AIMessage` Usage info is available on `.output.usage_metadata` once the stream has finished. !!! note "Output shape is always v1 content blocks" `.output.content` is always a list of v1 protocol blocks (text, reasoning, tool_call, image, …), regardless of the underlying model's `output_version` setting. That attribute only controls the legacy `stream()` / `astream()` / `invoke()` paths; `ChatModelStream` is built on the content-block protocol and emits v1 shapes by construction. Raw event iteration:: for event in stream: print(event) # MessagesData dicts """ _text_proj: SyncTextProjection _reasoning_proj: SyncTextProjection _tool_calls_proj: SyncProjection def __init__( # noqa: D107 self, *, namespace: list[str] | None = None, node: str | None = None, message_id: str | None = None, ) -> None: super().__init__(namespace=namespace, node=node, message_id=message_id) # Projections — created eagerly self._text_proj = SyncTextProjection() self._reasoning_proj = SyncTextProjection() self._tool_calls_proj = SyncProjection() # Pull callback (set by bind_pump or set_request_more) self._ensure_started: Callable[[], None] | None = None self._request_more: Callable[[], bool] | None = None # -- Pump/pull wiring -------------------------------------------------- def bind_pump(self, pump_one: Callable[[], bool]) -> None: """Bind a pump for standalone streaming. Delegates to `set_request_more`. Used by `BaseChatModel.stream_events(version="v3")`. """ self.set_request_more(pump_one) def set_start(self, cb: Callable[[], None] | None) -> None: """Install a lazy-start callback on this stream and its projections.""" self._ensure_started = cb self._text_proj.set_start(cb) self._reasoning_proj.set_start(cb) self._tool_calls_proj.set_start(cb) def set_request_more(self, cb: Callable[[], bool]) -> None: """Set the pull callback on this stream and all its projections. Used by langgraph's `GraphRunStream._wire_request_more` to connect the shared graph pump. """ self._request_more = cb self._text_proj.set_request_more(cb) self._reasoning_proj.set_request_more(cb) self._tool_calls_proj.set_request_more(cb) # -- Public projections ------------------------------------------------ @property def text(self) -> SyncTextProjection: """Text content — iterable of `str` deltas, `str()` for full.""" return self._text_proj @property def reasoning(self) -> SyncTextProjection: """Reasoning content — same interface as :attr:`text`.""" return self._reasoning_proj @property def tool_calls(self) -> SyncProjection: """Tool calls — iterable of `ToolCallChunk` deltas. `.get()` returns finalized `list[ToolCall]`. """ return self._tool_calls_proj @property def output(self) -> AIMessage: """Assembled `AIMessage` — blocks until the stream finishes.""" self._drain() if self._error is not None: raise self._error if self._output_message is None: msg = "Stream finished without producing a message" raise RuntimeError(msg) return self._output_message # -- Raw event iteration (replay buffer) ------------------------------- def __iter__(self) -> Iterator[MessagesData]: """Iterate raw protocol events with replay-buffer semantics.""" if self._ensure_started is not None: self._ensure_started() cursor = 0 while True: if cursor < len(self._events): yield self._events[cursor] cursor += 1 elif self._error is not None: raise self._error elif self._done: return elif self._request_more is not None: while cursor >= len(self._events) and not self._done: if not self._request_more(): break if cursor >= len(self._events): if self._error is not None: raise self._error return else: return # -- Internal helpers -------------------------------------------------- def _drain(self) -> None: """Pull all remaining events until done.""" if self._done: return if self._ensure_started is not None: self._ensure_started() if self._request_more is not None: while not self._done: if not self._request_more(): break # --------------------------------------------------------------------------- # Async stream # --------------------------------------------------------------------------- class AsyncChatModelStream(_ChatModelStreamBase): """Asynchronous per-message streaming object for a single LLM response. Returned by `BaseChatModel.astream_events(version="v3")`. Content-block events are fed into this object by a background producer task. Projections: - `.text` — async iterable of text deltas; awaitable for full text - `.reasoning` — async iterable of reasoning deltas; awaitable - `.tool_calls` — async iterable of `ToolCallChunk` deltas; awaitable for `list[ToolCall]` - `.output` — awaitable for assembled `AIMessage` Usage info is available on `.output.usage_metadata` once the stream has finished. !!! note "Output shape is always v1 content blocks" The assembled message's content is always a list of v1 protocol blocks, regardless of the model's `output_version` setting — see `ChatModelStream` for the full rationale. The stream itself is awaitable (`msg = await stream`) and async-iterable (`async for event in stream`). """ _text_proj: AsyncProjection _reasoning_proj: AsyncProjection _tool_calls_proj: AsyncProjection def __init__( # noqa: D107 self, *, namespace: list[str] | None = None, node: str | None = None, message_id: str | None = None, ) -> None: super().__init__(namespace=namespace, node=node, message_id=message_id) self._text_proj = AsyncProjection() self._reasoning_proj = AsyncProjection() self._tool_calls_proj = AsyncProjection() self._output_proj = AsyncProjection() self._events_proj = AsyncProjection() self._ensure_started: Callable[[], Awaitable[None]] | None = None self._producer_task: asyncio.Task[None] | None = None # Teardown callback invoked by `aclose()` only when the producer # task was cancelled before its body ran (so the normal # `_produce` CancelledError handler — which fires # `on_llm_error` — never executed). Set by `astream_events(version="v3")`. self._on_aclose_fail: Callable[[BaseException], Awaitable[None]] | None = None # -- Pump/pull wiring (async) ------------------------------------------ def set_arequest_more(self, cb: Callable[[], Awaitable[bool]] | None) -> None: """Fan the async pump callback out to every projection. Used by langgraph's `AsyncGraphRunStream._wire_arequest_more` so cursors on `stream.text`, `stream.reasoning`, etc. can drive the shared graph pump when their buffer is empty. Args: cb: Async no-arg callable returning `True` when a new event was produced, `False` when the source is exhausted. Pass `None` to unwire. """ for proj in ( self._text_proj, self._reasoning_proj, self._tool_calls_proj, self._output_proj, self._events_proj, ): proj.set_arequest_more(cb) def set_start(self, cb: Callable[[], Awaitable[None]] | None) -> None: """Install a lazy-start callback on this stream and its projections.""" self._ensure_started = cb for proj in ( self._text_proj, self._reasoning_proj, self._tool_calls_proj, self._output_proj, self._events_proj, ): proj.set_start(cb) # -- Public projections ------------------------------------------------ @property def text(self) -> AsyncProjection: """Text content — async iterable of deltas, awaitable for full.""" return self._text_proj @property def reasoning(self) -> AsyncProjection: """Reasoning content — same interface as :attr:`text`.""" return self._reasoning_proj @property def tool_calls(self) -> AsyncProjection: """Tool calls — async iterable, awaitable for finalized list.""" return self._tool_calls_proj @property def output(self) -> AsyncProjection: """Assembled `AIMessage` — awaitable.""" return self._output_proj def __await__(self) -> Generator[Any, None, AIMessage]: """Await the assembled `AIMessage` and full producer lifecycle. The producer task is awaited after the output projection resolves so that post-stream work (notably `on_llm_end` callbacks) has run by the time the caller's `await` returns. """ return self._await_full().__await__() async def _await_full(self) -> AIMessage: if self._ensure_started is not None: await self._ensure_started() message: AIMessage = await self._output_proj if self._producer_task is not None: await self._producer_task return message def __aiter__(self) -> _AsyncProjectionIterator: """Iterate raw protocol events asynchronously.""" return _AsyncProjectionIterator(self._events_proj) # -- Cleanup ----------------------------------------------------------- async def aclose(self) -> None: """Cancel the background producer task and release resources. If a consumer cancels mid-stream or decides to stop iterating early, the producer task keeps pumping the provider HTTP call to completion because `asyncio.Task` has no implicit link to its awaiter. Call this method to cancel the producer explicitly; the stream transitions to an errored state with `CancelledError`. If the stream has already produced a message successfully (for example, after `await stream.output`), the producer may still be running post-stream work such as `on_llm_end` callbacks. In that case `aclose()` awaits the task rather than cancelling it — turning a successful run into a cancelled one would drop the end callback and corrupt tracing. Idempotent: safe to call multiple times, including after the stream has finished normally. Also invoked by the async context manager protocol on `__aexit__`. """ if self._ensure_started is not None and self._producer_task is None: await self._ensure_started() task = self._producer_task if task is None: return if task.done() and self._done: return we_cancelled = not (self._output_message is not None and self._error is None) if we_cancelled and not task.done(): task.cancel() # Wait for the task via a linked `Future`, not by awaiting the # task directly. Awaiting the task would raise `CancelledError` # in two indistinguishable cases: (1) the task we just cancelled # completed, (2) our caller cancelled us. `asyncio.Task.cancelling()` # disambiguates on 3.11+ but doesn't exist on 3.10. # # The `done_future` resolves with `None` whenever the task # finishes (any reason). It is not a `Task` itself, so its # `await` only raises when our caller is cancelled — giving us # a portable, unambiguous signal to propagate. if not task.done(): loop = asyncio.get_running_loop() done_future: asyncio.Future[None] = loop.create_future() def _link(_: asyncio.Task[None]) -> None: if not done_future.done(): done_future.set_result(None) task.add_done_callback(_link) try: await done_future finally: task.remove_done_callback(_link) # If the task was cancelled before `_produce` ran (e.g. # `astream_events(version="v3")` immediately followed by `aclose()`), the stream # never reached `_produce`'s CancelledError handler — its # projections are still pending and no end-of-lifecycle callback # has fired. Resolve both here so callers of `await stream.output` # don't hang and tracing sees a matching end event. if we_cancelled and not self._done: cancel_exc = asyncio.CancelledError() self.fail(cancel_exc) teardown = self._on_aclose_fail if teardown is not None: with contextlib.suppress(Exception): await teardown(cancel_exc) async def __aenter__(self) -> Self: """Enter the async context — returns self.""" return self async def __aexit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: object, ) -> None: """Exit the async context — cancels the producer via `aclose()`.""" del exc_type, exc, tb await self.aclose() # -- Internal API (extend base to drive async projections) ------------- def _record_event(self, event: Mapping[str, Any]) -> None: """Record event and push to async event replay projection.""" super()._record_event(event) self._events_proj.push(cast("MessagesData", event)) def _finish(self, data: MessageFinishData) -> None: """Finish base projections and async-only projections.""" super()._finish(data) self._output_proj.complete(self._output_message) self._events_proj.complete(self._events) def fail(self, error: BaseException) -> None: """Fail base projections and async-only projections.""" super().fail(error) self._output_proj.fail(error) self._events_proj.fail(error) __all__ = [ "AsyncChatModelStream", "AsyncProjection", "ChatModelStream", "SyncProjection", "SyncTextProjection", ]