refactor(langgraph): scope on_stream_event to StreamMessagesHandlerV2

Moves `on_stream_event` off the v1 `StreamMessagesHandler` base and
onto the v2 subclass. Content-block events are a v2-only concept, so
forwarding them only when the v2 handler is attached keeps the
messages channel's shape predictable for v1 callers: plain
`graph.stream(stream_mode="messages")` now ignores `on_stream_event`
entirely, even if a node explicitly calls `model.stream_v2()` on a
v1-flagged run. Dedupe of the returned AIMessage still works in that
case via `_find_and_emit_messages` / `on_chain_end`.

Also makes the v2 handler's `on_llm_new_token` override an explicit
pass-through with a comment rather than `return None`, so it reads as
an intentional no-op instead of a missing return value.
This commit is contained in:
Nick Hollon
2026-04-17 16:22:50 -04:00
parent acaa767542
commit 0f2f66fc8f
+45 -39
View File
@@ -155,37 +155,6 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
stream_metadata["tags"] = filtered_tags
self.metadata[run_id] = (ns, stream_metadata)
def on_stream_event(
self,
event: dict[str, Any],
*,
run_id: UUID,
parent_run_id: UUID | None = None,
tags: list[str] | None = None,
**kwargs: Any,
) -> Any:
"""Forward a protocol event from `stream_v2` as a messages stream part.
Fires once per `MessagesData` event (`message-start`, per-block
`content-block-*`, `message-finish`). The transformer layer
correlates events back to a single `ChatModelStream` via
`metadata["run_id"]` — attached here so the v1
`stream_mode="messages"` output (which emits
`(AIMessageChunk, metadata)` via `on_llm_new_token`) keeps its
original metadata shape.
"""
if meta := self.metadata.get(run_id):
# Record message_id on message-start so on_chain_end's
# dedupe skips the finalized AIMessage the node returns
# (otherwise the messages projection double-counts: once
# from streaming, once from the chain output).
if event.get("event") == "message-start":
msg_id = event.get("message_id")
if msg_id:
self.seen.add(msg_id)
v2_meta = {**meta[1], "run_id": str(run_id)}
self.stream((meta[0], "messages", (event, v2_meta)))
def on_llm_new_token(
self,
token: str,
@@ -320,13 +289,50 @@ class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler
tags: list[str] | None = None,
**kwargs: Any,
) -> Any:
"""Suppress v1 chunk emission on the messages channel.
"""Intentional no-op — v1 chunks are not used on v2-flagged runs.
The v2 marker already steers `invoke` to the event generator,
so `on_llm_new_token` should not fire under normal routing.
This override guards against any caller (e.g. a node that
calls `model.stream()` directly, which still fires the v1
callback) leaking AIMessageChunks onto a v2-flagged messages
stream.
The v2 marker already steers `invoke` to the event generator, so
`on_llm_new_token` should not fire under normal routing. This
override stays a pass-through (no call to `super()`) to make
the intent explicit and to guard against any caller (e.g. a
node that calls `model.stream()` directly, which still fires
the v1 callback) leaking AIMessageChunks onto a v2-flagged
messages stream.
"""
return None
# Intentionally empty: v2 handler does not forward v1 chunks.
def on_stream_event(
self,
event: dict[str, Any],
*,
run_id: UUID,
parent_run_id: UUID | None = None,
tags: list[str] | None = None,
**kwargs: Any,
) -> Any:
"""Forward a protocol event from `stream_v2` as a messages stream part.
Fires once per `MessagesData` event (`message-start`, per-block
`content-block-*`, `message-finish`). The transformer layer
correlates events back to a single `ChatModelStream` via
`metadata["run_id"]` — attached here so the v1
`stream_mode="messages"` output (which emits
`(AIMessageChunk, metadata)` via `on_llm_new_token`) keeps its
original metadata shape.
Lives on the v2 handler rather than the v1 base: content-block
events are a v2-only concept, and forwarding them only when the
v2 handler is attached keeps the message channel's shape
predictable for v1 callers.
"""
if meta := self.metadata.get(run_id):
# Record message_id on message-start so on_chain_end's
# dedupe skips the finalized AIMessage the node returns
# (otherwise the messages projection double-counts: once
# from streaming, once from the chain output).
if event.get("event") == "message-start":
msg_id = event.get("message_id")
if msg_id:
self.seen.add(msg_id)
v2_meta = {**meta[1], "run_id": str(run_id)}
self.stream((meta[0], "messages", (event, v2_meta)))