379 lines
13 KiB
Python
379 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
|
|
from dataclasses import fields, is_dataclass
|
|
from typing import (
|
|
Any,
|
|
TypeVar,
|
|
cast,
|
|
)
|
|
from uuid import UUID, uuid4
|
|
|
|
from langchain_core.callbacks import BaseCallbackHandler
|
|
from langchain_core.messages import BaseMessage
|
|
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
|
|
from pydantic import BaseModel
|
|
|
|
from langgraph._internal._constants import NS_SEP
|
|
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
|
|
from langgraph.pregel.protocol import StreamChunk
|
|
from langgraph.types import Command
|
|
|
|
try:
|
|
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
|
except ImportError:
|
|
_StreamingCallbackHandler = object # type: ignore
|
|
|
|
try:
|
|
from langchain_core.tracers._streaming import _V2StreamingCallbackHandler
|
|
except ImportError:
|
|
_V2StreamingCallbackHandler = object # type: ignore
|
|
|
|
T = TypeVar("T")
|
|
Meta = tuple[tuple[str, ...], dict[str, Any]]
|
|
|
|
|
|
def _state_values(obj: Any) -> Sequence[Any]:
|
|
"""Extract top-level field values from a state object (dict, BaseModel, or dataclass)."""
|
|
if isinstance(obj, dict):
|
|
return list(obj.values())
|
|
elif isinstance(obj, BaseModel):
|
|
return [getattr(obj, k) for k in type(obj).model_fields]
|
|
elif is_dataclass(obj) and not isinstance(obj, type):
|
|
return [getattr(obj, f.name) for f in fields(obj)]
|
|
return ()
|
|
|
|
|
|
class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
|
"""A callback handler that implements stream_mode=messages.
|
|
|
|
Collects messages from:
|
|
(1) chat model stream events; and
|
|
(2) node outputs.
|
|
"""
|
|
|
|
run_inline = True
|
|
"""We want this callback to run in the main thread to avoid order/locking issues."""
|
|
|
|
def __init__(
|
|
self,
|
|
stream: Callable[[StreamChunk], None],
|
|
subgraphs: bool,
|
|
*,
|
|
parent_ns: tuple[str, ...] | None = None,
|
|
) -> None:
|
|
"""Configure the handler to stream messages from LLMs and nodes.
|
|
|
|
Args:
|
|
stream: A callable that takes a StreamChunk and emits it.
|
|
subgraphs: Whether to emit messages from subgraphs.
|
|
parent_ns: The namespace where the handler was created.
|
|
We keep track of this namespace to allow calls to subgraphs that
|
|
were explicitly requested as a stream with `messages` mode
|
|
configured.
|
|
|
|
Example:
|
|
parent_ns is used to handle scenarios where the subgraph is explicitly
|
|
streamed with `stream_mode="messages"`.
|
|
|
|
```python
|
|
def parent_graph_node():
|
|
# This node is in the parent graph.
|
|
async for event in some_subgraph(..., stream_mode="messages"):
|
|
do something with event # <-- these events will be emitted
|
|
return ...
|
|
|
|
parent_graph.invoke(subgraphs=False)
|
|
```
|
|
"""
|
|
self.stream = stream
|
|
self.subgraphs = subgraphs
|
|
self.metadata: dict[UUID, Meta] = {}
|
|
self.seen: set[int | str] = set()
|
|
self.parent_ns = parent_ns
|
|
|
|
def _emit(self, meta: Meta, message: BaseMessage, *, dedupe: bool = False) -> None:
|
|
if dedupe and message.id in self.seen:
|
|
return
|
|
else:
|
|
if message.id is None:
|
|
message.id = str(uuid4())
|
|
self.seen.add(message.id)
|
|
self.stream((meta[0], "messages", (message, meta[1])))
|
|
|
|
def _find_and_emit_messages(self, meta: Meta, response: Any) -> None:
|
|
if isinstance(response, BaseMessage):
|
|
self._emit(meta, response, dedupe=True)
|
|
elif isinstance(response, Sequence):
|
|
for value in response:
|
|
if isinstance(value, BaseMessage):
|
|
self._emit(meta, value, dedupe=True)
|
|
else:
|
|
for value in _state_values(response):
|
|
if isinstance(value, BaseMessage):
|
|
self._emit(meta, value, dedupe=True)
|
|
elif isinstance(value, Sequence):
|
|
for item in value:
|
|
if isinstance(item, BaseMessage):
|
|
self._emit(meta, item, dedupe=True)
|
|
|
|
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
|
|
|
|
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_tags := [t for t in tags if not t.startswith("seq:step")]:
|
|
metadata["tags"] = filtered_tags
|
|
self.metadata[run_id] = (ns, metadata)
|
|
|
|
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
|
|
if meta := self.metadata.get(run_id):
|
|
self._emit(meta, chunk.message)
|
|
|
|
def on_llm_end(
|
|
self,
|
|
response: LLMResult,
|
|
*,
|
|
run_id: UUID,
|
|
parent_run_id: UUID | None = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
if meta := self.metadata.get(run_id):
|
|
if response.generations and response.generations[0]:
|
|
gen = response.generations[0][0]
|
|
if isinstance(gen, ChatGeneration):
|
|
self._emit(meta, gen.message, dedupe=True)
|
|
self.metadata.pop(run_id, None)
|
|
|
|
def on_llm_error(
|
|
self,
|
|
error: BaseException,
|
|
*,
|
|
run_id: UUID,
|
|
parent_run_id: UUID | None = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.metadata.pop(run_id, None)
|
|
|
|
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)
|
|
for value in _state_values(inputs):
|
|
if isinstance(value, BaseMessage):
|
|
if value.id is not None:
|
|
self.seen.add(value.id)
|
|
elif isinstance(value, Sequence) and not isinstance(value, str):
|
|
for item in value:
|
|
if isinstance(item, BaseMessage):
|
|
if item.id is not None:
|
|
self.seen.add(item.id)
|
|
|
|
def on_chain_end(
|
|
self,
|
|
response: Any,
|
|
*,
|
|
run_id: UUID,
|
|
parent_run_id: UUID | None = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
if meta := self.metadata.pop(run_id, None):
|
|
# Handle Command node updates
|
|
if isinstance(response, Command):
|
|
self._find_and_emit_messages(meta, response.update)
|
|
# Handle list of Command updates
|
|
elif isinstance(response, Sequence) and any(
|
|
isinstance(value, Command) for value in response
|
|
):
|
|
for value in response:
|
|
if isinstance(value, Command):
|
|
self._find_and_emit_messages(meta, value.update)
|
|
else:
|
|
self._find_and_emit_messages(meta, value)
|
|
# Handle basic updates / streaming
|
|
else:
|
|
self._find_and_emit_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)
|
|
|
|
|
|
class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler):
|
|
"""v2 variant of `StreamMessagesHandler`.
|
|
|
|
Declaring `_V2StreamingCallbackHandler` as a base flips
|
|
`BaseChatModel.invoke` to route through `_stream_chat_model_events`
|
|
(firing `on_stream_event`) instead of `_stream` (firing
|
|
`on_llm_new_token`). Inherits `on_stream_event` from the parent,
|
|
which forwards protocol events onto the messages stream channel.
|
|
|
|
Pregel attaches this class instead of the v1 handler only when
|
|
`StreamingHandler` opts in via the internal
|
|
`CONFIG_KEY_STREAM_MESSAGES_V2` config key; direct
|
|
`graph.stream(stream_mode="messages")` callers keep the v1
|
|
AIMessageChunk shape.
|
|
"""
|
|
|
|
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:
|
|
"""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 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.
|
|
"""
|
|
# Intentionally empty: v2 handler does not forward v1 chunks.
|
|
|
|
def __init__(
|
|
self,
|
|
stream: Callable[[StreamChunk], None],
|
|
subgraphs: bool,
|
|
*,
|
|
parent_ns: tuple[str, ...] | None = None,
|
|
) -> None:
|
|
super().__init__(stream, subgraphs, parent_ns=parent_ns)
|
|
self._streamed_run_ids: set[UUID] = set()
|
|
|
|
def on_llm_end(
|
|
self,
|
|
response: LLMResult,
|
|
*,
|
|
run_id: UUID,
|
|
parent_run_id: UUID | None = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
if meta := self.metadata.get(run_id):
|
|
if response.generations and response.generations[0]:
|
|
gen = response.generations[0][0]
|
|
if isinstance(gen, ChatGeneration):
|
|
if run_id in self._streamed_run_ids:
|
|
if gen.message.id is None:
|
|
gen.message.id = str(uuid4())
|
|
self.seen.add(gen.message.id)
|
|
else:
|
|
self._emit(meta, gen.message, dedupe=True)
|
|
self._streamed_run_ids.discard(run_id)
|
|
self.metadata.pop(run_id, None)
|
|
|
|
def on_llm_error(
|
|
self,
|
|
error: BaseException,
|
|
*,
|
|
run_id: UUID,
|
|
parent_run_id: UUID | None = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self._streamed_run_ids.discard(run_id)
|
|
super().on_llm_error(
|
|
error,
|
|
run_id=run_id,
|
|
parent_run_id=parent_run_id,
|
|
**kwargs,
|
|
)
|
|
|
|
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_events(version="v3")` 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":
|
|
self._streamed_run_ids.add(run_id)
|
|
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)))
|