diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index dedfd9ab7..1d168e88d 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -60,6 +60,11 @@ class GraphRunStream: Sync iteration is caller-driven, so a cursor that catches up to the buffer's tail needs a way to ask the graph for more events. + + Also calls `_bind_pump` on any transformer that exposes it, so + that transformers producing ChatModelStream objects (e.g. + MessagesTransformer) can wire the pull callback on each stream as + it's created. """ mux._events._request_more = self._pump_next for value in mux.extensions.values(): @@ -67,6 +72,9 @@ class GraphRunStream: value._request_more = self._pump_next elif isinstance(value, StreamChannel): value._log._request_more = self._pump_next + for transformer in mux._transformers: + if hasattr(transformer, "_bind_pump"): + transformer._bind_pump(self._pump_next) def _pump_next(self) -> bool: """Pull one event from the graph and push it through the mux. diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index 3fe7c5ea9..57909247c 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -1,10 +1,21 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any, cast + +from langchain_core.language_models._compat_bridge import message_to_events +from langchain_core.language_models.chat_model_stream import ( + AsyncChatModelStream, + ChatModelStream, +) +from langchain_core.messages import AIMessageChunk, BaseMessage +from langchain_protocol.protocol import MessagesData from langgraph.stream._event_log import EventLog from langgraph.stream._types import ProtocolEvent, StreamTransformer +if TYPE_CHECKING: + from collections.abc import Callable + class ValuesTransformer(StreamTransformer): """Capture values events as an iterable of state snapshots. @@ -51,34 +62,142 @@ class ValuesTransformer(StreamTransformer): class MessagesTransformer(StreamTransformer): - """Pass through raw (chunk, metadata) tuples from messages events. + """Capture messages events as ChatModelStream objects. - This is the same shape as today's `stream_mode="messages"` output. - A follow-on PR will replace this with a richer transformer that - produces ChatModelStream objects using the protocol handler. + The messages projection yields one `ChatModelStream` (or + `AsyncChatModelStream`) per LLM call. Consumers iterate + `run.messages` to get stream handles, then use each handle's typed + projections (`.text`, `.reasoning`, `.tool_calls`, `.usage`, + `.output`) for per-message content. - Only root-namespace messages events are captured; tokens emitted - from subgraphs are dropped from the `messages` projection. Consumers - that need subgraph tokens should iterate the raw event stream or - register a custom transformer. + Two input shapes are handled (via `params["data"] = (payload, + metadata)` from `StreamMessagesHandler`): - Native transformer — projection keys are exposed as direct - attributes on the run stream (e.g. `run.messages`). + 1. Protocol event (dict with `"event"` key) — emitted by + `stream_v2()` / `astream_v2()` via the `on_stream_event` + callback. Routed to an existing `ChatModelStream` by + `metadata["run_id"]`. A `message-start` event creates a new + stream; `message-finish` closes it. + 2. Whole `AIMessage` — emitted from `on_chain_end` when a node + returns a finalized message. Replayed as a synthetic protocol + event lifecycle via `message_to_events`, then the + already-complete stream is pushed to the log. + + V1 `AIMessageChunk` tuples (from `on_llm_new_token`) are not + streamed into this projection: chat models that want to populate + `run.messages` with content-block streaming must use + `stream_v2()` / `astream_v2()`. Models called via the legacy + `stream()` method still surface their final `AIMessage` via + `on_chain_end` when a node returns it as state. + + Only root-namespace events are captured; tokens from subgraphs are + dropped. Consumers that need subgraph tokens should iterate the raw + event stream or register a custom transformer. + + Native transformer — the `messages` projection is exposed as a + direct attribute on the run stream. """ _native = True def __init__(self) -> None: - self._log: EventLog[tuple[Any, dict[str, Any]]] = EventLog() + self._log: EventLog[ChatModelStream] = EventLog() + # Correlate protocol events back to a ChatModelStream by run_id + # (attached to the event's metadata by StreamMessagesHandler). + self._by_run: dict[str, ChatModelStream] = {} + self._pump_fn: Callable[[], bool] | None = None def init(self) -> dict[str, Any]: return {"messages": self._log} + def _bind_pump(self, fn: Callable[[], bool]) -> None: + """Wire the sync pull callback. Called by GraphRunStream._wire_request_more.""" + self._pump_fn = fn + + def _make_stream( + self, + *, + namespace: list[str], + node: str | None, + message_id: str | None, + ) -> ChatModelStream: + """Create a ChatModelStream (sync) or AsyncChatModelStream (async).""" + if self._pump_fn is not None: + stream: ChatModelStream = ChatModelStream( + namespace=namespace, + node=node, + message_id=message_id, + ) + stream.set_request_more(self._pump_fn) + else: + stream = AsyncChatModelStream( + namespace=namespace, + node=node, + message_id=message_id, + ) + return stream + def process(self, event: ProtocolEvent) -> bool: if event["method"] != "messages": return True params = event["params"] if params["namespace"]: return True - self._log.push(params["data"]) + + payload, metadata = params["data"] + node: str | None = metadata.get("langgraph_node") + run_id = str(metadata.get("run_id", "")) if metadata else "" + + if isinstance(payload, dict) and "event" in payload: + self._route_protocol_event( + cast("MessagesData", payload), run_id=run_id, node=node + ) + elif isinstance(payload, BaseMessage) and not isinstance( + payload, AIMessageChunk + ): + self._route_whole_message(payload, node=node) + # Legacy AIMessageChunk tuples (from on_llm_new_token) are ignored; + # v1 streaming callers must switch to stream_v2() to populate this + # projection. + return True + + def _route_protocol_event( + self, + event: MessagesData, + *, + run_id: str, + node: str | None, + ) -> None: + event_type = event.get("event") + if event_type == "message-start": + message_id = event.get("message_id") + stream = self._make_stream( + namespace=[], + node=node, + message_id=str(message_id) if message_id is not None else None, + ) + self._by_run[run_id] = stream + self._log.push(stream) + stream.dispatch(event) + elif run_id in self._by_run: + stream = self._by_run[run_id] + stream.dispatch(event) + if event_type == "message-finish": + del self._by_run[run_id] + + def _route_whole_message(self, message: BaseMessage, *, node: str | None) -> None: + stream = self._make_stream(namespace=[], node=node, message_id=message.id) + for evt in message_to_events(message, message_id=message.id): + stream.dispatch(evt) + self._log.push(stream) + + def finalize(self) -> None: + """Clear any routing state — streams close themselves via `message-finish`.""" + self._by_run.clear() + + def fail(self, err: BaseException) -> None: + """Propagate run error to any streams still open when the graph fails.""" + for stream in list(self._by_run.values()): + stream.fail(err) + self._by_run.clear() diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_streaming_handler.py index d4797e4c1..261881e20 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_streaming_handler.py @@ -807,22 +807,40 @@ class TestValuesTransformer: class TestMessagesTransformer: def test_captures_root_messages(self) -> None: + """Protocol-event lifecycle produces a ChatModelStream in the log.""" t = MessagesTransformer() t.init() t._log._bind(is_async=False) + t._bind_pump(lambda: False) - t.process(_event("messages", ("chunk", {"meta": True}))) + meta = {"langgraph_node": "llm", "run_id": "run-1"} + for evt in ( + {"event": "message-start", "role": "ai", "message_id": "run-1"}, + {"event": "message-finish", "reason": "stop"}, + ): + t.process(_event("messages", (evt, meta))) t._log.close() items = list(t._log) assert len(items) == 1 - assert items[0] == ("chunk", {"meta": True}) + # Items in the messages log are ChatModelStream objects, not raw + # tuples — the content-block-centric projection. + assert hasattr(items[0], "dispatch") + assert items[0].message_id == "run-1" def test_ignores_non_root_namespace(self) -> None: t = MessagesTransformer() t.init() t._log._bind(is_async=False) + t._bind_pump(lambda: False) - t.process(_event("messages", ("chunk", {}), namespace=["sub"])) + meta = {"langgraph_node": "llm", "run_id": "run-1"} + t.process( + _event( + "messages", + ({"event": "message-start", "message_id": "run-1"}, meta), + namespace=["sub"], + ) + ) t._log.close() assert list(t._log) == []