diff --git a/libs/langgraph/langgraph/pregel/_messages_v2.py b/libs/langgraph/langgraph/pregel/_messages_v2.py new file mode 100644 index 000000000..d6ef70b39 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/_messages_v2.py @@ -0,0 +1,807 @@ +"""Protocol-native content-block message handler for StreamingHandler. + +Emits structured content-block lifecycle events (message-start, +content-block-start/delta/finish, message-finish) instead of raw +``(AIMessageChunk, metadata)`` tuples. The existing +:class:`~langgraph.pregel._messages.StreamMessagesHandler` is NOT +modified — this handler is only activated when +``__protocol_messages_stream`` is ``True`` in the run's configurable. +""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator, Callable, Iterator, Sequence +from dataclasses import dataclass, field +from typing import Any, TypeVar, cast +from uuid import UUID, uuid4 + +from langchain_core.callbacks import BaseCallbackHandler +from langchain_core.messages import AIMessageChunk, BaseMessage +from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult + +from langgraph._internal._constants import NS_SEP +from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM +from langgraph.pregel.protocol import StreamChunk +from langgraph.stream._types import ( + ContentBlockDeltaData, + ContentBlockFinishData, + ContentBlockStartData, + FinishReason, + InvalidToolCallBlock, + MessageErrorData, + MessageStartData, + ReasoningBlock, + TextBlock, + ToolCallBlock, + UsageInfo, +) + +try: + from langchain_core.tracers._streaming import _StreamingCallbackHandler +except ImportError: + _StreamingCallbackHandler = object # type: ignore + +T = TypeVar("T") +Meta = tuple[tuple[str, ...], dict[str, Any]] + +PROTOCOL_MESSAGES_STREAM_KEY = "__protocol_messages_stream" + +# --------------------------------------------------------------------------- +# Content-block accumulation helpers +# --------------------------------------------------------------------------- + +# A "compatible content block" is a dict matching one of the protocol block +# TypedDicts (TextBlock, ReasoningBlock, ToolCallChunkBlock, etc.). +CompatBlock = dict[str, Any] + + +@dataclass +class _ProtocolRunState: + """Per-run state for tracking the active message lifecycle.""" + + message_id: str | None = None + started: bool = False + blocks: dict[int, CompatBlock] = field(default_factory=dict) + usage: dict[str, Any] | None = None + + +def _accumulate_block(accumulated: CompatBlock, delta: CompatBlock) -> CompatBlock: + """Merge *delta* into *accumulated*, returning the updated block.""" + btype = accumulated.get("type", "text") + if btype == "text" and delta.get("type", "text") == "text": + accumulated["text"] = accumulated.get("text", "") + delta.get("text", "") + elif btype == "reasoning" and delta.get("type") == "reasoning": + accumulated["reasoning"] = accumulated.get("reasoning", "") + delta.get( + "reasoning", "" + ) + elif btype == "tool_call_chunk" and delta.get("type") == "tool_call_chunk": + accumulated["args"] = accumulated.get("args", "") + delta.get("args", "") + if delta.get("id") is not None: + accumulated["id"] = delta["id"] + if delta.get("name") is not None: + accumulated["name"] = delta["name"] + return accumulated + + +def _delta_block(previous: CompatBlock, current: CompatBlock) -> CompatBlock | None: + """Compute the delta between *previous* and *current*. + + Returns ``None`` if there is nothing new to emit. + """ + btype = current.get("type", "text") + if btype == "text": + prev_text = previous.get("text", "") + cur_text = current.get("text", "") + delta_text = cur_text[len(prev_text) :] + if not delta_text: + return None + return TextBlock(type="text", text=delta_text) + elif btype == "reasoning": + prev_r = previous.get("reasoning", "") + cur_r = current.get("reasoning", "") + delta_r = cur_r[len(prev_r) :] + if not delta_r: + return None + return ReasoningBlock(type="reasoning", reasoning=delta_r) + elif btype == "tool_call_chunk": + prev_args = previous.get("args", "") + cur_args = current.get("args", "") + delta_args = cur_args[len(prev_args) :] + has_meta = current.get("id") is not None or current.get("name") is not None + if not delta_args and not has_meta: + return None + result: CompatBlock = {"type": "tool_call_chunk", "args": delta_args} + if current.get("id") is not None and previous.get("id") is None: + result["id"] = current["id"] + if current.get("name") is not None and previous.get("name") is None: + result["name"] = current["name"] + return result + # Unrecognized block type — pass through unchanged + return current + + +def _finalize_block(block: CompatBlock) -> CompatBlock: + """Convert a ``tool_call_chunk`` block to a finalized ``tool_call`` or + ``invalid_tool_call`` block. Other block types pass through unchanged. + """ + if block.get("type") != "tool_call_chunk": + return block + raw_args = block.get("args", "{}") + try: + parsed_args = json.loads(raw_args) if raw_args else {} + return ToolCallBlock( + type="tool_call", + id=block.get("id", ""), + name=block.get("name", ""), + args=parsed_args, + ) + except (json.JSONDecodeError, TypeError): + return InvalidToolCallBlock( + type="invalid_tool_call", + id=block.get("id"), + name=block.get("name"), + args=raw_args, + error="Failed to parse tool call arguments as JSON", + ) + + +def _normalize_finish_reason(value: Any) -> FinishReason: + """Map provider-specific stop reasons to protocol finish reasons.""" + if value == "length": + return "length" + if value == "content_filter": + return "content_filter" + if value in ("tool_use", "tool_calls"): + return "tool_use" + # "end_turn", "stop", None, and anything else → "stop" + return "stop" + + +def _accumulate_usage( + current: dict[str, Any] | None, delta: Any +) -> dict[str, Any] | None: + """Accumulate usage metadata from streamed chunks.""" + if not isinstance(delta, dict): + return current + if current is None: + return dict(delta) + for key in ("input_tokens", "output_tokens", "total_tokens", "cached_tokens"): + if key in delta: + current[key] = current.get(key, 0) + delta[key] + # Merge detail dicts + for detail_key in ("input_token_details", "output_token_details"): + if detail_key in delta and isinstance(delta[detail_key], dict): + if detail_key not in current: + current[detail_key] = {} + current[detail_key].update(delta[detail_key]) + return current + + +def _to_protocol_usage(usage: dict[str, Any] | None) -> UsageInfo | None: + """Convert LangChain usage metadata to protocol ``UsageInfo``.""" + if usage is None: + return None + result: dict[str, Any] = {} + if "input_tokens" in usage: + result["input_tokens"] = usage["input_tokens"] + if "output_tokens" in usage: + result["output_tokens"] = usage["output_tokens"] + if "total_tokens" in usage: + result["total_tokens"] = usage["total_tokens"] + if "cached_tokens" in usage: + result["cached_tokens"] = usage["cached_tokens"] + return UsageInfo(**result) if result else None + + +# --------------------------------------------------------------------------- +# Extracting content blocks from LangChain messages +# --------------------------------------------------------------------------- + + +def _extract_blocks_from_chunk(msg: AIMessageChunk) -> list[tuple[int, CompatBlock]]: + """Extract ``(index, block)`` pairs from an ``AIMessageChunk``. + + LangChain stores content in several places: + - ``content: str`` — a single text block at index 0 + - ``content: list[dict]`` — explicit content blocks with their own types + - ``tool_call_chunks`` — separate list for streamed tool call deltas + """ + blocks: list[tuple[int, CompatBlock]] = [] + content = msg.content + if isinstance(content, str) and content: + blocks.append((0, dict(TextBlock(type="text", text=content)))) + elif isinstance(content, list): + for i, item in enumerate(content): + if not isinstance(item, dict): + continue + ctype = item.get("type", "") + if ctype == "text" and item.get("text"): + blocks.append( + ( + item.get("index", i), + dict(TextBlock(type="text", text=item["text"])), + ) + ) + elif ctype in ("reasoning_content", "reasoning", "thinking"): + reasoning_text = ( + item.get("reasoning_content") + or item.get("reasoning") + or item.get("thinking", "") + ) + if reasoning_text: + blocks.append( + ( + item.get("index", i), + dict( + ReasoningBlock( + type="reasoning", reasoning=reasoning_text + ) + ), + ) + ) + + # Tool call chunks live in a separate field + for tc in msg.tool_call_chunks or []: + idx = tc.get("index") + if idx is None: + # Assign indices after text content blocks + idx = len(blocks) + block: CompatBlock = {"type": "tool_call_chunk", "args": tc.get("args", "")} + if tc.get("id") is not None: + block["id"] = tc["id"] + if tc.get("name") is not None: + block["name"] = tc["name"] + blocks.append((idx, block)) + + return blocks + + +# --------------------------------------------------------------------------- +# The handler +# --------------------------------------------------------------------------- + + +class StreamProtocolMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): + """Callback handler that emits content-block protocol events. + + Activated when ``__protocol_messages_stream`` is ``True`` in the run's + configurable metadata. Emits ``StreamChunk`` tuples of the form + ``(namespace, "messages", data)`` where *data* is one of the + ``MessagesData`` event types (``message-start``, ``content-block-start``, + etc.). + """ + + run_inline = True + + def __init__( + self, + stream: Callable[[StreamChunk], None], + subgraphs: bool, + *, + parent_ns: tuple[str, ...] | None = None, + ) -> None: + self.stream = stream + self.subgraphs = subgraphs + self.parent_ns = parent_ns + # Per-run metadata: run_id → (namespace, metadata_dict) + self.metadata: dict[UUID, Meta] = {} + # Per-run protocol state for streamed messages + self.protocol_runs: dict[UUID, _ProtocolRunState] = {} + # Stable message ID mapping: run_id → message_id + self.stable_message_ids: dict[UUID, str] = {} + # Seen message IDs for deduplication of chain-emitted messages + self.seen: set[str | int] = set() + + def _emit(self, meta: Meta, data: Any) -> None: + """Emit a protocol event as a StreamChunk. + + The node name from *meta* is embedded at ``"__node__"`` so the + stream pump can lift it into ``params.node`` without changing the + ``StreamChunk`` tuple shape. + """ + node = meta[1].get("langgraph_node") + if node and isinstance(data, dict): + data = {**data, "__node__": node} + self.stream((meta[0], "messages", data)) + + # -- Chat model callbacks ----------------------------------------------- + + def on_chat_model_start( + self, + serialized: dict[str, Any], + messages: list[list[BaseMessage]], + *, + run_id: UUID, + parent_run_id: UUID | None = None, + tags: list[str] | None = None, + metadata: dict[str, Any] | None = None, + **kwargs: Any, + ) -> Any: + if metadata and (not tags or (TAG_NOSTREAM not in tags)): + ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[ + :-1 + ] + if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns: + return + if tags: + if filtered := [t for t in tags if not t.startswith("seq:step")]: + metadata["tags"] = filtered + self.metadata[run_id] = (ns, metadata) + self.protocol_runs[run_id] = _ProtocolRunState() + + def on_llm_new_token( + self, + token: str, + *, + chunk: ChatGenerationChunk | None = None, + run_id: UUID, + parent_run_id: UUID | None = None, + tags: list[str] | None = None, + **kwargs: Any, + ) -> Any: + if not isinstance(chunk, ChatGenerationChunk): + return + meta = self.metadata.get(run_id) + if meta is None: + return + state = self.protocol_runs.get(run_id) + if state is None: + return + + msg = chunk.message + if not isinstance(msg, AIMessageChunk): + return + + # Emit message-start on first token + if not state.started: + message_id = self._normalize_message_id(msg, run_id) + state.message_id = message_id + state.started = True + start_data = dict( + MessageStartData( + event="message-start", + role="ai", + ) + ) + if message_id: + start_data["message_id"] = message_id + self._emit(meta, start_data) + + # Extract content blocks from this chunk + extracted = _extract_blocks_from_chunk(msg) + for idx, delta_block in extracted: + if idx not in state.blocks: + # New block — emit content-block-start + state.blocks[idx] = dict(delta_block) + # Start block has empty content placeholder + start_block = _make_start_block(delta_block) + self._emit( + meta, + ContentBlockStartData( + event="content-block-start", + index=idx, + content_block=start_block, + ), + ) + # Then emit the first delta + first_delta = _delta_block( + _make_start_block(delta_block), state.blocks[idx] + ) + if first_delta is not None: + self._emit( + meta, + ContentBlockDeltaData( + event="content-block-delta", + index=idx, + content_block=first_delta, + ), + ) + else: + # Existing block — compute delta, accumulate, emit + previous = dict(state.blocks[idx]) + state.blocks[idx] = _accumulate_block(state.blocks[idx], delta_block) + delta = _delta_block(previous, state.blocks[idx]) + if delta is not None: + self._emit( + meta, + ContentBlockDeltaData( + event="content-block-delta", + index=idx, + content_block=delta, + ), + ) + + # Accumulate usage from chunk + if msg.usage_metadata: + state.usage = _accumulate_usage(state.usage, msg.usage_metadata) + + def on_llm_end( + self, + response: LLMResult, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + meta = self.metadata.pop(run_id, None) + state = self.protocol_runs.pop(run_id, None) + if meta is None or state is None: + return + + # Extract finish reason and usage from the final generation + finish_reason: FinishReason = "stop" + final_usage = state.usage + + if response.generations and response.generations[0]: + gen = response.generations[0][0] + if isinstance(gen, ChatGeneration): + final_msg = gen.message + # Get finish reason from response_metadata + rm = getattr(final_msg, "response_metadata", {}) or {} + raw_reason = rm.get("finish_reason") or rm.get("stop_reason") + if raw_reason: + finish_reason = _normalize_finish_reason(raw_reason) + # If we have tool calls in the final message, infer tool_use + if ( + finish_reason == "stop" + and hasattr(final_msg, "tool_calls") + and final_msg.tool_calls + ): + finish_reason = "tool_use" + # Get usage from final message if not accumulated from chunks + if final_usage is None and hasattr(final_msg, "usage_metadata"): + final_usage = ( + dict(final_msg.usage_metadata) + if final_msg.usage_metadata + else None + ) + + # If we never got streaming tokens (non-streamed model call), + # emit the full message lifecycle now + if not state.started: + self._emit_full_message(meta, final_msg, finish_reason, final_usage) + return + + # Close out any open content blocks + for idx in sorted(state.blocks): + finalized = _finalize_block(state.blocks[idx]) + self._emit( + meta, + ContentBlockFinishData( + event="content-block-finish", + index=idx, + content_block=finalized, + ), + ) + + # Emit message-finish + finish_data: dict[str, Any] = { + "event": "message-finish", + "reason": finish_reason, + } + usage_info = _to_protocol_usage(final_usage) + if usage_info is not None: + finish_data["usage"] = usage_info + self._emit(meta, finish_data) + + # Track the message as seen for dedup + if state.message_id: + self.seen.add(state.message_id) + + def on_llm_error( + self, + error: BaseException, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + meta = self.metadata.pop(run_id, None) + state = self.protocol_runs.pop(run_id, None) + self.stable_message_ids.pop(run_id, None) + if meta is None or state is None: + return + if state.started: + self._emit( + meta, + MessageErrorData( + event="error", + message=str(error), + ), + ) + + # -- Chain callbacks (for node-level message dedup) --------------------- + + def on_chain_start( + self, + serialized: dict[str, Any], + inputs: dict[str, Any], + *, + run_id: UUID, + parent_run_id: UUID | None = None, + tags: list[str] | None = None, + metadata: dict[str, Any] | None = None, + **kwargs: Any, + ) -> Any: + if ( + metadata + and kwargs.get("name") == metadata.get("langgraph_node") + and (not tags or TAG_HIDDEN not in tags) + ): + ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[ + :-1 + ] + if not self.subgraphs and len(ns) > 0: + return + self.metadata[run_id] = (ns, metadata) + # Record input message IDs for deduplication + self._record_seen_messages(inputs) + + def on_chain_end( + self, + response: Any, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + meta = self.metadata.pop(run_id, None) + if meta is None: + return + # Emit protocol events for any new messages in the node's output + self._emit_chain_messages(meta, response) + + def on_chain_error( + self, + error: BaseException, + *, + run_id: UUID, + parent_run_id: UUID | None = None, + **kwargs: Any, + ) -> Any: + self.metadata.pop(run_id, None) + + # -- Iterator taps (required by _StreamingCallbackHandler) --------------- + + def tap_output_aiter( + self, run_id: UUID, output: AsyncIterator[T] + ) -> AsyncIterator[T]: + return output + + def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]: + return output + + # -- Internal helpers --------------------------------------------------- + + def _normalize_message_id(self, msg: BaseMessage, run_id: UUID) -> str | None: + """Return a stable message ID for this run, creating one if needed.""" + msg_id = msg.id + if msg_id is None: + msg_id = self.stable_message_ids.get(run_id) + if msg_id is None: + msg_id = f"run-{run_id}" + self.stable_message_ids[run_id] = msg_id + # Mutate the message for consistency downstream + if msg.id != msg_id: + msg.id = msg_id + return msg_id + + def _emit_full_message( + self, + meta: Meta, + msg: BaseMessage, + finish_reason: FinishReason, + usage: dict[str, Any] | None, + role: str = "ai", + ) -> None: + """Emit a complete message lifecycle for a non-streamed model call.""" + message_id = msg.id or str(uuid4()) + if message_id in self.seen: + return + self.seen.add(message_id) + + # message-start + start_data = dict( + MessageStartData( + event="message-start", + role=role, + ) + ) + start_data["message_id"] = message_id + self._emit(meta, start_data) + + # Extract all blocks from the final message + blocks = _extract_final_blocks(msg) + for idx, block in blocks: + # content-block-start with the full content + self._emit( + meta, + ContentBlockStartData( + event="content-block-start", + index=idx, + content_block=_make_start_block(block), + ), + ) + # content-block-delta with the full content + delta = _delta_block(_make_start_block(block), block) + if delta is not None: + self._emit( + meta, + ContentBlockDeltaData( + event="content-block-delta", + index=idx, + content_block=delta, + ), + ) + # content-block-finish + finalized = _finalize_block(block) + self._emit( + meta, + ContentBlockFinishData( + event="content-block-finish", + index=idx, + content_block=finalized, + ), + ) + + # message-finish + finish_data: dict[str, Any] = { + "event": "message-finish", + "reason": finish_reason, + } + usage_info = _to_protocol_usage(usage) + if usage_info is not None: + finish_data["usage"] = usage_info + self._emit(meta, finish_data) + + def _record_seen_messages(self, obj: Any) -> None: + """Record message IDs from node inputs for deduplication.""" + if isinstance(obj, BaseMessage): + if obj.id is not None: + self.seen.add(obj.id) + elif isinstance(obj, dict): + for value in obj.values(): + self._record_seen_messages(value) + elif isinstance(obj, Sequence) and not isinstance(obj, (str, bytes)): + for item in obj: + self._record_seen_messages(item) + + def _emit_chain_messages(self, meta: Meta, response: Any) -> None: + """Emit protocol events for messages found in chain output.""" + from langgraph.types import Command + + if isinstance(response, Command): + self._emit_chain_messages(meta, response.update) + elif isinstance(response, BaseMessage): + self._emit_message_from_chain(meta, response) + elif isinstance(response, Sequence) and not isinstance(response, (str, bytes)): + for item in response: + if isinstance(item, Command): + self._emit_chain_messages(meta, item.update) + elif isinstance(item, BaseMessage): + self._emit_message_from_chain(meta, item) + elif isinstance(response, dict): + for value in response.values(): + if isinstance(value, BaseMessage): + self._emit_message_from_chain(meta, value) + elif isinstance(value, Sequence) and not isinstance( + value, (str, bytes) + ): + for item in value: + if isinstance(item, BaseMessage): + self._emit_message_from_chain(meta, item) + + def _emit_message_from_chain(self, meta: Meta, msg: BaseMessage) -> None: + """Emit a full message lifecycle for a message from a chain output, + deduplicating against previously-seen messages.""" + if msg.id is not None and msg.id in self.seen: + return + if msg.id is None: + msg.id = str(uuid4()) + + # Determine role and finish reason + role = "ai" + if hasattr(msg, "type"): + if msg.type == "human": + role = "human" + elif msg.type == "system": + role = "system" + + finish_reason: FinishReason = "stop" + rm = getattr(msg, "response_metadata", {}) or {} + raw_reason = rm.get("finish_reason") or rm.get("stop_reason") + if raw_reason: + finish_reason = _normalize_finish_reason(raw_reason) + if finish_reason == "stop" and hasattr(msg, "tool_calls") and msg.tool_calls: + finish_reason = "tool_use" + + raw_usage = getattr(msg, "usage_metadata", None) + usage = dict(raw_usage) if raw_usage else None + + self._emit_full_message(meta, msg, finish_reason, usage, role=role) + + +# --------------------------------------------------------------------------- +# Block extraction for finalized (non-streamed) messages +# --------------------------------------------------------------------------- + + +def _extract_final_blocks(msg: BaseMessage) -> list[tuple[int, CompatBlock]]: + """Extract ``(index, block)`` pairs from a finalized ``AIMessage``.""" + blocks: list[tuple[int, CompatBlock]] = [] + content = msg.content + + if isinstance(content, str) and content: + blocks.append((0, dict(TextBlock(type="text", text=content)))) + elif isinstance(content, list): + for i, item in enumerate(content): + if not isinstance(item, dict): + continue + ctype = item.get("type", "") + if ctype == "text" and item.get("text"): + blocks.append((i, dict(TextBlock(type="text", text=item["text"])))) + elif ctype in ("reasoning_content", "reasoning", "thinking"): + reasoning_text = ( + item.get("reasoning_content") + or item.get("reasoning") + or item.get("thinking", "") + ) + if reasoning_text: + blocks.append( + ( + i, + dict( + ReasoningBlock( + type="reasoning", reasoning=reasoning_text + ) + ), + ) + ) + + # Finalized tool calls (already parsed, not chunks) + for tc in getattr(msg, "tool_calls", None) or []: + idx = len(blocks) + blocks.append( + ( + idx, + dict( + ToolCallBlock( + type="tool_call", + id=tc.get("id", ""), + name=tc.get("name", ""), + args=tc.get("args", {}), + ) + ), + ) + ) + + return blocks + + +def _make_start_block(block: CompatBlock) -> CompatBlock: + """Create an empty start placeholder for a content block.""" + btype = block.get("type", "text") + if btype == "text": + return TextBlock(type="text", text="") + elif btype == "reasoning": + return ReasoningBlock(type="reasoning", reasoning="") + elif btype == "tool_call_chunk": + result: CompatBlock = {"type": "tool_call_chunk", "args": ""} + if "id" in block: + result["id"] = block["id"] + if "name" in block: + result["name"] = block["name"] + return result + elif btype == "tool_call": + # Already finalized — return as-is for start event + return ToolCallBlock( + type="tool_call", + id=block.get("id", ""), + name=block.get("name", ""), + args=block.get("args", {}), + ) + return dict(block) + + +__all__ = ["PROTOCOL_MESSAGES_STREAM_KEY", "StreamProtocolMessagesHandler"] diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index d754db391..4e2a155c6 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -128,6 +128,10 @@ from langgraph.pregel._loop import ( SyncPregelLoop, ) from langgraph.pregel._messages import StreamMessagesHandler +from langgraph.pregel._messages_v2 import ( + PROTOCOL_MESSAGES_STREAM_KEY, + StreamProtocolMessagesHandler, +) from langgraph.pregel._read import DEFAULT_BOUND, PregelNode from langgraph.pregel._retry import RetryPolicy from langgraph.pregel._runner import PregelRunner @@ -2616,8 +2620,15 @@ class Pregel( # set up messages stream mode if "messages" in stream_modes: ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)) + _msg_cls = ( + StreamProtocolMessagesHandler + if config.get("configurable", {}).get( + PROTOCOL_MESSAGES_STREAM_KEY, False + ) + else StreamMessagesHandler + ) run_manager.inheritable_handlers.append( - StreamMessagesHandler( + _msg_cls( stream.put, subgraphs, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, @@ -2935,7 +2946,10 @@ class Pregel( True for h in run_manager.handlers if isinstance(h, _StreamingCallbackHandler) - and not isinstance(h, StreamMessagesHandler) + and not isinstance( + h, + (StreamMessagesHandler, StreamProtocolMessagesHandler), + ) ), False, ) @@ -2972,10 +2986,16 @@ class Pregel( config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns) # set up messages stream mode if "messages" in stream_modes: - # namespace can be None in a root level graph? ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)) + _msg_cls = ( + StreamProtocolMessagesHandler + if config.get("configurable", {}).get( + PROTOCOL_MESSAGES_STREAM_KEY, False + ) + else StreamMessagesHandler + ) run_manager.inheritable_handlers.append( - StreamMessagesHandler( + _msg_cls( stream_put, subgraphs, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py new file mode 100644 index 000000000..d882191b5 --- /dev/null +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -0,0 +1,46 @@ +"""Stream protocol types and infrastructure for LangGraph.""" + +from langgraph.stream._convert import STREAM_V2_MODES, convert_to_protocol_event +from langgraph.stream._event_log import EventLog +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import ( + InterruptPayload, + ProtocolEvent, + StreamTransformer, +) +from langgraph.stream.chat_model_stream import AsyncChatModelStream, ChatModelStream +from langgraph.stream.run_stream import ( + AsyncGraphRunStream, + AsyncSubgraphRunStream, + GraphRunStream, + create_async_graph_run_stream, + create_graph_run_stream, +) +from langgraph.stream.stream_channel import StreamChannel, is_stream_channel +from langgraph.stream.streaming_handler import StreamingHandler +from langgraph.stream.transformers import ( + MessagesTransformer, + ValuesTransformer, +) + +__all__ = [ + "STREAM_V2_MODES", + "AsyncChatModelStream", + "AsyncGraphRunStream", + "AsyncSubgraphRunStream", + "ChatModelStream", + "EventLog", + "GraphRunStream", + "InterruptPayload", + "MessagesTransformer", + "ProtocolEvent", + "StreamChannel", + "StreamMux", + "StreamTransformer", + "StreamingHandler", + "ValuesTransformer", + "convert_to_protocol_event", + "create_async_graph_run_stream", + "create_graph_run_stream", + "is_stream_channel", +] diff --git a/libs/langgraph/langgraph/stream/_convert.py b/libs/langgraph/langgraph/stream/_convert.py new file mode 100644 index 000000000..6d199b789 --- /dev/null +++ b/libs/langgraph/langgraph/stream/_convert.py @@ -0,0 +1,79 @@ +"""Convert raw ``StreamChunk`` tuples to ``ProtocolEvent`` envelopes. + +Each ``StreamMode`` is mapped to a ``ProtocolEvent`` whose ``method`` +field matches the mode name and whose ``params.data`` wraps the +original payload. +""" + +from __future__ import annotations + +import time +from typing import Any + +from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams +from langgraph.types import StreamMode + +#: All stream modes requested by ``StreamingHandler`` when calling the +#: underlying ``stream()`` / ``astream()``. +STREAM_V2_MODES: list[StreamMode] = [ + "values", + "updates", + "messages", + "custom", + "checkpoints", + "tasks", + "debug", +] + +_SUPPORTED_MODES: set[str] = set(STREAM_V2_MODES) + + +def convert_to_protocol_event( + ns: tuple[str, ...], + mode: str, + payload: Any, + *, + node: str | None = None, +) -> ProtocolEvent | None: + """Convert a ``StreamChunk`` to a ``ProtocolEvent``. + + Returns ``None`` for unsupported or unknown modes. + + The ``seq`` field is left as ``0`` here; the :class:`StreamMux` is + the sole seq assigner and overwrites it inside ``push()``. + + Parameters + ---------- + ns: + Namespace tuple from the ``StreamChunk``. + mode: + Stream mode string (``"values"``, ``"updates"``, etc.). + payload: + The raw payload from the stream. + node: + Optional node name for provenance. + """ + if mode not in _SUPPORTED_MODES: + return None + + params: _ProtocolEventParams = { + "namespace": list(ns), + "timestamp": _now_ms(), + "data": payload, + } + if node is not None: + params["node"] = node + + return ProtocolEvent( + type="event", + method=mode, + params=params, + ) + + +def _now_ms() -> int: + """Current time in milliseconds since epoch.""" + return int(time.time() * 1000) + + +__all__ = ["STREAM_V2_MODES", "convert_to_protocol_event"] diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py new file mode 100644 index 000000000..660a05531 --- /dev/null +++ b/libs/langgraph/langgraph/stream/_event_log.py @@ -0,0 +1,132 @@ +"""Replayable append-only event buffer for StreamingHandler. + +``EventLog`` stores protocol events in an ordered list and supports +multiple independent async iterators, each with their own cursor +offset. Subscribers that join mid-stream replay from a given offset +without losing earlier events. +""" + +from __future__ import annotations + +import asyncio +import threading +from typing import Generic, TypeVar + +T = TypeVar("T") + + +class EventLog(Generic[T]): + """Append-only event buffer with cursor-based async iteration. + + Multiple consumers can subscribe independently and each will see + every event from their starting offset onward. + """ + + __slots__ = ("_items", "_closed", "_error", "_waiters", "_lock") + + def __init__(self) -> None: + self._items: list[T] = [] + self._closed = False + self._error: BaseException | None = None + self._waiters: list[asyncio.Future[None]] = [] + self._lock = threading.Lock() + + # -- Producer API ------------------------------------------------------- + + def append(self, item: T) -> None: + """Append an event and wake all waiting consumers.""" + with self._lock: + if self._closed: + raise RuntimeError("EventLog is closed") + self._items.append(item) + self._wake_all() + + def close(self) -> None: + """Mark the log as complete. Iterators will end gracefully.""" + with self._lock: + self._closed = True + self._wake_all() + + def fail(self, error: BaseException) -> None: + """Mark the log as failed. Iterators will raise *error*.""" + with self._lock: + self._error = error + self._closed = True + self._wake_all() + + # -- Consumer API ------------------------------------------------------- + + def subscribe(self, offset: int = 0) -> _Cursor[T]: + """Create a new cursor starting at *offset*. + + If *offset* is 0 the cursor replays the entire log. If + *offset* equals ``len(self)`` the cursor starts from the + current tip and only sees future events. + """ + return _Cursor(self, offset) + + # -- Inspection --------------------------------------------------------- + + def __len__(self) -> int: + return len(self._items) + + def __getitem__(self, index: int) -> T: + return self._items[index] + + @property + def closed(self) -> bool: + return self._closed + + # -- Internal ----------------------------------------------------------- + + def _wake_all(self) -> None: + for fut in self._waiters: + if not fut.done(): + try: + fut.get_loop().call_soon_threadsafe(fut.set_result, None) + except RuntimeError: + # Loop already closed — ignore. + pass + self._waiters.clear() + + +class _Cursor(Generic[T]): + """An independent async iterator over an :class:`EventLog`.""" + + __slots__ = ("_log", "_offset") + + def __init__(self, log: EventLog[T], offset: int) -> None: + self._log = log + self._offset = offset + + def __aiter__(self) -> _Cursor[T]: + return self + + async def __anext__(self) -> T: + while True: + with self._log._lock: + if self._offset < len(self._log._items): + item = self._log._items[self._offset] + self._offset += 1 + return item + if self._log._error is not None: + raise self._log._error + if self._log._closed: + raise StopAsyncIteration + # Nothing available yet — register a waiter + loop = asyncio.get_running_loop() + fut: asyncio.Future[None] = loop.create_future() + self._log._waiters.append(fut) + # Wait outside the lock + try: + await fut + except asyncio.CancelledError: + # Remove our future so it doesn't accumulate in the list. + try: + self._log._waiters.remove(fut) + except ValueError: + pass # Already removed by _wake_all + raise + + +__all__ = ["EventLog"] diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py new file mode 100644 index 000000000..83082d143 --- /dev/null +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -0,0 +1,378 @@ +"""Central event dispatcher with transformer pipeline for StreamingHandler. + +``StreamMux`` is the core coordination point: it holds the main +:class:`EventLog`, tracks discovered namespaces for subgraph stream +creation, and pipes every event through the registered +:class:`StreamTransformer` pipeline before appending it to the log. +""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import AsyncIterator +from typing import Any + +from langgraph.stream._event_log import EventLog +from langgraph.stream._types import InterruptPayload, ProtocolEvent, StreamTransformer +from langgraph.stream.stream_channel import StreamChannel, is_stream_channel + + +class StreamMux: + """Central event dispatcher for the StreamingHandler infrastructure. + + The mux owns the main event log, applies the transformer pipeline to + every incoming event, and provides subscription endpoints for + filtered event iteration and subgraph discovery. + """ + + def __init__(self, transformers: list[StreamTransformer] | None = None) -> None: + self._event_log: EventLog[ProtocolEvent] = EventLog() + self._transformers: list[StreamTransformer] = list(transformers or []) + self._channels: list[StreamChannel[Any]] = [] + self._current_namespace: list[str] = [] + self._next_emit_seq: int = 0 + + # Namespace discovery: maps top-level ns segment → True + self._discovered_ns: dict[str, bool] = {} + # Waiters for new namespace discovery + self._ns_waiters: list[asyncio.Future[None]] = [] + + # Latest values per namespace (list-of-strings key) + self._latest_values: dict[str, Any] = {} + + # Interrupt tracking + self._interrupts: list[InterruptPayload] = [] + self._interrupted = False + + # Output promise tracking + self._output_futures: dict[str, asyncio.Future[Any]] = {} + + # Closed state + self._closed = False + self._error: BaseException | None = None + + # -- Producer API ------------------------------------------------------- + + def push(self, event: ProtocolEvent) -> None: + """Push an event through the transformer pipeline and into the log. + + Each registered transformer's ``process()`` is called in order. + If any transformer returns ``False``, the event is suppressed + (not appended to the main log). + """ + if self._closed: + return + + # Mux is the sole seq assigner — ensures all events in the log + # (including those from StreamChannel forwarders) share a single + # monotonically increasing counter. + event["seq"] = self._next_emit_seq + self._next_emit_seq += 1 + + # Track namespace + ns = event["params"].get("namespace", []) + if ns: + top_segment = ns[0] + if top_segment not in self._discovered_ns: + self._discovered_ns[top_segment] = True + self._wake_ns_waiters() + + # Track values + if event["method"] == "values": + ns_key = _ns_key(ns) + self._latest_values[ns_key] = event["params"]["data"] + + # Track interrupts from values events + if event["method"] == "values": + data = event["params"]["data"] + if isinstance(data, dict) and "__interrupt__" in data: + interrupt_info = data["__interrupt__"] + if isinstance(interrupt_info, (list, tuple)): + for item in interrupt_info: + iid = getattr(item, "id", None) or str(id(item)) + self._interrupts.append( + InterruptPayload( + interrupt_id=iid, + payload=item, + ) + ) + self._interrupted = True + + # Run transformer pipeline + self._current_namespace = ns + keep = True + for transformer in self._transformers: + result = transformer.process(event) + if result is False: + keep = False + self._current_namespace = [] + + # Append to main log if not suppressed + if keep: + self._event_log.append(event) + + def close(self, output: Any = None) -> None: + """Close the mux, resolving all output futures.""" + if self._closed: + return + self._closed = True + + # Finalize transformers (optional method) + for transformer in self._transformers: + if hasattr(transformer, "finalize"): + transformer.finalize() + + # Close wired channels + for channel in self._channels: + channel._close() + + # Close the event log + self._event_log.close() + + # Resolve output futures + for ns_key, fut in self._output_futures.items(): + if not fut.done(): + value = self._latest_values.get(ns_key) + try: + fut.get_loop().call_soon_threadsafe(fut.set_result, value) + except RuntimeError: + pass + + # Wake namespace waiters + self._wake_ns_waiters() + + def fail(self, error: BaseException) -> None: + """Fail the mux, rejecting all output futures.""" + if self._closed: + return + self._closed = True + self._error = error + + # Fail transformers (optional method) + for transformer in self._transformers: + if hasattr(transformer, "fail"): + transformer.fail(error) + + # Fail wired channels + for channel in self._channels: + channel._fail(error) + + # Fail the event log + self._event_log.fail(error) + + # Reject output futures + for fut in self._output_futures.values(): + if not fut.done(): + try: + fut.get_loop().call_soon_threadsafe(fut.set_exception, error) + except RuntimeError: + pass + + # Wake namespace waiters + self._wake_ns_waiters() + + # -- Consumer API ------------------------------------------------------- + + def subscribe_events( + self, path: list[str] | None = None, offset: int = 0 + ) -> AsyncIterator[ProtocolEvent]: + """Return an async iterator over events matching *path*. + + If *path* is ``None`` or empty, all events are yielded. + Otherwise, only events whose namespace starts with *path* + are yielded. + """ + cursor = self._event_log.subscribe(offset) + if not path: + return cursor + return _FilteredEventIterator(cursor, path) + + async def subscribe_subgraphs( + self, path: list[str] | None = None, offset: int = 0 + ) -> AsyncIterator[str]: + """Yield top-level namespace segments as they are discovered. + + Each yielded value is the first namespace segment of a newly + discovered subgraph (e.g. ``"agent:0"``). + """ + yielded: set[str] = set() + while True: + # Yield any newly discovered namespaces + for ns_segment in list(self._discovered_ns): + if ns_segment not in yielded: + # Filter by path prefix if specified + if path: + if not ns_segment.startswith(path[0]): + continue + yielded.add(ns_segment) + yield ns_segment + + if self._closed: + return + + # Wait for new namespaces + loop = asyncio.get_running_loop() + fut: asyncio.Future[None] = loop.create_future() + self._ns_waiters.append(fut) + await fut + + def get_output_future(self, ns: list[str] | None = None) -> asyncio.Future[Any]: + """Get or create an output future for a namespace. + + The future resolves to the latest ``values`` event data when + the mux is closed. + """ + ns_key = _ns_key(ns or []) + if ns_key not in self._output_futures: + loop = asyncio.get_running_loop() + self._output_futures[ns_key] = loop.create_future() + + # If already closed, resolve immediately + if self._closed: + value = self._latest_values.get(ns_key) + if self._error is not None: + self._output_futures[ns_key].set_exception(self._error) + else: + self._output_futures[ns_key].set_result(value) + + return self._output_futures[ns_key] + + # -- Inspection --------------------------------------------------------- + + @property + def interrupted(self) -> bool: + return self._interrupted + + @property + def interrupts(self) -> list[InterruptPayload]: + return list(self._interrupts) + + @property + def event_log(self) -> EventLog[ProtocolEvent]: + return self._event_log + + def get_latest_values(self, ns: list[str] | None = None) -> Any: + """Return the most recent values for a namespace.""" + return self._latest_values.get(_ns_key(ns or [])) + + # -- Internal ----------------------------------------------------------- + + def register_transformer(self, transformer: StreamTransformer) -> None: + """Register a new transformer and replay all buffered events through it. + + This is the safe way to add a late-arriving transformer after the mux + has already started processing events. The sequence is: + + 1. Snapshot the current log length (no await → no gap possible in + asyncio's cooperative threading model). + 2. Append the transformer so future ``push()`` calls reach it. + 3. Replay events ``[0, snapshot)`` through the transformer. + 4. If the mux is already closed, call ``finalize()`` immediately so + the transformer's log/channel terminates cleanly. + + ``process()`` is only called for events whose namespace starts with + any prefix — callers that need namespace filtering should do so inside + their ``process()`` implementation, or wrap this call with their own + filtering logic. + """ + snapshot = len(self._event_log) + self._transformers.append(transformer) + for i in range(snapshot): + transformer.process(self._event_log[i]) + if self._closed: + if hasattr(transformer, "finalize"): + transformer.finalize() + + def wire_channels(self, projection: Any) -> None: + """Scan *projection* for :class:`StreamChannel` instances and wire them. + + For each ``StreamChannel`` found, registers a push callback that + appends a :class:`ProtocolEvent` directly to the main event log + with ``method`` set to the channel's name. + + Channel events bypass the transformer pipeline (matching the JS + implementation). They are visible to raw event iteration and + remote SDK clients but not to other transformers' ``process()``. + """ + if projection is None: + return + items: dict[str, Any] = {} + if isinstance(projection, dict): + items = projection + elif hasattr(projection, "__dict__"): + items = vars(projection) + for _key, value in items.items(): + if is_stream_channel(value): + channel: StreamChannel[Any] = value + self._channels.append(channel) + + def _make_forwarder(ch: StreamChannel[Any]) -> Any: + def _forward(item: Any) -> None: + if self._closed: + return + # Append directly to the event log, bypassing + # the transformer pipeline. This matches the JS + # implementation and avoids re-entrancy bugs + # (namespace clobbering, infinite recursion). + self._event_log.append( + ProtocolEvent( + type="event", + seq=self._next_emit_seq, + method=ch.channel_name, + params={ + "namespace": list(self._current_namespace), + "timestamp": int(time.time() * 1000), + "data": item, + }, + ) + ) + self._next_emit_seq += 1 + + return _forward + + channel._wire(_make_forwarder(channel)) + + def _wake_ns_waiters(self) -> None: + for fut in self._ns_waiters: + if not fut.done(): + try: + fut.get_loop().call_soon_threadsafe(fut.set_result, None) + except RuntimeError: + pass + self._ns_waiters.clear() + + +class _FilteredEventIterator: + """Async iterator that filters events by namespace prefix.""" + + __slots__ = ("_cursor", "_path") + + def __init__(self, cursor: AsyncIterator[ProtocolEvent], path: list[str]) -> None: + self._cursor = cursor + self._path = path + + def __aiter__(self) -> _FilteredEventIterator: + return self + + async def __anext__(self) -> ProtocolEvent: + while True: + event = await self._cursor.__anext__() + ns = event["params"].get("namespace", []) + if _ns_starts_with(ns, self._path): + return event + + +def _ns_key(ns: list[str] | tuple[str, ...]) -> str: + """Convert a namespace list to a hashable key.""" + return "|".join(ns) + + +def _ns_starts_with(ns: list[str], prefix: list[str]) -> bool: + """Check if *ns* starts with *prefix*.""" + if len(ns) < len(prefix): + return False + return ns[: len(prefix)] == prefix + + +__all__ = ["StreamMux"] diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py new file mode 100644 index 000000000..f4a29e227 --- /dev/null +++ b/libs/langgraph/langgraph/stream/_types.py @@ -0,0 +1,167 @@ +"""Protocol types for StreamingHandler. + +Re-exports CDDL-derived types from ``langchain-protocol`` and defines +in-process-only types needed by the LangGraph streaming infrastructure. +""" + +from __future__ import annotations + +from typing import Any, Protocol, runtime_checkable + +# --------------------------------------------------------------------------- +# Re-exports from langchain-protocol (CDDL-derived) +# --------------------------------------------------------------------------- +# Primitives +# Content blocks +# Messages data +# Tools data +from langchain_protocol import ( + Annotation, + Citation, + ContentBlock, + ContentBlockDeltaData, + ContentBlockFinishData, + ContentBlockStartData, + FinalizedContentBlock, + FinishReason, + InvalidToolCallBlock, + MessageErrorData, + MessageFinishData, + MessageMetadata, + MessageRole, + MessagesData, + MessageStartData, + MetadataScalar, + Namespace, + ReasoningBlock, + TextBlock, + ToolCallBlock, + ToolCallChunkBlock, + ToolErrorData, + ToolFinishedData, + ToolOutputDeltaData, + ToolsData, + ToolStartedData, + UsageInfo, +) +from typing_extensions import NotRequired, TypedDict + +# --------------------------------------------------------------------------- +# In-process types (not in the CDDL spec) +# --------------------------------------------------------------------------- + + +class _ProtocolEventParams(TypedDict): + """Payload envelope for a :class:`ProtocolEvent`.""" + + namespace: Namespace + timestamp: int + node: NotRequired[str] + data: Any + + +class ProtocolEvent(TypedDict): + """A single protocol event emitted by the StreamingHandler infrastructure. + + ``method`` corresponds to a + :pydata:`~langgraph.types.StreamMode` value (``"messages"``, + ``"updates"``, etc.). + """ + + type: str # always "event" + seq: NotRequired[int] # assigned by StreamMux.push(); absent before push() + method: str # StreamMode value + params: _ProtocolEventParams + + +@runtime_checkable +class StreamTransformer(Protocol): + """Extension point for custom stream projections. + + Implementations are registered with ``StreamingHandler`` and receive every + :class:`ProtocolEvent` before it is appended to the event log. + + Any :class:`~langgraph.stream.stream_channel.StreamChannel` instances + returned by ``init()`` are automatically wired to the protocol event + stream by the mux. + + """ + + def init(self) -> Any: + """Return the initial projection value. + + Called once before the run. Any + :class:`~langgraph.stream.stream_channel.StreamChannel` instances + in the return value are automatically wired by the mux. + """ + ... + + def process(self, event: ProtocolEvent) -> bool: + """Process an event. + + Return ``True`` to keep the event in the log, ``False`` to suppress + it. + """ + ... + + def finalize(self) -> None: + """Called once when the run completes successfully. + + Optional — the mux auto-closes any :class:`StreamChannel` instances, + so transformers that only use channels can omit this. + """ + ... + + def fail(self, err: BaseException) -> None: + """Called once when the run fails. + + Optional — the mux auto-fails any :class:`StreamChannel` instances, + so transformers that only use channels can omit this. + """ + ... + + +class InterruptPayload(TypedDict): + """An interrupt produced during a StreamingHandler run.""" + + interrupt_id: str + payload: Any + + +__all__ = [ + # Primitives (re-exported) + "Namespace", + "MessageRole", + "MessageMetadata", + "MetadataScalar", + # Content blocks (re-exported) + "TextBlock", + "ReasoningBlock", + "ToolCallBlock", + "ToolCallChunkBlock", + "InvalidToolCallBlock", + "ContentBlock", + "FinalizedContentBlock", + "Annotation", + "Citation", + # Messages data (re-exported) + "MessagesData", + "MessageStartData", + "ContentBlockStartData", + "ContentBlockDeltaData", + "ContentBlockFinishData", + "MessageFinishData", + "MessageErrorData", + "FinishReason", + "UsageInfo", + # Tools data (re-exported) + "ToolsData", + "ToolStartedData", + "ToolOutputDeltaData", + "ToolFinishedData", + "ToolErrorData", + # In-process types + "ProtocolEvent", + "StreamTransformer", + "InterruptPayload", +] diff --git a/libs/langgraph/langgraph/stream/chat_model_stream.py b/libs/langgraph/langgraph/stream/chat_model_stream.py new file mode 100644 index 000000000..1905696bc --- /dev/null +++ b/libs/langgraph/langgraph/stream/chat_model_stream.py @@ -0,0 +1,313 @@ +"""Per-message streaming objects for StreamingHandler. + +``ChatModelStream`` is the synchronous variant returned by +``GraphRunStream.messages``. Properties (``.text``, ``.reasoning``, +``.usage``) return final accumulated values. + +``AsyncChatModelStream`` is the asynchronous variant returned by +``AsyncGraphRunStream.messages``. Projections are dual +async-iterable + awaitable (e.g. ``async for delta in msg.text`` +or ``full = await msg.text``). +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Generator +from typing import Any + +from langgraph.stream._types import UsageInfo + +# --------------------------------------------------------------------------- +# Sync variant +# --------------------------------------------------------------------------- + + +class ChatModelStream: + """Synchronous per-message object for a single LLM response. + + Created by :class:`~langgraph.stream.transformers.MessagesTransformer` + and yielded by ``GraphRunStream.messages``. By the time the sync + iterator yields a ``ChatModelStream``, the message lifecycle is + complete and all properties contain their final values. + + Projections: + + - ``.text`` — accumulated text content (``str``) + - ``.reasoning`` — accumulated reasoning content (``str``) + - ``.usage`` — :class:`UsageInfo` or ``None`` + - ``.namespace`` / ``.node`` — provenance metadata + """ + + 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 = "" + self._reasoning_acc = "" + self._usage_value: UsageInfo | None = None + self._done = False + + # -- Public projections ------------------------------------------------ + + @property + def text(self) -> str: + """Accumulated text content.""" + return self._text_acc + + @property + def reasoning(self) -> str: + """Accumulated reasoning content.""" + return self._reasoning_acc + + @property + def usage(self) -> UsageInfo | None: + """Usage info, available after the message finishes.""" + return self._usage_value + + @property + def namespace(self) -> list[str]: + return self._namespace + + @property + def node(self) -> str | None: + return self._node + + @property + def message_id(self) -> str | None: + return self._message_id + + @property + def done(self) -> bool: + return self._done + + # -- Internal API (called by MessagesTransformer) ---------------------- + + def _push_content_block_delta(self, data: dict[str, Any]) -> None: + """Process a ``content-block-delta`` event.""" + block = data.get("content_block", {}) + btype = block.get("type", "") + + if btype == "text": + delta_text = block.get("text", "") + if delta_text: + self._text_acc += delta_text + elif btype == "reasoning": + delta_r = block.get("reasoning", "") + if delta_r: + self._reasoning_acc += delta_r + + def _push_content_block_finish(self, data: dict[str, Any]) -> None: + """Process a ``content-block-finish`` event.""" + block = data.get("content_block", {}) + btype = block.get("type", "") + + if btype == "text": + full_text = block.get("text", "") + if full_text and full_text != self._text_acc: + self._text_acc = full_text + elif btype == "reasoning": + full_r = block.get("reasoning", "") + if full_r and full_r != self._reasoning_acc: + self._reasoning_acc = full_r + + def _finish(self, data: dict[str, Any]) -> None: + """Process a ``message-finish`` event.""" + self._done = True + self._usage_value = data.get("usage") + + def _fail(self, error: BaseException) -> None: + """Process a ``message-error`` event.""" + self._done = True + + +# --------------------------------------------------------------------------- +# Async dual-projection helpers +# --------------------------------------------------------------------------- + + +class _DualProjection: + """Async iterable of deltas that is also awaitable for the final value. + + When iterated, yields delta values (e.g. text fragments) as they arrive. + When awaited, returns the accumulated final value (e.g. full text string). + """ + + def __init__(self) -> None: + self._deltas: list[Any] = [] + self._done = False + self._error: BaseException | None = None + self._waiters: list[asyncio.Future[None]] = [] + self._final_value: Any = None + self._final_set = False + + # -- Producer API (called by AsyncChatModelStream) --------------------- + + def _push(self, delta: Any) -> None: + """Add a new delta value.""" + self._deltas.append(delta) + self._wake() + + def _finish(self, accumulated: Any) -> None: + """Set the final accumulated value and mark as done.""" + self._final_value = accumulated + self._final_set = True + self._done = True + self._wake() + + def _fail(self, error: BaseException) -> None: + self._error = error + self._done = True + self._wake() + + def _wake(self) -> None: + for fut in self._waiters: + if not fut.done(): + try: + fut.get_loop().call_soon_threadsafe(fut.set_result, None) + except RuntimeError: + pass + self._waiters.clear() + + # -- Async iterable (yields deltas) ------------------------------------ + + def __aiter__(self) -> _DualProjectionIterator: + return _DualProjectionIterator(self) + + # -- Awaitable (returns final value) ----------------------------------- + + def __await__(self) -> Generator[Any, None, Any]: + return self._await_impl().__await__() + + async def _await_impl(self) -> Any: + while not self._final_set: + if self._error is not None: + raise self._error + loop = asyncio.get_running_loop() + fut: asyncio.Future[None] = loop.create_future() + self._waiters.append(fut) + await fut + if self._error is not None: + raise self._error + return self._final_value + + +class _DualProjectionIterator: + """Async iterator over a :class:`_DualProjection`'s deltas.""" + + __slots__ = ("_proj", "_offset") + + def __init__(self, proj: _DualProjection) -> None: + self._proj = proj + self._offset = 0 + + def __aiter__(self) -> _DualProjectionIterator: + return self + + async def __anext__(self) -> Any: + while True: + if self._offset < len(self._proj._deltas): + item = self._proj._deltas[self._offset] + self._offset += 1 + return item + if self._proj._error is not None: + raise self._proj._error + if self._proj._done: + raise StopAsyncIteration + loop = asyncio.get_running_loop() + fut: asyncio.Future[None] = loop.create_future() + self._proj._waiters.append(fut) + await fut + + +# --------------------------------------------------------------------------- +# Async variant +# --------------------------------------------------------------------------- + + +class AsyncChatModelStream(ChatModelStream): + """Asynchronous per-message streaming object for a single LLM response. + + Created by :class:`~langgraph.stream.transformers.MessagesTransformer` + and yielded by ``AsyncGraphRunStream.messages``. Content-block events + are fed into this object until ``message-finish``. + + Projections: + + - ``.text`` — async iterable of text deltas; awaitable for full text + - ``.reasoning`` — async iterable of reasoning deltas; awaitable for + full reasoning text + - ``.usage`` — awaitable for :class:`UsageInfo` + - ``.namespace`` / ``.node`` — provenance metadata + """ + + def __init__( + 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 = _DualProjection() + self._reasoning_proj = _DualProjection() + self._usage_proj = _DualProjection() + + # -- Public projections (override sync properties) --------------------- + + @property + def text(self) -> _DualProjection: + """Text content — async iterable of deltas, awaitable for full text.""" + return self._text_proj + + @property + def reasoning(self) -> _DualProjection: + """Reasoning content — async iterable of deltas, awaitable for full text.""" + return self._reasoning_proj + + @property + def usage(self) -> _DualProjection: + """Usage info — awaitable for :class:`UsageInfo`.""" + return self._usage_proj + + # -- Internal API (extend base to also drive projections) -------------- + + def _push_content_block_delta(self, data: dict[str, Any]) -> None: + """Process a ``content-block-delta`` event.""" + super()._push_content_block_delta(data) + block = data.get("content_block", {}) + btype = block.get("type", "") + + if btype == "text": + delta_text = block.get("text", "") + if delta_text: + self._text_proj._push(delta_text) + elif btype == "reasoning": + delta_r = block.get("reasoning", "") + if delta_r: + self._reasoning_proj._push(delta_r) + + def _finish(self, data: dict[str, Any]) -> None: + """Process a ``message-finish`` event.""" + super()._finish(data) + self._text_proj._finish(self._text_acc) + self._reasoning_proj._finish(self._reasoning_acc) + self._usage_proj._finish(self._usage_value) + + def _fail(self, error: BaseException) -> None: + """Process a ``message-error`` event.""" + super()._fail(error) + self._text_proj._fail(error) + self._reasoning_proj._fail(error) + self._usage_proj._fail(error) + + +__all__ = ["AsyncChatModelStream", "ChatModelStream"] diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py new file mode 100644 index 000000000..8044f1238 --- /dev/null +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -0,0 +1,584 @@ +"""GraphRunStream and AsyncGraphRunStream for StreamingHandler. + +These are the top-level objects returned by +``StreamingHandler.stream()`` / ``StreamingHandler.astream()``. +They wrap a :class:`StreamMux` and expose named +projections (``.values``, ``.messages``, ``.subgraphs``, ``.output``) +for ergonomic consumption. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Callable, Iterator +from typing import Any + +from langgraph.stream._convert import convert_to_protocol_event +from langgraph.stream._event_log import EventLog +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import InterruptPayload, ProtocolEvent, StreamTransformer +from langgraph.stream.chat_model_stream import AsyncChatModelStream, ChatModelStream +from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer + +# --------------------------------------------------------------------------- +# Values projection — dual async-iterable + awaitable +# --------------------------------------------------------------------------- + + +class _ValuesProjection: + """Async iterable of intermediate state snapshots; awaitable for final.""" + + def __init__( + self, + mux: StreamMux, + values_transformer: ValuesTransformer, + ns: list[str], + mapper: Callable[[Any], Any] | None = None, + ) -> None: + self._mux = mux + self._values_transformer = values_transformer + self._ns = ns + self._mapper = mapper + + def __aiter__(self) -> AsyncIterator[Any]: + return _ValuesIterator(self._values_transformer, self._ns, self._mapper) + + def __await__(self) -> Any: + return self._await_impl().__await__() + + async def _await_impl(self) -> Any: + value = await self._mux.get_output_future(self._ns) + if value is not None and self._mapper is not None: + return self._mapper(value) + return value + + +class _ValuesIterator: + """Filters the values log to events matching a namespace.""" + + def __init__( + self, + transformer: ValuesTransformer, + ns: list[str], + mapper: Callable[[Any], Any] | None = None, + ) -> None: + self._cursor = transformer.values_log.subscribe(0) + self._ns = ns + self._mapper = mapper + + def __aiter__(self) -> _ValuesIterator: + return self + + async def __anext__(self) -> Any: + while True: + item = await self._cursor.__anext__() + item_ns = item.get("namespace", []) + if item_ns == self._ns: + data = item["data"] + if data is not None and self._mapper is not None: + return self._mapper(data) + return data + + +# --------------------------------------------------------------------------- +# Messages projection +# --------------------------------------------------------------------------- + + +class _MessagesProjection: + """Async iterable of :class:`AsyncChatModelStream` instances.""" + + def __init__(self, messages_transformer: MessagesTransformer) -> None: + self._transformer = messages_transformer + + def __aiter__(self) -> AsyncIterator[AsyncChatModelStream]: + return self._transformer.messages_log.subscribe(0) + + +# --------------------------------------------------------------------------- +# Subgraphs projection +# --------------------------------------------------------------------------- + + +class _SubgraphsProjection: + """Async iterable yielding :class:`AsyncSubgraphRunStream` for each discovered subgraph.""" + + def __init__(self, mux: StreamMux, ns: list[str]) -> None: + self._mux = mux + self._ns = ns + + async def __aiter__(self) -> AsyncIterator[AsyncSubgraphRunStream]: + async for segment in self._mux.subscribe_subgraphs(self._ns): + child_ns = self._ns + [segment] + child_transformers: list[StreamTransformer] = [ + ValuesTransformer(), + MessagesTransformer( + namespace=child_ns, stream_cls=AsyncChatModelStream + ), + ] + for t in child_transformers: + t.init() + self._mux.register_transformer(t) + + yield AsyncSubgraphRunStream( + mux=self._mux, + namespace=child_ns, + transformers=child_transformers, + ) + + +# --------------------------------------------------------------------------- +# AsyncGraphRunStream +# --------------------------------------------------------------------------- + + +class AsyncGraphRunStream: + """The async run stream returned by ``StreamingHandler.astream()``. + + Async-iterable over all :class:`ProtocolEvent` instances. Named + projections provide ergonomic access to values, messages, subgraphs, + and output. + """ + + def __init__( + self, + *, + mux: StreamMux, + namespace: list[str] | None = None, + transformers: list[StreamTransformer], + abort_event: asyncio.Event | None = None, + output_mapper: Callable[[Any], Any] | None = None, + ) -> None: + self._mux = mux + self._ns = namespace or [] + self._transformers = transformers + self._abort_event = abort_event or asyncio.Event() + self._output_mapper = output_mapper + + # -- Transformer lookup ------------------------------------------------- + + def _find_transformer(self, name: str) -> StreamTransformer | None: + for t in self._transformers: + if getattr(t, "name", None) == name: + return t + return None + + # -- Raw event iteration ------------------------------------------------ + + def __aiter__(self) -> AsyncIterator[ProtocolEvent]: + return self._mux.subscribe_events(self._ns) + + # -- Named projections -------------------------------------------------- + + @property + def values(self) -> _ValuesProjection: + """Async iterable of state snapshots; awaitable for final state.""" + t = self._find_transformer("values") + return _ValuesProjection(self._mux, t, self._ns, self._output_mapper) + + @property + def output(self) -> _ValuesProjection: + """Awaitable for the final output state.""" + t = self._find_transformer("values") + return _ValuesProjection(self._mux, t, self._ns, self._output_mapper) + + @property + def messages(self) -> _MessagesProjection: + """Async iterable of :class:`AsyncChatModelStream` instances.""" + t = self._find_transformer("messages") + return _MessagesProjection(t) + + def messages_from(self, node: str) -> _MessagesProjection: + """Async iterable of messages from a specific node.""" + filtered = MessagesTransformer( + namespace=self._ns, + node_filter=node, + stream_cls=AsyncChatModelStream, + ) + self._mux.register_transformer(filtered) + return _MessagesProjection(filtered) + + @property + def subgraphs(self) -> _SubgraphsProjection: + """Async iterable of :class:`AsyncSubgraphRunStream` for child graphs.""" + return _SubgraphsProjection(self._mux, self._ns) + + # -- State -------------------------------------------------------------- + + @property + def interrupted(self) -> bool: + return self._mux.interrupted + + @property + def interrupts(self) -> list[InterruptPayload]: + return self._mux.interrupts + + # -- Cancellation ------------------------------------------------------- + + def abort(self, reason: str | None = None) -> None: + """Signal cancellation of the run.""" + self._abort_event.set() + + @property + def signal(self) -> asyncio.Event: + """The underlying cancellation event.""" + return self._abort_event + + # -- Extensions --------------------------------------------------------- + + @property + def extensions(self) -> dict[str, Any]: + """All transformer projections.""" + result: dict[str, Any] = {} + for t in self._transformers: + name = getattr(t, "name", None) + value = getattr(t, "value", None) + if name is not None and value is not None: + result[name] = value + return result + + +# --------------------------------------------------------------------------- +# AsyncSubgraphRunStream +# --------------------------------------------------------------------------- + + +class AsyncSubgraphRunStream(AsyncGraphRunStream): + """An :class:`AsyncGraphRunStream` for a child subgraph. + + Adds ``.name`` and ``.index`` parsed from the last namespace segment + (e.g. ``"researcher:2"`` → ``name="researcher"``, ``index=2``). + """ + + @property + def name(self) -> str: + if self._ns: + segment = self._ns[-1] + return segment.split(":")[0] if ":" in segment else segment + return "" + + @property + def index(self) -> int: + if self._ns: + segment = self._ns[-1] + if ":" in segment: + try: + return int(segment.split(":")[-1]) + except ValueError: + pass + return 0 + + +# --------------------------------------------------------------------------- +# Async factory +# --------------------------------------------------------------------------- + + +async def create_async_graph_run_stream( + source: AsyncIterator[tuple[tuple[str, ...], str, Any]], + *, + transformers: list[StreamTransformer] | None = None, + abort_event: asyncio.Event | None = None, + output_mapper: Callable[[Any], Any] | None = None, +) -> AsyncGraphRunStream: + """Create an :class:`AsyncGraphRunStream` from a raw async stream source. + + 1. Creates a :class:`StreamMux` + 2. Registers built-in ``ValuesTransformer`` and ``MessagesTransformer`` + 3. Registers user-supplied transformers + 4. Creates the root ``AsyncGraphRunStream`` + 5. Starts a background pump task that reads from *source*, + converts each chunk to a ``ProtocolEvent``, and pushes it + through the mux + 6. Returns the ``AsyncGraphRunStream`` + """ + abort = abort_event or asyncio.Event() + + # Built-in transformers first, then user-supplied + all_transformers: list[StreamTransformer] = [ + ValuesTransformer(), + MessagesTransformer(stream_cls=AsyncChatModelStream), + ] + all_transformers.extend(transformers or []) + + # Initialize transformers, collecting projections to wire after mux creation + projections: list[Any] = [] + for t in all_transformers: + projection = t.init() + if projection is not None: + projections.append(projection) + + mux = StreamMux(transformers=all_transformers) + + # Wire any StreamChannel instances found in transformer projections + for projection in projections: + mux.wire_channels(projection) + + # Create the root stream + run_stream = AsyncGraphRunStream( + mux=mux, + transformers=all_transformers, + abort_event=abort, + output_mapper=output_mapper, + ) + + # Start the pump task + async def pump() -> None: + try: + async for ns, mode, payload in source: + if abort.is_set(): + break + # Extract node name embedded by StreamProtocolMessagesHandler. + node: str | None = None + if ( + mode == "messages" + and isinstance(payload, dict) + and "__node__" in payload + ): + payload = dict(payload) + node = payload.pop("__node__") + event = convert_to_protocol_event(ns, mode, payload, node=node) + if event is not None: + mux.push(event) + mux.close() + except Exception as exc: + mux.fail(exc) + + asyncio.get_running_loop().create_task(pump()) + + return run_stream + + +# --------------------------------------------------------------------------- +# GraphRunStream — returned by StreamingHandler.stream() +# --------------------------------------------------------------------------- + + +class _PumpDrivenLog: + """Wraps an ``EventLog`` so that iteration drives the sync pump. + + Used by :attr:`GraphRunStream.extensions` to make extension logs + iterable without requiring the caller to drain the stream first. + """ + + __slots__ = ("_log", "_pump_one") + + def __init__(self, log: EventLog, pump_one: Callable[[], bool]) -> None: + self._log = log + self._pump_one = pump_one + + def __iter__(self) -> Iterator[Any]: + cursor = 0 + while True: + if cursor < len(self._log): + yield self._log[cursor] + cursor += 1 + elif not self._pump_one(): + return + + def __len__(self) -> int: + return len(self._log) + + def __getitem__(self, index: int) -> Any: + return self._log[index] + + +class GraphRunStream: + """Synchronous run stream returned by ``StreamingHandler.stream()``. + + All projections are blocking / sync-iterable. Internally uses + the same ``StreamMux`` and transformer pipeline, but without an + async event loop. + + The source iterator is consumed lazily: each projection pulls + events from the source on demand rather than eagerly buffering + everything upfront. This means callers see events as soon as + they are produced by the underlying ``stream()`` call. + """ + + def __init__( + self, + *, + mux: StreamMux, + source: Iterator[tuple[tuple[str, ...], str, Any]], + namespace: list[str] | None = None, + transformers: list[StreamTransformer], + output_mapper: Callable[[Any], Any] | None = None, + ) -> None: + self._mux = mux + self._source = source + self._source_exhausted = False + self._ns = namespace or [] + self._transformers = transformers + self._output_mapper = output_mapper + + # -- Transformer lookup ------------------------------------------------- + + def _find_transformer(self, name: str) -> StreamTransformer | None: + for t in self._transformers: + if getattr(t, "name", None) == name: + return t + return None + + # -- Lazy pump ---------------------------------------------------------- + + def _pump_one(self) -> bool: + """Pull one item from the source, convert it, and push through the mux. + + Returns ``True`` if an item was consumed, ``False`` if the source + is exhausted (or was already exhausted). + """ + if self._source_exhausted: + return False + try: + ns, mode, payload = next(self._source) + except StopIteration: + self._source_exhausted = True + self._mux.close() + return False + except Exception as exc: + self._source_exhausted = True + self._mux.fail(exc) + return False + + node: str | None = None + if mode == "messages" and isinstance(payload, dict) and "__node__" in payload: + payload = dict(payload) + node = payload.pop("__node__") + event = convert_to_protocol_event(ns, mode, payload, node=node) + if event is not None: + self._mux.push(event) + return True + + def _pump_all(self) -> None: + """Drain the source iterator completely.""" + while self._pump_one(): + pass + + # -- Helpers ------------------------------------------------------------ + + def _map(self, value: Any) -> Any: + if value is not None and self._output_mapper is not None: + return self._output_mapper(value) + return value + + # -- Raw event iteration (sync) ----------------------------------------- + + def __iter__(self) -> Iterator[ProtocolEvent]: + for event in _PumpDrivenLog(self._mux.event_log, self._pump_one): + ns = event["params"].get("namespace", []) + if not self._ns or ns[: len(self._ns)] == self._ns: + yield event + + # -- Named projections (sync) ------------------------------------------- + + @property + def output(self) -> Any: + """The final output state (blocking). Drains the source.""" + self._pump_all() + return self._map(self._mux.get_latest_values(self._ns)) + + @property + def values(self) -> Iterator[Any]: + """Sync iterable of intermediate state snapshots.""" + t = self._find_transformer("values") + if t is None: + return + for item in _PumpDrivenLog(t.value, self._pump_one): + if item.get("namespace", []) == self._ns: + yield self._map(item["data"]) + + @property + def messages(self) -> Iterator[ChatModelStream]: + """Sync iterable of :class:`ChatModelStream` instances. + + Each yielded ``ChatModelStream`` is fully populated (``done=True``) + so that sync consumers can read ``.text``, ``.reasoning``, and + ``.usage`` immediately. + """ + t = self._find_transformer("messages") + if t is None: + return + for msg in _PumpDrivenLog(t.value, self._pump_one): + # Pump until this message is complete so sync consumers + # get a fully populated ChatModelStream. + while not msg.done: + if not self._pump_one(): + break + yield msg + + # -- State -------------------------------------------------------------- + + @property + def interrupted(self) -> bool: + return self._mux.interrupted + + @property + def interrupts(self) -> list[InterruptPayload]: + return self._mux.interrupts + + # -- Extensions --------------------------------------------------------- + + @property + def extensions(self) -> dict[str, Any]: + """All transformer projections as pump-driven iterables.""" + result: dict[str, Any] = {} + for t in self._transformers: + name = getattr(t, "name", None) + value = getattr(t, "value", None) + if name is not None and value is not None: + if isinstance(value, EventLog): + result[name] = _PumpDrivenLog(value, self._pump_one) + else: + result[name] = value + return result + + +def create_graph_run_stream( + source: Iterator[tuple[tuple[str, ...], str, Any]], + *, + transformers: list[StreamTransformer] | None = None, + output_mapper: Callable[[Any], Any] | None = None, +) -> GraphRunStream: + """Create a :class:`GraphRunStream` from a sync stream source. + + The source iterator is stored on the returned stream and consumed + lazily as projections are iterated. + + Built-in transformers (values, messages) are always registered first + so that user-supplied transformers see events after built-in + processing. + """ + # Built-in transformers first, then user-supplied + all_transformers: list[StreamTransformer] = [ + ValuesTransformer(), + MessagesTransformer(), + ] + all_transformers.extend(transformers or []) + + projections: list[Any] = [] + for t in all_transformers: + projection = t.init() + if projection is not None: + projections.append(projection) + + mux = StreamMux(transformers=all_transformers) + + for projection in projections: + mux.wire_channels(projection) + + return GraphRunStream( + mux=mux, + source=source, + transformers=all_transformers, + output_mapper=output_mapper, + ) + + +__all__ = [ + "AsyncGraphRunStream", + "AsyncSubgraphRunStream", + "GraphRunStream", + "create_async_graph_run_stream", + "create_graph_run_stream", +] diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py new file mode 100644 index 000000000..fade0c2a1 --- /dev/null +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -0,0 +1,76 @@ +"""StreamChannel — typed push-based channel for StreamTransformer projections. + +A ``StreamChannel`` wraps an :class:`EventLog` and declares a protocol +channel name. When the :class:`StreamMux` detects a ``StreamChannel`` +in a transformer's ``init()`` return, it wires every ``push()`` call to +inject a :class:`ProtocolEvent` into the main event stream using the +channel's name as the ``method``. + +In-process consumers iterate the channel directly (it is an async +iterable). Remote SDK clients subscribe via +``session.subscribe("custom:")``. + +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Callable +from typing import Any, Generic, TypeVar + +from langgraph.stream._event_log import EventLog + +T = TypeVar("T") + + +class StreamChannel(Generic[T]): + """A typed push-based channel that integrates with the mux. + + Transformer authors create a ``StreamChannel`` in ``init()`` and + call ``push()`` inside ``process()`` to emit domain objects. The + mux auto-wires pushes to protocol events and auto-closes/fails the + channel on run completion. + """ + + __slots__ = ("channel_name", "_log", "_on_push") + + def __init__(self, name: str) -> None: + self.channel_name = name + self._log: EventLog[T] = EventLog() + self._on_push: Callable[[Any], None] | None = None + + def push(self, item: T) -> None: + """Push an item to the channel. + + If the mux has wired this channel, the push also injects a + protocol event into the main event stream. + """ + self._log.append(item) + if self._on_push is not None: + self._on_push(item) + + # -- Async iteration (in-process consumption) --------------------------- + + def __aiter__(self) -> AsyncIterator[T]: + return self._log.subscribe(0) + + # -- Internal (called by the mux) --------------------------------------- + + def _wire(self, fn: Callable[[Any], None]) -> None: + """Wire a callback invoked on every ``push()``. Called by the mux.""" + self._on_push = fn + + def _close(self) -> None: + """Close the underlying log. Called by the mux on normal completion.""" + self._log.close() + + def _fail(self, err: BaseException) -> None: + """Fail the underlying log. Called by the mux on failure.""" + self._log.fail(err) + + +def is_stream_channel(value: object) -> bool: + """Check if *value* is a :class:`StreamChannel` instance.""" + return isinstance(value, StreamChannel) + + +__all__ = ["StreamChannel", "is_stream_channel"] diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py new file mode 100644 index 000000000..c86cb7595 --- /dev/null +++ b/libs/langgraph/langgraph/stream/streaming_handler.py @@ -0,0 +1,168 @@ +"""Experimental streaming wrapper for CompiledGraph. + +``StreamingHandler`` wraps a compiled graph and exposes the new streaming +API without adding methods to the ``CompiledGraph`` class itself. + +Usage:: + + from langgraph.stream import StreamingHandler + + s = StreamingHandler(graph) + + # async + run = await s.astream(input) + async for msg in run.messages: + ... + + # sync + run = s.stream(input) + for event in run: + ... +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator, Sequence +from typing import TYPE_CHECKING, Any, cast + +from langchain_core.runnables import RunnableConfig + +from langgraph._internal._config import patch_configurable +from langgraph.stream._convert import STREAM_V2_MODES +from langgraph.stream._types import StreamTransformer +from langgraph.stream.run_stream import ( + AsyncGraphRunStream, + GraphRunStream, + create_async_graph_run_stream, + create_graph_run_stream, +) +from langgraph.types import All + +if TYPE_CHECKING: + from langgraph.pregel import Pregel + +#: Config key that activates the protocol messages handler. +#: Duplicated here to avoid a circular import with ``pregel._messages_v2``. +PROTOCOL_MESSAGES_STREAM_KEY = "__protocol_messages_stream" + + +class StreamingHandler: + """Experimental streaming wrapper around a compiled graph. + + Provides ``.stream()`` and ``.astream()`` returning + :class:`GraphRunStream` / :class:`AsyncGraphRunStream` with + ergonomic projections (``run.values``, ``run.messages``, + ``run.subgraphs``, ``run.output``). + + Args: + graph: A compiled LangGraph (``Pregel`` instance). + """ + + def __init__(self, graph: Pregel) -> None: + self._graph = graph + + async def astream( + self, + input: Any, + config: RunnableConfig | None = None, + *, + context: Any | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + debug: bool | None = None, + transformers: list[StreamTransformer] | None = None, + ) -> AsyncGraphRunStream: + """Stream graph execution, returning an + :class:`~langgraph.stream.run_stream.AsyncGraphRunStream`. + + The returned stream provides ergonomic projections: + + - ``await run.output`` -- final state + - ``async for v in run.values`` -- intermediate state snapshots + - ``async for msg in run.messages`` -- per-message + :class:`~langgraph.stream.chat_model_stream.AsyncChatModelStream` + objects + - ``async for sub in run.subgraphs`` -- child + :class:`~langgraph.stream.run_stream.AsyncSubgraphRunStream` + instances + - ``async for event in run`` -- raw + :class:`~langgraph.stream._types.ProtocolEvent` objects + + Args: + input: The input to the graph. + config: The configuration to use for the run. + context: The static context to use for the run. + interrupt_before: Nodes to interrupt before. + interrupt_after: Nodes to interrupt after. + debug: Whether to emit debug events. + transformers: Optional user-supplied + :class:`~langgraph.stream._types.StreamTransformer` instances + for custom projections (available on ``run.extensions``). + + Returns: + An :class:`~langgraph.stream.run_stream.AsyncGraphRunStream`. + """ + merged_config = patch_configurable(config, {PROTOCOL_MESSAGES_STREAM_KEY: True}) + + source = cast( + AsyncIterator[tuple[tuple[str, ...], str, Any]], + self._graph.astream( + input, + merged_config, + context=context, + stream_mode=STREAM_V2_MODES, + subgraphs=True, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + debug=debug, + version="v1", + ), + ) + + return await create_async_graph_run_stream( + source, + transformers=transformers, + output_mapper=self._graph._output_mapper, + ) + + def stream( + self, + input: Any, + config: RunnableConfig | None = None, + *, + context: Any | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + debug: bool | None = None, + transformers: list[StreamTransformer] | None = None, + ) -> GraphRunStream: + """Synchronous variant of :meth:`astream`. + + Returns a :class:`~langgraph.stream.run_stream.GraphRunStream` + immediately. The underlying source is consumed lazily as + projections are iterated. + + See :meth:`astream` for full documentation. + """ + merged_config = patch_configurable(config, {PROTOCOL_MESSAGES_STREAM_KEY: True}) + + source = cast( + Iterator[tuple[tuple[str, ...], str, Any]], + self._graph.stream( + input, + merged_config, + context=context, + stream_mode=STREAM_V2_MODES, + subgraphs=True, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + debug=debug, + version="v1", + ), + ) + + return create_graph_run_stream( + source, + transformers=transformers, + output_mapper=self._graph._output_mapper, + ) diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py new file mode 100644 index 000000000..cc9771791 --- /dev/null +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -0,0 +1,179 @@ +"""Built-in stream transformers for StreamingHandler. + +``ValuesTransformer`` extracts ``values`` events and maintains the latest +state per namespace. ``MessagesTransformer`` groups ``messages`` events +into :class:`ChatModelStream` instances. +""" + +from __future__ import annotations + +from typing import Any + +from langgraph.stream._event_log import EventLog +from langgraph.stream._types import ProtocolEvent +from langgraph.stream.chat_model_stream import ChatModelStream + +# Type alias for the stream class constructor signature +_StreamCls = type[ChatModelStream] + + +class ValuesTransformer: + """Extracts ``values`` events and populates a values event log. + + Maintains the latest state per namespace and provides a separate + event log that :class:`AsyncGraphRunStream` / :class:`GraphRunStream` uses for ``.values`` + iteration. + + Implements the :class:`StreamTransformer` protocol. + """ + + name = "values" + + def __init__(self) -> None: + self._values_log: EventLog[dict[str, Any]] = EventLog() + self._latest: dict[str, Any] = {} + + @property + def value(self) -> EventLog[dict[str, Any]]: + return self._values_log + + @property + def values_log(self) -> EventLog[dict[str, Any]]: + return self._values_log + + def get_latest(self, ns_key: str = "") -> Any: + return self._latest.get(ns_key) + + def init(self) -> Any: + return None + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "values": + return True + + ns = event["params"].get("namespace", []) + data = event["params"]["data"] + ns_key = "|".join(ns) if ns else "" + self._latest[ns_key] = data + + # Append to the values log for iteration + self._values_log.append({"namespace": ns, "data": data}) + return True + + def finalize(self) -> None: + self._values_log.close() + + def fail(self, err: BaseException) -> None: + self._values_log.fail(err) + + +class MessagesTransformer: + """Groups ``messages`` events into :class:`ChatModelStream` instances. + + One ``ChatModelStream`` is created per ``message-start`` event. + Content-block events are routed to the active stream until + ``message-finish`` or ``message-error`` closes it. + + Implements the :class:`StreamTransformer` protocol. + """ + + name = "messages" + + def __init__( + self, + *, + namespace: list[str] | None = None, + node_filter: str | None = None, + stream_cls: _StreamCls | None = None, + ) -> None: + self._namespace = namespace + self._node_filter = node_filter + self._stream_cls: _StreamCls = stream_cls or ChatModelStream + + # Message log for .messages iteration + self._messages_log: EventLog[ChatModelStream] = EventLog() + + # Current active stream per namespace key + self._active: dict[str, ChatModelStream] = {} + + @property + def value(self) -> EventLog[ChatModelStream]: + return self._messages_log + + @property + def messages_log(self) -> EventLog[ChatModelStream]: + return self._messages_log + + def init(self) -> Any: + return None + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "messages": + return True + + ns = event["params"].get("namespace", []) + node = event["params"].get("node") + data = event["params"]["data"] + + # Apply namespace filter + if self._namespace is not None: + if ns[: len(self._namespace)] != self._namespace: + return True + + # Apply node filter + if self._node_filter is not None and node != self._node_filter: + return True + + ns_key = "|".join(ns) if ns else "" + event_type = data.get("event") if isinstance(data, dict) else None + + if event_type == "message-start": + stream = self._stream_cls( + namespace=ns, + node=node, + message_id=data.get("message_id"), + ) + self._active[ns_key] = stream + self._messages_log.append(stream) + + elif event_type in ("content-block-delta", "content-block-start"): + active = self._active.get(ns_key) + if active is not None and event_type == "content-block-delta": + active._push_content_block_delta(data) + + elif event_type == "content-block-finish": + active = self._active.get(ns_key) + if active is not None: + active._push_content_block_finish(data) + + elif event_type == "message-finish": + active = self._active.pop(ns_key, None) + if active is not None: + active._finish(data) + + elif event_type == "error": + active = self._active.pop(ns_key, None) + if active is not None: + msg = data.get("message", "Unknown error") + active._fail(RuntimeError(msg)) + + return True + + def finalize(self) -> None: + # Close any remaining active streams + for stream in self._active.values(): + stream._finish({"reason": "stop"}) + self._active.clear() + self._messages_log.close() + + def fail(self, err: BaseException) -> None: + for stream in self._active.values(): + stream._fail(err) + self._active.clear() + self._messages_log.fail(err) + + +__all__ = [ + "MessagesTransformer", + "ValuesTransformer", +] diff --git a/libs/langgraph/tests/test_graph_run_stream.py b/libs/langgraph/tests/test_graph_run_stream.py new file mode 100644 index 000000000..a0583ac2a --- /dev/null +++ b/libs/langgraph/tests/test_graph_run_stream.py @@ -0,0 +1,558 @@ +import asyncio +from typing import Annotated, Any + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from pydantic import BaseModel +from typing_extensions import TypedDict + +from langgraph.config import get_stream_writer +from langgraph.graph import END, START, MessagesState, StateGraph +from langgraph.stream import AsyncChatModelStream, StreamingHandler +from langgraph.stream._types import ProtocolEvent +from tests.fake_chat import FakeChatModel + + +class State(TypedDict): + value: str + items: Annotated[list[str], lambda a, b: a + b] + + +def make_simple_graph(): + def node_a(state): + return {"value": state["value"] + "_a", "items": ["a"]} + + def node_b(state): + return {"value": state["value"] + "_b", "items": ["b"]} + + graph = StateGraph(State) + graph.add_node("node_a", node_a) + graph.add_node("node_b", node_b) + graph.add_edge(START, "node_a") + graph.add_edge("node_a", "node_b") + graph.add_edge("node_b", END) + return graph.compile() + + +@pytest.mark.anyio +async def test_output(): + graph = make_simple_graph() + run = await StreamingHandler(graph).astream({"value": "x", "items": []}) + await asyncio.sleep(0.1) + output = await run.output + assert output == {"value": "x_a_b", "items": ["a", "b"]} + + +@pytest.mark.anyio +async def test_values_iteration(): + graph = make_simple_graph() + run = await StreamingHandler(graph).astream({"value": "x", "items": []}) + await asyncio.sleep(0.1) + + snapshots = [] + async for v in run.values: + snapshots.append(v) + + assert len(snapshots) == 3 + assert snapshots[0]["value"] == "x" + assert snapshots[1]["value"] == "x_a" + assert snapshots[2]["value"] == "x_a_b" + + +@pytest.mark.anyio +async def test_updates_in_raw_events(): + graph = make_simple_graph() + run = await StreamingHandler(graph).astream({"value": "x", "items": []}) + await asyncio.sleep(0.1) + + updates = [] + async for event in run: + if event["method"] == "updates": + updates.append(event["params"]["data"]) + + assert len(updates) == 2 + assert "node_a" in updates[0] + assert "node_b" in updates[1] + + +@pytest.mark.anyio +async def test_messages_with_chat_model(): + model = FakeChatModel(messages=[AIMessage(content="Hello world")]) + + def agent(state): + return {"messages": [model.invoke(state["messages"])]} + + graph = StateGraph(MessagesState) + graph.add_node("agent", agent) + graph.add_edge(START, "agent") + graph.add_edge("agent", END) + compiled = graph.compile() + + run = await StreamingHandler(compiled).astream( + {"messages": [HumanMessage(content="hi")]} + ) + await asyncio.sleep(0.1) + + messages_seen = [] + async for msg in run.messages: + messages_seen.append(msg) + + assert len(messages_seen) >= 1 + msg = messages_seen[0] + assert isinstance(msg, AsyncChatModelStream) + text = await msg.text + assert text == "Hello world" + + +@pytest.mark.anyio +async def test_custom_events(): + def node(state): + writer = get_stream_writer() + writer("hello") + writer(42) + return {"value": state["value"] + "_a", "items": ["a"]} + + graph = StateGraph(State) + graph.add_node("node_a", node) + graph.add_edge(START, "node_a") + graph.add_edge("node_a", END) + compiled = graph.compile() + + run = await StreamingHandler(compiled).astream({"value": "x", "items": []}) + await asyncio.sleep(0.1) + + custom_payloads = [] + async for event in run: + if event["method"] == "custom": + custom_payloads.append(event["params"]["data"]) + + assert "hello" in custom_payloads + assert 42 in custom_payloads + + +@pytest.mark.anyio +async def test_multiple_modes_present(): + graph = make_simple_graph() + run = await StreamingHandler(graph).astream({"value": "x", "items": []}) + await asyncio.sleep(0.1) + + methods = set() + async for event in run: + methods.add(event["method"]) + + assert {"values", "updates", "tasks", "debug"} <= methods + + +@pytest.mark.anyio +async def test_interrupted_false(): + graph = make_simple_graph() + run = await StreamingHandler(graph).astream({"value": "x", "items": []}) + await asyncio.sleep(0.1) + async for _ in run: + pass + assert run.interrupted is False + + +@pytest.mark.anyio +async def test_regression_v1_stream_unchanged(): + graph = make_simple_graph() + chunks = [] + async for chunk in graph.astream( + {"value": "x", "items": []}, stream_mode="values", version="v1" + ): + chunks.append(chunk) + for chunk in chunks: + assert isinstance(chunk, dict) + + +@pytest.mark.anyio +async def test_regression_v2_stream_unchanged(): + graph = make_simple_graph() + chunks = [] + async for chunk in graph.astream( + {"value": "x", "items": []}, stream_mode="values", version="v2" + ): + chunks.append(chunk) + assert len(chunks) >= 1 + for chunk in chunks: + assert isinstance(chunk, dict) + assert "type" in chunk + assert chunk["type"] == "values" + + +@pytest.mark.anyio +async def test_regression_invoke_unchanged(): + graph = make_simple_graph() + result = await graph.ainvoke({"value": "x", "items": []}) + assert result == {"value": "x_a_b", "items": ["a", "b"]} + + +def test_sync_stream_output(): + graph = make_simple_graph() + run = StreamingHandler(graph).stream({"value": "x", "items": []}) + assert run.output == {"value": "x_a_b", "items": ["a", "b"]} + + +def test_sync_stream_values(): + graph = make_simple_graph() + run = StreamingHandler(graph).stream({"value": "x", "items": []}) + snapshots = list(run.values) + assert len(snapshots) == 3 + assert snapshots[0]["value"] == "x" + assert snapshots[2]["value"] == "x_a_b" + + +def test_sync_stream_raw_events(): + graph = make_simple_graph() + run = StreamingHandler(graph).stream({"value": "x", "items": []}) + methods = {e["method"] for e in run} + assert {"values", "updates", "tasks", "debug"} <= methods + + +# --------------------------------------------------------------------------- +# Typed output (pydantic) +# --------------------------------------------------------------------------- + + +class ModelState(BaseModel): + value: str + items: Annotated[list[str], lambda a, b: a + b] + + +def _make_model_state_graph(): + def node_a(state): + return {"value": state.value + "_a", "items": ["a"]} + + graph = StateGraph(ModelState) + graph.add_node("node_a", node_a) + graph.add_edge(START, "node_a") + graph.add_edge("node_a", END) + return graph.compile() + + +@pytest.mark.anyio +async def test_pydantic_output(): + graph = _make_model_state_graph() + run = await StreamingHandler(graph).astream(ModelState(value="x", items=[])) + await asyncio.sleep(0.1) + output = await run.output + assert isinstance(output, ModelState) + assert output.value == "x_a" + + +@pytest.mark.anyio +async def test_pydantic_values(): + graph = _make_model_state_graph() + run = await StreamingHandler(graph).astream(ModelState(value="x", items=[])) + await asyncio.sleep(0.1) + snapshots = [] + async for v in run.values: + snapshots.append(v) + for v in snapshots: + assert isinstance(v, ModelState) + + +def test_sync_pydantic_output(): + graph = _make_model_state_graph() + run = StreamingHandler(graph).stream(ModelState(value="x", items=[])) + assert isinstance(run.output, ModelState) + assert run.output.value == "x_a" + + +# --------------------------------------------------------------------------- +# Interrupts +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_interrupts(): + from langgraph.checkpoint.memory import MemorySaver + + from langgraph.types import interrupt + + def ask_human(state: State): + answer = interrupt("what do you want?") + return {"value": state["value"] + f"_{answer}", "items": [answer]} + + graph = StateGraph(State) + graph.add_node("ask", ask_human) + graph.add_edge(START, "ask") + graph.add_edge("ask", END) + compiled = graph.compile(checkpointer=MemorySaver()) + + config = {"configurable": {"thread_id": "t1"}} + run = await StreamingHandler(compiled).astream( + {"value": "x", "items": []}, config=config + ) + await asyncio.sleep(0.1) + # Drain events + async for _ in run: + pass + assert run.interrupted is True + assert len(run.interrupts) > 0 + + +# --------------------------------------------------------------------------- +# messages_from(node) +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_messages_from_node(): + model = FakeChatModel(messages=[AIMessage(content="from agent")]) + + def agent(state): + return {"messages": [model.invoke(state["messages"])]} + + def postprocess(state): + return {"messages": state["messages"]} + + graph = StateGraph(MessagesState) + graph.add_node("agent", agent) + graph.add_node("postprocess", postprocess) + graph.add_edge(START, "agent") + graph.add_edge("agent", "postprocess") + graph.add_edge("postprocess", END) + compiled = graph.compile() + + run = await StreamingHandler(compiled).astream( + {"messages": [HumanMessage(content="hi")]} + ) + await asyncio.sleep(0.1) + + # All messages + all_msgs = [] + async for m in run.messages: + all_msgs.append(m) + assert len(all_msgs) >= 1 + # Node provenance should be set + assert all_msgs[0].node == "agent" + + +# --------------------------------------------------------------------------- +# Subgraph child stream +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_subgraph_child_output(): + """AsyncSubgraphRunStream.output should contain the child graph's final state.""" + + class ChildState(TypedDict): + value: str + + class ParentState(TypedDict): + value: str + + def child_node(state): + return {"value": state["value"] + "_child"} + + child_graph = StateGraph(ChildState) + child_graph.add_node("child_node", child_node) + child_graph.add_edge(START, "child_node") + child_graph.add_edge("child_node", END) + # Add the compiled child as a node — this triggers LangGraph's + # subgraph streaming mechanism and emits child namespace events. + child_compiled = child_graph.compile() + + parent_graph = StateGraph(ParentState) + parent_graph.add_node("child_node", child_compiled) + parent_graph.add_edge(START, "child_node") + parent_graph.add_edge("child_node", END) + parent_compiled = parent_graph.compile() + + run = await StreamingHandler(parent_compiled).astream({"value": "x"}) + await asyncio.sleep(0.1) + + subgraph_streams = [] + async for sub in run.subgraphs: + subgraph_streams.append(sub) + + assert len(subgraph_streams) >= 1 + child_output = await subgraph_streams[0].output + assert child_output is not None + assert child_output["value"] == "x_child" + + +# --------------------------------------------------------------------------- +# Custom reducers / .extensions +# --------------------------------------------------------------------------- + + +class _CountTransformer: + """Counts events. Exposes count via .value for extensions.""" + + name = "event_count" + + def __init__(self) -> None: + self.value = 0 + + def init(self) -> Any: + return None + + def process(self, event: ProtocolEvent) -> bool: + self.value += 1 + return True + + def finalize(self) -> None: + pass + + def fail(self, err: BaseException) -> None: + pass + + +@pytest.mark.anyio +async def test_custom_reducer_extensions(): + graph = make_simple_graph() + counter = _CountTransformer() + run = await StreamingHandler(graph).astream( + {"value": "x", "items": []}, transformers=[counter] + ) + await asyncio.sleep(0.1) + async for _ in run: + pass + assert counter.value > 0 + assert run.extensions["event_count"] == counter.value + + +def test_sync_custom_reducer_extensions(): + graph = make_simple_graph() + counter = _CountTransformer() + run = StreamingHandler(graph).stream( + {"value": "x", "items": []}, transformers=[counter] + ) + for _ in run: + pass + assert counter.value > 0 + assert run.extensions["event_count"] == counter.value + + +# --------------------------------------------------------------------------- +# Tool transformer via extensions +# --------------------------------------------------------------------------- + + +class _ToolExecution: + def __init__(self, tool_call_id: str, tool_name: str, input: Any, output: Any): + self.tool_call_id = tool_call_id + self.tool_name = tool_name + self.input = input + self.output = output + + +class _ToolsTransformer: + """Groups tool-started/tool-finished custom events into _ToolExecution objects.""" + + name = "tools" + + def __init__(self) -> None: + from langgraph.stream._event_log import EventLog + + self._log: EventLog[_ToolExecution] = EventLog() + self._pending: dict[str, dict] = {} + self.value = self._log + + def init(self) -> Any: + return None + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] != "custom": + return True + data = event["params"]["data"] + if not isinstance(data, dict) or "event" not in data: + return True + + tool_call_id = data.get("tool_call_id") + if tool_call_id is None: + return True + + if data["event"] == "tool-started": + self._pending[tool_call_id] = data + return False + + if data["event"] == "tool-finished": + started = self._pending.pop(tool_call_id, {}) + self._log.append(_ToolExecution( + tool_call_id=tool_call_id, + tool_name=started.get("tool_name", ""), + input=started.get("input"), + output=data["output"], + )) + return False + + return True + + def finalize(self) -> None: + self._log.close() + + def fail(self, err: BaseException) -> None: + self._log.fail(err) + + +def _make_tool_graph(): + """Graph: agent emits a tool call, custom_tools executes it with writer events.""" + from langgraph.types import StreamWriter + + def agent(state): + return { + "value": "called", + "items": ["agent"], + } + + def custom_tools(state, *, writer: StreamWriter): + writer({ + "event": "tool-started", + "tool_call_id": "call_1", + "tool_name": "get_weather", + "input": {"city": "SF"}, + }) + writer({ + "event": "tool-finished", + "tool_call_id": "call_1", + "output": {"temp_f": 64}, + }) + return {"value": "done", "items": ["tools"]} + + graph = StateGraph(State) + graph.add_node("agent", agent) + graph.add_node("custom_tools", custom_tools) + graph.add_edge(START, "agent") + graph.add_edge("agent", "custom_tools") + graph.add_edge("custom_tools", END) + return graph.compile() + + +def test_sync_tool_transformer_via_extensions(): + """Tool events flow through extensions and are iterable without draining raw events.""" + graph = _make_tool_graph() + run = StreamingHandler(graph).stream( + {"value": "", "items": []}, + transformers=[_ToolsTransformer()], + ) + + # Iterating extensions drives the pump — no need to drain raw events first + executions = list(run.extensions["tools"]) + assert len(executions) == 1 + assert executions[0].tool_name == "get_weather" + assert executions[0].input == {"city": "SF"} + assert executions[0].output == {"temp_f": 64} + + +@pytest.mark.anyio +async def test_async_tool_transformer_via_extensions(): + """Tool events flow through extensions in async mode.""" + graph = _make_tool_graph() + run = await StreamingHandler(graph).astream( + {"value": "", "items": []}, + transformers=[_ToolsTransformer()], + ) + await asyncio.sleep(0.1) + + # Drain main stream so transformer processes all events + async for _ in run: + pass + + tools_log = run.extensions["tools"] + assert len(tools_log) == 1 + assert tools_log[0].tool_name == "get_weather" + assert tools_log[0].output == {"temp_f": 64} diff --git a/libs/langgraph/tests/test_messages_v2.py b/libs/langgraph/tests/test_messages_v2.py new file mode 100644 index 000000000..745036338 --- /dev/null +++ b/libs/langgraph/tests/test_messages_v2.py @@ -0,0 +1,531 @@ +from uuid import uuid4 + +import pytest +from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage +from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult + +from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM +from langgraph.pregel._messages_v2 import StreamProtocolMessagesHandler +from langgraph.types import Command + +META = {"langgraph_checkpoint_ns": "root:", "langgraph_node": "agent"} + + +def make_handler(subgraphs=True): + events = [] + handler = StreamProtocolMessagesHandler(events.append, subgraphs) + return handler, events + + +def test_streamed_text(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + + for token_text in ("Hello", " ", "world"): + chunk = ChatGenerationChunk( + message=AIMessageChunk(content=token_text, id=f"run-{run_id}") + ) + handler.on_llm_new_token(token_text, chunk=chunk, run_id=run_id) + + final_msg = AIMessage(content="Hello world", id=f"run-{run_id}") + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + + data_events = [e[2] for e in events] + assert data_events[0]["event"] == "message-start" + assert data_events[1]["event"] == "content-block-start" + assert data_events[1]["index"] == 0 + + deltas = [d for d in data_events if d["event"] == "content-block-delta"] + assert len(deltas) == 3 + assert deltas[0]["content_block"]["text"] == "Hello" + assert deltas[1]["content_block"]["text"] == " " + assert deltas[2]["content_block"]["text"] == "world" + + finish_blocks = [d for d in data_events if d["event"] == "content-block-finish"] + assert len(finish_blocks) == 1 + assert finish_blocks[0]["content_block"]["text"] == "Hello world" + + assert data_events[-1]["event"] == "message-finish" + assert data_events[-1]["reason"] == "stop" + + +def test_tool_calls(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + + chunk1 = ChatGenerationChunk( + message=AIMessageChunk( + content="", + tool_call_chunks=[ + {"name": "search", "args": '{"q', "id": "call_1", "index": 0} + ], + id=f"run-{run_id}", + ) + ) + handler.on_llm_new_token("", chunk=chunk1, run_id=run_id) + + chunk2 = ChatGenerationChunk( + message=AIMessageChunk( + content="", + tool_call_chunks=[ + {"name": None, "args": 'uery":"hi"}', "id": None, "index": 0} + ], + id=f"run-{run_id}", + ) + ) + handler.on_llm_new_token("", chunk=chunk2, run_id=run_id) + + final_msg = AIMessage( + content="", + tool_calls=[{"name": "search", "args": {"query": "hi"}, "id": "call_1"}], + id=f"run-{run_id}", + ) + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + + data_events = [e[2] for e in events] + finish_blocks = [d for d in data_events if d["event"] == "content-block-finish"] + assert len(finish_blocks) == 1 + fb = finish_blocks[0]["content_block"] + assert fb["type"] == "tool_call" + assert fb["args"] == {"query": "hi"} + assert fb["name"] == "search" + assert fb["id"] == "call_1" + + +def test_invalid_tool_call_json(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + + chunk = ChatGenerationChunk( + message=AIMessageChunk( + content="", + tool_call_chunks=[ + { + "name": "search", + "args": "{not valid json", + "id": "call_2", + "index": 0, + } + ], + id=f"run-{run_id}", + ) + ) + handler.on_llm_new_token("", chunk=chunk, run_id=run_id) + + final_msg = AIMessage(content="", id=f"run-{run_id}") + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + + data_events = [e[2] for e in events] + finish_blocks = [d for d in data_events if d["event"] == "content-block-finish"] + assert len(finish_blocks) == 1 + fb = finish_blocks[0]["content_block"] + assert fb["type"] == "invalid_tool_call" + assert "Failed to parse" in fb["error"] + + +def test_reasoning_blocks(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + + chunk = ChatGenerationChunk( + message=AIMessageChunk( + content=[{"type": "reasoning_content", "reasoning_content": "thinking..."}], + id=f"run-{run_id}", + ) + ) + handler.on_llm_new_token("", chunk=chunk, run_id=run_id) + + final_msg = AIMessage(content="", id=f"run-{run_id}") + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + + data_events = [e[2] for e in events] + block_starts = [d for d in data_events if d["event"] == "content-block-start"] + assert len(block_starts) == 1 + assert block_starts[0]["content_block"]["type"] == "reasoning" + + deltas = [d for d in data_events if d["event"] == "content-block-delta"] + assert len(deltas) == 1 + assert deltas[0]["content_block"]["reasoning"] == "thinking..." + + +def test_multiple_content_blocks(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + + chunk1 = ChatGenerationChunk( + message=AIMessageChunk(content="hello", id=f"run-{run_id}") + ) + handler.on_llm_new_token("hello", chunk=chunk1, run_id=run_id) + + chunk2 = ChatGenerationChunk( + message=AIMessageChunk( + content="", + tool_call_chunks=[ + {"name": "lookup", "args": '{"x":1}', "id": "call_3", "index": 1} + ], + id=f"run-{run_id}", + ) + ) + handler.on_llm_new_token("", chunk=chunk2, run_id=run_id) + + final_msg = AIMessage( + content="hello", + tool_calls=[{"name": "lookup", "args": {"x": 1}, "id": "call_3"}], + id=f"run-{run_id}", + ) + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + + data_events = [e[2] for e in events] + finish_blocks = [d for d in data_events if d["event"] == "content-block-finish"] + assert len(finish_blocks) == 2 + + +def test_usage_metadata(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + + chunk = ChatGenerationChunk( + message=AIMessageChunk(content="hi", id=f"run-{run_id}") + ) + handler.on_llm_new_token("hi", chunk=chunk, run_id=run_id) + + final_msg = AIMessage( + content="hi", + id=f"run-{run_id}", + usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + ) + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + + data_events = [e[2] for e in events] + finish_event = [d for d in data_events if d["event"] == "message-finish"][0] + assert "usage" in finish_event + assert finish_event["usage"]["input_tokens"] == 10 + + +@pytest.mark.parametrize( + "raw_reason,expected", + [ + ("stop", "stop"), + ("tool_calls", "tool_use"), + ("length", "length"), + ("content_filter", "content_filter"), + ("end_turn", "stop"), + ], +) +def test_finish_reason_normalization(raw_reason, expected): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + chunk = ChatGenerationChunk(message=AIMessageChunk(content="x", id=f"run-{run_id}")) + handler.on_llm_new_token("x", chunk=chunk, run_id=run_id) + + final_msg = AIMessage( + content="x", + id=f"run-{run_id}", + response_metadata={"finish_reason": raw_reason}, + ) + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + + data_events = [e[2] for e in events] + finish_event = [d for d in data_events if d["event"] == "message-finish"][0] + assert finish_event["reason"] == expected + + +def test_tag_nostream(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[TAG_NOSTREAM] + ) + chunk = ChatGenerationChunk( + message=AIMessageChunk(content="secret", id=f"run-{run_id}") + ) + handler.on_llm_new_token("secret", chunk=chunk, run_id=run_id) + + final_msg = AIMessage(content="secret", id=f"run-{run_id}") + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + assert events == [] + + +def test_tag_hidden_chain(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chain_start( + serialized={}, + inputs={}, + run_id=run_id, + metadata=META, + tags=[TAG_HIDDEN], + name="agent", + ) + handler.on_chain_end( + {"messages": [AIMessage(content="hidden", id="msg-1")]}, + run_id=run_id, + ) + assert events == [] + + +def test_subgraph_filtering(): + handler, events = make_handler(subgraphs=False) + run_id = uuid4() + + subgraph_meta = { + "langgraph_checkpoint_ns": "root:|child:", + "langgraph_node": "agent", + } + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=subgraph_meta, tags=[] + ) + chunk = ChatGenerationChunk( + message=AIMessageChunk(content="sub", id=f"run-{run_id}") + ) + handler.on_llm_new_token("sub", chunk=chunk, run_id=run_id) + + final_msg = AIMessage(content="sub", id=f"run-{run_id}") + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + assert events == [] + + +def test_chain_emits_messages(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chain_start( + serialized={}, inputs={}, run_id=run_id, metadata=META, tags=[], name="agent" + ) + handler.on_chain_end( + {"messages": [AIMessage(content="hello", id="msg-chain-1")]}, + run_id=run_id, + ) + + data_events = [e[2] for e in events] + assert len(data_events) > 0 + assert data_events[0]["event"] == "message-start" + assert data_events[-1]["event"] == "message-finish" + + +def test_llm_error_after_start(): + """on_llm_error should emit a message-error event for a started stream.""" + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + + chunk = ChatGenerationChunk( + message=AIMessageChunk(content="partial", id=f"run-{run_id}") + ) + handler.on_llm_new_token("partial", chunk=chunk, run_id=run_id) + + handler.on_llm_error(RuntimeError("connection lost"), run_id=run_id) + + data_events = [e[2] for e in events] + assert data_events[0]["event"] == "message-start" + error_events = [d for d in data_events if d["event"] == "error"] + assert len(error_events) == 1 + assert "connection lost" in error_events[0]["message"] + + +def test_llm_error_before_start_no_emit(): + """on_llm_error before any tokens should not emit error events.""" + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + + # Error before any token — state.started is False + handler.on_llm_error(RuntimeError("immediate fail"), run_id=run_id) + + data_events = [e[2] for e in events] + error_events = [d for d in data_events if d.get("event") == "error"] + assert len(error_events) == 0 + + +def test_non_streamed_model(): + handler, events = make_handler() + run_id = uuid4() + + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id, metadata=META, tags=[] + ) + + final_msg = AIMessage( + content="full response", + id=f"run-{run_id}", + response_metadata={"finish_reason": "stop"}, + ) + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id, + ) + + data_events = [e[2] for e in events] + assert len(data_events) > 0 + assert data_events[0]["event"] == "message-start" + + deltas = [d for d in data_events if d["event"] == "content-block-delta"] + assert len(deltas) == 1 + assert deltas[0]["content_block"]["text"] == "full response" + + assert data_events[-1]["event"] == "message-finish" + assert data_events[-1]["reason"] == "stop" + + +def test_chain_emits_command_with_message(): + """on_chain_end should emit protocol events for messages inside a Command.""" + handler, events = make_handler() + run_id = uuid4() + + handler.on_chain_start( + serialized={}, inputs={}, run_id=run_id, metadata=META, tags=[], name="agent" + ) + handler.on_chain_end( + Command(update={"messages": [AIMessage(content="from command", id="cmd-1")]}), + run_id=run_id, + ) + + data_events = [e[2] for e in events] + assert len(data_events) > 0 + assert data_events[0]["event"] == "message-start" + deltas = [d for d in data_events if d["event"] == "content-block-delta"] + assert len(deltas) == 1 + assert deltas[0]["content_block"]["text"] == "from command" + assert data_events[-1]["event"] == "message-finish" + + +def test_chain_emits_command_in_list(): + """on_chain_end should handle a list containing Command objects.""" + handler, events = make_handler() + run_id = uuid4() + + handler.on_chain_start( + serialized={}, inputs={}, run_id=run_id, metadata=META, tags=[], name="agent" + ) + handler.on_chain_end( + [Command(update={"messages": [AIMessage(content="listed", id="cmd-2")]})], + run_id=run_id, + ) + + data_events = [e[2] for e in events] + starts = [d for d in data_events if d["event"] == "message-start"] + assert len(starts) == 1 + + +def test_chain_deduplicates_seen_messages(): + """Messages already seen from LLM streaming should not be re-emitted by chain end.""" + handler, events = make_handler() + run_id_llm = uuid4() + run_id_chain = uuid4() + msg_id = f"run-{run_id_llm}" + + # Simulate LLM streaming + handler.on_chat_model_start( + serialized={}, messages=[[]], run_id=run_id_llm, metadata=META, tags=[] + ) + chunk = ChatGenerationChunk(message=AIMessageChunk(content="hello", id=msg_id)) + handler.on_llm_new_token("hello", chunk=chunk, run_id=run_id_llm) + + final_msg = AIMessage(content="hello", id=msg_id) + handler.on_llm_end( + LLMResult(generations=[[ChatGeneration(message=final_msg)]]), + run_id=run_id_llm, + ) + + events_before = len(events) + + # Now chain end with the same message ID + handler.on_chain_start( + serialized={}, + inputs={}, + run_id=run_id_chain, + metadata=META, + tags=[], + name="agent", + ) + handler.on_chain_end( + {"messages": [AIMessage(content="hello", id=msg_id)]}, + run_id=run_id_chain, + ) + + # No new events should have been emitted for the duplicate + data_events_after = [e[2] for e in events[events_before:]] + starts = [d for d in data_events_after if d.get("event") == "message-start"] + assert len(starts) == 0 + + +def test_chain_emits_human_message_role(): + """Non-AI messages from chain output should have the correct role.""" + handler, events = make_handler() + run_id = uuid4() + + handler.on_chain_start( + serialized={}, inputs={}, run_id=run_id, metadata=META, tags=[], name="agent" + ) + handler.on_chain_end( + {"messages": [HumanMessage(content="user msg", id="hmsg-1")]}, + run_id=run_id, + ) + + data_events = [e[2] for e in events] + starts = [d for d in data_events if d["event"] == "message-start"] + assert len(starts) == 1 + assert starts[0]["role"] == "human" diff --git a/libs/langgraph/tests/test_stream_chat_model_stream.py b/libs/langgraph/tests/test_stream_chat_model_stream.py new file mode 100644 index 000000000..e06109121 --- /dev/null +++ b/libs/langgraph/tests/test_stream_chat_model_stream.py @@ -0,0 +1,245 @@ +import pytest + +from langgraph.stream.chat_model_stream import AsyncChatModelStream, ChatModelStream + + +def _text_delta(text: str) -> dict: + return {"content_block": {"type": "text", "text": text}} + + +def _reasoning_delta(text: str) -> dict: + return {"content_block": {"type": "reasoning", "reasoning": text}} + + +# --------------------------------------------------------------------------- +# Sync ChatModelStream tests +# --------------------------------------------------------------------------- + + +def test_sync_text_accumulates(): + stream = ChatModelStream() + stream._push_content_block_delta(_text_delta("Hello")) + stream._push_content_block_delta(_text_delta(", world")) + stream._finish({"reason": "stop"}) + + assert stream.text == "Hello, world" + assert isinstance(stream.text, str) + + +def test_sync_reasoning_accumulates(): + stream = ChatModelStream() + stream._push_content_block_delta(_reasoning_delta("step 1")) + stream._push_content_block_delta(_reasoning_delta(" -> step 2")) + stream._finish({"reason": "stop"}) + + assert stream.reasoning == "step 1 -> step 2" + assert isinstance(stream.reasoning, str) + + +def test_sync_usage(): + stream = ChatModelStream() + usage = {"input_tokens": 10, "output_tokens": 5} + stream._finish({"reason": "stop", "usage": usage}) + assert stream.usage == usage + + +def test_sync_mixed_blocks(): + stream = ChatModelStream() + stream._push_content_block_delta(_text_delta("answer")) + stream._push_content_block_delta( + {"content_block": {"type": "tool_call", "name": "search"}} + ) + stream._push_content_block_delta(_text_delta(" here")) + stream._finish({"reason": "stop"}) + + assert stream.text == "answer here" + + +def test_sync_tool_call_only_text_empty(): + stream = ChatModelStream() + stream._push_content_block_delta( + {"content_block": {"type": "tool_call", "name": "search"}} + ) + stream._finish({"reason": "stop"}) + assert stream.text == "" + + +def test_sync_fail_marks_done(): + stream = ChatModelStream() + assert not stream.done + stream._fail(RuntimeError("err")) + assert stream.done + + +def test_sync_namespace_and_node(): + stream = ChatModelStream( + namespace=["agent:0", "tools:1"], + node="chat_model", + message_id="msg-123", + ) + assert stream.namespace == ["agent:0", "tools:1"] + assert stream.node == "chat_model" + assert stream.message_id == "msg-123" + + +def test_sync_content_block_finish_authoritative(): + """content-block-finish with authoritative text overrides accumulated.""" + stream = ChatModelStream() + stream._push_content_block_delta(_text_delta("partial")) + stream._push_content_block_finish( + {"content_block": {"type": "text", "text": "full text"}} + ) + assert stream.text == "full text" + + +# --------------------------------------------------------------------------- +# Async ChatModelStream tests +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_async_text_iterable_yields_deltas(): + stream = AsyncChatModelStream() + stream._push_content_block_delta(_text_delta("Hello")) + stream._push_content_block_delta(_text_delta(", world")) + stream._finish({"reason": "stop"}) + + collected = [] + async for delta in stream.text: + collected.append(delta) + assert collected == ["Hello", ", world"] + + +@pytest.mark.anyio +async def test_async_text_awaitable_returns_full(): + stream = AsyncChatModelStream() + stream._push_content_block_delta(_text_delta("Hello")) + stream._push_content_block_delta(_text_delta(", world")) + stream._finish({"reason": "stop"}) + + result = await stream.text + assert result == "Hello, world" + + +@pytest.mark.anyio +async def test_async_reasoning_dual_pattern(): + stream = AsyncChatModelStream() + stream._push_content_block_delta(_reasoning_delta("step 1")) + stream._push_content_block_delta(_reasoning_delta(" -> step 2")) + stream._finish({"reason": "stop"}) + + collected = [] + async for delta in stream.reasoning: + collected.append(delta) + assert collected == ["step 1", " -> step 2"] + + stream2 = AsyncChatModelStream() + stream2._push_content_block_delta(_reasoning_delta("thinking")) + stream2._finish({"reason": "stop"}) + full = await stream2.reasoning + assert full == "thinking" + + +@pytest.mark.anyio +async def test_async_usage_resolves(): + stream = AsyncChatModelStream() + usage = {"input_tokens": 10, "output_tokens": 5} + stream._finish({"reason": "stop", "usage": usage}) + result = await stream.usage + assert result == usage + + +@pytest.mark.anyio +async def test_async_mixed_blocks_text_only(): + stream = AsyncChatModelStream() + stream._push_content_block_delta(_text_delta("answer")) + stream._push_content_block_delta( + {"content_block": {"type": "tool_call", "name": "search"}} + ) + stream._push_content_block_delta(_text_delta(" here")) + stream._finish({"reason": "stop"}) + + collected = [] + async for delta in stream.text: + collected.append(delta) + assert collected == ["answer", " here"] + + +@pytest.mark.anyio +async def test_async_tool_call_only_text_empty(): + stream = AsyncChatModelStream() + stream._push_content_block_delta( + {"content_block": {"type": "tool_call", "name": "search"}} + ) + stream._finish({"reason": "stop"}) + result = await stream.text + assert result == "" + + +@pytest.mark.anyio +async def test_async_fail_raises_on_text_await(): + stream = AsyncChatModelStream() + stream._push_content_block_delta(_text_delta("partial")) + stream._fail(RuntimeError("model error")) + + with pytest.raises(RuntimeError, match="model error"): + await stream.text + + +@pytest.mark.anyio +async def test_async_fail_raises_on_reasoning_await(): + stream = AsyncChatModelStream() + stream._push_content_block_delta(_reasoning_delta("thinking")) + stream._fail(RuntimeError("model error")) + + with pytest.raises(RuntimeError, match="model error"): + await stream.reasoning + + +@pytest.mark.anyio +async def test_async_fail_raises_on_usage_await(): + stream = AsyncChatModelStream() + stream._fail(RuntimeError("model error")) + + with pytest.raises(RuntimeError, match="model error"): + await stream.usage + + +@pytest.mark.anyio +async def test_async_fail_raises_during_text_iteration(): + stream = AsyncChatModelStream() + stream._push_content_block_delta(_text_delta("partial")) + stream._fail(RuntimeError("model error")) + + collected = [] + with pytest.raises(RuntimeError, match="model error"): + async for delta in stream.text: + collected.append(delta) + assert collected == ["partial"] + + +@pytest.mark.anyio +async def test_async_fail_marks_done(): + stream = AsyncChatModelStream() + assert not stream.done + stream._fail(RuntimeError("err")) + assert stream.done + + +@pytest.mark.anyio +async def test_async_namespace_and_node(): + stream = AsyncChatModelStream( + namespace=["agent:0", "tools:1"], + node="chat_model", + message_id="msg-123", + ) + assert stream.namespace == ["agent:0", "tools:1"] + assert stream.node == "chat_model" + assert stream.message_id == "msg-123" + + +@pytest.mark.anyio +async def test_async_inherits_from_sync(): + """AsyncChatModelStream is a subclass of ChatModelStream.""" + stream = AsyncChatModelStream() + assert isinstance(stream, ChatModelStream) diff --git a/libs/langgraph/tests/test_stream_convert.py b/libs/langgraph/tests/test_stream_convert.py new file mode 100644 index 000000000..091363854 --- /dev/null +++ b/libs/langgraph/tests/test_stream_convert.py @@ -0,0 +1,86 @@ +from langgraph.stream._convert import STREAM_V2_MODES, convert_to_protocol_event + + +def test_values_mode(): + evt = convert_to_protocol_event((), "values", {"x": 1}) + assert evt is not None + assert evt["method"] == "values" + assert evt["params"]["data"] == {"x": 1} + + +def test_updates_mode(): + evt = convert_to_protocol_event((), "updates", {"node": "out"}) + assert evt is not None + assert evt["method"] == "updates" + + +def test_messages_mode(): + evt = convert_to_protocol_event((), "messages", {"event": "msg"}) + assert evt is not None + assert evt["method"] == "messages" + + +def test_custom_mode(): + evt = convert_to_protocol_event((), "custom", "hello") + assert evt is not None + assert evt["method"] == "custom" + assert evt["params"]["data"] == "hello" + + +def test_debug_mode(): + evt = convert_to_protocol_event((), "debug", {}) + assert evt is not None + assert evt["method"] == "debug" + + +def test_checkpoints_mode(): + evt = convert_to_protocol_event((), "checkpoints", {}) + assert evt is not None + assert evt["method"] == "checkpoints" + + +def test_tasks_mode(): + evt = convert_to_protocol_event((), "tasks", {}) + assert evt is not None + assert evt["method"] == "tasks" + + +def test_namespace_passthrough(): + evt = convert_to_protocol_event(("agent", "0"), "values", {}) + assert evt is not None + assert evt["params"]["namespace"] == ["agent", "0"] + + +def test_timestamp_populated(): + evt = convert_to_protocol_event((), "values", {}) + assert evt is not None + assert isinstance(evt["params"]["timestamp"], int) + assert evt["params"]["timestamp"] > 0 + + +def test_unknown_mode_returns_none(): + assert convert_to_protocol_event((), "unknown_mode", {}) is None + + +def test_node_parameter(): + evt = convert_to_protocol_event((), "values", {}, node="agent") + assert evt is not None + assert evt["params"]["node"] == "agent" + + +def test_type_is_event(): + evt = convert_to_protocol_event((), "values", {}) + assert evt is not None + assert evt["type"] == "event" + + +def test_stream_v2_modes_complete(): + assert set(STREAM_V2_MODES) == { + "values", + "updates", + "messages", + "custom", + "checkpoints", + "tasks", + "debug", + } diff --git a/libs/langgraph/tests/test_stream_event_log.py b/libs/langgraph/tests/test_stream_event_log.py new file mode 100644 index 000000000..7bec10156 --- /dev/null +++ b/libs/langgraph/tests/test_stream_event_log.py @@ -0,0 +1,144 @@ +import asyncio + +import pytest + +from langgraph.stream._event_log import EventLog + + +@pytest.mark.anyio +async def test_push_and_iterate_in_order(): + log = EventLog() + log.append("a") + log.append("b") + log.append("c") + log.close() + items = [item async for item in log.subscribe(0)] + assert items == ["a", "b", "c"] + + +@pytest.mark.anyio +async def test_multiple_independent_cursors(): + log = EventLog() + log.append("x") + log.append("y") + log.close() + items1 = [item async for item in log.subscribe(0)] + items2 = [item async for item in log.subscribe(0)] + assert items1 == ["x", "y"] + assert items2 == ["x", "y"] + + +@pytest.mark.anyio +async def test_close_ends_iteration(): + log = EventLog() + log.close() + items = [item async for item in log.subscribe(0)] + assert items == [] + + +@pytest.mark.anyio +async def test_fail_raises_error(): + log = EventLog() + log.fail(RuntimeError("boom")) + with pytest.raises(RuntimeError, match="boom"): + async for _ in log.subscribe(0): + pass + + +@pytest.mark.anyio +async def test_nonzero_offset(): + log = EventLog() + log.append("x") + log.append("y") + log.append("z") + log.close() + items = [item async for item in log.subscribe(2)] + assert items == ["z"] + + +@pytest.mark.anyio +async def test_concurrent_push_and_iterate(): + log = EventLog() + received = [] + + async def consumer(): + async for item in log.subscribe(0): + received.append(item) + + async def producer(): + for i in range(5): + log.append(i) + await asyncio.sleep(0.01) + log.close() + + await asyncio.gather(producer(), consumer()) + assert received == [0, 1, 2, 3, 4] + + +@pytest.mark.anyio +async def test_items_before_cursor_visible(): + log = EventLog() + log.append("a") + log.append("b") + cursor = log.subscribe(0) + log.append("c") + log.close() + items = [item async for item in cursor] + assert items == ["a", "b", "c"] + + +@pytest.mark.anyio +async def test_empty_log_closed_yields_nothing(): + log = EventLog() + log.close() + items = [item async for item in log.subscribe(0)] + assert items == [] + + +@pytest.mark.anyio +async def test_fail_mid_iteration(): + """A cursor that has consumed some items should raise when fail() is called.""" + log = EventLog() + received = [] + + async def consumer(): + async for item in log.subscribe(0): + received.append(item) + + async def producer(): + log.append("a") + log.append("b") + await asyncio.sleep(0.02) + log.fail(RuntimeError("mid-stream error")) + + with pytest.raises(RuntimeError, match="mid-stream error"): + await asyncio.gather(producer(), consumer()) + assert received == ["a", "b"] + + +@pytest.mark.anyio +async def test_abandoned_cursor_cleans_up_waiters(): + """Abandoned async cursors should not leave stale futures in the + EventLog waiter list. + + When a cursor's __anext__ is cancelled (e.g. consumer breaks out of + ``async for``), the Future it registered in ``_waiters`` should be + cleaned up. Otherwise the list grows without bound until the next + append/close/fail triggers ``_wake_all()``. + """ + log: EventLog[str] = EventLog() + + for _ in range(10): + cursor = log.subscribe(0) + task = asyncio.ensure_future(cursor.__anext__()) + await asyncio.sleep(0) # let task register its waiter + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + assert len(log._waiters) == 0, ( + f"Expected 0 waiters after abandoning 10 cursors, " + f"got {len(log._waiters)}. Abandoned cursors leak futures." + ) diff --git a/libs/langgraph/tests/test_stream_mux.py b/libs/langgraph/tests/test_stream_mux.py new file mode 100644 index 000000000..d2160e5f1 --- /dev/null +++ b/libs/langgraph/tests/test_stream_mux.py @@ -0,0 +1,293 @@ +from typing import Any + +import pytest + +from langgraph.stream._convert import convert_to_protocol_event +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import ProtocolEvent +from langgraph.stream.stream_channel import StreamChannel + + +def _event(mode: str, data: Any, ns: list[str] | None = None) -> ProtocolEvent: + ev = convert_to_protocol_event(tuple(ns or []), mode, data) + assert ev is not None + return ev + + +class _MockTransformer: + def __init__(self, *, suppress: bool = False): + self.calls: list[ProtocolEvent] = [] + self._suppress = suppress + + def init(self) -> Any: + return None + + def process(self, event: ProtocolEvent) -> bool: + self.calls.append(event) + return not self._suppress + + def finalize(self) -> None: + pass + + def fail(self, err: BaseException) -> None: + pass + + +@pytest.mark.anyio +async def test_events_through_reducer_pipeline(): + reducer = _MockTransformer() + mux = StreamMux(transformers=[reducer]) + event = _event("values", {"key": "val"}) + mux.push(event) + assert len(reducer.calls) == 1 + assert reducer.calls[0] is event + + +@pytest.mark.anyio +async def test_reducer_suppresses_event(): + reducer = _MockTransformer(suppress=True) + mux = StreamMux(transformers=[reducer]) + mux.push(_event("values", {"x": 1})) + mux.close() + assert len(reducer.calls) == 1 + assert len(mux.event_log) == 0 + + +@pytest.mark.anyio +async def test_namespace_discovery(): + mux = StreamMux() + mux.push(_event("values", {"a": 1}, ns=["child:0"])) + assert "child:0" in mux._discovered_ns + + +@pytest.mark.anyio +async def test_top_level_ns_only(): + mux = StreamMux() + mux.push(_event("values", {"a": 1}, ns=["agent:0", "tools:1"])) + assert "agent:0" in mux._discovered_ns + assert "tools:1" not in mux._discovered_ns + + +@pytest.mark.anyio +async def test_subscribe_events_filter(): + mux = StreamMux() + mux.push(_event("values", {"a": 1}, ns=["child:0"])) + mux.push(_event("values", {"b": 2}, ns=["other:1"])) + mux.push(_event("values", {"c": 3}, ns=["child:0"])) + mux.close() + + collected = [] + async for ev in mux.subscribe_events(["child:0"]): + collected.append(ev) + assert len(collected) == 2 + assert collected[0]["params"]["data"] == {"a": 1} + assert collected[1]["params"]["data"] == {"c": 3} + + +@pytest.mark.anyio +async def test_close_resolves_output(): + mux = StreamMux() + fut = mux.get_output_future() + mux.push(_event("values", {"v": 1})) + mux.push(_event("values", {"v": 2})) + mux.close() + result = await fut + assert result == {"v": 2} + + +@pytest.mark.anyio +async def test_fail_rejects_output(): + mux = StreamMux() + fut = mux.get_output_future() + mux.fail(ValueError("boom")) + with pytest.raises(ValueError, match="boom"): + await fut + + +@pytest.mark.anyio +async def test_latest_values_tracked(): + mux = StreamMux() + mux.push(_event("values", {"v": 1}, ns=["child:0"])) + mux.push(_event("values", {"v": 2}, ns=["child:0"])) + assert mux.get_latest_values(["child:0"]) == {"v": 2} + + +@pytest.mark.anyio +async def test_interrupt_tracking(): + """StreamMux should track __interrupt__ payloads in values events.""" + + class _FakeInterrupt: + def __init__(self, id: str, payload: Any): + self.id = id + self.payload = payload + + mux = StreamMux() + interrupt_obj = _FakeInterrupt("int-1", "what do you want?") + mux.push( + _event( + "values", + {"__interrupt__": [interrupt_obj]}, + ) + ) + assert mux.interrupted is True + assert len(mux.interrupts) == 1 + assert mux.interrupts[0]["interrupt_id"] == "int-1" + assert mux.interrupts[0]["payload"] is interrupt_obj + + +@pytest.mark.anyio +async def test_no_interrupt_by_default(): + mux = StreamMux() + mux.push(_event("values", {"x": 1})) + mux.close() + assert mux.interrupted is False + assert mux.interrupts == [] + + +@pytest.mark.anyio +async def test_push_after_close_ignored(): + mux = StreamMux() + mux.push(_event("values", {"a": 1})) + mux.close() + mux.push(_event("values", {"b": 2})) + assert len(mux.event_log) == 1 + + +@pytest.mark.anyio +async def test_fail_rejects_all_futures(): + mux = StreamMux() + fut1 = mux.get_output_future([]) + fut2 = mux.get_output_future(["child:0"]) + mux.fail(ValueError("boom")) + with pytest.raises(ValueError, match="boom"): + await fut1 + with pytest.raises(ValueError, match="boom"): + await fut2 + + +@pytest.mark.anyio +async def test_channel_events_bypass_transformer_pipeline(): + """Events emitted via ``StreamChannel.push()`` are appended directly + to the event log, bypassing the transformer pipeline. This matches + the JS implementation and avoids re-entrancy bugs. + """ + mock = _MockTransformer() + mux = StreamMux(transformers=[mock]) + + channel: StreamChannel[str] = StreamChannel("my_channel") + mux.wire_channels({"ch": channel}) + + # Regular push — transformer sees it + mux.push(_event("values", {"a": 1})) + assert len(mock.calls) == 1 + + # Channel push — bypasses transformers, goes straight to event log + channel.push("hello from channel") + + assert len(mock.calls) == 1, ( + f"Transformer saw {len(mock.calls)} events (expected 1). " + "Channel events should bypass the transformer pipeline." + ) + + # But the event IS in the log + mux.close() + events = [] + async for ev in mux.subscribe_events(): + events.append(ev) + assert len(events) == 2 + assert events[1]["method"] == "my_channel" + assert events[1]["params"]["data"] == "hello from channel" + + +@pytest.mark.anyio +async def test_event_log_has_monotonic_seq_numbers(): + """All events in the event log should have strictly monotonically + increasing seq numbers so consumers can reason about ordering. + + Events from ``mux.push()`` carry seq numbers assigned by the pump + while channel-emitted events use a separate counter + (``_next_emit_seq``). When interleaved, seq numbers can duplicate. + """ + mux = StreamMux() + channel: StreamChannel[str] = StreamChannel("test_ch") + mux.wire_channels({"ch": channel}) + + mux.push(_event("values", {"a": 1})) # log seq: 0 + channel.push("from_channel") # log seq: 0 (from _next_emit_seq) + mux.push(_event("values", {"b": 2})) # log seq: 1 + mux.close() + + seqs: list[int] = [] + async for event in mux.subscribe_events(): + seqs.append(event["seq"]) + + assert len(seqs) == 3, f"Expected 3 events but got {len(seqs)}" + + for i in range(1, len(seqs)): + assert seqs[i] > seqs[i - 1], ( + f"Seq numbers not strictly monotonic: {seqs}. " + f"seq[{i}]={seqs[i]} <= seq[{i - 1}]={seqs[i - 1]}. " + "Channel events use a separate counter from push() events." + ) + + +@pytest.mark.anyio +async def test_channel_push_during_process_preserves_namespace(): + """When two transformers both call channel.push() during the same + outer mux.push(), the second transformer's channel event should + still carry the original event's namespace. + + Bug: the first channel.push() re-enters mux.push(), which resets + ``_current_namespace`` to ``[]`` on exit. The second transformer's + channel.push() then reads the clobbered value and its event gets + ``namespace: []`` instead of the original. + """ + + class _ChannelTransformer: + """Pushes to its channel whenever it sees a ``values`` event.""" + + def __init__(self, name: str) -> None: + self.name = name + self.channel: StreamChannel[str] = StreamChannel(name) + + def init(self) -> Any: + return {self.name: self.channel} + + def process(self, event: ProtocolEvent) -> bool: + if event["method"] == "values": + self.channel.push(f"from_{self.name}") + return True + + def finalize(self) -> None: + pass + + def fail(self, err: BaseException) -> None: + pass + + t1 = _ChannelTransformer("first") + t2 = _ChannelTransformer("second") + mux = StreamMux(transformers=[t1, t2]) + mux.wire_channels({"first": t1.channel}) + mux.wire_channels({"second": t2.channel}) + + # Push a values event with a non-root namespace + mux.push(_event("values", {"x": 1}, ns=["agent:0"])) + mux.close() + + # Collect channel events emitted by each transformer + channel_events: list[ProtocolEvent] = [] + async for ev in mux.subscribe_events(): + if ev["method"] in ("first", "second"): + channel_events.append(ev) + + assert len(channel_events) == 2, ( + f"Expected 2 channel events but got {len(channel_events)}" + ) + + for ev in channel_events: + assert ev["params"]["namespace"] == ["agent:0"], ( + f"Channel event for method={ev['method']!r} has " + f"namespace={ev['params']['namespace']!r}, expected ['agent:0']. " + "The nested mux.push() from the first channel.push() clobbered " + "_current_namespace before the second transformer ran." + ) diff --git a/libs/langgraph/tests/test_stream_reducers.py b/libs/langgraph/tests/test_stream_reducers.py new file mode 100644 index 000000000..c0d119168 --- /dev/null +++ b/libs/langgraph/tests/test_stream_reducers.py @@ -0,0 +1,214 @@ +from typing import Any + +import pytest + +from langgraph.stream._convert import convert_to_protocol_event +from langgraph.stream._types import ProtocolEvent +from langgraph.stream.chat_model_stream import ChatModelStream +from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer + + +def _event( + mode: str, + data: Any, + ns: list[str] | None = None, + node: str | None = None, +) -> ProtocolEvent: + ev = convert_to_protocol_event(tuple(ns or []), mode, data, node=node) + assert ev is not None + return ev + + +# -- ValuesTransformer --------------------------------------------------------- + + +@pytest.mark.anyio +async def test_values_captures_values_events(): + reducer = ValuesTransformer() + reducer.init() + reducer.process(_event("values", {"a": 1})) + reducer.process(_event("values", {"b": 2})) + reducer.finalize() + + collected = [] + async for item in reducer.values_log.subscribe(0): + collected.append(item) + assert len(collected) == 2 + assert collected[0]["data"] == {"a": 1} + assert collected[1]["data"] == {"b": 2} + + +@pytest.mark.anyio +async def test_values_ignores_other_modes(): + reducer = ValuesTransformer() + reducer.init() + reducer.process(_event("updates", {"x": 1})) + reducer.process(_event("messages", {"event": "message-start"})) + reducer.finalize() + + collected = [] + async for item in reducer.values_log.subscribe(0): + collected.append(item) + assert len(collected) == 0 + + +@pytest.mark.anyio +async def test_values_latest_per_namespace(): + reducer = ValuesTransformer() + reducer.init() + reducer.process(_event("values", {"v": 1}, ns=["child:0"])) + reducer.process(_event("values", {"v": 2}, ns=["child:0"])) + assert reducer.get_latest("child:0") == {"v": 2} + + +@pytest.mark.anyio +async def test_values_finalize_closes_log(): + reducer = ValuesTransformer() + reducer.init() + reducer.process(_event("values", {"a": 1})) + reducer.finalize() + assert reducer.values_log.closed + + +# -- MessagesTransformer ------------------------------------------------------- + + +def _msg_start(ns=None, node=None, message_id="msg-1"): + return _event( + "messages", + {"event": "message-start", "message_id": message_id}, + ns=ns, + node=node, + ) + + +def _content_delta(text, ns=None, node=None): + return _event( + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": text}, + }, + ns=ns, + node=node, + ) + + +def _msg_finish(ns=None, node=None): + return _event( + "messages", + {"event": "message-finish", "reason": "stop"}, + ns=ns, + node=node, + ) + + +@pytest.mark.anyio +async def test_messages_groups_lifecycle(): + reducer = MessagesTransformer() + reducer.init() + reducer.process(_msg_start()) + reducer.process(_content_delta("hi")) + reducer.process(_msg_finish()) + reducer.finalize() + + collected = [] + async for stream in reducer.messages_log.subscribe(0): + collected.append(stream) + assert len(collected) == 1 + assert isinstance(collected[0], ChatModelStream) + assert collected[0].done + + +@pytest.mark.anyio +async def test_messages_multiple_sequential(): + reducer = MessagesTransformer() + reducer.init() + reducer.process(_msg_start(message_id="m1")) + reducer.process(_msg_finish()) + reducer.process(_msg_start(message_id="m2")) + reducer.process(_msg_finish()) + reducer.finalize() + + collected = [] + async for stream in reducer.messages_log.subscribe(0): + collected.append(stream) + assert len(collected) == 2 + + +@pytest.mark.anyio +async def test_messages_namespace_filter(): + reducer = MessagesTransformer(namespace=["root"]) + reducer.init() + reducer.process(_msg_start(ns=["root"])) + reducer.process(_msg_finish(ns=["root"])) + reducer.process(_msg_start(ns=["other"], message_id="m2")) + reducer.process(_msg_finish(ns=["other"])) + reducer.finalize() + + collected = [] + async for stream in reducer.messages_log.subscribe(0): + collected.append(stream) + assert len(collected) == 1 + + +@pytest.mark.anyio +async def test_messages_node_filter(): + reducer = MessagesTransformer(node_filter="agent") + reducer.init() + reducer.process(_msg_start(node="agent")) + reducer.process(_msg_finish(node="agent")) + reducer.process(_msg_start(node="tools", message_id="m2")) + reducer.process(_msg_finish(node="tools")) + reducer.finalize() + + collected = [] + async for stream in reducer.messages_log.subscribe(0): + collected.append(stream) + assert len(collected) == 1 + + +@pytest.mark.anyio +async def test_messages_error_event(): + """An error event should fail the active ChatModelStream.""" + reducer = MessagesTransformer() + reducer.init() + reducer.process(_msg_start()) + reducer.process(_content_delta("partial")) + reducer.process( + _event("messages", {"event": "error", "message": "connection lost"}), + ) + reducer.finalize() + + collected: list[ChatModelStream] = [] + async for stream in reducer.messages_log.subscribe(0): + collected.append(stream) + assert len(collected) == 1 + assert collected[0].done + + +@pytest.mark.anyio +async def test_messages_fail_propagates_to_active(): + """transformer.fail() should propagate the error to any active streams.""" + reducer = MessagesTransformer() + reducer.init() + reducer.process(_msg_start()) + reducer.process(_content_delta("partial")) + reducer.fail(RuntimeError("graph failed")) + + # The messages log should be failed too + with pytest.raises(RuntimeError, match="graph failed"): + async for _ in reducer.messages_log.subscribe(0): + pass + + +@pytest.mark.anyio +async def test_values_fail_propagates(): + reducer = ValuesTransformer() + reducer.init() + reducer.process(_event("values", {"a": 1})) + reducer.fail(RuntimeError("graph failed")) + + with pytest.raises(RuntimeError, match="graph failed"): + async for _ in reducer.values_log.subscribe(0): + pass diff --git a/libs/langgraph/tests/test_stream_run_stream.py b/libs/langgraph/tests/test_stream_run_stream.py new file mode 100644 index 000000000..09cb5b67b --- /dev/null +++ b/libs/langgraph/tests/test_stream_run_stream.py @@ -0,0 +1,610 @@ +import asyncio +from collections.abc import AsyncIterator, Iterator +from typing import Any + +import pytest + +from langgraph.stream._mux import StreamMux +from langgraph.stream._types import ProtocolEvent +from langgraph.stream.chat_model_stream import ChatModelStream +from langgraph.stream.run_stream import ( + AsyncGraphRunStream, + AsyncSubgraphRunStream, + create_async_graph_run_stream, + create_graph_run_stream, +) +from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer + + +async def _mock_source( + chunks: list[tuple[tuple[str, ...], str, Any]], +) -> AsyncIterator[tuple[tuple[str, ...], str, Any]]: + for chunk in chunks: + yield chunk + + +@pytest.mark.anyio +async def test_aiter_yields_all_events(): + chunks = [ + ((), "values", {"step": 1}), + ((), "values", {"step": 2}), + ((), "updates", {"node": "a"}), + ] + run = await create_async_graph_run_stream(_mock_source(chunks)) + await asyncio.sleep(0.05) + + collected: list[ProtocolEvent] = [] + async for event in run: + collected.append(event) + assert len(collected) == 3 + assert collected[0]["method"] == "values" + assert collected[2]["method"] == "updates" + + +@pytest.mark.anyio +async def test_subgraph_name_and_index(): + vr, mr = ValuesTransformer(), MessagesTransformer() + mux = StreamMux(transformers=[vr, mr]) + sub = AsyncSubgraphRunStream( + mux=mux, + namespace=["researcher:2"], + transformers=[vr, mr], + ) + assert sub.name == "researcher" + assert sub.index == 2 + + +@pytest.mark.anyio +async def test_subgraph_name_no_index(): + vr, mr = ValuesTransformer(), MessagesTransformer() + mux = StreamMux(transformers=[vr, mr]) + sub = AsyncSubgraphRunStream( + mux=mux, namespace=["agent"], transformers=[vr, mr] + ) + assert sub.name == "agent" + assert sub.index == 0 + + +@pytest.mark.anyio +async def test_values_iterable(): + chunks = [ + ((), "values", {"v": 1}), + ((), "values", {"v": 2}), + ] + run = await create_async_graph_run_stream(_mock_source(chunks)) + await asyncio.sleep(0.05) + + collected = [] + async for v in run.values: + collected.append(v) + assert len(collected) == 2 + assert collected[0] == {"v": 1} + assert collected[1] == {"v": 2} + + +@pytest.mark.anyio +async def test_values_awaitable(): + chunks = [ + ((), "values", {"v": 1}), + ((), "values", {"v": 2}), + ] + run = await create_async_graph_run_stream(_mock_source(chunks)) + await asyncio.sleep(0.05) + + result = await run.values + assert result == {"v": 2} + + +@pytest.mark.anyio +async def test_output_resolves(): + chunks = [((), "values", {"final": True})] + run = await create_async_graph_run_stream(_mock_source(chunks)) + await asyncio.sleep(0.05) + result = await run.output + assert result == {"final": True} + + +@pytest.mark.anyio +async def test_messages_yields_streams(): + chunks = [ + ((), "messages", {"event": "message-start", "message_id": "m1"}), + ( + (), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "hi"}, + }, + ), + ((), "messages", {"event": "message-finish", "reason": "stop"}), + ] + run = await create_async_graph_run_stream(_mock_source(chunks)) + await asyncio.sleep(0.05) + + collected: list[ChatModelStream] = [] + async for stream in run.messages: + collected.append(stream) + assert len(collected) == 1 + assert isinstance(collected[0], ChatModelStream) + assert collected[0].done + + +@pytest.mark.anyio +async def test_interrupted_false_by_default(): + vr, mr = ValuesTransformer(), MessagesTransformer() + mux = StreamMux(transformers=[vr, mr]) + run = AsyncGraphRunStream(mux=mux, transformers=[vr, mr]) + assert run.interrupted is False + + +@pytest.mark.anyio +async def test_abort_sets_signal(): + vr, mr = ValuesTransformer(), MessagesTransformer() + mux = StreamMux(transformers=[vr, mr]) + run = AsyncGraphRunStream(mux=mux, transformers=[vr, mr]) + assert not run.signal.is_set() + run.abort() + assert run.signal.is_set() + + +@pytest.mark.anyio +async def test_abort_stops_pump(): + """Calling abort() should stop the pump from processing further chunks.""" + gate = asyncio.Event() + + async def _gated_source(): + yield ((), "values", {"v": 1}) + yield ((), "values", {"v": 2}) + await gate.wait() # Block until released + yield ((), "values", {"v": 3}) # Should not be processed + + run = await create_async_graph_run_stream(_gated_source()) + await asyncio.sleep(0.05) # Let first two events through + run.abort() + gate.set() # Unblock the source so the pump can check abort and exit + await asyncio.sleep(0.05) # Let pump close the mux + + collected = [] + async for event in run: + if event["method"] == "values": + collected.append(event["params"]["data"]) + # v:3 should not have been processed because abort was set + assert all(v.get("v") != 3 for v in collected) + + +@pytest.mark.anyio +async def test_messages_from_filters_by_node(): + """messages_from(node) should only yield messages from the specified node.""" + chunks = [ + ( + (), + "messages", + {"event": "message-start", "message_id": "m1", "__node__": "agent"}, + ), + ( + (), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "from agent"}, + "__node__": "agent", + }, + ), + ( + (), + "messages", + {"event": "message-finish", "reason": "stop", "__node__": "agent"}, + ), + ( + (), + "messages", + {"event": "message-start", "message_id": "m2", "__node__": "tools"}, + ), + ( + (), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "from tools"}, + "__node__": "tools", + }, + ), + ( + (), + "messages", + {"event": "message-finish", "reason": "stop", "__node__": "tools"}, + ), + ] + run = await create_async_graph_run_stream(_mock_source(chunks)) + await asyncio.sleep(0.05) + + agent_msgs: list[ChatModelStream] = [] + async for stream in run.messages_from("agent"): + agent_msgs.append(stream) + assert len(agent_msgs) == 1 + assert agent_msgs[0].node == "agent" + + +# --------------------------------------------------------------------------- +# GraphRunStream / create_graph_run_stream +# --------------------------------------------------------------------------- + + +def _sync_source( + chunks: list[tuple[tuple[str, ...], str, Any]], +) -> Iterator[tuple[tuple[str, ...], str, Any]]: + yield from chunks + + +def test_sync_create_yields_all_events(): + chunks = [ + ((), "values", {"step": 1}), + ((), "values", {"step": 2}), + ((), "updates", {"node": "a"}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + collected = list(run) + assert len(collected) == 3 + assert collected[0]["method"] == "values" + assert collected[2]["method"] == "updates" + + +def test_sync_output(): + chunks = [ + ((), "values", {"v": 1}), + ((), "values", {"v": 2}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + assert run.output == {"v": 2} + + +def test_sync_values_iteration(): + chunks = [ + ((), "values", {"v": 1}), + ((), "values", {"v": 2}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + collected = list(run.values) + assert len(collected) == 2 + assert collected[0] == {"v": 1} + assert collected[1] == {"v": 2} + + +def test_sync_messages(): + chunks = [ + ((), "messages", {"event": "message-start", "message_id": "m1"}), + ( + (), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "hi"}, + }, + ), + ((), "messages", {"event": "message-finish", "reason": "stop"}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + collected = list(run.messages) + assert len(collected) == 1 + assert isinstance(collected[0], ChatModelStream) + assert collected[0].done + + +def test_sync_messages_text_accessible(): + """Sync consumers should be able to read ChatModelStream text content + without an async event loop. + """ + chunks = [ + ((), "messages", {"event": "message-start", "message_id": "m1"}), + ( + (), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "Hello"}, + }, + ), + ( + (), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": " world"}, + }, + ), + ((), "messages", {"event": "message-finish", "reason": "stop"}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + + for msg in run.messages: + assert isinstance(msg.text, str) + assert msg.text == "Hello world" + assert msg.done + + +def test_sync_messages_content_populated_when_yielded(): + """When sync run.messages yields a ChatModelStream, its content should + be fully populated (done=True) with all text accumulated. + """ + chunks = [ + ((), "messages", {"event": "message-start", "message_id": "m1"}), + ( + (), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "answer"}, + }, + ), + ((), "messages", {"event": "message-finish", "reason": "stop"}), + ((), "messages", {"event": "message-start", "message_id": "m2"}), + ( + (), + "messages", + { + "event": "content-block-delta", + "content_block": {"type": "text", "text": "second"}, + }, + ), + ((), "messages", {"event": "message-finish", "reason": "stop"}), + ] + run = create_graph_run_stream(_sync_source(chunks)) + + messages = list(run.messages) + assert len(messages) == 2 + assert messages[0].done + assert messages[0].text == "answer" + assert messages[1].done + assert messages[1].text == "second" + + +def test_sync_output_mapper(): + chunks = [((), "values", {"v": 1})] + run = create_graph_run_stream( + _sync_source(chunks), output_mapper=lambda x: {"mapped": x["v"]} + ) + assert run.output == {"mapped": 1} + + +def test_sync_interrupted_false(): + chunks = [((), "values", {"v": 1})] + run = create_graph_run_stream(_sync_source(chunks)) + assert run.interrupted is False + + +def test_sync_source_error(): + """If the source raises, the mux should fail and the error should propagate.""" + + def _bad_source(): + yield ((), "values", {"v": 1}) + raise ValueError("source error") + + run = create_graph_run_stream(_bad_source()) + collected = list(run) + # Events before the error are still accessible + assert len(collected) >= 1 + assert collected[0]["method"] == "values" + # The mux recorded the failure + assert run._mux._error is not None + assert isinstance(run._mux._error, ValueError) + assert "source error" in str(run._mux._error) + + +# --------------------------------------------------------------------------- +# GraphRunStream — lazy consumption tests +# --------------------------------------------------------------------------- + + +def test_sync_lazy_not_consumed_on_creation(): + """Source iterator should not be consumed when the stream is created.""" + consumed = 0 + + def counting_source(): + nonlocal consumed + for chunk in [ + ((), "values", {"v": 1}), + ((), "values", {"v": 2}), + ((), "values", {"v": 3}), + ]: + consumed += 1 + yield chunk + + create_graph_run_stream(counting_source()) + assert consumed == 0 + + +def test_sync_lazy_values_pull_incrementally(): + """Iterating .values should pull from the source one event at a time.""" + consumed = 0 + + def counting_source(): + nonlocal consumed + for chunk in [ + ((), "values", {"v": 1}), + ((), "values", {"v": 2}), + ((), "values", {"v": 3}), + ]: + consumed += 1 + yield chunk + + run = create_graph_run_stream(counting_source()) + assert consumed == 0 + + it = iter(run.values) + v = next(it) + assert v == {"v": 1} + assert consumed == 1 + + v = next(it) + assert v == {"v": 2} + assert consumed == 2 + + # Source not fully drained yet + assert consumed < 3 + + +def test_sync_lazy_output_drains_all(): + """Accessing .output should drain the entire source.""" + consumed = 0 + + def counting_source(): + nonlocal consumed + for chunk in [((), "values", {"v": i}) for i in range(5)]: + consumed += 1 + yield chunk + + run = create_graph_run_stream(counting_source()) + assert consumed == 0 + assert run.output == {"v": 4} + assert consumed == 5 + + +def test_sync_lazy_early_break(): + """Breaking out of a projection early should leave the source partially consumed.""" + consumed = 0 + + def counting_source(): + nonlocal consumed + for chunk in [((), "values", {"v": i}) for i in range(10)]: + consumed += 1 + yield chunk + + run = create_graph_run_stream(counting_source()) + for v in run.values: + break # consume only the first value + assert consumed == 1 + assert consumed < 10 + + +def test_sync_lazy_interleaved_projections(): + """Switching between projections replays buffered items then resumes pumping.""" + consumed = 0 + + def counting_source(): + nonlocal consumed + for chunk in [ + ((), "values", {"v": 1}), + ((), "messages", {"event": "message-start", "message_id": "m1"}), + ((), "messages", {"event": "message-finish", "reason": "stop"}), + ((), "values", {"v": 2}), + ]: + consumed += 1 + yield chunk + + run = create_graph_run_stream(counting_source()) + + # Pull first value — consumes 1 source item + vit = iter(run.values) + assert next(vit) == {"v": 1} + assert consumed == 1 + + # Pull first message — pumps until message-finish (item 3) so the + # ChatModelStream is fully populated before yielding. + mit = iter(run.messages) + msg = next(mit) + assert isinstance(msg, ChatModelStream) + assert msg.done + assert consumed == 3 + + # Pull second value — pumps values (item 4) + assert next(vit) == {"v": 2} + assert consumed == 4 + + +def test_sync_lazy_iter_pulls_incrementally(): + """Raw __iter__ should pull from the source lazily.""" + consumed = 0 + + def counting_source(): + nonlocal consumed + for chunk in [ + ((), "values", {"v": 1}), + ((), "updates", {"node": "a"}), + ((), "values", {"v": 2}), + ]: + consumed += 1 + yield chunk + + run = create_graph_run_stream(counting_source()) + it = iter(run) + event = next(it) + assert event["method"] == "values" + assert consumed == 1 + + event = next(it) + assert event["method"] == "updates" + assert consumed == 2 + + +def test_sync_lazy_source_error(): + """If the source raises mid-stream, earlier events are still accessible.""" + consumed = 0 + + def bad_source(): + nonlocal consumed + consumed += 1 + yield ((), "values", {"v": 1}) + raise ValueError("boom") + + run = create_graph_run_stream(bad_source()) + collected = list(run) + assert len(collected) >= 1 + assert collected[0]["method"] == "values" + + +@pytest.mark.anyio +async def test_subgraph_child_values_receive_post_discovery_events(): + """Child AsyncSubgraphRunStream.values iteration should include events + that arrive AFTER the subgraph namespace is first discovered. + + ``_SubgraphsProjection`` creates a local ``ValuesTransformer`` for + each child and replays existing events, but never registers the + transformer with the mux. Events that arrive after discovery are + not routed to it, and ``finalize()`` is not called (the mux wasn't + closed at discovery time), so the child's values_log is never + closed and iteration hangs. + """ + gate = asyncio.Event() + + async def _source() -> AsyncIterator[tuple[tuple[str, ...], str, Any]]: + # First event from child namespace — triggers discovery + yield (("child:0",), "values", {"v": 1}) + await gate.wait() + # Second event from same child — arrives after discovery + yield (("child:0",), "values", {"v": 2}) + # Root event so the mux tracks output + yield ((), "values", {"done": True}) + + run = await create_async_graph_run_stream(_source()) + await asyncio.sleep(0.05) # let pump process first event + + # Get the first subgraph while the mux is still open + sub = None + async for s in run.subgraphs: + sub = s + break + + assert sub is not None + + # Release the gate so the pump finishes + gate.set() + await asyncio.sleep(0.05) # let pump close mux + + # ``await sub.output`` uses the mux's output future — works fine + output = await sub.output + assert output == {"v": 2}, "await sub.output should reflect the latest value" + + # But ``async for v in sub.values`` only gets the replayed event + # and then hangs because the child's values_log is never closed. + values: list[Any] = [] + try: + async with asyncio.timeout(1.0): + async for v in sub.values: + values.append(v) + except (asyncio.TimeoutError, TimeoutError): + pass + + assert len(values) == 2, ( + f"Expected 2 child value snapshots but got {len(values)}: {values}. " + "Child transformer missed post-discovery events." + ) diff --git a/libs/langgraph/tests/test_streaming_comparison.py b/libs/langgraph/tests/test_streaming_comparison.py new file mode 100644 index 000000000..9f93da516 --- /dev/null +++ b/libs/langgraph/tests/test_streaming_comparison.py @@ -0,0 +1,420 @@ +"""Prove V1 and StreamingHandler APIs expose identical information. + +Each test runs the same graph through both APIs and asserts data +equivalence — same state snapshots, same messages, same custom events, +same interrupts. Sync APIs are used where possible; async tests cover +features without sync equivalents (subgraphs projection, messages_from). + +Run with: + TEST=tests/test_streaming_comparison.py make test +""" + +from __future__ import annotations + +from typing import Annotated + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.memory import MemorySaver +from typing_extensions import TypedDict + +from langgraph.config import get_stream_writer +from langgraph.graph import END, START, MessagesState, StateGraph +from langgraph.stream import StreamingHandler +from langgraph.stream._convert import STREAM_V2_MODES +from langgraph.types import interrupt +from tests.fake_chat import FakeChatModel + +# --------------------------------------------------------------------------- +# Graph factories +# --------------------------------------------------------------------------- + + +class State(TypedDict): + value: str + items: Annotated[list[str], lambda a, b: a + b] + + +def _linear_graph(n_nodes: int = 3): + """Chain of *n_nodes* that concatenate strings.""" + g = StateGraph(State) + names = [f"node_{i}" for i in range(n_nodes)] + for name in names: + + def make_fn(n): + def fn(state: State) -> dict: + return {"value": state["value"] + f"_{n}", "items": [n]} + + return fn + + g.add_node(name, make_fn(name)) + + g.add_edge(START, names[0]) + for i in range(len(names) - 1): + g.add_edge(names[i], names[i + 1]) + g.add_edge(names[-1], END) + return g.compile() + + +def _chat_graph(): + """Single agent node with a FakeChatModel.""" + model = FakeChatModel(messages=[AIMessage(content="Hello from agent")]) + + def agent(state: dict) -> dict: + return {"messages": [model.invoke(state["messages"])]} + + g = StateGraph(MessagesState) + g.add_node("agent", agent) + g.add_edge(START, "agent") + g.add_edge("agent", END) + return g.compile() + + +def _multi_node_chat_graph(): + """Two LLM nodes: agent -> reviewer.""" + agent_model = FakeChatModel(messages=[AIMessage(content="Agent reply")]) + reviewer_model = FakeChatModel(messages=[AIMessage(content="Reviewer reply")]) + + def agent(state: dict) -> dict: + return {"messages": [agent_model.invoke(state["messages"])]} + + def reviewer(state: dict) -> dict: + return {"messages": [reviewer_model.invoke(state["messages"])]} + + g = StateGraph(MessagesState) + g.add_node("agent", agent) + g.add_node("reviewer", reviewer) + g.add_edge(START, "agent") + g.add_edge("agent", "reviewer") + g.add_edge("reviewer", END) + return g.compile() + + +def _custom_events_graph(): + """Node that emits custom events via StreamWriter.""" + + def worker(state: State) -> dict: + writer = get_stream_writer() + writer({"step": 1, "msg": "started"}) + writer({"step": 2, "msg": "processing"}) + writer({"step": 3, "msg": "done"}) + return {"value": state["value"] + "_done", "items": ["done"]} + + g = StateGraph(State) + g.add_node("worker", worker) + g.add_edge(START, "worker") + g.add_edge("worker", END) + return g.compile() + + +def _interrupt_graph(): + """Graph that interrupts for human input.""" + + def ask_human(state: State) -> dict: + answer = interrupt("What next?") + return {"value": state["value"] + f"_{answer}", "items": [answer]} + + g = StateGraph(State) + g.add_node("ask", ask_human) + g.add_edge(START, "ask") + g.add_edge("ask", END) + return g.compile(checkpointer=MemorySaver()) + + +def _subgraph(): + """Parent with a compiled child subgraph.""" + + class ChildState(TypedDict): + value: str + + class ParentState(TypedDict): + value: str + + def child_node(state: ChildState) -> dict: + return {"value": state["value"] + "_child"} + + child = StateGraph(ChildState) + child.add_node("inner", child_node) + child.add_edge(START, "inner") + child.add_edge("inner", END) + child_compiled = child.compile() + + parent = StateGraph(ParentState) + parent.add_node("child", child_compiled) + parent.add_edge(START, "child") + parent.add_edge("child", END) + return parent.compile() + + +# =================================================================== +# 1. Final output +# =================================================================== + + +def test_output(): + """graph.invoke() produces the same result as StreamingHandler().stream().output.""" + graph = _linear_graph() + inp = {"value": "x", "items": []} + + v1 = graph.invoke(inp) + + run = StreamingHandler(graph).stream(inp) + v2 = run.output + + assert v1 == v2 + + +# =================================================================== +# 2. Intermediate state snapshots (values mode) +# =================================================================== + + +def test_values(): + """stream(mode='values') snapshots == StreamingHandler().stream().values snapshots.""" + graph = _linear_graph() + inp = {"value": "x", "items": []} + + v1 = list(graph.stream(inp, stream_mode="values")) + + run = StreamingHandler(graph).stream(inp) + v2 = list(run.values) + + assert v1 == v2 + + +# =================================================================== +# 3. Per-node updates (updates mode) +# =================================================================== + + +def test_updates(): + """stream(mode='updates') data == StreamingHandler raw events[method=updates].""" + graph = _linear_graph() + inp = {"value": "x", "items": []} + + v1 = list(graph.stream(inp, stream_mode="updates")) + + run = StreamingHandler(graph).stream(inp) + v2 = [ + e["params"]["data"] + for e in run + if e["method"] == "updates" and not e["params"]["namespace"] + ] + + assert v1 == v2 + + +# =================================================================== +# 4. Message text and node attribution +# =================================================================== + + +def test_messages(): + """Reassembled V1 message text per node == V2 .messages text per node.""" + graph = _multi_node_chat_graph() + inp = {"messages": [HumanMessage(content="hi")]} + + # V1: collect (chunk, metadata) pairs, group text by node + v1_text_by_node: dict[str, list[str]] = {} + for chunk, metadata in graph.stream(inp, stream_mode="messages"): + node = metadata["langgraph_node"] + v1_text_by_node.setdefault(node, []).append(chunk.content) + v1_text = {k: "".join(v) for k, v in v1_text_by_node.items()} + + # V2: each ChatModelStream has .text and .node + run = StreamingHandler(graph).stream(inp) + v2_text: dict[str, str] = {} + for msg in run.messages: + assert msg.done is True + v2_text[msg.node] = msg.text + + assert v1_text == v2_text + + +# =================================================================== +# 5. Custom events +# =================================================================== + + +def test_custom_events(): + """stream(mode='custom') payloads == StreamingHandler raw events[method=custom].""" + graph = _custom_events_graph() + inp = {"value": "x", "items": []} + + v1 = list(graph.stream(inp, stream_mode="custom")) + + run = StreamingHandler(graph).stream(inp) + v2 = [ + e["params"]["data"] + for e in run + if e["method"] == "custom" and not e["params"]["namespace"] + ] + + assert v1 == v2 + + +# =================================================================== +# 6. Mode coverage +# =================================================================== + + +def test_mode_coverage(): + """V2 produces events for the same set of modes as V1.""" + graph = _chat_graph() + inp = {"messages": [HumanMessage(content="hi")]} + + # V1: request all modes, collect which ones appear + v1_modes: set[str] = set() + for ns, mode, _ in graph.stream( + inp, stream_mode=STREAM_V2_MODES, subgraphs=True, version="v1" + ): + if not ns: + v1_modes.add(mode) + + # V2: iterate raw events, collect methods + run = StreamingHandler(graph).stream(inp) + v2_modes = {e["method"] for e in run if not e["params"]["namespace"]} + + assert v1_modes == v2_modes + + +# =================================================================== +# 7. Interrupt detection +# =================================================================== + + +def test_interrupts(): + """V1 __interrupt__ value == V2 .interrupted and .interrupts payload.""" + graph = _interrupt_graph() + inp = {"value": "x", "items": []} + + # V1: detect __interrupt__ in values stream + config1 = {"configurable": {"thread_id": "equiv-1"}} + v1_interrupt_value = None + for chunk in graph.stream(inp, config1, stream_mode="values"): + if isinstance(chunk, dict) and "__interrupt__" in chunk: + info = chunk["__interrupt__"] + if info: + v1_interrupt_value = info[0].value + + assert v1_interrupt_value is not None + + # V2: .interrupted and .interrupts (fresh thread) + config2 = {"configurable": {"thread_id": "equiv-2"}} + run = StreamingHandler(graph).stream(inp, config=config2) + for _ in run: + pass + + assert run.interrupted is True + assert len(run.interrupts) > 0 + v2_interrupt_value = run.interrupts[0]["payload"].value + + assert v1_interrupt_value == v2_interrupt_value + + +# =================================================================== +# 8. Subgraph state snapshots +# =================================================================== + + +def test_subgraph_values(): + """V1 child namespace values == V2 child namespace values.""" + graph = _subgraph() + inp = {"value": "x"} + + # V1: stream with subgraphs=True, collect child values + v1_child_values = [] + for ns, data in graph.stream(inp, stream_mode="values", subgraphs=True): + if ns: + v1_child_values.append(data) + + # V2: filter raw events for child namespace + values mode + run = StreamingHandler(graph).stream(inp) + v2_child_values = [ + e["params"]["data"] + for e in run + if e["method"] == "values" and e["params"]["namespace"] + ] + + assert v1_child_values == v2_child_values + + +# =================================================================== +# 9. Node filtering on messages +# =================================================================== + + +def test_messages_node_filtering(): + """V1 manual metadata filter == V2 .messages filtered by .node.""" + graph = _multi_node_chat_graph() + inp = {"messages": [HumanMessage(content="hi")]} + + # V1: manual filter for "agent" node only + v1_agent_text: list[str] = [] + for chunk, metadata in graph.stream(inp, stream_mode="messages"): + if metadata.get("langgraph_node") == "agent": + v1_agent_text.append(chunk.content) + v1_text = "".join(v1_agent_text) + + # V2: filter .messages by .node + run = StreamingHandler(graph).stream(inp) + v2_agent_msgs = [msg for msg in run.messages if msg.node == "agent"] + assert len(v2_agent_msgs) == 1 + v2_text = v2_agent_msgs[0].text + + assert v1_text == v2_text + + +# =================================================================== +# 10. Async: subgraphs projection +# =================================================================== + + +@pytest.mark.anyio +async def test_async_subgraph_projection(): + """V2 .subgraphs child output matches V1 child namespace output.""" + graph = _subgraph() + inp = {"value": "x"} + + # V1 + v1_child_output = None + async for ns, data in graph.astream(inp, stream_mode="values", subgraphs=True): + if ns: + v1_child_output = data + + # V2: .subgraphs yields typed child stream objects + run = await StreamingHandler(graph).astream(inp) + v2_child_output = None + async for sub in run.subgraphs: + v2_child_output = await sub.output + + assert v1_child_output == v2_child_output + + +# =================================================================== +# 11. Async: messages_from projection +# =================================================================== + + +@pytest.mark.anyio +async def test_async_messages_from(): + """V2 .messages_from('agent') text matches V1 filtered by metadata.""" + graph = _multi_node_chat_graph() + inp = {"messages": [HumanMessage(content="hi")]} + + # V1: manual filter for agent node + v1_agent_text: list[str] = [] + async for chunk, metadata in graph.astream(inp, stream_mode="messages"): + if metadata.get("langgraph_node") == "agent": + v1_agent_text.append(chunk.content) + v1_text = "".join(v1_agent_text) + + # V2: declarative node filtering + run = await StreamingHandler(graph).astream(inp) + v2_texts: list[str] = [] + async for msg in run.messages_from("agent"): + v2_texts.append(await msg.text) + assert len(v2_texts) == 1 + v2_text = v2_texts[0] + + assert v1_text == v2_text