mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 06:35:46 +02:00
Nick Hollon
acaa767542
When `StreamingHandler(graph).stream()` is used, content-block (v2) protocol events now flow through `stream_mode="messages"` for every `model.invoke()` call inside a node — with no node-level code changes. Adds `StreamMessagesHandlerV2`, a `StreamMessagesHandler` subclass that also inherits `_V2StreamingCallbackHandler` from langchain-core. The marker base flips `BaseChatModel.invoke` to drive the protocol event generator (firing `on_stream_event`) instead of `_stream` (firing `on_llm_new_token`). The handler inherits `on_stream_event` from the parent — events forward onto the messages channel unchanged — and overrides `on_llm_new_token` to no-op so a node calling `model.stream()` directly on a v2-flagged run can't leak AIMessageChunks onto the same channel. Opt-in is scoped to `StreamingHandler`: it merges a new internal `CONFIG_KEY_STREAM_MESSAGES_V2=True` into `config.configurable` before dispatching to `graph.stream` / `graph.astream`. Pregel reads the flag at handler-construction time in both sync and async stream paths and attaches the v2 subclass only when set. Direct `graph.stream(stream_mode="messages")` callers keep the v1 `(AIMessageChunk, metadata)` shape — confirmed by a regression test. Existing dedupe between the streamed v2 lifecycle and a node returning the same assembled `AIMessage` transfers for free: the handler populates `self.seen` from `message-start` events (via the inherited `on_stream_event` body), and `on_chain_end`'s `_find_and_emit_messages` already gates on `seen` — so an invoking node surfaces as exactly one `ChatModelStream`, not two. Test coverage in `tests/test_stream_messages_transformer.py`: - `TestEndToEndV2Invoke` — node calling `model.invoke()` produces a single `ChatModelStream` with the full v2 event lifecycle, text projection accumulates correctly, multi-node graphs produce one stream per model call, constructed-message nodes still replay via `message_to_events`, async mirror via `ainvoke` + `astream`. - `TestDirectMessagesModeStaysV1` — regression guard: direct `graph.stream(stream_mode="messages")` still yields AIMessageChunk tuples (not event dicts). - `TestStreamMessagesHandlerV2Unit` — direct unit test that the v2 handler's `on_llm_new_token` does not emit.