mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 13:17:52 +02:00
feat(langgraph): add streaming transformer infrastructure and tests (#7519)
This commit is contained in:
@@ -66,6 +66,9 @@ CONFIG_KEY_RUNTIME = sys.intern("__pregel_runtime")
|
||||
# holds a `Runtime` instance with context, store, stream writer, etc.
|
||||
CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map")
|
||||
# holds a mapping of task ns -> resume value for resuming tasks
|
||||
CONFIG_KEY_STREAM_MESSAGES_V2 = sys.intern("__pregel_stream_messages_v2")
|
||||
# when True, attach StreamMessagesHandlerV2 so content-block (v2) events
|
||||
# flow through stream_mode="messages"; set by StreamingHandler only.
|
||||
|
||||
# --- Other constants ---
|
||||
PUSH = sys.intern("__pregel_push")
|
||||
@@ -107,6 +110,7 @@ RESERVED = {
|
||||
CONFIG_KEY_CHECKPOINT_ID,
|
||||
CONFIG_KEY_CHECKPOINT_NS,
|
||||
CONFIG_KEY_RESUME_MAP,
|
||||
CONFIG_KEY_STREAM_MESSAGES_V2,
|
||||
# other constants
|
||||
PUSH,
|
||||
PULL,
|
||||
|
||||
@@ -1045,6 +1045,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
||||
interrupt_after: All | list[str] | None = None,
|
||||
debug: bool = False,
|
||||
name: str | None = None,
|
||||
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||||
) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]:
|
||||
"""Compiles the `StateGraph` into a `CompiledStateGraph` object.
|
||||
|
||||
@@ -1077,6 +1078,13 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
||||
interrupt_after: An optional list of node names to interrupt after.
|
||||
debug: A flag indicating whether to enable debug mode.
|
||||
name: The name to use for the compiled graph.
|
||||
transformers: Optional sequence of `StreamTransformer` classes or
|
||||
configured factories. Classes and factories are instantiated
|
||||
per run whenever `stream_v2` / `astream_v2` is called and are
|
||||
propagated to subgraph scopes. Custom factories should follow
|
||||
the standard `StreamTransformer` constructor shape by
|
||||
accepting `scope` as their first argument. Appended after the
|
||||
built-in stream transformers.
|
||||
|
||||
Returns:
|
||||
CompiledStateGraph: The compiled `StateGraph`.
|
||||
@@ -1159,6 +1167,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
||||
store=store,
|
||||
cache=cache,
|
||||
name=name or "LangGraph",
|
||||
stream_transformers=transformers,
|
||||
)
|
||||
compiled._serde_allowlist = serde_allowlist
|
||||
|
||||
|
||||
@@ -24,6 +24,11 @@ try:
|
||||
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]]
|
||||
|
||||
@@ -248,3 +253,126 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
**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_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":
|
||||
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)))
|
||||
|
||||
@@ -0,0 +1,268 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Callable, Iterator
|
||||
from contextvars import ContextVar, Token
|
||||
from typing import Any, TypeVar, cast
|
||||
from uuid import UUID
|
||||
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph.constants import TAG_NOSTREAM
|
||||
from langgraph.pregel.protocol import StreamChunk
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = object # type: ignore[assignment,misc]
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
ToolCallWriter = Callable[[Any], None]
|
||||
"""A closure bound to a single tool call that emits `tool-output-delta` events."""
|
||||
|
||||
_tool_call_writer: ContextVar[ToolCallWriter | None] = ContextVar(
|
||||
"langgraph_tool_call_writer", default=None
|
||||
)
|
||||
"""ContextVar holding the writer for the currently-executing tool call.
|
||||
|
||||
Set by `StreamToolCallHandler.on_tool_start` and reset on end/error.
|
||||
Read by `ToolRuntime.emit_output_delta` (in `langgraph.prebuilt`).
|
||||
"""
|
||||
|
||||
|
||||
class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
"""Callback handler that emits tool-call lifecycle events on the stream.
|
||||
|
||||
Fires on LangChain's `on_tool_*` callbacks and pushes to the `tools`
|
||||
stream mode. Emits `tool-started` / `tool-output-delta` /
|
||||
`tool-finished` / `tool-error` payloads keyed by `tool_call_id`.
|
||||
|
||||
While a tool is executing, this handler sets `_tool_call_writer` to a
|
||||
closure bound to that call's namespace and `tool_call_id`.
|
||||
`ToolRuntime.emit_output_delta` reads that ContextVar so tool bodies
|
||||
can stream partial output without threading the writer through their
|
||||
own signature.
|
||||
|
||||
Attached by `Pregel.stream` / `astream` when `"tools"` is in
|
||||
`stream_modes`. `run_inline = True` keeps event ordering
|
||||
deterministic.
|
||||
"""
|
||||
|
||||
run_inline = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stream: Callable[[StreamChunk], None],
|
||||
subgraphs: bool,
|
||||
*,
|
||||
parent_ns: tuple[str, ...] | None = None,
|
||||
) -> None:
|
||||
"""Configure the handler to stream tool-call events.
|
||||
|
||||
Args:
|
||||
stream: Callable that accepts a `StreamChunk` tuple
|
||||
`(namespace, mode, payload)` and enqueues it.
|
||||
subgraphs: Whether to emit events from tools called inside
|
||||
nested subgraphs. When False, only tools at the
|
||||
handler's own scope (`parent_ns`) emit.
|
||||
parent_ns: Namespace where the handler was attached.
|
||||
Mirrors the `StreamMessagesHandler` escape hatch:
|
||||
tools whose containing namespace equals `parent_ns`
|
||||
still emit even with `subgraphs=False`, so a node that
|
||||
explicitly streams a subgraph with `stream_mode="tools"`
|
||||
sees its own tools.
|
||||
"""
|
||||
self.stream = stream
|
||||
self.subgraphs = subgraphs
|
||||
self.parent_ns = parent_ns
|
||||
# run_id → (namespace, tool_call_id, ContextVar token)
|
||||
# `on_tool_end` does not receive `tool_call_id` in kwargs, so
|
||||
# we correlate by `run_id` which is present on every callback.
|
||||
self._run_to_call: dict[
|
||||
UUID, tuple[tuple[str, ...], str, Token[ToolCallWriter | None]]
|
||||
] = {}
|
||||
|
||||
def _ns_for_emit(
|
||||
self,
|
||||
metadata: dict[str, Any] | None,
|
||||
tags: list[str] | None,
|
||||
) -> tuple[str, ...] | None:
|
||||
"""Resolve the namespace this tool call should emit at, or `None` to skip.
|
||||
|
||||
Mirrors `StreamMessagesHandler.on_chat_model_start`'s namespace
|
||||
derivation: parses `langgraph_checkpoint_ns` (which ends with
|
||||
the `node_name:task_id` of the calling node), drops that
|
||||
trailing segment, and returns the containing subgraph's own
|
||||
namespace. Returns `None` when the call should be silently
|
||||
suppressed:
|
||||
|
||||
- `metadata` is missing — handler is attached to a context
|
||||
without Pregel routing info.
|
||||
- `TAG_NOSTREAM` is in `tags` — caller explicitly opted out.
|
||||
- Tool runs in a subgraph (`len(ns) > 0`) and the handler was
|
||||
attached with `subgraphs=False` and a different `parent_ns`
|
||||
than the call's containing subgraph.
|
||||
"""
|
||||
if not metadata:
|
||||
return None
|
||||
if tags and TAG_NOSTREAM in tags:
|
||||
return None
|
||||
nskey = metadata.get("langgraph_checkpoint_ns")
|
||||
if not nskey:
|
||||
ns: tuple[str, ...] = ()
|
||||
else:
|
||||
ns = tuple(cast(str, nskey).split(NS_SEP))[:-1]
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return None
|
||||
return ns
|
||||
|
||||
def _start(
|
||||
self,
|
||||
serialized: dict[str, Any] | None,
|
||||
input_str: str,
|
||||
*,
|
||||
run_id: UUID,
|
||||
metadata: dict[str, Any] | None,
|
||||
tags: list[str] | None,
|
||||
inputs: dict[str, Any] | None,
|
||||
kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
ns = self._ns_for_emit(metadata, tags)
|
||||
if ns is None:
|
||||
return
|
||||
tool_call_id = cast("str | None", kwargs.get("tool_call_id")) or str(run_id)
|
||||
tool_name = (
|
||||
(serialized or {}).get("name")
|
||||
or cast("str | None", kwargs.get("name"))
|
||||
or ""
|
||||
)
|
||||
|
||||
def writer(delta: Any) -> None:
|
||||
self.stream(
|
||||
(
|
||||
ns,
|
||||
"tools",
|
||||
{
|
||||
"event": "tool-output-delta",
|
||||
"tool_call_id": tool_call_id,
|
||||
"delta": delta,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
token = _tool_call_writer.set(writer)
|
||||
self._run_to_call[run_id] = (ns, tool_call_id, token)
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"event": "tool-started",
|
||||
"tool_call_id": tool_call_id,
|
||||
"tool_name": tool_name,
|
||||
}
|
||||
if inputs is not None:
|
||||
payload["input"] = inputs
|
||||
self.stream((ns, "tools", payload))
|
||||
|
||||
def _end(self, output: Any, *, run_id: UUID) -> None:
|
||||
info = self._run_to_call.pop(run_id, None)
|
||||
if info is None:
|
||||
return
|
||||
ns, tool_call_id, token = info
|
||||
self._reset_writer(token)
|
||||
self.stream(
|
||||
(
|
||||
ns,
|
||||
"tools",
|
||||
{
|
||||
"event": "tool-finished",
|
||||
"tool_call_id": tool_call_id,
|
||||
"output": output,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
def _error(self, error: BaseException, *, run_id: UUID) -> None:
|
||||
info = self._run_to_call.pop(run_id, None)
|
||||
if info is None:
|
||||
return
|
||||
ns, tool_call_id, token = info
|
||||
self._reset_writer(token)
|
||||
self.stream(
|
||||
(
|
||||
ns,
|
||||
"tools",
|
||||
{
|
||||
"event": "tool-error",
|
||||
"tool_call_id": tool_call_id,
|
||||
"message": str(error),
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
def tap_output_aiter(
|
||||
self, run_id: UUID, output: AsyncIterator[T]
|
||||
) -> AsyncIterator[T]:
|
||||
"""Pass-through — required by the `_StreamingCallbackHandler` protocol."""
|
||||
return output
|
||||
|
||||
def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]:
|
||||
"""Pass-through — sync counterpart to `tap_output_aiter`."""
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def _reset_writer(token: Token[ToolCallWriter | None]) -> None:
|
||||
# Token is invalid if `on_tool_end` runs in a different context
|
||||
# than `on_tool_start` (e.g. langchain may hand off to a thread
|
||||
# worker without copying the context). Swallow that case; the
|
||||
# ContextVar lifetime is bounded by the enclosing task anyway.
|
||||
try:
|
||||
_tool_call_writer.reset(token)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Sync callbacks
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def on_tool_start(
|
||||
self,
|
||||
serialized: dict[str, Any],
|
||||
input_str: str,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
inputs: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
self._start(
|
||||
serialized,
|
||||
input_str,
|
||||
run_id=run_id,
|
||||
metadata=metadata,
|
||||
tags=tags,
|
||||
inputs=inputs,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
def on_tool_end(
|
||||
self,
|
||||
output: Any,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
self._end(output, run_id=run_id)
|
||||
|
||||
def on_tool_error(
|
||||
self,
|
||||
error: BaseException,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
self._error(error, run_id=run_id)
|
||||
@@ -73,6 +73,7 @@ from langgraph._internal._constants import (
|
||||
CONFIG_KEY_RUNTIME,
|
||||
CONFIG_KEY_SEND,
|
||||
CONFIG_KEY_STREAM,
|
||||
CONFIG_KEY_STREAM_MESSAGES_V2,
|
||||
CONFIG_KEY_TASK_ID,
|
||||
CONFIG_KEY_THREAD_ID,
|
||||
ERROR,
|
||||
@@ -133,10 +134,14 @@ from langgraph.pregel._loop import (
|
||||
AsyncPregelLoop,
|
||||
SyncPregelLoop,
|
||||
)
|
||||
from langgraph.pregel._messages import StreamMessagesHandler
|
||||
from langgraph.pregel._messages import (
|
||||
StreamMessagesHandler,
|
||||
StreamMessagesHandlerV2,
|
||||
)
|
||||
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
|
||||
from langgraph.pregel._retry import RetryPolicy
|
||||
from langgraph.pregel._runner import PregelRunner
|
||||
from langgraph.pregel._tools import StreamToolCallHandler
|
||||
from langgraph.pregel._utils import get_new_channel_versions
|
||||
from langgraph.pregel._validate import validate_graph, validate_keys
|
||||
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
|
||||
@@ -148,6 +153,15 @@ from langgraph.runtime import (
|
||||
Runtime,
|
||||
ServerInfo,
|
||||
)
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import StreamTransformer
|
||||
from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream
|
||||
from langgraph.stream.transformers import (
|
||||
LifecycleTransformer,
|
||||
MessagesTransformer,
|
||||
SubgraphTransformer,
|
||||
ValuesTransformer,
|
||||
)
|
||||
from langgraph.types import (
|
||||
All,
|
||||
CachePolicy,
|
||||
@@ -340,6 +354,58 @@ class NodeBuilder:
|
||||
)
|
||||
|
||||
|
||||
def _collect_stream_modes(mux: Any) -> list[StreamMode]:
|
||||
"""Return the union of `required_stream_modes` across registered transformers.
|
||||
|
||||
Transformers declare the stream modes they need to function, and
|
||||
`stream_v2` asks the graph for exactly that union — no hardcoded
|
||||
default set. If zero transformers declare a given mode, the graph
|
||||
does not stream events for it.
|
||||
"""
|
||||
modes: set[StreamMode] = set()
|
||||
for transformer in mux._transformers:
|
||||
modes.update(
|
||||
cast(
|
||||
"tuple[StreamMode, ...]",
|
||||
getattr(transformer, "required_stream_modes", ()),
|
||||
)
|
||||
)
|
||||
return list(modes)
|
||||
|
||||
|
||||
def _normalize_stream_transformer_factories(
|
||||
specs: Sequence[Callable[[tuple[str, ...]], Any]] | None,
|
||||
) -> list[Callable[[tuple[str, ...]], Any]]:
|
||||
"""Normalize stream transformer specs to scoped factories.
|
||||
|
||||
A stream transformer spec is a callable that accepts
|
||||
`scope: tuple[str, ...]` and returns a fresh `StreamTransformer`.
|
||||
Transformer classes work when their constructor follows the same
|
||||
shape. Pre-built instances are rejected because they cannot be
|
||||
cloned into subgraph scopes.
|
||||
"""
|
||||
factories: list[Callable[[tuple[str, ...]], Any]] = []
|
||||
for spec in specs or ():
|
||||
if isinstance(spec, StreamTransformer):
|
||||
raise TypeError(
|
||||
"stream_v2 transformers must be scope-aware callables, "
|
||||
f"got pre-built instance {type(spec).__name__}. Pass the "
|
||||
"transformer class or a factory like "
|
||||
"`lambda scope: MyTransformer(scope, ...)`."
|
||||
)
|
||||
if not callable(spec):
|
||||
raise TypeError(
|
||||
"stream_v2 transformers must be scope-aware callables, "
|
||||
f"got {type(spec).__name__}."
|
||||
)
|
||||
|
||||
def factory(scope: tuple[str, ...], _spec: Callable[..., Any] = spec) -> Any:
|
||||
return _spec(scope)
|
||||
|
||||
factories.append(factory)
|
||||
return factories
|
||||
|
||||
|
||||
class Pregel(
|
||||
PregelProtocol[StateT, ContextT, InputT, OutputT],
|
||||
Generic[StateT, ContextT, InputT, OutputT],
|
||||
@@ -671,6 +737,7 @@ class Pregel(
|
||||
config: RunnableConfig | None = None,
|
||||
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
|
||||
name: str = "LangGraph",
|
||||
stream_transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||||
**deprecated_kwargs: Unpack[DeprecatedKwargs],
|
||||
) -> None:
|
||||
if (
|
||||
@@ -717,6 +784,9 @@ class Pregel(
|
||||
self.config = config
|
||||
self.trigger_to_nodes = trigger_to_nodes or {}
|
||||
self.name = name
|
||||
self.stream_transformers: tuple[Callable[[tuple[str, ...]], Any], ...] = tuple(
|
||||
stream_transformers or ()
|
||||
)
|
||||
self._serde_allowlist: set[tuple[str, ...]] | None = None
|
||||
if auto_validate:
|
||||
self.validate()
|
||||
@@ -2582,19 +2652,7 @@ class Pregel(
|
||||
stream = SyncQueue()
|
||||
|
||||
config = ensure_config(self.config, config)
|
||||
callback_manager = get_callback_manager_for_config(config)
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
run_manager = callback_manager.on_chain_start(
|
||||
None,
|
||||
input,
|
||||
name=config.get("run_name", self.get_name()),
|
||||
run_id=config.get("run_id"),
|
||||
)
|
||||
graph_callback_manager = get_sync_graph_callback_manager_for_config(
|
||||
config,
|
||||
run_id=run_manager.run_id,
|
||||
)
|
||||
run_manager = None
|
||||
try:
|
||||
# assign defaults
|
||||
(
|
||||
@@ -2615,6 +2673,36 @@ class Pregel(
|
||||
interrupt_after=interrupt_after,
|
||||
durability=durability,
|
||||
)
|
||||
callback_manager = get_callback_manager_for_config(config)
|
||||
if "messages" in stream_modes and version != "v2":
|
||||
# Strip any inherited v2 messages handler so a v1 stream
|
||||
# does not get routed through the content-block event
|
||||
# protocol. Leave v1 handlers in place — an outer
|
||||
# stream(stream_mode="messages", subgraphs=True) relies
|
||||
# on its inheritable handler to observe events emitted
|
||||
# by inner stream(stream_mode="messages") calls.
|
||||
callback_manager.handlers = [
|
||||
h
|
||||
for h in callback_manager.handlers
|
||||
if not isinstance(h, StreamMessagesHandlerV2)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not isinstance(h, StreamMessagesHandlerV2)
|
||||
]
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
run_manager = callback_manager.on_chain_start(
|
||||
None,
|
||||
input,
|
||||
name=config.get("run_name", self.get_name()),
|
||||
run_id=config.get("run_id"),
|
||||
)
|
||||
graph_callback_manager = get_sync_graph_callback_manager_for_config(
|
||||
config,
|
||||
run_id=run_manager.run_id,
|
||||
)
|
||||
if checkpointer is None and durability is not None:
|
||||
warnings.warn(
|
||||
"`durability` has no effect when no checkpointer is present.",
|
||||
@@ -2626,14 +2714,33 @@ class Pregel(
|
||||
# set up messages stream mode
|
||||
if "messages" in stream_modes:
|
||||
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
use_stream_messages_v2 = bool(
|
||||
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||||
)
|
||||
messages_handler_cls = (
|
||||
StreamMessagesHandlerV2
|
||||
if use_stream_messages_v2
|
||||
else StreamMessagesHandler
|
||||
)
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamMessagesHandler(
|
||||
messages_handler_cls(
|
||||
stream.put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||
)
|
||||
)
|
||||
|
||||
# set up tools stream mode
|
||||
if "tools" in stream_modes:
|
||||
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamToolCallHandler(
|
||||
stream.put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
|
||||
)
|
||||
)
|
||||
|
||||
# set up custom stream mode
|
||||
if "custom" in stream_modes:
|
||||
|
||||
@@ -2804,7 +2911,8 @@ class Pregel(
|
||||
# set final channel values as run output
|
||||
run_manager.on_chain_end(loop.output)
|
||||
except BaseException as e:
|
||||
run_manager.on_chain_error(e)
|
||||
if run_manager is not None:
|
||||
run_manager.on_chain_error(e)
|
||||
raise
|
||||
|
||||
@overload
|
||||
@@ -2944,33 +3052,7 @@ class Pregel(
|
||||
)
|
||||
|
||||
config = ensure_config(self.config, config)
|
||||
callback_manager = get_async_callback_manager_for_config(config)
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
run_manager = await callback_manager.on_chain_start(
|
||||
None,
|
||||
input,
|
||||
name=config.get("run_name", self.get_name()),
|
||||
run_id=config.get("run_id"),
|
||||
)
|
||||
graph_callback_manager = get_async_graph_callback_manager_for_config(
|
||||
config,
|
||||
run_id=run_manager.run_id,
|
||||
)
|
||||
# if running from astream_log() run each proc with streaming
|
||||
do_stream = (
|
||||
next(
|
||||
(
|
||||
True
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
and not isinstance(h, StreamMessagesHandler)
|
||||
),
|
||||
False,
|
||||
)
|
||||
if _StreamingCallbackHandler is not None
|
||||
else False
|
||||
)
|
||||
run_manager = None
|
||||
try:
|
||||
# assign defaults
|
||||
(
|
||||
@@ -2991,6 +3073,50 @@ class Pregel(
|
||||
interrupt_after=interrupt_after,
|
||||
durability=durability,
|
||||
)
|
||||
callback_manager = get_async_callback_manager_for_config(config)
|
||||
if "messages" in stream_modes and version != "v2":
|
||||
# Strip any inherited v2 messages handler so a v1 stream
|
||||
# does not get routed through the content-block event
|
||||
# protocol. Leave v1 handlers in place — an outer
|
||||
# astream(stream_mode="messages", subgraphs=True) relies
|
||||
# on its inheritable handler to observe events emitted
|
||||
# by inner astream(stream_mode="messages") calls.
|
||||
callback_manager.handlers = [
|
||||
h
|
||||
for h in callback_manager.handlers
|
||||
if not isinstance(h, StreamMessagesHandlerV2)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not isinstance(h, StreamMessagesHandlerV2)
|
||||
]
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
run_manager = await callback_manager.on_chain_start(
|
||||
None,
|
||||
input,
|
||||
name=config.get("run_name", self.get_name()),
|
||||
run_id=config.get("run_id"),
|
||||
)
|
||||
graph_callback_manager = get_async_graph_callback_manager_for_config(
|
||||
config,
|
||||
run_id=run_manager.run_id,
|
||||
)
|
||||
# if running from astream_log() run each proc with streaming
|
||||
do_stream = (
|
||||
next(
|
||||
(
|
||||
True
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
and not isinstance(h, StreamMessagesHandler)
|
||||
),
|
||||
False,
|
||||
)
|
||||
if _StreamingCallbackHandler is not None
|
||||
else False
|
||||
)
|
||||
if checkpointer is None and durability is not None:
|
||||
warnings.warn(
|
||||
"`durability` has no effect when no checkpointer is present.",
|
||||
@@ -3003,14 +3129,33 @@ class Pregel(
|
||||
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))
|
||||
use_stream_messages_v2 = bool(
|
||||
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||||
)
|
||||
messages_handler_cls = (
|
||||
StreamMessagesHandlerV2
|
||||
if use_stream_messages_v2
|
||||
else StreamMessagesHandler
|
||||
)
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamMessagesHandler(
|
||||
messages_handler_cls(
|
||||
stream_put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||
)
|
||||
)
|
||||
|
||||
# set up tools stream mode
|
||||
if "tools" in stream_modes:
|
||||
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamToolCallHandler(
|
||||
stream_put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
|
||||
)
|
||||
)
|
||||
|
||||
# set up custom stream mode
|
||||
def stream_writer(c: Any) -> None:
|
||||
aioloop.call_soon_threadsafe(
|
||||
@@ -3234,9 +3379,150 @@ class Pregel(
|
||||
# set final channel values as run output
|
||||
await run_manager.on_chain_end(loop.output)
|
||||
except BaseException as e:
|
||||
await asyncio.shield(run_manager.on_chain_error(e))
|
||||
if run_manager is not None:
|
||||
await asyncio.shield(run_manager.on_chain_error(e))
|
||||
raise
|
||||
|
||||
def stream_v2(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||||
) -> Any:
|
||||
"""Start a sync v2 streaming run driven by transformer projections.
|
||||
|
||||
Builds a `StreamMux` from the built-in `ValuesTransformer` /
|
||||
`MessagesTransformer`, this graph's compile-time
|
||||
`stream_transformers`, and any additional `transformers=`
|
||||
supplied at the call site. Returns a `GraphRunStream` that the
|
||||
caller drives by iterating any projection — no background
|
||||
thread.
|
||||
|
||||
Note:
|
||||
Nesting v1 `stream(stream_mode="messages")` inside a node
|
||||
of a `stream_v2` run is not fully supported. The outer v2
|
||||
messages handler is inheritable, so it sits in the inner
|
||||
chat model's callback chain; `BaseChatModel.invoke` then
|
||||
routes through the v2 event protocol and the inner v1
|
||||
messages handler does not see `on_llm_new_token` chunks.
|
||||
The inner stream still yields a finalized message via
|
||||
`on_llm_end`, but token-by-token output is lost. Use
|
||||
`stream_v2` for the inner graph as well, or call
|
||||
`chat_model.stream(...)` explicitly inside the node, to
|
||||
get token-level streaming.
|
||||
|
||||
Args:
|
||||
input: Graph input.
|
||||
config: Optional runnable config forwarded to the graph.
|
||||
interrupt_before: Nodes to interrupt before, if any.
|
||||
interrupt_after: Nodes to interrupt after, if any.
|
||||
transformers: Extra transformer classes or configured factories
|
||||
appended after compile-time `stream_transformers`. Factories
|
||||
are called as `factory(scope)` so they can propagate to
|
||||
subgraph scopes.
|
||||
|
||||
Returns:
|
||||
A `GraphRunStream` the caller iterates to drive the run.
|
||||
"""
|
||||
parent_ns = _resolve_parent_ns(self.config, config)
|
||||
compiled_factories = _normalize_stream_transformer_factories(
|
||||
self.stream_transformers
|
||||
)
|
||||
extra_factories = _normalize_stream_transformer_factories(transformers)
|
||||
mux = StreamMux(
|
||||
factories=[
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
*compiled_factories,
|
||||
*extra_factories,
|
||||
],
|
||||
scope=parent_ns,
|
||||
is_async=False,
|
||||
)
|
||||
values_t = cast(ValuesTransformer, mux.transformer_by_key("values"))
|
||||
graph_iter = iter(
|
||||
self.stream(
|
||||
input,
|
||||
patch_configurable(config, {CONFIG_KEY_STREAM_MESSAGES_V2: True}),
|
||||
stream_mode=_collect_stream_modes(mux),
|
||||
subgraphs=True,
|
||||
version="v2",
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
)
|
||||
)
|
||||
return GraphRunStream(graph_iter, mux, values_t)
|
||||
|
||||
async def astream_v2(
|
||||
self,
|
||||
input: InputT | Command | None,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||||
) -> Any:
|
||||
"""Async counterpart to `stream_v2`.
|
||||
|
||||
Returns an `AsyncGraphRunStream` whose projections can be awaited
|
||||
concurrently; each subscribed cursor drives the pump when its
|
||||
buffer is empty.
|
||||
|
||||
Note:
|
||||
Same nesting limitation as `stream_v2`: nesting v1
|
||||
`astream(stream_mode="messages")` inside a node of an
|
||||
`astream_v2` run drops `on_llm_new_token` chunks because
|
||||
the outer v2 handler reroutes `BaseChatModel.invoke`
|
||||
through the v2 event protocol. The inner stream still
|
||||
yields a finalized message at end-of-call. Use
|
||||
`astream_v2` for the inner graph as well, or call
|
||||
`chat_model.astream(...)` explicitly inside the node, to
|
||||
get token-level streaming.
|
||||
|
||||
Args:
|
||||
input: Graph input.
|
||||
config: Optional runnable config forwarded to the graph.
|
||||
interrupt_before: Nodes to interrupt before, if any.
|
||||
interrupt_after: Nodes to interrupt after, if any.
|
||||
transformers: Extra transformer classes or configured factories
|
||||
appended after compile-time `stream_transformers`. Factories
|
||||
are called as `factory(scope)` so they can propagate to
|
||||
subgraph scopes.
|
||||
"""
|
||||
parent_ns = _resolve_parent_ns(self.config, config)
|
||||
compiled_factories = _normalize_stream_transformer_factories(
|
||||
self.stream_transformers
|
||||
)
|
||||
extra_factories = _normalize_stream_transformer_factories(transformers)
|
||||
mux = StreamMux(
|
||||
factories=[
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
*compiled_factories,
|
||||
*extra_factories,
|
||||
],
|
||||
scope=parent_ns,
|
||||
is_async=True,
|
||||
)
|
||||
values_t = cast(ValuesTransformer, mux.transformer_by_key("values"))
|
||||
graph_aiter = self.astream(
|
||||
input,
|
||||
patch_configurable(config, {CONFIG_KEY_STREAM_MESSAGES_V2: True}),
|
||||
stream_mode=_collect_stream_modes(mux),
|
||||
subgraphs=True,
|
||||
version="v2",
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
).__aiter__()
|
||||
return AsyncGraphRunStream(graph_aiter, mux, values_t)
|
||||
|
||||
@overload
|
||||
def invoke(
|
||||
self,
|
||||
@@ -3712,6 +3998,24 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
|
||||
payload["values"] = mapper(payload["values"])
|
||||
|
||||
|
||||
def _resolve_parent_ns(
|
||||
graph_config: RunnableConfig | None, call_config: RunnableConfig | None
|
||||
) -> tuple[str, ...]:
|
||||
"""Return the checkpoint namespace the caller is running under.
|
||||
|
||||
`stream_v2` uses this to scope its native projections
|
||||
(`ValuesTransformer`, `MessagesTransformer`) to events emitted at
|
||||
the run's own level. A root call resolves to `()`; a call made
|
||||
from inside a node carries the outer graph's task namespace so the
|
||||
projection still matches its own root-level events.
|
||||
"""
|
||||
merged = ensure_config(graph_config, call_config)
|
||||
ns = merged.get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS)
|
||||
if not ns:
|
||||
return ()
|
||||
return tuple(ns.split(NS_SEP))
|
||||
|
||||
|
||||
def _build_server_info(
|
||||
config: RunnableConfig, parent_runtime: Runtime[Any]
|
||||
) -> ServerInfo | None:
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Streaming infrastructure for LangGraph.
|
||||
|
||||
Compile a graph with `transformers=[...]` and call `graph.stream_v2()` /
|
||||
`graph.astream_v2()` to drive a transformer pipeline that projects the
|
||||
graph's raw events into ergonomic per-channel streams.
|
||||
"""
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.run_stream import (
|
||||
AsyncGraphRunStream,
|
||||
AsyncSubgraphRunStream,
|
||||
GraphRunStream,
|
||||
SubgraphRunStream,
|
||||
)
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
from langgraph.stream.transformers import (
|
||||
LifecyclePayload,
|
||||
LifecycleTransformer,
|
||||
SubgraphStatus,
|
||||
SubgraphTransformer,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AsyncGraphRunStream",
|
||||
"AsyncSubgraphRunStream",
|
||||
"EventLog",
|
||||
"GraphRunStream",
|
||||
"LifecyclePayload",
|
||||
"LifecycleTransformer",
|
||||
"ProtocolEvent",
|
||||
"StreamChannel",
|
||||
"StreamTransformer",
|
||||
"SubgraphRunStream",
|
||||
"SubgraphStatus",
|
||||
"SubgraphTransformer",
|
||||
]
|
||||
@@ -0,0 +1,32 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any, cast
|
||||
|
||||
from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams
|
||||
from langgraph.types import StreamPart
|
||||
|
||||
|
||||
def convert_to_protocol_event(part: StreamPart) -> ProtocolEvent:
|
||||
"""Convert a v2 StreamPart to a ProtocolEvent.
|
||||
|
||||
Args:
|
||||
part: A stream part with keys `type`, `ns`, `data`, and
|
||||
optionally `interrupts` (present on values events).
|
||||
|
||||
Returns:
|
||||
The equivalent ProtocolEvent.
|
||||
"""
|
||||
part_dict = cast(dict[str, Any], part)
|
||||
params: _ProtocolEventParams = {
|
||||
"namespace": list(part_dict["ns"]),
|
||||
"timestamp": int(time.time() * 1000),
|
||||
"data": part_dict["data"],
|
||||
}
|
||||
if "interrupts" in part_dict:
|
||||
params["interrupts"] = part_dict["interrupts"]
|
||||
return {
|
||||
"type": "event",
|
||||
"method": part_dict["type"],
|
||||
"params": params,
|
||||
}
|
||||
@@ -0,0 +1,306 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class EventLog(Generic[T]):
|
||||
"""Single-consumer drainable queue for streaming events.
|
||||
|
||||
Items are popped off the front as the consumer advances — there is
|
||||
no retention beyond what's currently queued. A log accepts exactly
|
||||
one subscriber; a second `__iter__` / `__aiter__` call raises. Use
|
||||
`tee(n)` / `atee(n)` for fan-out.
|
||||
|
||||
Starts unbound — neither `__iter__` nor `__aiter__` is available
|
||||
until the StreamMux calls `_bind(is_async)`. After binding, only
|
||||
the matching iteration protocol works; the other raises `TypeError`.
|
||||
|
||||
Pump wiring (set by the run stream, not by `_bind`):
|
||||
- `_request_more`: sync pump callable, returns True if a new
|
||||
event was produced.
|
||||
- `_arequest_more`: async pump coroutine factory, same contract.
|
||||
|
||||
Memory is bounded by caller pace: both sync and async use caller-
|
||||
driven pumps, so each cursor advance produces at most one event.
|
||||
The only shape where a log can accumulate meaningfully is
|
||||
concurrent async consumers at unequal rates — a slow consumer's
|
||||
log grows while fast consumers drive the shared pump. That's the
|
||||
documented tradeoff for concurrent consumption; consume at similar
|
||||
rates or use a single consumer if memory matters.
|
||||
|
||||
Lazy-subscribe: `push` is a no-op when no subscriber has registered.
|
||||
Transformers still execute `process()` (so scalar state like
|
||||
`ValuesTransformer._latest` stays current); only the log append is
|
||||
skipped.
|
||||
"""
|
||||
|
||||
def __init__(self, maxlen: int | None = None) -> None:
|
||||
"""Initialize an empty, unbound log.
|
||||
|
||||
Args:
|
||||
maxlen: Accepted for forward compatibility; currently unused.
|
||||
The caller-driven pump bounds memory naturally for
|
||||
single-consumer use.
|
||||
|
||||
Raises:
|
||||
ValueError: If `maxlen` is not a positive integer or `None`.
|
||||
"""
|
||||
if maxlen is not None and maxlen <= 0:
|
||||
raise ValueError("EventLog maxlen must be a positive int or None")
|
||||
self._items: deque[T] = deque()
|
||||
self._maxlen: int | None = maxlen
|
||||
self._closed = False
|
||||
self._error: BaseException | None = None
|
||||
|
||||
# Binding state — None means unbound.
|
||||
self._is_async: bool | None = None
|
||||
|
||||
# Flipped on first __iter__ / __aiter__. Pre-subscription
|
||||
# pushes are silent no-ops.
|
||||
self._subscribed = False
|
||||
|
||||
# Pump wiring set by the run stream after bind.
|
||||
self._request_more: Callable[[], bool] | None = None
|
||||
self._arequest_more: Callable[[], Awaitable[bool]] | None = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Binding
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _bind(self, *, is_async: bool) -> None:
|
||||
"""Bind this log to sync or async mode.
|
||||
|
||||
Called by the StreamMux after transformer registration. Must be
|
||||
called exactly once before any iteration.
|
||||
|
||||
Args:
|
||||
is_async: True to enable async iteration, False for sync.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the log has already been bound.
|
||||
"""
|
||||
if self._is_async is not None:
|
||||
raise RuntimeError("EventLog is already bound")
|
||||
self._is_async = is_async
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Producer API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def push(self, item: T) -> None:
|
||||
"""Append an item. No-op when no subscriber is registered.
|
||||
|
||||
Non-blocking in both sync and async — matches v1's
|
||||
`put_nowait` producer shape. Memory is bounded by caller pace
|
||||
via the caller-driven pump.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the log is closed (and subscribed).
|
||||
"""
|
||||
if not self._subscribed:
|
||||
return
|
||||
if self._closed:
|
||||
raise RuntimeError("Cannot push to a closed EventLog")
|
||||
self._items.append(item)
|
||||
|
||||
def close(self) -> None:
|
||||
"""Mark the log as complete."""
|
||||
self._closed = True
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Mark the log as errored.
|
||||
|
||||
Args:
|
||||
err: The exception to surface to the subscriber.
|
||||
"""
|
||||
self._error = err
|
||||
self._closed = True
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Sync iteration (caller-driven pump)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def __iter__(self) -> Iterator[T]:
|
||||
"""Subscribe and return a sync cursor. Can be called only once.
|
||||
|
||||
Raises:
|
||||
TypeError: If the log is unbound or bound to async mode.
|
||||
RuntimeError: If the log already has a subscriber.
|
||||
"""
|
||||
if self._is_async is None:
|
||||
raise TypeError(
|
||||
"EventLog has not been bound yet. "
|
||||
"Register the transformer with a StreamMux first."
|
||||
)
|
||||
if self._is_async:
|
||||
raise TypeError(
|
||||
"This EventLog is bound to async mode — use 'async for' instead."
|
||||
)
|
||||
if self._subscribed:
|
||||
raise RuntimeError(
|
||||
"EventLog already has a subscriber; use .tee(n) for fan-out."
|
||||
)
|
||||
self._subscribed = True
|
||||
return self._sync_cursor()
|
||||
|
||||
def _sync_cursor(self) -> Iterator[T]:
|
||||
while True:
|
||||
if self._items:
|
||||
yield self._items.popleft()
|
||||
elif self._closed:
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return
|
||||
elif self._request_more is not None:
|
||||
if not self._request_more():
|
||||
if not self._items and not self._closed:
|
||||
return
|
||||
else:
|
||||
return
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Async iteration (caller-driven pump)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[T]:
|
||||
"""Subscribe and return an async cursor. Can be called only once.
|
||||
|
||||
Raises:
|
||||
TypeError: If the log is unbound or bound to sync mode.
|
||||
RuntimeError: If the log already has a subscriber.
|
||||
"""
|
||||
if self._is_async is None:
|
||||
raise TypeError(
|
||||
"EventLog has not been bound yet. "
|
||||
"Register the transformer with a StreamMux first."
|
||||
)
|
||||
if not self._is_async:
|
||||
raise TypeError("This EventLog is bound to sync mode — use 'for' instead.")
|
||||
if self._subscribed:
|
||||
raise RuntimeError(
|
||||
"EventLog already has a subscriber; use .atee(n) for fan-out."
|
||||
)
|
||||
self._subscribed = True
|
||||
return self._async_cursor()
|
||||
|
||||
async def _async_cursor(self) -> AsyncIterator[T]:
|
||||
while True:
|
||||
if self._items:
|
||||
yield self._items.popleft()
|
||||
elif self._closed:
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return
|
||||
elif self._arequest_more is not None:
|
||||
if not await self._arequest_more():
|
||||
if not self._items and not self._closed:
|
||||
return
|
||||
else:
|
||||
return
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Fan-out via tee
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def tee(self, n: int = 2) -> tuple[Iterator[T], ...]:
|
||||
"""Subscribe and return `n` independent sync iterators.
|
||||
|
||||
Each branch has its own buffer; items pulled from the
|
||||
underlying cursor are copied into every branch. Branches are
|
||||
naturally bounded by caller pace since the sync pump is
|
||||
caller-driven.
|
||||
|
||||
Args:
|
||||
n: Number of branches to create. Must be >= 1.
|
||||
|
||||
Returns:
|
||||
A tuple of `n` iterators over the same underlying stream.
|
||||
|
||||
Raises:
|
||||
TypeError: If the log is unbound or bound to async mode.
|
||||
RuntimeError: If the log already has a subscriber.
|
||||
ValueError: If `n` < 1.
|
||||
"""
|
||||
if n < 1:
|
||||
raise ValueError("tee() requires n >= 1")
|
||||
source = self.__iter__()
|
||||
buffers: list[deque[T]] = [deque() for _ in range(n)]
|
||||
exhausted = [False]
|
||||
|
||||
def branch(i: int) -> Iterator[T]:
|
||||
buf = buffers[i]
|
||||
while True:
|
||||
if buf:
|
||||
yield buf.popleft()
|
||||
elif exhausted[0]:
|
||||
return
|
||||
else:
|
||||
try:
|
||||
item = next(source)
|
||||
except StopIteration:
|
||||
exhausted[0] = True
|
||||
return
|
||||
for b in buffers:
|
||||
b.append(item)
|
||||
|
||||
return tuple(branch(i) for i in range(n))
|
||||
|
||||
def atee(self, n: int = 2) -> tuple[AsyncIterator[T], ...]:
|
||||
"""Subscribe and return `n` independent async iterators.
|
||||
|
||||
Caller-driven fan-out: each branch's `__anext__` either pops
|
||||
from its own buffer or, under a shared `asyncio.Lock`, pulls
|
||||
one item from the underlying cursor and distributes it to
|
||||
every branch's buffer.
|
||||
|
||||
Args:
|
||||
n: Number of branches to create. Must be >= 1.
|
||||
|
||||
Returns:
|
||||
A tuple of `n` async iterators over the same underlying
|
||||
stream.
|
||||
|
||||
Raises:
|
||||
TypeError: If the log is unbound or bound to sync mode.
|
||||
RuntimeError: If the log already has a subscriber.
|
||||
ValueError: If `n` < 1.
|
||||
"""
|
||||
if n < 1:
|
||||
raise ValueError("atee() requires n >= 1")
|
||||
source = self.__aiter__()
|
||||
buffers: list[deque[T]] = [deque() for _ in range(n)]
|
||||
exhausted = [False]
|
||||
error: list[BaseException | None] = [None]
|
||||
lock = asyncio.Lock()
|
||||
|
||||
async def branch(i: int) -> AsyncIterator[T]:
|
||||
buf = buffers[i]
|
||||
while True:
|
||||
if buf:
|
||||
yield buf.popleft()
|
||||
continue
|
||||
if exhausted[0]:
|
||||
if error[0] is not None:
|
||||
raise error[0]
|
||||
return
|
||||
async with lock:
|
||||
if buf or exhausted[0]:
|
||||
continue
|
||||
try:
|
||||
item = await source.__anext__()
|
||||
except StopAsyncIteration:
|
||||
exhausted[0] = True
|
||||
continue
|
||||
except Exception as e:
|
||||
error[0] = e
|
||||
exhausted[0] = True
|
||||
continue
|
||||
for b in buffers:
|
||||
b.append(item)
|
||||
|
||||
return tuple(branch(i) for i in range(n))
|
||||
@@ -0,0 +1,510 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._types import (
|
||||
ProtocolEvent,
|
||||
StreamTransformer,
|
||||
transformer_requires_async,
|
||||
)
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
|
||||
TransformerFactory = Callable[["tuple[str, ...]"], StreamTransformer]
|
||||
"""Factory that builds a scoped transformer for a mux.
|
||||
|
||||
Called once per `StreamMux` with the mux's scope (typically `()` for
|
||||
the root). Standard transformer classes accept a single positional
|
||||
scope argument, so the class itself is a valid factory. User
|
||||
transformers can close over their config:
|
||||
`lambda scope: MyTransformer(scope, foo=...)`.
|
||||
"""
|
||||
|
||||
|
||||
class StreamMux:
|
||||
"""Central event dispatcher for the streaming infrastructure.
|
||||
|
||||
Owns the main event log and routes events through a transformer
|
||||
pipeline. StreamChannels discovered in transformer projections are
|
||||
auto-wired so that every `push()` also injects a `ProtocolEvent`
|
||||
into the main log.
|
||||
|
||||
Pass `is_async=True` when the mux will be consumed via async
|
||||
iteration (`handler.astream()`). All EventLog and StreamChannel
|
||||
instances discovered during registration are automatically bound
|
||||
to the matching mode.
|
||||
|
||||
Attributes:
|
||||
extensions: Merged projection dict across all registered
|
||||
transformers. Treat as read-only — mutations won't be
|
||||
reflected back in individual transformers' state.
|
||||
native_keys: Projection keys contributed by transformers with
|
||||
`_native = True`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
*,
|
||||
is_async: bool = False,
|
||||
factories: list[TransformerFactory] | None = None,
|
||||
scope: tuple[str, ...] = (),
|
||||
_assign_seq: bool = True,
|
||||
) -> None:
|
||||
"""Initialize the mux and register transformers in order.
|
||||
|
||||
Callers pass either `transformers` (pre-built instances) or
|
||||
`factories` (callables producing fresh instances per mux). Each
|
||||
transformer's `init()` is called, projections are merged into
|
||||
`extensions`, `_native` keys are recorded in `native_keys`, and
|
||||
any EventLog / StreamChannel instances are bound and wired.
|
||||
|
||||
Args:
|
||||
transformers: Already-built transformer instances. Registered
|
||||
only on this mux — they are NOT cloned into child
|
||||
mini-muxes built by `_make_child`. Use `factories` for
|
||||
transformers that should propagate to nested scopes.
|
||||
is_async: True for async dispatch (`apush` / `aclose` /
|
||||
`afail`), False for the sync path.
|
||||
factories: One-argument callables `(scope) -> StreamTransformer`.
|
||||
Called once with this mux's `scope` here, and cloned
|
||||
again per child scope by `_make_child` so each
|
||||
sub-mux gets fresh instances.
|
||||
scope: The namespace the mux operates within. The root mux
|
||||
is `()`.
|
||||
_assign_seq: Internal flag for child muxes. Root muxes assign
|
||||
monotonic `seq` numbers when appending to their main event
|
||||
log; child muxes share forwarded event objects and must not
|
||||
mutate their envelopes.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If any transformer requires an async run but
|
||||
the mux is in sync mode.
|
||||
TypeError: If a transformer's `init()` doesn't return a dict.
|
||||
ValueError: If transformers' projection keys collide.
|
||||
"""
|
||||
self.is_async = is_async
|
||||
self.scope: tuple[str, ...] = scope
|
||||
self._assign_seq = _assign_seq
|
||||
self._events: EventLog[ProtocolEvent] = EventLog()
|
||||
self._events._bind(is_async=is_async)
|
||||
self._transformers: list[StreamTransformer] = []
|
||||
self._channels: list[StreamChannel[Any]] = []
|
||||
self._logs: list[EventLog[Any]] = []
|
||||
self._seq = 0
|
||||
|
||||
self.extensions: dict[str, Any] = {}
|
||||
self.native_keys: set[str] = set()
|
||||
self._projection_owners: dict[str, str] = {}
|
||||
self._transformer_by_key: dict[str, StreamTransformer] = {}
|
||||
|
||||
# Stored only when constructed from factories — used by
|
||||
# `_make_child` to clone the transformer pipeline at a deeper
|
||||
# scope. Pre-built transformers can't be cloned, so a mux
|
||||
# built with `transformers=` rejects child construction.
|
||||
self._factories: list[TransformerFactory] | None = (
|
||||
list(factories) if factories is not None else None
|
||||
)
|
||||
self._pump_fn: Callable[[], bool] | None = None
|
||||
self._apump_fn: Callable[[], Awaitable[bool]] | None = None
|
||||
|
||||
# Factories run first (they propagate to child mini-muxes
|
||||
# via `_make_child`), then any pre-built `transformers=`
|
||||
# instances are registered as root-only — they aren't cloned
|
||||
# for child scopes.
|
||||
if factories is not None:
|
||||
for factory in factories:
|
||||
self._register(factory(scope))
|
||||
for transformer in transformers or ():
|
||||
self._register(transformer)
|
||||
|
||||
def transformer_by_key(self, key: str) -> StreamTransformer | None:
|
||||
"""Return the transformer that contributed `key` to the projection."""
|
||||
return self._transformer_by_key.get(key)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Pump wiring + mini-mux nesting
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def bind_pump(self, fn: Callable[[], bool]) -> None:
|
||||
"""Wire the sync pull callback onto every projection in this mux.
|
||||
|
||||
Records the pump on the mux so child mini-muxes built by
|
||||
`_make_child` can inherit it. Propagates to:
|
||||
- the main event log (`self._events`)
|
||||
- every projection EventLog and StreamChannel in `extensions`
|
||||
- any registered transformer that exposes `_bind_pump` (e.g.
|
||||
`MessagesTransformer` so `ChatModelStream` instances drive the
|
||||
shared pump from their cursors)
|
||||
"""
|
||||
self._pump_fn = fn
|
||||
self._events._request_more = fn
|
||||
for value in self.extensions.values():
|
||||
if isinstance(value, EventLog):
|
||||
value._request_more = fn
|
||||
elif isinstance(value, StreamChannel):
|
||||
value._log._request_more = fn
|
||||
for transformer in self._transformers:
|
||||
bind = getattr(transformer, "_bind_pump", None)
|
||||
if bind is not None:
|
||||
bind(fn)
|
||||
|
||||
def bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:
|
||||
"""Async counterpart to `bind_pump`."""
|
||||
self._apump_fn = fn
|
||||
self._events._arequest_more = fn
|
||||
for value in self.extensions.values():
|
||||
if isinstance(value, EventLog):
|
||||
value._arequest_more = fn
|
||||
elif isinstance(value, StreamChannel):
|
||||
value._log._arequest_more = fn
|
||||
for transformer in self._transformers:
|
||||
abind = getattr(transformer, "_bind_apump", None)
|
||||
if abind is not None:
|
||||
abind(fn)
|
||||
|
||||
def _make_child(self, scope: tuple[str, ...]) -> StreamMux:
|
||||
"""Build a mini-mux with the same factories scoped to `scope`.
|
||||
|
||||
Used by `SubgraphTransformer` to attach a fresh transformer
|
||||
pipeline to each discovered subgraph handle. The child mux
|
||||
inherits the current pump bindings (so cursors on its
|
||||
projection logs drive the root pump), carries the same factory
|
||||
list forward to any grandchild subgraphs, and does not assign
|
||||
`seq` numbers so forwarded events can be shared without
|
||||
mutating their envelope.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the mux was not constructed with
|
||||
`factories=`. Mini-muxes require factories so each scope
|
||||
gets its own fresh transformer instances.
|
||||
"""
|
||||
if self._factories is None:
|
||||
raise RuntimeError(
|
||||
"StreamMux._make_child requires the mux to be constructed "
|
||||
"with `factories=`; pre-built transformers can't be "
|
||||
"cloned to a new scope."
|
||||
)
|
||||
child = StreamMux(
|
||||
factories=self._factories,
|
||||
is_async=self.is_async,
|
||||
scope=scope,
|
||||
_assign_seq=False,
|
||||
)
|
||||
if self._pump_fn is not None:
|
||||
child.bind_pump(self._pump_fn)
|
||||
if self._apump_fn is not None:
|
||||
child.bind_apump(self._apump_fn)
|
||||
return child
|
||||
|
||||
def _register(self, transformer: StreamTransformer) -> None:
|
||||
"""Register a single transformer.
|
||||
|
||||
Calls `transformer.init()`, stores the transformer for event
|
||||
processing, binds any EventLog or StreamChannel instances in
|
||||
the projection, and merges the projection into `extensions`.
|
||||
"""
|
||||
if transformer_requires_async(transformer) and not self.is_async:
|
||||
raise RuntimeError(
|
||||
f"{type(transformer).__name__} requires an async run — "
|
||||
"it overrides aprocess/afinalize/afail or sets "
|
||||
"requires_async=True. Use astream(), not stream()."
|
||||
)
|
||||
projection = transformer.init()
|
||||
if not isinstance(projection, dict):
|
||||
raise TypeError(
|
||||
f"StreamTransformer.init() must return a dict, "
|
||||
f"got {type(projection).__name__}"
|
||||
)
|
||||
conflicts = set(projection) & set(self.extensions)
|
||||
if conflicts:
|
||||
attributions = ", ".join(
|
||||
f"{key!r} (owned by {self._projection_owners[key]})"
|
||||
for key in sorted(conflicts)
|
||||
)
|
||||
raise ValueError(
|
||||
f"Transformer {type(transformer).__name__} returned "
|
||||
f"projection keys that conflict with already-registered "
|
||||
f"keys: {attributions}"
|
||||
)
|
||||
is_native = bool(getattr(transformer, "_native", False))
|
||||
self._transformers.append(transformer)
|
||||
self._bind_and_wire(projection, native=is_native)
|
||||
self.extensions.update(projection)
|
||||
owner_name = type(transformer).__name__
|
||||
for key in projection:
|
||||
self._projection_owners[key] = owner_name
|
||||
self._transformer_by_key[key] = transformer
|
||||
if is_native:
|
||||
self.native_keys.update(projection.keys())
|
||||
transformer._on_register(self)
|
||||
|
||||
def push(self, event: ProtocolEvent) -> None:
|
||||
"""Route an event through all transformers, then append to the main log.
|
||||
|
||||
Each transformer's `process()` is called in registration order.
|
||||
If any transformer returns False, the event is suppressed from
|
||||
the main log, but transformers that already saw it keep their
|
||||
side effects.
|
||||
|
||||
On the root mux, `seq` is assigned right before an event enters
|
||||
the main log, not before the transformer pipeline runs. This
|
||||
ensures that events auto-forwarded from StreamChannels during
|
||||
`process()` get earlier seq numbers than the original event,
|
||||
preserving monotonic ordering in the root log. Child muxes do
|
||||
not assign `seq`, so subgraph forwarding can share event objects
|
||||
without mutating their envelopes.
|
||||
|
||||
Args:
|
||||
event: The protocol event to dispatch.
|
||||
"""
|
||||
keep = True
|
||||
for transformer in self._transformers:
|
||||
if not transformer.process(event):
|
||||
keep = False
|
||||
if keep:
|
||||
if self._assign_seq:
|
||||
self._seq += 1
|
||||
event["seq"] = self._seq
|
||||
self._events.push(event)
|
||||
|
||||
def close(self) -> None:
|
||||
"""Finalize all transformers, close all projections and the main log.
|
||||
|
||||
EventLogs and StreamChannels discovered in transformer
|
||||
projections are auto-closed after `finalize()` runs —
|
||||
transformers don't need to close them manually. If any
|
||||
transformer's `finalize()` raises, the remaining transformers,
|
||||
projections, and the main log are still closed; the first error
|
||||
is re-raised after cleanup completes.
|
||||
|
||||
Raises:
|
||||
BaseException: The first error raised by a transformer's
|
||||
`finalize()`, re-raised after cleanup finishes.
|
||||
"""
|
||||
first_error: BaseException | None = None
|
||||
for transformer in self._transformers:
|
||||
try:
|
||||
transformer.finalize()
|
||||
except BaseException as e:
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
for log in self._logs:
|
||||
if not log._closed:
|
||||
log.close()
|
||||
for ch in self._channels:
|
||||
if not ch._log._closed:
|
||||
ch._close()
|
||||
self._events.close()
|
||||
if first_error is not None:
|
||||
raise first_error
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Fail all transformers, projections, and the main log.
|
||||
|
||||
EventLogs and StreamChannels discovered in transformer
|
||||
projections are auto-failed — transformers don't need to fail
|
||||
them manually. If any transformer's `fail()` raises, the
|
||||
remaining transformers, projections, and the main log are still
|
||||
failed.
|
||||
|
||||
Args:
|
||||
err: The exception that ended the run.
|
||||
"""
|
||||
for transformer in self._transformers:
|
||||
try:
|
||||
transformer.fail(err)
|
||||
except BaseException:
|
||||
pass
|
||||
for log in self._logs:
|
||||
if not log._closed:
|
||||
log.fail(err)
|
||||
for ch in self._channels:
|
||||
if not ch._log._closed:
|
||||
ch._fail(err)
|
||||
self._events.fail(err)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Async dispatch
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def apush(self, event: ProtocolEvent) -> None:
|
||||
"""Dispatch an event on the async lane.
|
||||
|
||||
Awaits each transformer's `aprocess` in registration order
|
||||
before appending to the main log. A slow `aprocess` serializes
|
||||
the pipeline by design — that's the guarantee that lets a later
|
||||
transformer (or a synchronous consumer) see the result of the
|
||||
async work. For decoupled work, use `schedule()` from inside
|
||||
`process` / `aprocess` instead.
|
||||
|
||||
The main log append is a non-blocking `push` — matching v1's
|
||||
`put_nowait` shape. The root mux assigns `seq`; child muxes do
|
||||
not, so forwarded subgraph events can be shared without copying.
|
||||
Memory is bounded by caller pace via the caller-driven pump; see
|
||||
`EventLog` for the full tradeoff story.
|
||||
|
||||
Args:
|
||||
event: The protocol event to dispatch.
|
||||
"""
|
||||
keep = True
|
||||
for transformer in self._transformers:
|
||||
if not await transformer.aprocess(event):
|
||||
keep = False
|
||||
if keep:
|
||||
if self._assign_seq:
|
||||
self._seq += 1
|
||||
event["seq"] = self._seq
|
||||
self._events.push(event)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Finalize on the async lane.
|
||||
|
||||
Awaits every task started via `StreamTransformer.schedule()`
|
||||
across all transformers, then calls `afinalize()` on each,
|
||||
then auto-closes logs, channels, and the main event log.
|
||||
|
||||
If any scheduled task raised under `on_error="raise"`, or any
|
||||
transformer's `afinalize` raises, the exception propagates.
|
||||
The caller (the pump) handles it by routing into `afail`.
|
||||
|
||||
Raises:
|
||||
BaseException: The first scheduled-task or `afinalize`
|
||||
error, re-raised after cleanup.
|
||||
"""
|
||||
pending = self._collect_scheduled_tasks()
|
||||
if pending:
|
||||
results = await asyncio.gather(*pending, return_exceptions=True)
|
||||
first_err = next(
|
||||
(
|
||||
r
|
||||
for r in results
|
||||
if isinstance(r, BaseException)
|
||||
and not isinstance(r, asyncio.CancelledError)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if first_err is not None:
|
||||
raise first_err
|
||||
|
||||
first_error: BaseException | None = None
|
||||
for transformer in self._transformers:
|
||||
try:
|
||||
await transformer.afinalize()
|
||||
except BaseException as e:
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
for log in self._logs:
|
||||
if not log._closed:
|
||||
log.close()
|
||||
for ch in self._channels:
|
||||
if not ch._log._closed:
|
||||
ch._close()
|
||||
self._events.close()
|
||||
if first_error is not None:
|
||||
raise first_error
|
||||
|
||||
async def afail(self, err: BaseException) -> None:
|
||||
"""Fail on the async lane.
|
||||
|
||||
Cancels every scheduled task across all transformers, awaits
|
||||
them to completion, then runs each transformer's `afail` hook
|
||||
and auto-fails logs, channels, and the main event log.
|
||||
|
||||
Args:
|
||||
err: The exception that ended the run.
|
||||
"""
|
||||
pending = self._collect_scheduled_tasks()
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
if pending:
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
|
||||
for transformer in self._transformers:
|
||||
try:
|
||||
await transformer.afail(err)
|
||||
except BaseException:
|
||||
pass
|
||||
for log in self._logs:
|
||||
if not log._closed:
|
||||
log.fail(err)
|
||||
for ch in self._channels:
|
||||
if not ch._log._closed:
|
||||
ch._fail(err)
|
||||
if not self._events._closed:
|
||||
self._events.fail(err)
|
||||
|
||||
def _collect_scheduled_tasks(self) -> list[asyncio.Task[Any]]:
|
||||
"""Return a snapshot of in-flight tasks scheduled via transformers."""
|
||||
return [
|
||||
task
|
||||
for transformer in self._transformers
|
||||
for task in getattr(transformer, "_stream_scheduled_tasks", ())
|
||||
if not task.done()
|
||||
]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Binding and StreamChannel auto-wiring
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _bind_and_wire(
|
||||
self, projection: dict[str, Any], *, native: bool = False
|
||||
) -> None:
|
||||
"""Bind and wire EventLog / StreamChannel instances in a projection.
|
||||
|
||||
Args:
|
||||
projection: The projection dict returned by a transformer's
|
||||
`init()`.
|
||||
native: True when the owning transformer is `_native`.
|
||||
Channels owned by a native transformer use the channel
|
||||
name directly as the protocol method; user-defined
|
||||
channels are prefixed with `custom:`.
|
||||
"""
|
||||
for value in projection.values():
|
||||
if isinstance(value, StreamChannel):
|
||||
value._bind(is_async=self.is_async)
|
||||
self._channels.append(value)
|
||||
method = value.name if native else f"custom:{value.name}"
|
||||
|
||||
def _make_forward(method_name: str) -> Callable[[Any], None]:
|
||||
def _forward(item: Any) -> None:
|
||||
self._forward(method_name, item)
|
||||
|
||||
return _forward
|
||||
|
||||
value._wire(_make_forward(method))
|
||||
elif isinstance(value, EventLog):
|
||||
value._bind(is_async=self.is_async)
|
||||
self._logs.append(value)
|
||||
|
||||
def _forward(self, method: str, item: Any) -> None:
|
||||
"""Inject a ProtocolEvent for a StreamChannel push.
|
||||
|
||||
Forwarded events bypass the transformer pipeline to avoid
|
||||
infinite recursion (a transformer that pushes to a channel
|
||||
during `process()` would re-trigger itself). These events are
|
||||
visible in this mux's main event log but are not passed through
|
||||
transformers' `process()` methods. Only the root mux assigns
|
||||
`seq` to forwarded channel events.
|
||||
|
||||
Args:
|
||||
method: The full protocol method (already with or without
|
||||
the `custom:` prefix; resolved by `_bind_and_wire`).
|
||||
item: The payload pushed onto the channel.
|
||||
"""
|
||||
event: ProtocolEvent = {
|
||||
"type": "event",
|
||||
"method": method,
|
||||
"params": {
|
||||
"namespace": [],
|
||||
"timestamp": int(time.time() * 1000),
|
||||
"data": item,
|
||||
},
|
||||
}
|
||||
if self._assign_seq:
|
||||
self._seq += 1
|
||||
event["seq"] = self._seq
|
||||
self._events.push(event)
|
||||
@@ -0,0 +1,312 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Coroutine
|
||||
from typing import Any, ClassVar, Literal
|
||||
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class _ProtocolEventParams(TypedDict):
|
||||
"""Parameters for a protocol event.
|
||||
|
||||
`timestamp` is wall-clock milliseconds since the epoch and can go
|
||||
backwards across NTP adjustments — use `ProtocolEvent.seq` for
|
||||
ordering.
|
||||
"""
|
||||
|
||||
namespace: list[str]
|
||||
timestamp: int
|
||||
data: Any
|
||||
interrupts: NotRequired[tuple[Any, ...]]
|
||||
|
||||
|
||||
class ProtocolEvent(TypedDict):
|
||||
"""A protocol event emitted by the streaming infrastructure.
|
||||
|
||||
Wraps a raw stream part (values, messages, custom, etc.) in a uniform
|
||||
envelope with a monotonic sequence number assigned by the root StreamMux.
|
||||
Consumers that need a total order across root events should use `seq`, not
|
||||
`params.timestamp` (which is wall-clock and not monotonic).
|
||||
"""
|
||||
|
||||
type: Literal["event"]
|
||||
eventId: NotRequired[str]
|
||||
seq: NotRequired[int]
|
||||
method: str # StreamMode value: "values", "messages", "custom", etc.
|
||||
params: _ProtocolEventParams
|
||||
|
||||
|
||||
class StreamTransformer(ABC):
|
||||
"""Extension point for custom stream projections.
|
||||
|
||||
Transformers observe protocol events flowing through the StreamMux and
|
||||
build typed derived projections (EventLogs, StreamChannels, promises,
|
||||
etc.).
|
||||
|
||||
Set `_native = True` on a transformer to have its projection keys
|
||||
exposed as direct attributes on the run stream (in addition to
|
||||
appearing in `run.extensions`).
|
||||
|
||||
Subclasses must implement `init` and override at least one of
|
||||
`process` / `aprocess`. The `finalize` / `afinalize` and `fail` /
|
||||
`afail` hooks are optional — the default implementations are no-ops.
|
||||
EventLog and StreamChannel instances in the projection dict are
|
||||
auto-closed / auto-failed by the mux, so most transformers don't
|
||||
need `finalize` or `fail` at all.
|
||||
|
||||
Transformers that need async work pick the async lane by:
|
||||
|
||||
1. Overriding `aprocess` (and optionally `afinalize` / `afail`), or
|
||||
2. Calling `self.schedule(coro)` from inside a sync `process`, or
|
||||
3. Setting `requires_async = True` explicitly.
|
||||
|
||||
The mux detects these cases at registration and raises if they're
|
||||
used under sync `stream()` — they only work under `astream()`.
|
||||
|
||||
Use `aprocess` when the pump must wait for async work before the
|
||||
next transformer sees the event (e.g. PII redaction that mutates
|
||||
`event` in place). Use `schedule()` for decoupled async work whose
|
||||
result lands on an independent projection (e.g. async moderation
|
||||
scoring, cost lookup, external tracing).
|
||||
|
||||
Attributes:
|
||||
scope: Namespace the transformer operates within — `()` for the
|
||||
root mux. Set at construction from the mux's scope (each
|
||||
factory is called as `factory(scope)`).
|
||||
requires_async: Explicit opt-in for transformers that need a
|
||||
running event loop but don't override any async method (for
|
||||
example, transformers that call `schedule()` from a sync
|
||||
`process`). The mux also auto-detects the async lane when
|
||||
`aprocess`, `afinalize`, or `afail` is overridden.
|
||||
supports_sync: Set True only for transformers that override
|
||||
async-lane hooks while still fully supporting the sync lane.
|
||||
Such transformers may be registered under `stream()`.
|
||||
required_stream_modes: Stream modes the graph must emit for
|
||||
this transformer to have anything to process. Computed as
|
||||
the union across all registered transformers to determine
|
||||
which modes a `stream_v2` run requests from the graph.
|
||||
Empty tuple means the transformer consumes only synthetic
|
||||
events (or is purely passive).
|
||||
"""
|
||||
|
||||
requires_async: ClassVar[bool] = False
|
||||
supports_sync: ClassVar[bool] = False
|
||||
required_stream_modes: ClassVar[tuple[str, ...]] = ()
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
"""Initialize the transformer with its mux's scope.
|
||||
|
||||
Args:
|
||||
scope: The namespace tuple the owning mux is scoped to.
|
||||
`()` for the root. Factories receive this at
|
||||
construction time (`factory(scope)` in `StreamMux`).
|
||||
"""
|
||||
self.scope: tuple[str, ...] = scope
|
||||
|
||||
@abstractmethod
|
||||
def init(self) -> dict[str, Any]:
|
||||
"""Return the projection dict.
|
||||
|
||||
Keys become entries in `run.extensions`. If the transformer has
|
||||
`_native = True`, keys are also set as direct attributes on the
|
||||
run stream.
|
||||
|
||||
StreamChannel instances in the return value are automatically
|
||||
wired by the StreamMux for protocol event auto-forwarding.
|
||||
"""
|
||||
...
|
||||
|
||||
def _on_register(self, mux: Any) -> None:
|
||||
"""Called by `StreamMux._register` after this transformer is wired in.
|
||||
|
||||
Default is a no-op. Override to capture a reference to the
|
||||
owning mux — needed for transformers that build mini-muxes
|
||||
via `mux._make_child(...)` (e.g. `SubgraphTransformer`).
|
||||
"""
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
"""Handle an event on the sync lane.
|
||||
|
||||
Called for every event before it is appended to the main event
|
||||
log. Subclasses must override either `process` or `aprocess`.
|
||||
The default raises so a missing override fails loudly rather
|
||||
than silently passing every event through.
|
||||
|
||||
Args:
|
||||
event: The protocol event to observe.
|
||||
|
||||
Returns:
|
||||
True to keep the event in the main log, False to suppress it.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} must override process() or aprocess()"
|
||||
)
|
||||
|
||||
async def aprocess(self, event: ProtocolEvent) -> bool:
|
||||
"""Handle an event on the async lane.
|
||||
|
||||
The mux awaits this before dispatching to the next transformer,
|
||||
so a slow `aprocess` serializes the pipeline. Use it only when
|
||||
a later transformer — or a consumer reading the event
|
||||
synchronously — must see the result of the async work (e.g.
|
||||
PII redaction that mutates `event` in place).
|
||||
|
||||
The default delegates to `process`, so purely-sync transformers
|
||||
run unchanged under `astream()`.
|
||||
|
||||
Args:
|
||||
event: The protocol event to observe.
|
||||
|
||||
Returns:
|
||||
True to keep the event in the main log, False to suppress it.
|
||||
"""
|
||||
return self.process(event)
|
||||
|
||||
def finalize(self) -> None:
|
||||
"""Called when the run ends normally (sync lane).
|
||||
|
||||
Override to close EventLogs, resolve promises, or perform other
|
||||
teardown. StreamChannel instances are auto-closed by the mux.
|
||||
"""
|
||||
|
||||
async def afinalize(self) -> None:
|
||||
"""Called when the run ends normally (async lane).
|
||||
|
||||
By the time this runs, the mux has already awaited every task
|
||||
started via `schedule()`, so EventLogs can be closed here
|
||||
without a last-task-wins race.
|
||||
|
||||
The default delegates to `finalize`.
|
||||
"""
|
||||
self.finalize()
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Called when the run ends with an error (sync lane).
|
||||
|
||||
Override to fail EventLogs, reject promises, or perform other
|
||||
teardown. StreamChannel instances are auto-failed by the mux.
|
||||
|
||||
Args:
|
||||
err: The exception that ended the run.
|
||||
"""
|
||||
|
||||
async def afail(self, err: BaseException) -> None:
|
||||
"""Called when the run ends with an error (async lane).
|
||||
|
||||
The mux cancels and awaits every task started via `schedule()`
|
||||
before calling this, so cleanup doesn't race with in-flight work.
|
||||
|
||||
The default delegates to `fail`.
|
||||
|
||||
Args:
|
||||
err: The exception that ended the run.
|
||||
"""
|
||||
self.fail(err)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Scheduled async work
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def schedule(
|
||||
self,
|
||||
coro: Coroutine[Any, Any, Any],
|
||||
*,
|
||||
on_error: Literal["log", "raise"] = "log",
|
||||
) -> asyncio.Task[Any]:
|
||||
"""Schedule a coroutine tied to this transformer's lifecycle.
|
||||
|
||||
The mux holds the task reference, awaits all scheduled tasks
|
||||
during `aclose()` before calling `afinalize()`, and cancels
|
||||
them on `afail()`. Authors don't need to track tasks or
|
||||
implement the last-task-closes-the-log dance.
|
||||
|
||||
Requires a running event loop — call only under `astream()`.
|
||||
Set `requires_async = True` on the class so registration under
|
||||
sync `stream()` fails fast with a clear message.
|
||||
|
||||
Args:
|
||||
coro: The coroutine to run. Its lifecycle is owned by the
|
||||
mux from this point on.
|
||||
on_error: `"log"` (default) catches and logs any exception
|
||||
the coroutine raises, so a single failure doesn't tear
|
||||
down the run. `"raise"` lets the exception propagate
|
||||
when the mux joins pendings, converting the close path
|
||||
into the fail path.
|
||||
|
||||
Returns:
|
||||
The asyncio Task. Authors rarely need to await it directly
|
||||
— consumers read results from whatever projection the
|
||||
coroutine pushes into.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If called without a running event loop (i.e.
|
||||
under sync `stream()` rather than `astream()`).
|
||||
"""
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
raise RuntimeError(
|
||||
f"{type(self).__name__}.schedule() requires a running "
|
||||
"event loop; this transformer must run under astream(), "
|
||||
"not stream(). Set requires_async=True on the class so "
|
||||
"this fails at registration rather than at first event."
|
||||
) from None
|
||||
|
||||
wrapped = self._wrap_scheduled(coro) if on_error == "log" else coro
|
||||
task = asyncio.create_task(wrapped)
|
||||
tasks = self._scheduled_task_set()
|
||||
tasks.add(task)
|
||||
task.add_done_callback(tasks.discard)
|
||||
return task
|
||||
|
||||
@staticmethod
|
||||
async def _wrap_scheduled(coro: Coroutine[Any, Any, Any]) -> Any:
|
||||
try:
|
||||
return await coro
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except BaseException:
|
||||
_logger.exception("Scheduled StreamTransformer task failed")
|
||||
|
||||
def _scheduled_task_set(self) -> set[asyncio.Task[Any]]:
|
||||
"""Return the lazily-allocated task set.
|
||||
|
||||
Avoids requiring subclasses to call `super().__init__()`.
|
||||
"""
|
||||
tasks: set[asyncio.Task[Any]] | None = getattr(
|
||||
self, "_stream_scheduled_tasks", None
|
||||
)
|
||||
if tasks is None:
|
||||
tasks = set()
|
||||
self._stream_scheduled_tasks = tasks
|
||||
return tasks
|
||||
|
||||
|
||||
def transformer_requires_async(transformer: StreamTransformer) -> bool:
|
||||
"""Return True if the transformer needs a running event loop.
|
||||
|
||||
A transformer requires async if it explicitly opts in
|
||||
(`requires_async = True`) or overrides any of the async-lane methods
|
||||
(`aprocess`, `afinalize`, `afail`) without also declaring that it
|
||||
supports the sync lane.
|
||||
|
||||
Args:
|
||||
transformer: The transformer to inspect.
|
||||
|
||||
Returns:
|
||||
True if the transformer cannot run under sync `stream()`.
|
||||
"""
|
||||
if transformer.requires_async:
|
||||
return True
|
||||
if transformer.supports_sync:
|
||||
return False
|
||||
cls = type(transformer)
|
||||
for name in ("aprocess", "afinalize", "afail"):
|
||||
if getattr(cls, name) is not getattr(StreamTransformer, name):
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,521 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping
|
||||
from types import MappingProxyType, TracebackType
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from langgraph.stream._convert import convert_to_protocol_event
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import ProtocolEvent
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langgraph.stream.transformers import SubgraphStatus, ValuesTransformer
|
||||
|
||||
|
||||
def _drive_until_done(pump: Callable[[], bool]) -> None:
|
||||
"""Call the sync pump until it returns False."""
|
||||
while pump():
|
||||
pass
|
||||
|
||||
|
||||
async def _adrive_until_done(pump: Callable[[], Awaitable[bool]]) -> None:
|
||||
"""Call the async pump until it returns False."""
|
||||
while await pump():
|
||||
pass
|
||||
|
||||
|
||||
class GraphRunStream:
|
||||
"""Sync run stream with caller-driven pumping.
|
||||
|
||||
The caller's iteration on any projection (`values`, `messages`,
|
||||
raw events, or `output`) drives the graph forward. No background
|
||||
thread is used — the caller's `for` loop is the pump.
|
||||
|
||||
Projections are single-consumer — iterating `run.values` twice
|
||||
raises. Use `projection.tee(n)` if you genuinely need fan-out.
|
||||
|
||||
All transformer projections live in `extensions`. Native transformer
|
||||
projections (those with `_native = True`) are also set as direct
|
||||
attributes on this instance (e.g. `run.values`, `run.messages`).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
graph_iter: Iterator[Any] | None,
|
||||
mux: StreamMux,
|
||||
values_transformer: ValuesTransformer,
|
||||
*,
|
||||
wire_pump: bool = True,
|
||||
) -> None:
|
||||
"""Initialize the run stream.
|
||||
|
||||
Args:
|
||||
graph_iter: Pull-based iterator over the graph's stream,
|
||||
or `None` for nested run streams whose pump is driven
|
||||
by an outer run (e.g. `SubgraphRunStream`).
|
||||
mux: The StreamMux owning projections and the main log.
|
||||
values_transformer: The built-in values transformer
|
||||
providing `output` / `interrupted` / `interrupts`.
|
||||
wire_pump: When True (default), bind `_pump_next` as the
|
||||
mux's pump callable. Subclasses that inherit a parent
|
||||
pump via `StreamMux._make_child` should pass False to
|
||||
preserve the parent binding.
|
||||
"""
|
||||
self._graph_iter = graph_iter
|
||||
self._mux = mux
|
||||
self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
|
||||
self._values_transformer = values_transformer
|
||||
self._exhausted = False
|
||||
for key in mux.native_keys:
|
||||
setattr(self, key, mux.extensions[key])
|
||||
if wire_pump:
|
||||
self._wire_request_more(mux)
|
||||
|
||||
def _wire_request_more(self, mux: StreamMux) -> None:
|
||||
"""Wire the sync pull callback through the mux.
|
||||
|
||||
Routing through `mux.bind_pump` (rather than walking
|
||||
projections directly here) lets child mini-muxes built by
|
||||
`mux._make_child(...)` inherit the same pump callable, so
|
||||
cursors on a subgraph handle's projections drive the root
|
||||
pump just like cursors on `run.values` do.
|
||||
"""
|
||||
mux.bind_pump(self._pump_next)
|
||||
|
||||
def _pump_next(self) -> bool:
|
||||
"""Pull one event from the graph and push it through the mux.
|
||||
|
||||
Returns:
|
||||
True if an event was pulled, False if the graph is exhausted
|
||||
or has raised. Always False when constructed with
|
||||
`graph_iter=None` (the run is driven by an outer pump).
|
||||
"""
|
||||
if self._exhausted or self._graph_iter is None:
|
||||
return False
|
||||
try:
|
||||
part = next(self._graph_iter)
|
||||
self._mux.push(convert_to_protocol_event(part))
|
||||
return True
|
||||
except StopIteration:
|
||||
self._mux.close()
|
||||
self._exhausted = True
|
||||
return False
|
||||
except Exception as e:
|
||||
self._mux.fail(e)
|
||||
self._exhausted = True
|
||||
return False
|
||||
|
||||
def abort(self) -> None:
|
||||
"""Stop the run early.
|
||||
|
||||
Closes the mux and marks the stream exhausted. The graph
|
||||
iterator is dropped; any in-flight nodes see the closure on
|
||||
their next yield point. Idempotent.
|
||||
"""
|
||||
if self._exhausted:
|
||||
return
|
||||
self._exhausted = True
|
||||
try:
|
||||
self._mux.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def __enter__(self) -> GraphRunStream:
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> None:
|
||||
self.abort()
|
||||
|
||||
@property
|
||||
def output(self) -> dict[str, Any] | None:
|
||||
"""Drive the run to completion and return the final state."""
|
||||
_drive_until_done(self._pump_next)
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._latest
|
||||
|
||||
@property
|
||||
def interrupted(self) -> bool:
|
||||
"""Drive the run to completion, then return whether it was
|
||||
interrupted.
|
||||
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
_drive_until_done(self._pump_next)
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._interrupted
|
||||
|
||||
@property
|
||||
def interrupts(self) -> list[Any]:
|
||||
"""Drive the run to completion, then return interrupt payloads.
|
||||
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
_drive_until_done(self._pump_next)
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._interrupts
|
||||
|
||||
def __iter__(self) -> Iterator[ProtocolEvent]:
|
||||
"""Subscribe to the main event log and iterate protocol events."""
|
||||
return iter(self._mux._events)
|
||||
|
||||
def interleave(self, *names: str) -> Iterator[tuple[str, Any]]:
|
||||
"""Iterate multiple projections round-robin, yielding ``(name, item)``.
|
||||
|
||||
Each turn advances one projection's cursor; when a cursor's buffer
|
||||
is empty, pulling from it drives the pump once, which fans out to
|
||||
every subscribed projection log. Projections whose items aren't
|
||||
consumed on this turn sit in their own buffers only until the next
|
||||
turn reaches them, bounding memory by the skew between projection
|
||||
rates rather than letting any single log grow to the full run
|
||||
length.
|
||||
|
||||
Projections are exhausted independently; a projection that finishes
|
||||
early drops out of the rotation while others continue. The overall
|
||||
iterator ends once all named projections are done.
|
||||
|
||||
Args:
|
||||
*names: Projection keys to interleave. Must match keys in
|
||||
``extensions``.
|
||||
|
||||
Yields:
|
||||
``(name, item)`` tuples in round-robin order across the named
|
||||
projections.
|
||||
|
||||
Raises:
|
||||
KeyError: If a name doesn't match a registered projection.
|
||||
|
||||
Example:
|
||||
```python
|
||||
for name, item in run.interleave("messages", "values"):
|
||||
if name == "messages":
|
||||
print("msg:", item)
|
||||
else:
|
||||
print("val:", item)
|
||||
```
|
||||
"""
|
||||
cursors: dict[str, Iterator[Any]] = {
|
||||
name: iter(self.extensions[name]) for name in names
|
||||
}
|
||||
done: set[str] = set()
|
||||
while len(done) < len(cursors):
|
||||
for name, cursor in cursors.items():
|
||||
if name in done:
|
||||
continue
|
||||
try:
|
||||
item = next(cursor)
|
||||
except StopIteration:
|
||||
done.add(name)
|
||||
continue
|
||||
yield (name, item)
|
||||
|
||||
|
||||
class AsyncGraphRunStream:
|
||||
"""Async run stream with caller-driven pumping.
|
||||
|
||||
Async iteration on any projection drives the graph forward — there
|
||||
is no background task. Concurrent consumers share a single-flight
|
||||
pump via an `asyncio.Lock`, so each awaiting cursor contributes one
|
||||
event per acquisition. Backpressure comes from the logs: when a
|
||||
subscribed log's buffer reaches `maxlen`, `apush` awaits the
|
||||
subscriber to drain, which holds back the pump and paces the graph.
|
||||
|
||||
Projections are single-consumer — a second `aiter(run.values)`
|
||||
raises. Use `projection.tee(n)` for fan-out.
|
||||
|
||||
Use as an async context manager to guarantee clean shutdown on
|
||||
early exit:
|
||||
|
||||
```python
|
||||
async with await handler.astream(input) as run:
|
||||
async for msg in run.messages:
|
||||
...
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
graph_aiter: AsyncIterator[Any] | None,
|
||||
mux: StreamMux,
|
||||
values_transformer: ValuesTransformer,
|
||||
*,
|
||||
wire_pump: bool = True,
|
||||
) -> None:
|
||||
"""Initialize the async run stream.
|
||||
|
||||
Args:
|
||||
graph_aiter: Async iterator over the graph's stream, or
|
||||
`None` for nested run streams whose pump is driven by
|
||||
an outer run (e.g. `AsyncSubgraphRunStream`).
|
||||
mux: The StreamMux owning projections and the main log.
|
||||
values_transformer: The built-in values transformer
|
||||
providing `output` / `interrupted` / `interrupts`.
|
||||
wire_pump: When True (default), bind `_apump_next` as the
|
||||
mux's async pump callable. Subclasses that inherit a
|
||||
parent pump via `StreamMux._make_child` should pass
|
||||
False to preserve the parent binding.
|
||||
"""
|
||||
self._graph_aiter = graph_aiter
|
||||
self._mux = mux
|
||||
self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
|
||||
self._values_transformer = values_transformer
|
||||
self._exhausted = False
|
||||
self._pump_cond = asyncio.Condition()
|
||||
self._pumping = False
|
||||
for key in mux.native_keys:
|
||||
setattr(self, key, mux.extensions[key])
|
||||
if wire_pump:
|
||||
self._wire_arequest_more(mux)
|
||||
|
||||
def _wire_arequest_more(self, mux: StreamMux) -> None:
|
||||
"""Wire the async pull callback through the mux.
|
||||
|
||||
Mirrors `_wire_request_more`: routing through
|
||||
`mux.bind_apump` lets child mini-muxes inherit the pump
|
||||
callable so cursors on subgraph handles drive the root
|
||||
pump.
|
||||
"""
|
||||
mux.bind_apump(self._apump_next)
|
||||
|
||||
async def _apump_next(self) -> bool:
|
||||
"""Drive one pump step, or wait for the active pumper to drive one.
|
||||
|
||||
"Take-a-number" semantics: at most one task at a time calls
|
||||
`graph_aiter.__anext__()` (asyncio iterators can't be advanced
|
||||
concurrently). Other callers wait on a Condition that the
|
||||
active pumper notifies after each step. This lets a "passive"
|
||||
consumer — one whose projection's buffer is being filled by the
|
||||
active pumper's push — wake up as soon as its data lands,
|
||||
instead of queueing on the pump and only observing its data one
|
||||
graph event late.
|
||||
|
||||
`except Exception` is intentional — `CancelledError` and other
|
||||
`BaseException` subclasses propagate, matching asyncio's
|
||||
cancellation contract.
|
||||
|
||||
Returns:
|
||||
True if a pump step completed (by this task or another),
|
||||
False if the graph is exhausted.
|
||||
"""
|
||||
async with self._pump_cond:
|
||||
if self._exhausted or self._graph_aiter is None:
|
||||
return False
|
||||
if self._pumping:
|
||||
# Another task is pumping; wait for its progress signal.
|
||||
await self._pump_cond.wait()
|
||||
return not self._exhausted
|
||||
self._pumping = True
|
||||
|
||||
try:
|
||||
try:
|
||||
part = await self._graph_aiter.__anext__()
|
||||
await self._mux.apush(convert_to_protocol_event(part))
|
||||
return True
|
||||
except StopAsyncIteration:
|
||||
self._exhausted = True
|
||||
await self._mux.aclose()
|
||||
return False
|
||||
except Exception as e:
|
||||
self._exhausted = True
|
||||
await self._mux.afail(e)
|
||||
return False
|
||||
finally:
|
||||
async with self._pump_cond:
|
||||
self._pumping = False
|
||||
self._pump_cond.notify_all()
|
||||
|
||||
async def abort(self) -> None:
|
||||
"""Stop the run early.
|
||||
|
||||
Marks the stream exhausted, wakes any pump-waiters, and closes
|
||||
the mux. Any `apush` blocked on backpressure wakes and returns
|
||||
without appending. Idempotent.
|
||||
"""
|
||||
async with self._pump_cond:
|
||||
if self._exhausted:
|
||||
return
|
||||
self._exhausted = True
|
||||
self._pump_cond.notify_all()
|
||||
try:
|
||||
await self._mux.aclose()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def __aenter__(self) -> AsyncGraphRunStream:
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> None:
|
||||
await self.abort()
|
||||
|
||||
async def output(self) -> dict[str, Any] | None:
|
||||
"""Drive the run to completion and return the final state.
|
||||
|
||||
Methods (not properties) on the async lane so `run.output`
|
||||
without `await` raises at type-check time instead of silently
|
||||
yielding a coroutine object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
output = await run.output()
|
||||
```
|
||||
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
await _adrive_until_done(self._apump_next)
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._latest
|
||||
|
||||
async def interrupted(self) -> bool:
|
||||
"""Drive the run to completion and return whether it was
|
||||
interrupted.
|
||||
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
await _adrive_until_done(self._apump_next)
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._interrupted
|
||||
|
||||
async def interrupts(self) -> list[Any]:
|
||||
"""Drive the run to completion and return interrupt payloads.
|
||||
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
"""
|
||||
await _adrive_until_done(self._apump_next)
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._interrupts
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
|
||||
"""Subscribe to the main event log and iterate protocol events."""
|
||||
return self._mux._events.__aiter__()
|
||||
|
||||
|
||||
class _SubgraphRunStreamMixin:
|
||||
"""Subgraph metadata + parent-pump delegation shared by both lanes.
|
||||
|
||||
Inherits from `GraphRunStream` (or `AsyncGraphRunStream`) with
|
||||
`graph_iter=None` + `wire_pump=False` — the mini-mux is driven
|
||||
by the parent's pump (inherited via `StreamMux._make_child`), and
|
||||
the handle never pulls upstream itself. Pump-driving methods
|
||||
delegate to the parent pump so `handle.output` and friends drive
|
||||
the root run.
|
||||
|
||||
Subclasses set the parent pump function captured at construction
|
||||
(`_parent_pump_fn` / `_parent_apump_fn`) and override
|
||||
`_pump_next` / `_apump_next` to delegate to it.
|
||||
|
||||
Status is updated in place by `SubgraphTransformer`. Iterate
|
||||
`run.subgraphs` to receive handles as subgraphs spawn, then
|
||||
drill into projections inside the loop body **before** the next
|
||||
pump cycle — same lazy-subscribe constraint as root projections.
|
||||
"""
|
||||
|
||||
path: tuple[str, ...]
|
||||
graph_name: str | None
|
||||
trigger_call_id: str | None
|
||||
status: SubgraphStatus
|
||||
error: str | None
|
||||
_seen_terminal: bool
|
||||
|
||||
|
||||
class SubgraphRunStream(GraphRunStream, _SubgraphRunStreamMixin):
|
||||
"""Sync handle for a discovered subgraph (extends `GraphRunStream`)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mux: StreamMux,
|
||||
values_transformer: ValuesTransformer,
|
||||
*,
|
||||
path: tuple[str, ...],
|
||||
graph_name: str | None = None,
|
||||
trigger_call_id: str | None = None,
|
||||
) -> None:
|
||||
# Capture the parent-inherited pump before super().__init__
|
||||
# touches anything; we delegate to it from `_pump_next`.
|
||||
self._parent_pump_fn: Callable[[], bool] | None = mux._pump_fn
|
||||
super().__init__(
|
||||
graph_iter=None,
|
||||
mux=mux,
|
||||
values_transformer=values_transformer,
|
||||
wire_pump=False,
|
||||
)
|
||||
self.path = path
|
||||
self.graph_name = graph_name
|
||||
self.trigger_call_id = trigger_call_id
|
||||
self.status = "started"
|
||||
self.error = None
|
||||
self._seen_terminal = False
|
||||
|
||||
def _pump_next(self) -> bool:
|
||||
"""Delegate to the parent's pump.
|
||||
|
||||
Cursors on this handle's projections call here when their
|
||||
buffers empty. Driving the parent fans events into our
|
||||
mini-mux, transparently advancing the whole run.
|
||||
"""
|
||||
if (
|
||||
self._exhausted
|
||||
or self._seen_terminal
|
||||
or self._mux._events._closed
|
||||
or self._parent_pump_fn is None
|
||||
):
|
||||
return False
|
||||
return self._parent_pump_fn()
|
||||
|
||||
|
||||
class AsyncSubgraphRunStream(AsyncGraphRunStream, _SubgraphRunStreamMixin):
|
||||
"""Async handle for a discovered subgraph (extends `AsyncGraphRunStream`)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mux: StreamMux,
|
||||
values_transformer: ValuesTransformer,
|
||||
*,
|
||||
path: tuple[str, ...],
|
||||
graph_name: str | None = None,
|
||||
trigger_call_id: str | None = None,
|
||||
) -> None:
|
||||
self._parent_apump_fn: Callable[[], Awaitable[bool]] | None = mux._apump_fn
|
||||
super().__init__(
|
||||
graph_aiter=None,
|
||||
mux=mux,
|
||||
values_transformer=values_transformer,
|
||||
wire_pump=False,
|
||||
)
|
||||
self.path = path
|
||||
self.graph_name = graph_name
|
||||
self.trigger_call_id = trigger_call_id
|
||||
self.status = "started"
|
||||
self.error = None
|
||||
self._seen_terminal = False
|
||||
|
||||
async def _apump_next(self) -> bool:
|
||||
"""Delegate to the parent's async pump."""
|
||||
if (
|
||||
self._exhausted
|
||||
or self._seen_terminal
|
||||
or self._mux._events._closed
|
||||
or self._parent_apump_fn is None
|
||||
):
|
||||
return False
|
||||
return await self._parent_apump_fn()
|
||||
@@ -0,0 +1,113 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Callable, Iterator
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class StreamChannel(Generic[T]):
|
||||
"""A named projection channel with optional protocol auto-forwarding.
|
||||
|
||||
Wraps an event log and declares a protocol channel name. When the
|
||||
StreamMux detects a StreamChannel in a transformer's `init()`
|
||||
return value, it automatically wires every `push()` to inject a
|
||||
`ProtocolEvent` into the main event stream using the channel's
|
||||
name as the method.
|
||||
|
||||
Auto-forwarded events bypass the transformer pipeline — other
|
||||
transformers' `process()` / `aprocess()` methods do not see
|
||||
`custom:<name>` events produced by a channel push. This prevents a
|
||||
transformer that pushes to its own channel during `process()` from
|
||||
re-triggering itself, but it also means filter- or tap-style
|
||||
transformers cannot observe channel output from peer transformers.
|
||||
Consumers that need that should iterate the main event stream.
|
||||
|
||||
In-process consumers iterate the channel directly (`for item in ch`
|
||||
or `async for item in ch`). Remote SDK clients subscribe via
|
||||
`session.subscribe("custom:<channelName>")`.
|
||||
|
||||
Like EventLog, a StreamChannel starts unbound. The mux calls
|
||||
`_bind(is_async)` during registration so the correct iteration
|
||||
protocol is available by the time user code sees it.
|
||||
|
||||
Lifecycle (`_close` / `_fail`) is managed by the mux — transformers
|
||||
using only StreamChannels don't need `finalize` or `fail` hooks.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, *, maxlen: int | None = None) -> None:
|
||||
"""Initialize the channel with an empty inner log.
|
||||
|
||||
Args:
|
||||
name: The protocol channel name used for auto-forwarded
|
||||
events. Surfaced on the wire as `custom:<name>` for
|
||||
user-defined transformers, or as `<name>` for channels
|
||||
owned by a native transformer (`_native = True`). The
|
||||
StreamMux derives the prefix from the owning
|
||||
transformer at registration time.
|
||||
maxlen: Optional retention cap on the inner EventLog. See
|
||||
`EventLog.__init__` for semantics.
|
||||
"""
|
||||
self.name = name
|
||||
self._log: EventLog[T] = EventLog(maxlen=maxlen)
|
||||
self._wire_fn: Callable[[T], None] | None = None
|
||||
|
||||
def _bind(self, *, is_async: bool) -> None:
|
||||
"""Bind the underlying event log to sync or async mode.
|
||||
|
||||
Args:
|
||||
is_async: True for async iteration, False for sync.
|
||||
"""
|
||||
self._log._bind(is_async=is_async)
|
||||
|
||||
def push(self, item: T) -> None:
|
||||
"""Append an item to the log and auto-forward if wired.
|
||||
|
||||
Args:
|
||||
item: The item to push.
|
||||
"""
|
||||
self._log.push(item)
|
||||
if self._wire_fn is not None:
|
||||
self._wire_fn(item)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Mux lifecycle hooks (not called by transformers directly)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _wire(self, fn: Callable[[T], None]) -> None:
|
||||
"""Install the auto-forward callback (called by StreamMux)."""
|
||||
self._wire_fn = fn
|
||||
|
||||
def _close(self) -> None:
|
||||
"""Close the underlying log (called by StreamMux on run end)."""
|
||||
self._log.close()
|
||||
|
||||
def _fail(self, err: BaseException) -> None:
|
||||
"""Fail the underlying log (called by StreamMux on run error)."""
|
||||
self._log.fail(err)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Iteration — delegates to the inner event log (multi-cursor)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def __iter__(self) -> Iterator[T]:
|
||||
return iter(self._log)
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[T]:
|
||||
return self._log.__aiter__()
|
||||
|
||||
def tee(self, n: int = 2) -> tuple[Iterator[T], ...]:
|
||||
"""Fan out the channel into `n` independent sync iterators.
|
||||
|
||||
Delegates to the underlying EventLog's `tee()`.
|
||||
"""
|
||||
return self._log.tee(n)
|
||||
|
||||
def atee(self, n: int = 2) -> tuple[AsyncIterator[T], ...]:
|
||||
"""Fan out the channel into `n` independent async iterators.
|
||||
|
||||
Delegates to the underlying EventLog's `atee()`.
|
||||
"""
|
||||
return self._log.atee(n)
|
||||
@@ -0,0 +1,748 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Literal, 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 typing_extensions import NotRequired, TypedDict
|
||||
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.run_stream import AsyncSubgraphRunStream, SubgraphRunStream
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from langgraph.stream._mux import StreamMux
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ValuesTransformer(StreamTransformer):
|
||||
"""Capture values events as a drainable stream of state snapshots.
|
||||
|
||||
Keeps `_latest` / `_interrupted` / `_interrupts` as scalar state
|
||||
regardless of whether the log has a subscriber — so `run.output()`
|
||||
and `run.interrupted` work without forcing the caller to iterate
|
||||
`run.values`. Log pushes are silent no-ops when unsubscribed.
|
||||
|
||||
Native transformer — projection keys are exposed as direct
|
||||
attributes on the run stream (e.g. `run.values`).
|
||||
|
||||
Only values events at the run's own level are captured; snapshots
|
||||
from deeper subgraphs are left in the main event log but excluded
|
||||
from the projection. "Own level" is defined by `scope`, which
|
||||
`stream_v2` / `astream_v2` populate from the caller's checkpoint
|
||||
namespace so that a nested `stream_v2` call still sees its own
|
||||
root snapshots.
|
||||
"""
|
||||
|
||||
_native = True
|
||||
required_stream_modes = ("values",)
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._log: EventLog[dict[str, Any]] = EventLog()
|
||||
self._latest: dict[str, Any] | None = None
|
||||
self._interrupted = False
|
||||
self._interrupts: list[Any] = []
|
||||
# Cached as a list once for cheap equality with the protocol
|
||||
# event's `namespace` field, which is `list[str]`.
|
||||
self._scope_list: list[str] = list(scope)
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"values": self._log}
|
||||
|
||||
@property
|
||||
def error(self) -> BaseException | None:
|
||||
"""The error that ended the run, or `None` if it succeeded.
|
||||
|
||||
Set by the mux when it auto-fails the projection log.
|
||||
"""
|
||||
return self._log._error
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] != "values":
|
||||
return True
|
||||
params = event["params"]
|
||||
if params["namespace"] != self._scope_list:
|
||||
return True
|
||||
self._latest = params["data"]
|
||||
interrupts = params.get("interrupts", ())
|
||||
if interrupts:
|
||||
self._interrupted = True
|
||||
self._interrupts.extend(interrupts)
|
||||
self._log.push(params["data"])
|
||||
return True
|
||||
|
||||
|
||||
class MessagesTransformer(StreamTransformer):
|
||||
"""Capture messages events as ChatModelStream objects.
|
||||
|
||||
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.
|
||||
|
||||
Two input shapes are handled (via `params["data"] = (payload,
|
||||
metadata)` from `StreamMessagesHandler`):
|
||||
|
||||
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 events at the run's own level are projected; tokens from
|
||||
deeper subgraphs are left in the main event log but excluded from
|
||||
`.messages`. "Own level" is defined by `scope`, which
|
||||
`stream_v2` / `astream_v2` populate from the caller's checkpoint
|
||||
namespace so that a `stream_v2` call inside a node still sees its
|
||||
own root chat model streams on `.messages`. 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
|
||||
required_stream_modes = ("messages",)
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
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
|
||||
self._apump_fn: Callable[[], Awaitable[bool]] | None = None
|
||||
# Cached as a list once for cheap equality with the protocol
|
||||
# event's `namespace` field, which is `list[str]`.
|
||||
self._scope_list: list[str] = list(scope)
|
||||
|
||||
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 _bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:
|
||||
"""Wire the async pull callback.
|
||||
|
||||
Called by `AsyncGraphRunStream._wire_arequest_more` so each
|
||||
`AsyncChatModelStream` this transformer creates can drive the
|
||||
shared graph pump from its projection cursors.
|
||||
"""
|
||||
self._apump_fn = fn
|
||||
|
||||
def _make_stream(
|
||||
self,
|
||||
*,
|
||||
namespace: list[str],
|
||||
node: str | None,
|
||||
message_id: str | None,
|
||||
) -> ChatModelStream:
|
||||
"""Create a ChatModelStream (sync) or AsyncChatModelStream (async).
|
||||
|
||||
Wires whichever pump is bound. Prefers the async pump so nested
|
||||
iteration under `AsyncGraphRunStream` drives the graph forward
|
||||
without a background task. The unwired fallback (no pump bound)
|
||||
is used by unit tests that dispatch events manually.
|
||||
"""
|
||||
if self._apump_fn is not None:
|
||||
astream = AsyncChatModelStream(
|
||||
namespace=namespace,
|
||||
node=node,
|
||||
message_id=message_id,
|
||||
)
|
||||
astream.set_arequest_more(self._apump_fn)
|
||||
return astream
|
||||
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)
|
||||
return stream
|
||||
return AsyncChatModelStream(
|
||||
namespace=namespace,
|
||||
node=node,
|
||||
message_id=message_id,
|
||||
)
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] != "messages":
|
||||
return True
|
||||
params = event["params"]
|
||||
if params["namespace"] != self._scope_list:
|
||||
return True
|
||||
|
||||
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()
|
||||
|
||||
|
||||
SubgraphStatus = Literal["started", "completed", "failed", "interrupted"]
|
||||
|
||||
|
||||
def _parse_ns_segment(segment: str) -> tuple[str, str | None]:
|
||||
"""Split a namespace segment into `(graph_name, trigger_call_id)`.
|
||||
|
||||
Segments are formatted `node_name:task_id` by `prepare_next_tasks`.
|
||||
Returns `(segment, None)` if no `:` is present.
|
||||
"""
|
||||
name, sep, task_id = segment.partition(":")
|
||||
return name, task_id if sep else None
|
||||
|
||||
|
||||
class LifecyclePayload(TypedDict, total=False):
|
||||
"""Payload of a lifecycle event surfaced on the `lifecycle` channel.
|
||||
|
||||
Auto-forwarded as `lifecycle` protocol events (no `custom:` prefix
|
||||
because `LifecycleTransformer` is a native transformer) so remote
|
||||
SDK clients receive the same data in-process consumers see via
|
||||
`run.lifecycle`.
|
||||
"""
|
||||
|
||||
event: SubgraphStatus
|
||||
namespace: list[str]
|
||||
graph_name: NotRequired[str]
|
||||
trigger_call_id: NotRequired[str]
|
||||
error: NotRequired[str]
|
||||
|
||||
|
||||
class _TasksLifecycleBase(StreamTransformer):
|
||||
"""Shared bookkeeping for `tasks`-event-driven lifecycle inference.
|
||||
|
||||
Both `LifecycleTransformer` (wire-serializable channel) and
|
||||
`SubgraphTransformer` (in-process navigation handles) discover
|
||||
subgraphs by watching the same `tasks` stream — `started` on the
|
||||
first event at a tracked namespace, terminal status when the
|
||||
parent's `TaskResultPayload` arrives. Centralizing the dispatch
|
||||
+ open-set bookkeeping here keeps the inference rules from
|
||||
drifting between the two surfaces.
|
||||
|
||||
Subclasses provide three template-method hooks:
|
||||
|
||||
- `_should_track(ns)` — scope filter (e.g. multi-depth vs
|
||||
direct-children-only).
|
||||
- `_on_started(ns, graph_name, trigger_call_id)` — first sighting
|
||||
action (push payload / build handle / etc.). Called once per
|
||||
discovered namespace.
|
||||
- `_on_terminal(ns, status, error)` — terminal action (push
|
||||
terminal payload / mark handle status). Called once per
|
||||
tracked namespace at result time, or via `finalize` / `fail`
|
||||
sweeps if no parent result arrived.
|
||||
|
||||
Tasks events are suppressed from the main event log (`process`
|
||||
returns False) — they're folded into whichever projection the
|
||||
subclass populates; consumers iterating the raw protocol stream
|
||||
see the higher-level view.
|
||||
"""
|
||||
|
||||
required_stream_modes = ("tasks",)
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._seen: set[tuple[str, ...]] = set()
|
||||
# Maps tracked namespace -> task_id of the parent task whose
|
||||
# `TaskResultPayload` will close it.
|
||||
self._open: dict[tuple[str, ...], str] = {}
|
||||
|
||||
# --- Template-method hooks (subclass overrides) ---
|
||||
|
||||
def _should_track(self, ns: tuple[str, ...]) -> bool:
|
||||
"""Scope filter — return True iff `ns` is in this transformer's region."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _on_started(
|
||||
self,
|
||||
ns: tuple[str, ...],
|
||||
graph_name: str | None,
|
||||
trigger_call_id: str | None,
|
||||
) -> None:
|
||||
"""Fired once per discovered namespace (first observed task event)."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _on_terminal(
|
||||
self,
|
||||
ns: tuple[str, ...],
|
||||
status: SubgraphStatus,
|
||||
error: str | None,
|
||||
) -> None:
|
||||
"""Fired once per tracked namespace when its parent's result arrives,
|
||||
or via finalize/fail safety-net sweeps.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
# --- Dispatch + bookkeeping (shared) ---
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] != "tasks":
|
||||
return True
|
||||
ns = tuple(event["params"]["namespace"])
|
||||
data = event["params"]["data"]
|
||||
if "result" in data:
|
||||
self._handle_task_result(ns, data)
|
||||
else:
|
||||
self._handle_task_start(ns)
|
||||
# Tasks events are folded into the synthesized projections;
|
||||
# suppress from the main event log so iterators don't double-see
|
||||
# the same information in two shapes.
|
||||
return False
|
||||
|
||||
def _handle_task_start(self, ns: tuple[str, ...]) -> None:
|
||||
if not self._should_track(ns) or ns in self._seen:
|
||||
return
|
||||
self._seen.add(ns)
|
||||
graph_name, trigger_call_id = _parse_ns_segment(ns[-1])
|
||||
self._on_started(ns, graph_name or None, trigger_call_id)
|
||||
if trigger_call_id is not None:
|
||||
self._open[ns] = trigger_call_id
|
||||
|
||||
def _pop_terminal_transitions(
|
||||
self, ns: tuple[str, ...], data: dict[str, Any]
|
||||
) -> list[tuple[tuple[str, ...], SubgraphStatus, str | None]]:
|
||||
"""Return and remove tracked children closed by this task result."""
|
||||
result_id = data.get("id")
|
||||
if not result_id:
|
||||
return []
|
||||
transitions: list[tuple[tuple[str, ...], SubgraphStatus, str | None]] = []
|
||||
for child_ns, parent_task_id in list(self._open.items()):
|
||||
if child_ns[:-1] != ns or parent_task_id != result_id:
|
||||
continue
|
||||
status, error = _terminal_from_result(data)
|
||||
transitions.append((child_ns, status, error))
|
||||
del self._open[child_ns]
|
||||
return transitions
|
||||
|
||||
def _handle_task_result(self, ns: tuple[str, ...], data: dict[str, Any]) -> None:
|
||||
for child_ns, status, error in self._pop_terminal_transitions(ns, data):
|
||||
self._on_terminal(child_ns, status, error)
|
||||
|
||||
def finalize(self) -> None:
|
||||
"""Emit `completed` for any tracked namespace still open at run end."""
|
||||
for ns in list(self._open):
|
||||
self._on_terminal(ns, "completed", None)
|
||||
self._open.clear()
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Emit `failed` / `interrupted` for any tracked namespace still open."""
|
||||
is_interrupt = isinstance(err, GraphInterrupt)
|
||||
status: SubgraphStatus = "interrupted" if is_interrupt else "failed"
|
||||
error_str = None if is_interrupt else str(err)
|
||||
for ns in list(self._open):
|
||||
self._on_terminal(ns, status, error_str)
|
||||
self._open.clear()
|
||||
|
||||
|
||||
def _status_from_exception(err: BaseException) -> tuple[SubgraphStatus, str | None]:
|
||||
"""Map a run exception to a subgraph terminal status and error string."""
|
||||
if isinstance(err, GraphInterrupt):
|
||||
return "interrupted", None
|
||||
return "failed", str(err)
|
||||
|
||||
|
||||
def _terminal_from_result(
|
||||
payload: dict[str, Any],
|
||||
) -> tuple[SubgraphStatus, str | None]:
|
||||
"""Map a `TaskResultPayload` to a `(status, error)` pair.
|
||||
|
||||
Order matters: a result with both `error` and `interrupts` prefers
|
||||
the interrupt classification, since `GraphInterrupt` manifests as
|
||||
a populated `interrupts` list, not as `error`.
|
||||
"""
|
||||
if payload.get("interrupts"):
|
||||
return "interrupted", None
|
||||
error = payload.get("error")
|
||||
if error:
|
||||
return "failed", str(error)
|
||||
return "completed", None
|
||||
|
||||
|
||||
class LifecycleTransformer(_TasksLifecycleBase):
|
||||
"""Surface subgraph lifecycle as `lifecycle` protocol events.
|
||||
|
||||
Pushes `LifecyclePayload` to a `StreamChannel` named `lifecycle`.
|
||||
The channel is auto-forwarded by the mux so payloads land in the
|
||||
main event log under `method = "lifecycle"` (native transformer —
|
||||
no `custom:` prefix) — visible to remote SDK clients over the
|
||||
wire and to in-process consumers via `run.lifecycle`.
|
||||
|
||||
Tracks subgraphs at every depth strictly below the transformer's
|
||||
scope, so a graph → subgraph → subgraph chain produces lifecycle
|
||||
events for both nested levels in a flat stream.
|
||||
|
||||
Native transformer — projection key `lifecycle` is exposed as
|
||||
`run.lifecycle`.
|
||||
"""
|
||||
|
||||
_native = True
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._channel: StreamChannel[LifecyclePayload] = StreamChannel("lifecycle")
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"lifecycle": self._channel}
|
||||
|
||||
def _should_track(self, ns: tuple[str, ...]) -> bool:
|
||||
depth = len(self.scope)
|
||||
return len(ns) > depth and ns[:depth] == self.scope
|
||||
|
||||
def _on_started(
|
||||
self,
|
||||
ns: tuple[str, ...],
|
||||
graph_name: str | None,
|
||||
trigger_call_id: str | None,
|
||||
) -> None:
|
||||
if trigger_call_id is None:
|
||||
# Without a task id we can't correlate a parent-result
|
||||
# event back to this namespace — skip the started payload
|
||||
# and rely on finalize/fail to close.
|
||||
return
|
||||
payload: LifecyclePayload = {"event": "started", "namespace": list(ns)}
|
||||
if graph_name:
|
||||
payload["graph_name"] = graph_name
|
||||
payload["trigger_call_id"] = trigger_call_id
|
||||
self._channel.push(payload)
|
||||
|
||||
def _on_terminal(
|
||||
self,
|
||||
ns: tuple[str, ...],
|
||||
status: SubgraphStatus,
|
||||
error: str | None,
|
||||
) -> None:
|
||||
payload: LifecyclePayload = {"event": status, "namespace": list(ns)}
|
||||
if error is not None:
|
||||
payload["error"] = error
|
||||
self._channel.push(payload)
|
||||
|
||||
|
||||
class SubgraphTransformer(_TasksLifecycleBase):
|
||||
"""Discover subgraph invocations as in-process navigation handles.
|
||||
|
||||
Per discovered direct-child subgraph, builds a `SubgraphRunStream`
|
||||
(or `AsyncSubgraphRunStream`) wrapping a child mini-mux scoped to
|
||||
the subgraph's namespace. Consumers iterate `run.subgraphs` to
|
||||
receive handles, then drill into `handle.values` / `handle.messages`
|
||||
/ `handle.subgraphs` (recursive grandchildren) / `handle.lifecycle`.
|
||||
|
||||
Each mini-mux owns its own scope and uses its own
|
||||
`SubgraphTransformer` to discover its direct children, so
|
||||
grandchildren live on the child handle — never on the root's
|
||||
`subgraphs` log. Forwarding events into the matching child mini-mux
|
||||
is what keeps the child's projections populated.
|
||||
|
||||
Native transformer — `subgraphs` is exposed as `run.subgraphs`.
|
||||
"""
|
||||
|
||||
_native = True
|
||||
supports_sync = True
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._log: EventLog[SubgraphRunStream | AsyncSubgraphRunStream] = EventLog()
|
||||
self._handles: dict[
|
||||
tuple[str, ...], SubgraphRunStream | AsyncSubgraphRunStream
|
||||
] = {}
|
||||
self._mux: StreamMux | None = None
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"subgraphs": self._log}
|
||||
|
||||
def _on_register(self, mux: Any) -> None:
|
||||
self._mux = mux
|
||||
|
||||
def _should_track(self, ns: tuple[str, ...]) -> bool:
|
||||
# Direct children only — grandchildren are picked up by the
|
||||
# child mini-mux's own SubgraphTransformer.
|
||||
depth = len(self.scope)
|
||||
return len(ns) == depth + 1 and ns[:depth] == self.scope
|
||||
|
||||
def _on_started(
|
||||
self,
|
||||
ns: tuple[str, ...],
|
||||
graph_name: str | None,
|
||||
trigger_call_id: str | None,
|
||||
) -> None:
|
||||
if self._mux is None:
|
||||
return
|
||||
try:
|
||||
child_mux = self._mux._make_child(ns)
|
||||
except RuntimeError:
|
||||
# Mux wasn't built from factories — no mini-mux navigation
|
||||
# available. Skip; LifecycleTransformer still tracks the
|
||||
# subgraph via the flat event stream.
|
||||
return
|
||||
values_t = child_mux.transformer_by_key("values")
|
||||
if not isinstance(values_t, ValuesTransformer):
|
||||
return
|
||||
handle_cls = AsyncSubgraphRunStream if child_mux.is_async else SubgraphRunStream
|
||||
handle = handle_cls(
|
||||
mux=child_mux,
|
||||
values_transformer=values_t,
|
||||
path=ns,
|
||||
graph_name=graph_name,
|
||||
trigger_call_id=trigger_call_id,
|
||||
)
|
||||
self._handles[ns] = handle
|
||||
self._log.push(handle)
|
||||
|
||||
def _on_terminal(
|
||||
self,
|
||||
ns: tuple[str, ...],
|
||||
status: SubgraphStatus,
|
||||
error: str | None,
|
||||
) -> None:
|
||||
handle = self._handles.get(ns)
|
||||
if handle is None or not self._mark_terminal(handle, status, error):
|
||||
return
|
||||
self._close_or_fail_handle(handle, status, error)
|
||||
|
||||
async def _aon_terminal(
|
||||
self,
|
||||
ns: tuple[str, ...],
|
||||
status: SubgraphStatus,
|
||||
error: str | None,
|
||||
) -> None:
|
||||
handle = self._handles.get(ns)
|
||||
if handle is None or not self._mark_terminal(handle, status, error):
|
||||
return
|
||||
await self._aclose_or_fail_handle(handle, status, error)
|
||||
|
||||
def _mark_terminal(
|
||||
self,
|
||||
handle: SubgraphRunStream | AsyncSubgraphRunStream,
|
||||
status: SubgraphStatus,
|
||||
error: str | None,
|
||||
) -> bool:
|
||||
"""Mark a handle terminal once. Returns True on first transition."""
|
||||
if handle._seen_terminal:
|
||||
return False
|
||||
handle.status = status
|
||||
if error is not None and handle.error is None:
|
||||
handle.error = error
|
||||
handle._seen_terminal = True
|
||||
return True
|
||||
|
||||
def _close_or_fail_handle(
|
||||
self,
|
||||
handle: SubgraphRunStream | AsyncSubgraphRunStream,
|
||||
status: SubgraphStatus,
|
||||
error: str | None,
|
||||
) -> None:
|
||||
if handle._mux is None or handle._mux._events._closed:
|
||||
return
|
||||
if status == "failed":
|
||||
handle._mux.fail(RuntimeError(error or "Subgraph failed"))
|
||||
else:
|
||||
handle._mux.close()
|
||||
|
||||
async def _aclose_or_fail_handle(
|
||||
self,
|
||||
handle: SubgraphRunStream | AsyncSubgraphRunStream,
|
||||
status: SubgraphStatus,
|
||||
error: str | None,
|
||||
) -> None:
|
||||
if handle._mux is None or handle._mux._events._closed:
|
||||
return
|
||||
if status == "failed":
|
||||
await handle._mux.afail(RuntimeError(error or "Subgraph failed"))
|
||||
else:
|
||||
await handle._mux.aclose()
|
||||
|
||||
def _child_mux_for_event(self, event: ProtocolEvent) -> StreamMux | None:
|
||||
ns = tuple(event["params"]["namespace"])
|
||||
depth = len(self.scope)
|
||||
if len(ns) < depth + 1:
|
||||
return None
|
||||
handle = self._handles.get(ns[: depth + 1])
|
||||
if handle is None or handle._mux is None or handle._mux._events._closed:
|
||||
return None
|
||||
return handle._mux
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
# Discover / update terminal status before forwarding so a
|
||||
# `started` handle exists by the time the child mini-mux sees
|
||||
# its own first event.
|
||||
keep = super().process(event)
|
||||
child_mux = self._child_mux_for_event(event)
|
||||
if child_mux is not None:
|
||||
child_mux.push(event)
|
||||
return keep
|
||||
|
||||
async def aprocess(self, event: ProtocolEvent) -> bool:
|
||||
# Async counterpart to `process`: repeat the tasks bookkeeping
|
||||
# here instead of delegating to `process`, so child mini-muxes
|
||||
# receive events through their async lane.
|
||||
if event["method"] == "tasks":
|
||||
ns = tuple(event["params"]["namespace"])
|
||||
data = event["params"]["data"]
|
||||
if "result" in data:
|
||||
for child_ns, status, error in self._pop_terminal_transitions(ns, data):
|
||||
await self._aon_terminal(child_ns, status, error)
|
||||
else:
|
||||
self._handle_task_start(ns)
|
||||
keep = False
|
||||
else:
|
||||
keep = True
|
||||
child_mux = self._child_mux_for_event(event)
|
||||
if child_mux is not None:
|
||||
await child_mux.apush(event)
|
||||
return keep
|
||||
|
||||
def _complete_open_handles(self) -> BaseException | None:
|
||||
first_error: BaseException | None = None
|
||||
for ns in list(self._open):
|
||||
try:
|
||||
self._on_terminal(ns, "completed", None)
|
||||
except BaseException as e:
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
self._open.clear()
|
||||
for handle in self._handles.values():
|
||||
if self._mark_terminal(handle, "completed", None):
|
||||
try:
|
||||
self._close_or_fail_handle(handle, "completed", None)
|
||||
except BaseException as e:
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
return first_error
|
||||
|
||||
async def _acomplete_open_handles(self) -> BaseException | None:
|
||||
first_error: BaseException | None = None
|
||||
for ns in list(self._open):
|
||||
try:
|
||||
await self._aon_terminal(ns, "completed", None)
|
||||
except BaseException as e:
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
self._open.clear()
|
||||
for handle in self._handles.values():
|
||||
if self._mark_terminal(handle, "completed", None):
|
||||
try:
|
||||
await self._aclose_or_fail_handle(handle, "completed", None)
|
||||
except BaseException as e:
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
return first_error
|
||||
|
||||
def finalize(self) -> None:
|
||||
first_error = self._complete_open_handles()
|
||||
if first_error is not None:
|
||||
raise first_error
|
||||
|
||||
async def afinalize(self) -> None:
|
||||
first_error = await self._acomplete_open_handles()
|
||||
if first_error is not None:
|
||||
raise first_error
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
status, error_str = _status_from_exception(err)
|
||||
self._open.clear()
|
||||
for handle in self._handles.values():
|
||||
self._mark_terminal(handle, status, error_str)
|
||||
if handle._mux is not None and not handle._mux._events._closed:
|
||||
try:
|
||||
handle._mux.fail(err)
|
||||
except Exception:
|
||||
_logger.warning(
|
||||
"Error failing subgraph mini-mux at %s; "
|
||||
"subscribers may not see the terminal error.",
|
||||
handle.path,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
async def afail(self, err: BaseException) -> None:
|
||||
status, error_str = _status_from_exception(err)
|
||||
self._open.clear()
|
||||
for handle in self._handles.values():
|
||||
self._mark_terminal(handle, status, error_str)
|
||||
if handle._mux is not None and not handle._mux._events._closed:
|
||||
try:
|
||||
await handle._mux.afail(err)
|
||||
except Exception:
|
||||
_logger.warning(
|
||||
"Error failing subgraph mini-mux at %s; "
|
||||
"subscribers may not see the terminal error.",
|
||||
handle.path,
|
||||
exc_info=True,
|
||||
)
|
||||
@@ -24,7 +24,7 @@ classifiers = [
|
||||
'Programming Language :: Python :: 3.13',
|
||||
]
|
||||
dependencies = [
|
||||
"langchain-core>=1.3.0,<2",
|
||||
"langchain-core>=1.3.2,<2",
|
||||
"langgraph-checkpoint>=2.1.0,<5.0.0",
|
||||
"langgraph-sdk>=0.3.0,<0.4.0",
|
||||
"langgraph-prebuilt>=1.0.12,<1.1.0",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,401 @@
|
||||
"""Tests for LifecycleTransformer.
|
||||
|
||||
Consumes the `tasks` stream mode and emits subgraph lifecycle payloads
|
||||
on the `lifecycle` channel for both in-process iteration via
|
||||
`run.lifecycle` and wire delivery via `custom:lifecycle` protocol
|
||||
events. Most tests dispatch synthetic protocol events through a
|
||||
`StreamMux` to keep the inference logic isolated; the end-of-file
|
||||
group exercises the path through real graphs (multi-depth
|
||||
discovery, nested `stream_v2` calls with non-empty `parent_ns`).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
import time
|
||||
from typing import Annotated, Any
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph._internal._constants import CONF, CONFIG_KEY_CHECKPOINT_NS
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream.transformers import (
|
||||
LifecyclePayload,
|
||||
LifecycleTransformer,
|
||||
)
|
||||
|
||||
TS = int(time.time() * 1000)
|
||||
|
||||
|
||||
def _tasks_start(
|
||||
namespace: list[str],
|
||||
*,
|
||||
task_id: str,
|
||||
name: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a `tasks` ProtocolEvent carrying a TaskPayload (start)."""
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "tasks",
|
||||
"params": {
|
||||
"namespace": namespace,
|
||||
"timestamp": TS,
|
||||
"data": {
|
||||
"id": task_id,
|
||||
"name": name,
|
||||
"input": None,
|
||||
"triggers": [],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _tasks_result(
|
||||
namespace: list[str],
|
||||
*,
|
||||
task_id: str,
|
||||
name: str,
|
||||
error: str | None = None,
|
||||
interrupts: list[dict[str, Any]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a `tasks` ProtocolEvent carrying a TaskResultPayload (finish)."""
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "tasks",
|
||||
"params": {
|
||||
"namespace": namespace,
|
||||
"timestamp": TS,
|
||||
"data": {
|
||||
"id": task_id,
|
||||
"name": name,
|
||||
"error": error,
|
||||
"interrupts": interrupts or [],
|
||||
"result": {},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _arm(mux: StreamMux) -> None:
|
||||
"""Force projection logs to accept pushes (skip lazy-subscribe gate).
|
||||
|
||||
`EventLog.push` is a no-op until a subscriber attaches. Tests that
|
||||
inspect `_items` directly need the gate flipped before any event
|
||||
is dispatched.
|
||||
"""
|
||||
mux._events._subscribed = True
|
||||
for transformer in mux._transformers:
|
||||
if isinstance(transformer, LifecycleTransformer):
|
||||
transformer._channel._log._subscribed = True
|
||||
|
||||
|
||||
def _drain_lifecycle(mux: StreamMux) -> list[LifecyclePayload]:
|
||||
"""Snapshot the lifecycle channel's underlying log."""
|
||||
transformer = mux.transformer_by_key("lifecycle")
|
||||
assert isinstance(transformer, LifecycleTransformer)
|
||||
return list(transformer._channel._log._items)
|
||||
|
||||
|
||||
def _build_lifecycle_mux(*, scope: tuple[str, ...] = ()) -> StreamMux:
|
||||
mux = StreamMux([LifecycleTransformer(scope=scope)], is_async=False)
|
||||
_arm(mux)
|
||||
return mux
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LifecycleTransformer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_started_emitted_on_first_direct_child_task() -> None:
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc123"], task_id="t1", name="tool"))
|
||||
|
||||
[payload] = _drain_lifecycle(mux)
|
||||
assert payload["event"] == "started"
|
||||
assert payload["namespace"] == ["agent:abc123"]
|
||||
assert payload["graph_name"] == "agent"
|
||||
assert payload["trigger_call_id"] == "abc123"
|
||||
|
||||
|
||||
def test_started_dedup_on_repeat_namespace() -> None:
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="a"))
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t2", name="b"))
|
||||
|
||||
payloads = _drain_lifecycle(mux)
|
||||
assert [p["event"] for p in payloads] == ["started"]
|
||||
|
||||
|
||||
def test_grandchild_namespace_discovered() -> None:
|
||||
"""Subgraphs at any depth below scope are tracked, not just direct children."""
|
||||
mux = _build_lifecycle_mux()
|
||||
# First-seen task at length-2 ns means a 2nd-level subgraph started.
|
||||
mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t1", name="x"))
|
||||
|
||||
[payload] = _drain_lifecycle(mux)
|
||||
assert payload["event"] == "started"
|
||||
assert payload["namespace"] == ["agent:abc", "tool:def"]
|
||||
|
||||
|
||||
def test_nested_chain_emits_started_at_each_depth() -> None:
|
||||
"""A graph → subgraph → subgraph chain produces a started event per level."""
|
||||
mux = _build_lifecycle_mux()
|
||||
# Subgraph1 starts emitting tasks (events tagged with its own ns).
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
# Subgraph1 invokes subgraph2; subgraph2's first task event arrives.
|
||||
mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t2", name="deep"))
|
||||
|
||||
payloads = _drain_lifecycle(mux)
|
||||
assert [p["namespace"] for p in payloads] == [
|
||||
["agent:abc"],
|
||||
["agent:abc", "tool:def"],
|
||||
]
|
||||
assert all(p["event"] == "started" for p in payloads)
|
||||
|
||||
|
||||
def test_nested_chain_emits_completed_at_each_depth() -> None:
|
||||
"""Each subgraph in a nested chain closes when its parent task result arrives."""
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t2", name="deep"))
|
||||
|
||||
# Subgraph2's owning task (id=def, inside subgraph1) finishes.
|
||||
mux.push(_tasks_result(["agent:abc"], task_id="def", name="tool"))
|
||||
# Subgraph1's owning task (id=abc, at root) finishes.
|
||||
mux.push(_tasks_result([], task_id="abc", name="agent"))
|
||||
|
||||
payloads = _drain_lifecycle(mux)
|
||||
events = [(p["event"], p["namespace"]) for p in payloads]
|
||||
assert events == [
|
||||
("started", ["agent:abc"]),
|
||||
("started", ["agent:abc", "tool:def"]),
|
||||
("completed", ["agent:abc", "tool:def"]),
|
||||
("completed", ["agent:abc"]),
|
||||
]
|
||||
|
||||
|
||||
def test_completed_on_parent_task_result() -> None:
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(_tasks_result([], task_id="abc", name="agent"))
|
||||
|
||||
events = [p["event"] for p in _drain_lifecycle(mux)]
|
||||
assert events == ["started", "completed"]
|
||||
|
||||
|
||||
def test_failed_on_parent_task_result_with_error() -> None:
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(_tasks_result([], task_id="abc", name="agent", error="boom"))
|
||||
|
||||
payloads = _drain_lifecycle(mux)
|
||||
assert [p["event"] for p in payloads] == ["started", "failed"]
|
||||
assert payloads[1]["error"] == "boom"
|
||||
|
||||
|
||||
def test_interrupted_on_parent_task_result_with_interrupts() -> None:
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(
|
||||
_tasks_result(
|
||||
[],
|
||||
task_id="abc",
|
||||
name="agent",
|
||||
interrupts=[{"value": "pause"}],
|
||||
)
|
||||
)
|
||||
|
||||
payloads = _drain_lifecycle(mux)
|
||||
assert [p["event"] for p in payloads] == ["started", "interrupted"]
|
||||
|
||||
|
||||
def test_interrupt_takes_precedence_over_error() -> None:
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(
|
||||
_tasks_result(
|
||||
[],
|
||||
task_id="abc",
|
||||
name="agent",
|
||||
error="should-be-suppressed",
|
||||
interrupts=[{"value": "pause"}],
|
||||
)
|
||||
)
|
||||
|
||||
last = _drain_lifecycle(mux)[-1]
|
||||
assert last["event"] == "interrupted"
|
||||
assert "error" not in last
|
||||
|
||||
|
||||
def test_finalize_completes_open_subgraphs() -> None:
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
mux.close()
|
||||
payloads = _drain_lifecycle(mux)
|
||||
assert [p["event"] for p in payloads] == ["started", "completed"]
|
||||
|
||||
|
||||
def test_fail_emits_interrupted_for_graph_interrupt() -> None:
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
mux.fail(GraphInterrupt())
|
||||
payloads = _drain_lifecycle(mux)
|
||||
assert [p["event"] for p in payloads] == ["started", "interrupted"]
|
||||
assert "error" not in payloads[1]
|
||||
|
||||
|
||||
def test_fail_emits_failed_for_other_exceptions() -> None:
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
mux.fail(RuntimeError("boom"))
|
||||
payloads = _drain_lifecycle(mux)
|
||||
assert [p["event"] for p in payloads] == ["started", "failed"]
|
||||
assert payloads[1]["error"] == "boom"
|
||||
|
||||
|
||||
def test_unrelated_methods_pass_through() -> None:
|
||||
"""Non-`tasks` events are not consumed and don't emit lifecycle."""
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(
|
||||
{
|
||||
"type": "event",
|
||||
"method": "values",
|
||||
"params": {"namespace": ["agent:abc"], "timestamp": TS, "data": {}},
|
||||
}
|
||||
)
|
||||
assert _drain_lifecycle(mux) == []
|
||||
|
||||
|
||||
def test_scoped_transformer_filters_outside_scope_but_tracks_all_depths() -> None:
|
||||
"""Scope filters the prefix; subgraphs at any depth below scope are tracked."""
|
||||
mux = _build_lifecycle_mux(scope=("agent:abc",))
|
||||
# Root-level task — out of scope (no shared prefix).
|
||||
mux.push(_tasks_start(["other:1"], task_id="t1", name="other"))
|
||||
# Direct child of agent:abc — in scope.
|
||||
mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t2", name="tool"))
|
||||
# Grandchild of agent:abc — also in scope, tracked at its own depth.
|
||||
mux.push(
|
||||
_tasks_start(["agent:abc", "tool:def", "deep:ghi"], task_id="t3", name="deep")
|
||||
)
|
||||
|
||||
payloads = _drain_lifecycle(mux)
|
||||
assert [p["namespace"] for p in payloads] == [
|
||||
["agent:abc", "tool:def"],
|
||||
["agent:abc", "tool:def", "deep:ghi"],
|
||||
]
|
||||
|
||||
|
||||
def test_required_stream_modes_declared() -> None:
|
||||
assert LifecycleTransformer.required_stream_modes == ("tasks",)
|
||||
|
||||
|
||||
def test_protocol_event_method_is_native() -> None:
|
||||
"""Native transformer — auto-forwarded events use `lifecycle`, not `custom:lifecycle`."""
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
methods = {evt["method"] for evt in mux._events._items}
|
||||
assert "lifecycle" in methods
|
||||
assert "custom:lifecycle" not in methods
|
||||
|
||||
|
||||
def test_tasks_events_suppressed_from_main_log() -> None:
|
||||
"""Tasks events are folded into lifecycle and don't appear on the main log."""
|
||||
mux = _build_lifecycle_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(_tasks_result([], task_id="abc", name="agent"))
|
||||
|
||||
methods = [evt["method"] for evt in mux._events._items]
|
||||
assert "tasks" not in methods
|
||||
# Lifecycle events did make it through, though.
|
||||
assert "lifecycle" in methods
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end: real graphs through stream_v2
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
value: str
|
||||
items: Annotated[list[str], operator.add]
|
||||
|
||||
|
||||
def _passthrough(state: _State) -> dict[str, Any]:
|
||||
return {"value": state["value"] + "!", "items": ["x"]}
|
||||
|
||||
|
||||
def _make_two_level_nested() -> Any:
|
||||
"""Build outer → middle → inner. Three Pregel instances, two nesting levels."""
|
||||
inner_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
inner_b.add_node("inner_node", _passthrough)
|
||||
inner_b.add_edge(START, "inner_node")
|
||||
inner_b.add_edge("inner_node", END)
|
||||
inner = inner_b.compile()
|
||||
|
||||
middle_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
middle_b.add_node("inner", inner)
|
||||
middle_b.add_edge(START, "inner")
|
||||
middle_b.add_edge("inner", END)
|
||||
middle = middle_b.compile()
|
||||
|
||||
outer_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
outer_b.add_node("middle", middle)
|
||||
outer_b.add_edge(START, "middle")
|
||||
outer_b.add_edge("middle", END)
|
||||
return outer_b.compile()
|
||||
|
||||
|
||||
def test_stream_v2_real_graph_emits_lifecycle_at_each_depth() -> None:
|
||||
"""Outer graph with two nested subgraphs surfaces lifecycle for both."""
|
||||
graph = _make_two_level_nested()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
# Iterating the projection drives the pump and drains synthesized
|
||||
# lifecycle events at the same time.
|
||||
payloads = list(run.lifecycle)
|
||||
# Each subgraph instance produces a started + a terminal event. Two
|
||||
# nested instances, so four payloads total in some interleaving.
|
||||
by_event = {p["event"] for p in payloads}
|
||||
assert "started" in by_event
|
||||
assert "completed" in by_event
|
||||
# Two distinct namespaces — direct child of root, and grandchild.
|
||||
namespaces = {tuple(p["namespace"]) for p in payloads}
|
||||
direct_children = {ns for ns in namespaces if len(ns) == 1}
|
||||
grandchildren = {ns for ns in namespaces if len(ns) == 2}
|
||||
assert direct_children, f"expected a level-1 lifecycle namespace, got {namespaces}"
|
||||
assert grandchildren, f"expected a level-2 lifecycle namespace, got {namespaces}"
|
||||
# Every direct-child namespace has a matching grandchild whose path extends it.
|
||||
for parent in direct_children:
|
||||
assert any(gc[: len(parent)] == parent for gc in grandchildren), (
|
||||
f"grandchild does not extend parent {parent}: {grandchildren}"
|
||||
)
|
||||
|
||||
|
||||
def test_stream_v2_with_nested_parent_ns_scopes_lifecycle() -> None:
|
||||
"""When `stream_v2` is called with a non-empty checkpoint_ns in config,
|
||||
`_resolve_parent_ns` returns that namespace and the registered
|
||||
`LifecycleTransformer` is constructed with `scope=parent_ns`. This
|
||||
exercises the path that exists today purely for nested-stream_v2
|
||||
callers; the test simulates such a caller by injecting a
|
||||
checkpoint_ns into the config.
|
||||
"""
|
||||
graph = _make_two_level_nested()
|
||||
config = {CONF: {CONFIG_KEY_CHECKPOINT_NS: "outer:abc"}}
|
||||
run = graph.stream_v2({"value": "x", "items": []}, config=config)
|
||||
|
||||
payloads = list(run.lifecycle)
|
||||
# Every emitted lifecycle namespace must extend the caller's scope —
|
||||
# nothing at root-level, nothing under a sibling prefix.
|
||||
for p in payloads:
|
||||
ns = tuple(p["namespace"])
|
||||
assert ns[:1] == ("outer:abc",), (
|
||||
f"namespace {ns} not within scoped prefix ('outer:abc',)"
|
||||
)
|
||||
@@ -0,0 +1,872 @@
|
||||
"""Tests for MessagesTransformer: protocol event routing, whole-message fallback,
|
||||
legacy v1 chunk filtering, and end-to-end via stream_v2 / astream_v2."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.language_models import GenericFakeChatModel
|
||||
from langchain_core.language_models.chat_model_stream import (
|
||||
AsyncChatModelStream,
|
||||
ChatModelStream,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import MessagesState, StateGraph
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream.run_stream import GraphRunStream
|
||||
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
|
||||
|
||||
TS = int(time.time() * 1000)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _proto_event(
|
||||
event: dict[str, Any],
|
||||
*,
|
||||
run_id: str = "run-1",
|
||||
node: str = "llm",
|
||||
) -> dict[str, Any]:
|
||||
"""Build a messages ProtocolEvent carrying a protocol event dict (v2 path)."""
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "messages",
|
||||
"params": {
|
||||
"namespace": [],
|
||||
"timestamp": TS,
|
||||
"data": (event, {"langgraph_node": node, "run_id": run_id}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _v1_chunk(
|
||||
text: str,
|
||||
msg_id: str = "msg-1",
|
||||
*,
|
||||
finish: bool = False,
|
||||
node: str = "llm",
|
||||
) -> dict[str, Any]:
|
||||
"""Build a messages ProtocolEvent carrying a v1 AIMessageChunk tuple."""
|
||||
rm: dict[str, Any] = {"finish_reason": "stop"} if finish else {}
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "messages",
|
||||
"params": {
|
||||
"namespace": [],
|
||||
"timestamp": TS,
|
||||
"data": (
|
||||
AIMessageChunk(content=text, id=msg_id, response_metadata=rm),
|
||||
{"langgraph_node": node},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _whole_msg(
|
||||
text: str,
|
||||
msg_id: str = "msg-10",
|
||||
*,
|
||||
node: str = "node",
|
||||
) -> dict[str, Any]:
|
||||
"""Build a messages ProtocolEvent carrying a completed AIMessage."""
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "messages",
|
||||
"params": {
|
||||
"namespace": [],
|
||||
"timestamp": TS,
|
||||
"data": (AIMessage(content=text, id=msg_id), {"langgraph_node": node}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _make_sync_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]:
|
||||
t = MessagesTransformer()
|
||||
log: EventLog[ChatModelStream] = t.init()["messages"]
|
||||
log._bind(is_async=False)
|
||||
# Subscribe up front so pushes during process() are retained.
|
||||
log._subscribed = True
|
||||
t._bind_pump(lambda: False)
|
||||
return t, log
|
||||
|
||||
|
||||
def _make_async_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]:
|
||||
t = MessagesTransformer()
|
||||
log: EventLog[ChatModelStream] = t.init()["messages"]
|
||||
log._bind(is_async=True)
|
||||
log._subscribed = True
|
||||
return t, log
|
||||
|
||||
|
||||
def _lifecycle(
|
||||
*, text: str = "hello world", message_id: str = "run-1"
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Produce a valid protocol event lifecycle: start, delta, finish."""
|
||||
half = len(text) // 2
|
||||
first, second = text[:half], text[half:]
|
||||
return [
|
||||
{"event": "message-start", "role": "ai", "message_id": message_id},
|
||||
{
|
||||
"event": "content-block-start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
},
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": first},
|
||||
},
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": second},
|
||||
},
|
||||
{
|
||||
"event": "content-block-finish",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": text},
|
||||
},
|
||||
{"event": "message-finish", "reason": "stop"},
|
||||
]
|
||||
|
||||
|
||||
def _simple_graph():
|
||||
def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
model = GenericFakeChatModel(messages=iter(["hello world"]))
|
||||
stream = model.stream_v2(state["messages"])
|
||||
return {"messages": stream.output}
|
||||
|
||||
return (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Protocol event routing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestProtocolEventRouting:
|
||||
def test_message_start_creates_stream(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(
|
||||
_proto_event(
|
||||
{"event": "message-start", "role": "ai", "message_id": "run-1"},
|
||||
run_id="run-1",
|
||||
)
|
||||
)
|
||||
log.close()
|
||||
(stream,) = list(log._items)
|
||||
assert isinstance(stream, ChatModelStream)
|
||||
assert stream.message_id == "run-1"
|
||||
|
||||
def test_full_lifecycle_yields_done_stream(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
for evt in _lifecycle(text="hello world"):
|
||||
t.process(_proto_event(evt, run_id="run-1"))
|
||||
log.close()
|
||||
(stream,) = list(log._items)
|
||||
assert stream.done
|
||||
assert stream.output.text == "hello world"
|
||||
|
||||
def test_message_finish_cleans_up_routing(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
for evt in _lifecycle():
|
||||
t.process(_proto_event(evt, run_id="run-1"))
|
||||
assert t._by_run == {}
|
||||
|
||||
def test_events_without_prior_start_are_ignored(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(
|
||||
_proto_event(
|
||||
{
|
||||
"event": "content-block-delta",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": "orphan"},
|
||||
},
|
||||
run_id="unknown",
|
||||
)
|
||||
)
|
||||
log.close()
|
||||
assert list(log._items) == []
|
||||
|
||||
def test_concurrent_streams_routed_by_run_id(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
life_a = _lifecycle(text="aaaa", message_id="run-a")
|
||||
life_b = _lifecycle(text="bbbb", message_id="run-b")
|
||||
for a, b in zip(life_a, life_b):
|
||||
t.process(_proto_event(a, run_id="run-a"))
|
||||
t.process(_proto_event(b, run_id="run-b"))
|
||||
log.close()
|
||||
streams = list(log._items)
|
||||
assert len(streams) == 2
|
||||
by_id = {s.message_id: s for s in streams}
|
||||
assert by_id["run-a"].output.text == "aaaa"
|
||||
assert by_id["run-b"].output.text == "bbbb"
|
||||
|
||||
def test_text_deltas_accumulated_on_stream(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
for evt in _lifecycle(text="abcdef"):
|
||||
t.process(_proto_event(evt))
|
||||
log.close()
|
||||
(stream,) = list(log._items)
|
||||
assert "".join(stream._text_proj._deltas) == "abcdef"
|
||||
|
||||
def test_stream_pushed_on_message_start_not_finish(self) -> None:
|
||||
# Consumer can see the stream before message-finish arrives.
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(
|
||||
_proto_event(
|
||||
{"event": "message-start", "role": "ai", "message_id": "run-1"},
|
||||
run_id="run-1",
|
||||
)
|
||||
)
|
||||
assert len(log._items) == 1
|
||||
|
||||
def test_node_metadata_set_on_stream(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(
|
||||
_proto_event(
|
||||
{"event": "message-start", "role": "ai", "message_id": "run-1"},
|
||||
run_id="run-1",
|
||||
node="my_llm",
|
||||
)
|
||||
)
|
||||
(stream,) = list(log._items)
|
||||
assert stream.node == "my_llm"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Whole-message fallback
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestWholeMessageFallback:
|
||||
def test_whole_ai_message_produces_complete_stream(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(_whole_msg("the full answer"))
|
||||
log.close()
|
||||
(stream,) = list(log._items)
|
||||
assert stream.done
|
||||
assert stream.output.text == "the full answer"
|
||||
|
||||
def test_whole_message_has_full_lifecycle(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(_whole_msg("full"))
|
||||
log.close()
|
||||
(stream,) = list(log._items)
|
||||
assert [e["event"] for e in stream._events] == [
|
||||
"message-start",
|
||||
"content-block-start",
|
||||
"content-block-delta",
|
||||
"content-block-finish",
|
||||
"message-finish",
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Filtering
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFiltering:
|
||||
def test_non_messages_events_pass_through(self) -> None:
|
||||
t, _ = _make_sync_transformer()
|
||||
assert (
|
||||
t.process(
|
||||
{
|
||||
"type": "event",
|
||||
"method": "values",
|
||||
"params": {"namespace": [], "timestamp": TS, "data": {"x": 1}},
|
||||
}
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_subgraph_namespace_dropped(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(
|
||||
{
|
||||
"type": "event",
|
||||
"method": "messages",
|
||||
"params": {
|
||||
"namespace": ["subgraph"],
|
||||
"timestamp": TS,
|
||||
"data": (
|
||||
{"event": "message-start", "message_id": "run-x"},
|
||||
{"run_id": "run-x"},
|
||||
),
|
||||
},
|
||||
}
|
||||
)
|
||||
log.close()
|
||||
assert list(log._items) == []
|
||||
|
||||
def test_legacy_v1_chunks_ignored(self) -> None:
|
||||
# v1 AIMessageChunk tuples (from on_llm_new_token) are not streamed
|
||||
# into this projection; callers must migrate to stream_v2.
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(_v1_chunk("hello"))
|
||||
t.process(_v1_chunk(" world", finish=True))
|
||||
log.close()
|
||||
assert list(log._items) == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lifecycle: fail / finalize
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLifecycle:
|
||||
def test_fail_propagates_to_open_streams(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(
|
||||
_proto_event(
|
||||
{"event": "message-start", "message_id": "run-1"}, run_id="run-1"
|
||||
)
|
||||
)
|
||||
streams = list(log._items)
|
||||
err = RuntimeError("graph died")
|
||||
t.fail(err)
|
||||
assert t._by_run == {}
|
||||
assert streams[0]._error is err
|
||||
|
||||
def test_finalize_clears_routing_state(self) -> None:
|
||||
t, _ = _make_sync_transformer()
|
||||
t.process(
|
||||
_proto_event(
|
||||
{"event": "message-start", "message_id": "run-1"}, run_id="run-1"
|
||||
)
|
||||
)
|
||||
assert "run-1" in t._by_run
|
||||
t.finalize()
|
||||
assert t._by_run == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async mode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAsyncMode:
|
||||
def test_async_mode_creates_async_stream(self) -> None:
|
||||
t, log = _make_async_transformer()
|
||||
for evt in _lifecycle(text="async stream"):
|
||||
t.process(_proto_event(evt))
|
||||
assert isinstance(list(log._items)[0], AsyncChatModelStream)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_text_projection_yields_deltas(self) -> None:
|
||||
t, log = _make_async_transformer()
|
||||
for evt in _lifecycle(text="hello world"):
|
||||
t.process(_proto_event(evt))
|
||||
(stream,) = list(log._items)
|
||||
assert isinstance(stream, AsyncChatModelStream)
|
||||
assert "".join([d async for d in stream.text]) == "hello world"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_output_awaitable(self) -> None:
|
||||
t, log = _make_async_transformer()
|
||||
for evt in _lifecycle(text="async"):
|
||||
t.process(_proto_event(evt))
|
||||
(stream,) = list(log._items)
|
||||
assert (await stream.output).text == "async"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GraphRunStream integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestWireRequestMore:
|
||||
def test_bind_pump_called_on_wire(self) -> None:
|
||||
values_t = ValuesTransformer()
|
||||
messages_t = MessagesTransformer()
|
||||
mux = StreamMux([values_t, messages_t], is_async=False)
|
||||
|
||||
assert messages_t._pump_fn is None
|
||||
run = GraphRunStream(iter([]), mux, values_t)
|
||||
assert messages_t._pump_fn is not None
|
||||
assert messages_t._pump_fn() is False
|
||||
assert run._exhausted
|
||||
|
||||
def test_created_streams_have_request_more(self) -> None:
|
||||
values_t = ValuesTransformer()
|
||||
messages_t = MessagesTransformer()
|
||||
mux = StreamMux([values_t, messages_t], is_async=False)
|
||||
GraphRunStream(iter([]), mux, values_t)
|
||||
|
||||
log: EventLog[ChatModelStream] = mux.extensions["messages"]
|
||||
log._subscribed = True
|
||||
for evt in _lifecycle():
|
||||
messages_t.process(_proto_event(evt))
|
||||
|
||||
(stream,) = list(log._items)
|
||||
assert stream._request_more is messages_t._pump_fn
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end via StreamMux
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestViaMux:
|
||||
def _make_mux(
|
||||
self,
|
||||
) -> tuple[MessagesTransformer, StreamMux, EventLog[ChatModelStream]]:
|
||||
t = MessagesTransformer()
|
||||
v = ValuesTransformer()
|
||||
mux = StreamMux([v, t], is_async=False)
|
||||
t._bind_pump(lambda: False)
|
||||
log: EventLog[ChatModelStream] = mux.extensions["messages"]
|
||||
log._subscribed = True
|
||||
return t, mux, log
|
||||
|
||||
def test_streaming_via_mux(self) -> None:
|
||||
t, mux, log = self._make_mux()
|
||||
for evt in _lifecycle(text="mux stream"):
|
||||
mux.push(_proto_event(evt))
|
||||
mux.close()
|
||||
(stream,) = list(log._items)
|
||||
assert stream.output.text == "mux stream"
|
||||
|
||||
def test_whole_message_via_mux(self) -> None:
|
||||
t, mux, log = self._make_mux()
|
||||
mux.push(_whole_msg("result"))
|
||||
mux.close()
|
||||
(stream,) = list(log._items)
|
||||
assert stream.output.text == "result"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_streaming_via_mux(self) -> None:
|
||||
t = MessagesTransformer()
|
||||
v = ValuesTransformer()
|
||||
mux = StreamMux([v, t], is_async=True)
|
||||
log: EventLog[ChatModelStream] = mux.extensions["messages"]
|
||||
log._subscribed = True
|
||||
|
||||
for evt in _lifecycle(text="async mux"):
|
||||
await mux.apush(_proto_event(evt))
|
||||
|
||||
(stream,) = list(log._items)
|
||||
assert (await stream.output).text == "async mux"
|
||||
await mux.aclose()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end: graph → stream_v2 → run.messages (node calls stream_v2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEndToEnd:
|
||||
"""stream_v2 path: node calls model.stream_v2() explicitly."""
|
||||
|
||||
def test_node_calling_stream_v2_populates_messages(self) -> None:
|
||||
model = GenericFakeChatModel(messages=iter(["hello world"]))
|
||||
|
||||
def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
stream = model.stream_v2(state["messages"])
|
||||
return {"messages": stream.output}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = graph.stream_v2({"messages": "hi"})
|
||||
(stream,) = list(run.messages)
|
||||
assert isinstance(stream, ChatModelStream)
|
||||
assert stream.output.text == "hello world"
|
||||
|
||||
def test_node_stream_v2_text_deltas_iterate(self) -> None:
|
||||
"""Consumer can iterate `.text` on the streamed message in real time."""
|
||||
model = GenericFakeChatModel(messages=iter(["streamed answer"]))
|
||||
|
||||
def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
stream = model.stream_v2(state["messages"])
|
||||
return {"messages": stream.output}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = graph.stream_v2({"messages": "go"})
|
||||
(stream,) = list(run.messages)
|
||||
assert "".join(stream.text) == "streamed answer"
|
||||
|
||||
def test_non_llm_message_returned_from_node(self) -> None:
|
||||
"""Whole-message fallback: node returns a finalized AIMessage directly."""
|
||||
|
||||
def return_message(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": AIMessage(content="hardcoded", id="msg-abc")}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("return_message", return_message)
|
||||
.add_edge(START, "return_message")
|
||||
.add_edge("return_message", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = graph.stream_v2({"messages": "hi"})
|
||||
(stream,) = list(run.messages)
|
||||
assert stream.output.text == "hardcoded"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_node_calling_astream_v2(self) -> None:
|
||||
model = GenericFakeChatModel(messages=iter(["async answer"]))
|
||||
|
||||
async def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
stream = await model.astream_v2(state["messages"])
|
||||
return {"messages": await stream}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = await graph.astream_v2({"messages": "hi"})
|
||||
streams = [s async for s in run.messages]
|
||||
assert len(streams) == 1
|
||||
assert isinstance(streams[0], AsyncChatModelStream)
|
||||
assert (await streams[0].output).text == "async answer"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_nested_async_iteration_yields_text_deltas(self) -> None:
|
||||
"""Inner stream.text drives the shared graph pump via the async pump binding."""
|
||||
import asyncio
|
||||
|
||||
model = GenericFakeChatModel(messages=iter(["hello world"]))
|
||||
|
||||
async def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
stream = await model.astream_v2(state["messages"])
|
||||
return {"messages": await stream}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = await graph.astream_v2({"messages": "hi"})
|
||||
|
||||
async def consume() -> list[str]:
|
||||
collected: list[str] = []
|
||||
async for stream in run.messages:
|
||||
async for delta in stream.text:
|
||||
collected.append(delta)
|
||||
return collected
|
||||
|
||||
assert "".join(await asyncio.wait_for(consume(), timeout=2.0)) == "hello world"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end: graph → stream_v2 → run.messages (node calls invoke)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEndToEndV2Invoke:
|
||||
"""Auto-routing path: stream_v2 injects CONFIG_KEY_STREAM_MESSAGES_V2,
|
||||
causing BaseChatModel to drive the v2 protocol event generator even for
|
||||
model.invoke()."""
|
||||
|
||||
def _graph(self, model):
|
||||
def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": model.invoke(state["messages"])}
|
||||
|
||||
return (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
def test_invoke_populates_messages(self) -> None:
|
||||
run = self._graph(
|
||||
GenericFakeChatModel(messages=iter(["hello world"]))
|
||||
).stream_v2({"messages": "hi"})
|
||||
(stream,) = list(run.messages)
|
||||
assert isinstance(stream, ChatModelStream)
|
||||
assert stream.output.text == "hello world"
|
||||
|
||||
def test_invoke_emits_protocol_events(self) -> None:
|
||||
"""Iterating the stream yields the full v2 lifecycle, not v1 chunks."""
|
||||
run = self._graph(
|
||||
GenericFakeChatModel(messages=iter(["streamed answer"]))
|
||||
).stream_v2({"messages": "go"})
|
||||
(stream,) = list(run.messages)
|
||||
|
||||
events = list(stream)
|
||||
event_types = [e.get("event") for e in events]
|
||||
assert "message-start" in event_types
|
||||
assert "content-block-start" in event_types
|
||||
assert "content-block-delta" in event_types
|
||||
assert "content-block-finish" in event_types
|
||||
assert "message-finish" in event_types
|
||||
# Sanity: every event is a dict carrying an "event" key — not an
|
||||
# AIMessageChunk tuple from the v1 path.
|
||||
for event in events:
|
||||
assert isinstance(event, dict)
|
||||
assert "event" in event
|
||||
# Typed projection still assembles the final text.
|
||||
assert stream.output.text == "streamed answer"
|
||||
|
||||
def test_invoke_text_deltas_iterate(self) -> None:
|
||||
run = self._graph(
|
||||
GenericFakeChatModel(messages=iter(["delta streaming works"]))
|
||||
).stream_v2({"messages": "hi"})
|
||||
(stream,) = list(run.messages)
|
||||
assert "".join(stream.text) == "delta streaming works"
|
||||
|
||||
def test_invoke_two_nodes_two_streams(self) -> None:
|
||||
model_a = GenericFakeChatModel(messages=iter(["alpha"]))
|
||||
model_b = GenericFakeChatModel(messages=iter(["beta"]))
|
||||
|
||||
def node_a(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": model_a.invoke(state["messages"])}
|
||||
|
||||
def node_b(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": model_b.invoke(state["messages"])}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("node_a", node_a)
|
||||
.add_node("node_b", node_b)
|
||||
.add_edge(START, "node_a")
|
||||
.add_edge("node_a", "node_b")
|
||||
.add_edge("node_b", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
streams = list(graph.stream_v2({"messages": "hi"}).messages)
|
||||
assert len(streams) == 2
|
||||
assert {s.output.text for s in streams} == {"alpha", "beta"}
|
||||
|
||||
def test_invoke_plus_constructed_message_two_streams(self) -> None:
|
||||
"""Live-streamed node + constructed-message node → two ChatModelStreams."""
|
||||
model = GenericFakeChatModel(messages=iter(["live stream"]))
|
||||
|
||||
def streaming_node(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": model.invoke(state["messages"])}
|
||||
|
||||
def constructed_node(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": [AIMessage(content="hardcoded", id="constructed-1")]}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("streaming_node", streaming_node)
|
||||
.add_node("constructed_node", constructed_node)
|
||||
.add_edge(START, "streaming_node")
|
||||
.add_edge("streaming_node", "constructed_node")
|
||||
.add_edge("constructed_node", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = graph.stream_v2({"messages": "hi"})
|
||||
streams = list(run.messages)
|
||||
assert len(streams) == 2
|
||||
assert streams[0].node == "streaming_node"
|
||||
assert streams[0].output.text == "live stream"
|
||||
assert streams[1].node == "constructed_node"
|
||||
assert streams[1].output.text == "hardcoded"
|
||||
assert streams[1].message_id == "constructed-1"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_ainvoke_populates_messages(self) -> None:
|
||||
model = GenericFakeChatModel(messages=iter(["async invoke"]))
|
||||
|
||||
async def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": await model.ainvoke(state["messages"])}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = await graph.astream_v2({"messages": "hi"})
|
||||
streams = [s async for s in run.messages]
|
||||
assert len(streams) == 1
|
||||
assert isinstance(streams[0], AsyncChatModelStream)
|
||||
assert (await streams[0].output).text == "async invoke"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression: direct stream_mode="messages" must stay v1
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDirectMessagesModeStaysV1:
|
||||
def test_direct_graph_stream_messages_yields_ai_message_chunks(self) -> None:
|
||||
"""graph.stream(stream_mode="messages") must not leak v2 event dicts —
|
||||
the v2 flag is only injected by stream_v2 / astream_v2."""
|
||||
model = GenericFakeChatModel(messages=iter(["legacy path"]))
|
||||
|
||||
def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": model.invoke(state["messages"])}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
parts = list(graph.stream({"messages": "hi"}, stream_mode="messages"))
|
||||
assert parts, "expected stream_mode='messages' to emit tuples"
|
||||
for payload, _metadata in parts:
|
||||
assert isinstance(payload, AIMessageChunk)
|
||||
assert (
|
||||
"".join(p[0].content for p in parts if isinstance(p[0].content, str))
|
||||
== "legacy path"
|
||||
)
|
||||
|
||||
def test_nested_graph_stream_messages_stays_v1_under_outer_stream_v2(self) -> None:
|
||||
"""An outer `stream_v2()` run must not flip an inner direct
|
||||
`stream_mode="messages"` call onto the v2 event protocol."""
|
||||
model = GenericFakeChatModel(messages=iter(["nested legacy path"]))
|
||||
|
||||
def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": model.invoke(state["messages"])}
|
||||
|
||||
inner = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
class OuterState(TypedDict, total=False):
|
||||
saw_only_chunks: bool
|
||||
first_payload_type: str
|
||||
text: str
|
||||
|
||||
def call_subgraph(state: OuterState, config: RunnableConfig) -> dict[str, Any]:
|
||||
parts = list(
|
||||
inner.stream(
|
||||
{"messages": "hi"},
|
||||
config,
|
||||
stream_mode="messages",
|
||||
)
|
||||
)
|
||||
assert parts
|
||||
payloads = [payload for payload, _metadata in parts]
|
||||
return {
|
||||
"saw_only_chunks": all(
|
||||
isinstance(payload, AIMessageChunk) for payload in payloads
|
||||
),
|
||||
"first_payload_type": type(payloads[0]).__name__,
|
||||
"text": "".join(
|
||||
payload.content
|
||||
for payload in payloads
|
||||
if isinstance(payload, AIMessageChunk)
|
||||
and isinstance(payload.content, str)
|
||||
),
|
||||
}
|
||||
|
||||
outer = (
|
||||
StateGraph(OuterState)
|
||||
.add_node("call_subgraph", call_subgraph)
|
||||
.add_edge(START, "call_subgraph")
|
||||
.add_edge("call_subgraph", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
result = outer.stream_v2({}).output
|
||||
|
||||
assert result is not None
|
||||
assert result["saw_only_chunks"] is True
|
||||
assert result["first_payload_type"] == "AIMessageChunk"
|
||||
assert result["text"] == "nested legacy path"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# StreamMessagesHandlerV2 unit
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStreamMessagesHandlerV2Unit:
|
||||
def test_on_llm_new_token_is_noop(self) -> None:
|
||||
"""v2 handler must not emit v1 chunks even when on_llm_new_token fires."""
|
||||
from uuid import uuid4
|
||||
|
||||
from langchain_core.outputs import ChatGenerationChunk
|
||||
|
||||
from langgraph.pregel._messages import StreamMessagesHandlerV2
|
||||
|
||||
emitted: list[Any] = []
|
||||
handler = StreamMessagesHandlerV2(emitted.append, subgraphs=False)
|
||||
run_id = uuid4()
|
||||
handler.metadata[run_id] = ((), {"langgraph_node": "x"})
|
||||
|
||||
handler.on_llm_new_token(
|
||||
"hello",
|
||||
chunk=ChatGenerationChunk(message=AIMessageChunk(content="hello")),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
assert emitted == []
|
||||
|
||||
def test_on_llm_end_dedupes_when_final_message_id_differs(self) -> None:
|
||||
"""A streamed v2 message should not be emitted again from the final
|
||||
AIMessage fallback when its final id does not match `message-start`."""
|
||||
from uuid import uuid4
|
||||
|
||||
from langchain_core.outputs import ChatGeneration, LLMResult
|
||||
|
||||
from langgraph.pregel._messages import StreamMessagesHandlerV2
|
||||
|
||||
emitted: list[Any] = []
|
||||
handler = StreamMessagesHandlerV2(emitted.append, subgraphs=False)
|
||||
run_id = uuid4()
|
||||
handler.metadata[run_id] = ((), {"langgraph_node": "x"})
|
||||
|
||||
handler.on_stream_event(
|
||||
{"event": "message-start", "message_id": "stream-msg-1"},
|
||||
run_id=run_id,
|
||||
)
|
||||
handler.on_llm_end(
|
||||
LLMResult(
|
||||
generations=[
|
||||
[
|
||||
ChatGeneration(
|
||||
message=AIMessage(content="hello", id="final-msg-1")
|
||||
)
|
||||
]
|
||||
]
|
||||
),
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
assert len(emitted) == 1
|
||||
@@ -0,0 +1,865 @@
|
||||
"""Tests for SubgraphTransformer.
|
||||
|
||||
Subscribes to `tasks` events and produces in-process `SubgraphRunStream`
|
||||
handles backed by mini-muxes (built via `StreamMux._make_child`). The
|
||||
synthetic-event tests isolate the inference / mini-mux wiring; the
|
||||
real-graph tests exercise the end-to-end navigation path through
|
||||
`stream_v2`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from functools import partial
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.pregel.main import _normalize_stream_transformer_factories
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.run_stream import (
|
||||
AsyncGraphRunStream,
|
||||
AsyncSubgraphRunStream,
|
||||
GraphRunStream,
|
||||
SubgraphRunStream,
|
||||
)
|
||||
from langgraph.stream.transformers import (
|
||||
LifecycleTransformer,
|
||||
MessagesTransformer,
|
||||
SubgraphTransformer,
|
||||
ValuesTransformer,
|
||||
)
|
||||
|
||||
TS = int(time.time() * 1000)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _tasks_start(
|
||||
namespace: list[str],
|
||||
*,
|
||||
task_id: str,
|
||||
name: str,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "tasks",
|
||||
"params": {
|
||||
"namespace": namespace,
|
||||
"timestamp": TS,
|
||||
"data": {
|
||||
"id": task_id,
|
||||
"name": name,
|
||||
"input": None,
|
||||
"triggers": [],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _tasks_result(
|
||||
namespace: list[str],
|
||||
*,
|
||||
task_id: str,
|
||||
name: str,
|
||||
error: str | None = None,
|
||||
interrupts: list[dict[str, Any]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "tasks",
|
||||
"params": {
|
||||
"namespace": namespace,
|
||||
"timestamp": TS,
|
||||
"data": {
|
||||
"id": task_id,
|
||||
"name": name,
|
||||
"error": error,
|
||||
"interrupts": interrupts or [],
|
||||
"result": {},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _native_factories() -> list[Any]:
|
||||
"""Mirror the factory list `Pregel.stream_v2` registers."""
|
||||
return [
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
]
|
||||
|
||||
|
||||
def _stream_part(
|
||||
method: str,
|
||||
namespace: tuple[str, ...],
|
||||
data: Any,
|
||||
) -> dict[str, Any]:
|
||||
return {"type": method, "ns": namespace, "data": data}
|
||||
|
||||
|
||||
async def _astream_parts(*parts: dict[str, Any]) -> AsyncIterator[dict[str, Any]]:
|
||||
for part in parts:
|
||||
yield part
|
||||
|
||||
|
||||
def _arm(mux: StreamMux) -> None:
|
||||
"""Pre-subscribe every projection in the mux so synthetic pushes accumulate.
|
||||
|
||||
Real consumer code subscribes by iterating the projection; tests
|
||||
inspect `_items` directly, so the lazy-subscribe gate has to be
|
||||
flipped manually before any synthetic events are pushed.
|
||||
"""
|
||||
mux._events._subscribed = True
|
||||
for value in mux.extensions.values():
|
||||
if hasattr(value, "_subscribed"): # EventLog
|
||||
value._subscribed = True
|
||||
elif hasattr(value, "_log"): # StreamChannel
|
||||
value._log._subscribed = True
|
||||
|
||||
|
||||
def _arm_recursive(mux: StreamMux) -> None:
|
||||
"""Arm `mux` and every mini-mux currently held by SubgraphTransformer handles.
|
||||
|
||||
Mini-muxes are created during `mux.push(...)` when a new direct
|
||||
child is discovered. Tests must call this after each push that
|
||||
might have created a new mini-mux so subsequent pushes' projection
|
||||
side effects accumulate (rather than dropping silently against an
|
||||
unsubscribed log).
|
||||
"""
|
||||
_arm(mux)
|
||||
for handle in _subgraph_transformer(mux)._handles.values():
|
||||
if handle._mux is not None:
|
||||
_arm_recursive(handle._mux)
|
||||
|
||||
|
||||
def _build_root_mux(*, scope: tuple[str, ...] = ()) -> StreamMux:
|
||||
mux = StreamMux(
|
||||
factories=_native_factories(),
|
||||
scope=scope,
|
||||
is_async=False,
|
||||
)
|
||||
_arm(mux)
|
||||
return mux
|
||||
|
||||
|
||||
def _subgraph_transformer(mux: StreamMux) -> SubgraphTransformer:
|
||||
transformer = mux.transformer_by_key("subgraphs")
|
||||
assert isinstance(transformer, SubgraphTransformer)
|
||||
return transformer
|
||||
|
||||
|
||||
def _drain_subgraphs(mux: StreamMux) -> list[SubgraphRunStream]:
|
||||
return list(_subgraph_transformer(mux)._log._items)
|
||||
|
||||
|
||||
def _child_mux(handle: SubgraphRunStream | AsyncSubgraphRunStream) -> StreamMux:
|
||||
assert handle._mux is not None
|
||||
return handle._mux
|
||||
|
||||
|
||||
def _event_items(mux: StreamMux) -> list[ProtocolEvent]:
|
||||
return list(mux._events._items)
|
||||
|
||||
|
||||
def _lifecycle_payloads(mux: StreamMux) -> list[dict[str, Any]]:
|
||||
lifecycle_t = mux.transformer_by_key("lifecycle")
|
||||
assert isinstance(lifecycle_t, LifecycleTransformer)
|
||||
return list(lifecycle_t._channel._log._items)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Synthetic-event tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_handle_created_on_first_direct_child_task() -> None:
|
||||
mux = _build_root_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
[handle] = _drain_subgraphs(mux)
|
||||
assert handle.path == ("agent:abc",)
|
||||
assert handle.graph_name == "agent"
|
||||
assert handle.trigger_call_id == "abc"
|
||||
assert handle.status == "started"
|
||||
_child_mux(handle) # mini-mux backed
|
||||
|
||||
|
||||
def test_handle_status_completes_on_parent_result() -> None:
|
||||
mux = _build_root_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(_tasks_result([], task_id="abc", name="agent"))
|
||||
|
||||
[handle] = _drain_subgraphs(mux)
|
||||
assert handle.status == "completed"
|
||||
assert handle.error is None
|
||||
|
||||
|
||||
def test_handle_status_failed_with_error() -> None:
|
||||
mux = _build_root_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(_tasks_result([], task_id="abc", name="agent", error="boom"))
|
||||
|
||||
[handle] = _drain_subgraphs(mux)
|
||||
assert handle.status == "failed"
|
||||
assert handle.error == "boom"
|
||||
|
||||
|
||||
def test_handle_status_interrupted() -> None:
|
||||
mux = _build_root_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(
|
||||
_tasks_result(
|
||||
[],
|
||||
task_id="abc",
|
||||
name="agent",
|
||||
interrupts=[{"value": "pause"}],
|
||||
)
|
||||
)
|
||||
|
||||
[handle] = _drain_subgraphs(mux)
|
||||
assert handle.status == "interrupted"
|
||||
|
||||
|
||||
def test_grandchild_discovered_via_child_mini_mux() -> None:
|
||||
"""Each mini-mux owns its own scope; grandchildren live on the child handle."""
|
||||
mux = _build_root_mux()
|
||||
# Direct child started — creates the mini-mux.
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
# Pre-subscribe the freshly-created mini-mux so subsequent
|
||||
# forwarded events land on its projections (consumer would
|
||||
# subscribe naturally by iterating handle.subgraphs, but the
|
||||
# test inspects `_items` directly).
|
||||
_arm_recursive(mux)
|
||||
# Grandchild's first task event flows down into the child mini-mux.
|
||||
mux.push(_tasks_start(["agent:abc", "tool:def"], task_id="t2", name="deep"))
|
||||
|
||||
[child_handle] = _drain_subgraphs(mux)
|
||||
assert child_handle.path == ("agent:abc",)
|
||||
# The grandchild appears on the CHILD'S subgraphs projection.
|
||||
grandchildren = list(child_handle.subgraphs._items)
|
||||
assert len(grandchildren) == 1
|
||||
assert grandchildren[0].path == ("agent:abc", "tool:def")
|
||||
|
||||
|
||||
def test_finalize_completes_open_handles() -> None:
|
||||
mux = _build_root_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
mux.close()
|
||||
[handle] = _drain_subgraphs(mux)
|
||||
assert handle.status == "completed"
|
||||
|
||||
|
||||
def test_fail_marks_open_handles_interrupted_for_graph_interrupt() -> None:
|
||||
mux = _build_root_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
mux.fail(GraphInterrupt())
|
||||
[handle] = _drain_subgraphs(mux)
|
||||
assert handle.status == "interrupted"
|
||||
|
||||
|
||||
def test_fail_marks_open_handles_failed_for_other_errors() -> None:
|
||||
mux = _build_root_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
mux.fail(RuntimeError("boom"))
|
||||
[handle] = _drain_subgraphs(mux)
|
||||
assert handle.status == "failed"
|
||||
assert handle.error == "boom"
|
||||
|
||||
|
||||
def test_child_mux_requires_factories() -> None:
|
||||
"""A mux constructed only from `transformers=` can't clone factories."""
|
||||
transformer = SubgraphTransformer()
|
||||
mux = StreamMux(transformers=[transformer], is_async=False)
|
||||
with pytest.raises(RuntimeError, match="factories"):
|
||||
mux._make_child(("anything",))
|
||||
|
||||
|
||||
def test_subgraph_and_lifecycle_agree_on_terminal_status() -> None:
|
||||
"""Both transformers consume the same tasks signal — no drift."""
|
||||
mux = _build_root_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(_tasks_result([], task_id="abc", name="agent", error="boom"))
|
||||
|
||||
[handle] = _drain_subgraphs(mux)
|
||||
payloads = _lifecycle_payloads(mux)
|
||||
assert handle.status == "failed"
|
||||
assert payloads[-1]["event"] == "failed"
|
||||
assert handle.error == payloads[-1]["error"]
|
||||
|
||||
|
||||
def test_required_stream_modes_declared() -> None:
|
||||
assert SubgraphTransformer.required_stream_modes == ("tasks",)
|
||||
|
||||
|
||||
def test_tasks_events_suppressed_from_main_log() -> None:
|
||||
"""Tasks events are folded into discovery and don't appear on the main log."""
|
||||
mux = _build_root_mux()
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
mux.push(_tasks_result([], task_id="abc", name="agent"))
|
||||
|
||||
methods = [evt["method"] for evt in _event_items(mux)]
|
||||
assert "tasks" not in methods
|
||||
|
||||
|
||||
class _ChildEventObserver(StreamTransformer):
|
||||
"""Records child-scope event identity without mutating it."""
|
||||
|
||||
records: list[tuple[tuple[str, ...], int, int, bool]] = []
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if self.scope and event["method"] == "values":
|
||||
self.records.append(
|
||||
(
|
||||
self.scope,
|
||||
id(event),
|
||||
id(event["params"]["data"]),
|
||||
"seq" in event,
|
||||
)
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def test_child_forwarding_reuses_event_without_assigning_seq() -> None:
|
||||
_ChildEventObserver.records = []
|
||||
mux = StreamMux(
|
||||
factories=[
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
_ChildEventObserver,
|
||||
],
|
||||
is_async=False,
|
||||
)
|
||||
_arm(mux)
|
||||
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
data = {"x": 1}
|
||||
event: ProtocolEvent = {
|
||||
"type": "event",
|
||||
"method": "values",
|
||||
"params": {
|
||||
"namespace": ["agent:abc"],
|
||||
"timestamp": TS,
|
||||
"data": data,
|
||||
},
|
||||
}
|
||||
mux.push(event)
|
||||
|
||||
assert _ChildEventObserver.records == [(("agent:abc",), id(event), id(data), False)]
|
||||
[root_event] = [evt for evt in _event_items(mux) if evt["method"] == "values"]
|
||||
assert root_event is event
|
||||
assert "seq" in root_event
|
||||
|
||||
|
||||
class _AsyncProbeTransformer(StreamTransformer):
|
||||
"""Async-only transformer used to verify mini-mux async dispatch."""
|
||||
|
||||
required_stream_modes = ("tasks",)
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self.seen: list[tuple[str, ...]] = []
|
||||
self.finalized = False
|
||||
self.failed: BaseException | None = None
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"async_probe": self}
|
||||
|
||||
async def aprocess(self, event: ProtocolEvent) -> bool:
|
||||
self.seen.append(tuple(event["params"]["namespace"]))
|
||||
return True
|
||||
|
||||
async def afinalize(self) -> None:
|
||||
self.finalized = True
|
||||
|
||||
async def afail(self, err: BaseException) -> None:
|
||||
self.failed = err
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_child_mini_mux_uses_async_lane() -> None:
|
||||
mux = StreamMux(
|
||||
factories=[
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
_AsyncProbeTransformer,
|
||||
],
|
||||
is_async=True,
|
||||
)
|
||||
await mux.apush(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
handle = _subgraph_transformer(mux)._handles[("agent:abc",)]
|
||||
assert isinstance(handle, AsyncSubgraphRunStream)
|
||||
probe = _child_mux(handle).transformer_by_key("async_probe")
|
||||
assert isinstance(probe, _AsyncProbeTransformer)
|
||||
assert probe.seen == [("agent:abc",)]
|
||||
|
||||
await mux.apush(_tasks_result([], task_id="abc", name="agent"))
|
||||
assert probe.finalized is True
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_child_mini_mux_fail_uses_async_lane() -> None:
|
||||
mux = StreamMux(
|
||||
factories=[
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
_AsyncProbeTransformer,
|
||||
],
|
||||
is_async=True,
|
||||
)
|
||||
await mux.apush(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
|
||||
|
||||
handle = _subgraph_transformer(mux)._handles[("agent:abc",)]
|
||||
probe = _child_mux(handle).transformer_by_key("async_probe")
|
||||
assert isinstance(probe, _AsyncProbeTransformer)
|
||||
|
||||
err = RuntimeError("boom")
|
||||
await mux.afail(err)
|
||||
assert probe.failed is err
|
||||
|
||||
|
||||
class _StandardCtorTransformer(StreamTransformer):
|
||||
"""Transformer class that inherits the standard scoped constructor."""
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"standard_ctor": self}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class _ScopedTransformer(StreamTransformer):
|
||||
"""Transformer class that uses the inherited scoped construction."""
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"scoped": self}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class _ConfigurableFactoryTransformer(StreamTransformer):
|
||||
"""Transformer built by a configured per-scope factory."""
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = (), *, label: str) -> None:
|
||||
super().__init__(scope)
|
||||
self.label = label
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"configurable": self}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class _ChildExploder(StreamTransformer):
|
||||
"""Raise from child mini-muxes to verify errors propagate upstream."""
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if self.scope and event["method"] == "values":
|
||||
raise RuntimeError("child boom")
|
||||
return True
|
||||
|
||||
|
||||
class _ChildFinalizeExploder(StreamTransformer):
|
||||
"""Raise from child mini-mux finalization."""
|
||||
|
||||
supports_sync = True
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
return True
|
||||
|
||||
def finalize(self) -> None:
|
||||
if self.scope:
|
||||
raise RuntimeError("child finalize boom")
|
||||
|
||||
async def afinalize(self) -> None:
|
||||
if self.scope:
|
||||
raise RuntimeError("child afinalize boom")
|
||||
|
||||
|
||||
def test_normalize_transformer_factories_supports_scoped_classes() -> None:
|
||||
factories = _normalize_stream_transformer_factories(
|
||||
[_StandardCtorTransformer, _ScopedTransformer]
|
||||
)
|
||||
|
||||
standard_ctor = factories[0](("child",))
|
||||
scoped = factories[1](("child",))
|
||||
assert isinstance(standard_ctor, _StandardCtorTransformer)
|
||||
assert standard_ctor.scope == ("child",)
|
||||
assert isinstance(scoped, _ScopedTransformer)
|
||||
assert scoped.scope == ("child",)
|
||||
|
||||
|
||||
def test_normalize_transformer_factories_supports_configured_factories() -> None:
|
||||
factories = _normalize_stream_transformer_factories(
|
||||
[partial(_ConfigurableFactoryTransformer, label="configured")]
|
||||
)
|
||||
|
||||
built = factories[0](("child",))
|
||||
assert isinstance(built, _ConfigurableFactoryTransformer)
|
||||
assert built.label == "configured"
|
||||
assert built.scope == ("child",)
|
||||
|
||||
|
||||
def test_normalize_transformer_factories_rejects_instances() -> None:
|
||||
with pytest.raises(TypeError, match="pre-built instance"):
|
||||
_normalize_stream_transformer_factories([_StandardCtorTransformer()])
|
||||
|
||||
|
||||
def test_child_forwarding_errors_fail_sync_run() -> None:
|
||||
mux = StreamMux(
|
||||
factories=[
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
_ChildExploder,
|
||||
],
|
||||
is_async=False,
|
||||
)
|
||||
values_t = mux.transformer_by_key("values")
|
||||
assert isinstance(values_t, ValuesTransformer)
|
||||
run = GraphRunStream(
|
||||
iter(
|
||||
[
|
||||
_stream_part(
|
||||
"tasks",
|
||||
("agent:abc",),
|
||||
{
|
||||
"id": "t1",
|
||||
"name": "tool",
|
||||
"input": None,
|
||||
"triggers": [],
|
||||
},
|
||||
),
|
||||
_stream_part("values", ("agent:abc",), {"x": 1}),
|
||||
]
|
||||
),
|
||||
mux,
|
||||
values_t,
|
||||
)
|
||||
|
||||
handle = next(iter(run.subgraphs))
|
||||
assert handle.path == ("agent:abc",)
|
||||
with pytest.raises(RuntimeError, match="child boom"):
|
||||
_ = run.output
|
||||
assert run._mux._events._error is not None
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_child_forwarding_errors_fail_async_run() -> None:
|
||||
mux = StreamMux(
|
||||
factories=[
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
_ChildExploder,
|
||||
],
|
||||
is_async=True,
|
||||
)
|
||||
values_t = mux.transformer_by_key("values")
|
||||
assert isinstance(values_t, ValuesTransformer)
|
||||
run = AsyncGraphRunStream(
|
||||
_astream_parts(
|
||||
_stream_part(
|
||||
"tasks",
|
||||
("agent:abc",),
|
||||
{
|
||||
"id": "t1",
|
||||
"name": "tool",
|
||||
"input": None,
|
||||
"triggers": [],
|
||||
},
|
||||
),
|
||||
_stream_part("values", ("agent:abc",), {"x": 1}),
|
||||
),
|
||||
mux,
|
||||
values_t,
|
||||
)
|
||||
|
||||
handle = await run.subgraphs.__aiter__().__anext__()
|
||||
assert handle.path == ("agent:abc",)
|
||||
with pytest.raises(RuntimeError, match="child boom"):
|
||||
await run.output()
|
||||
assert run._mux._events._error is not None
|
||||
|
||||
|
||||
def test_child_finalize_errors_propagate_to_sync_run() -> None:
|
||||
mux = StreamMux(
|
||||
factories=[
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
_ChildFinalizeExploder,
|
||||
],
|
||||
is_async=False,
|
||||
)
|
||||
values_t = mux.transformer_by_key("values")
|
||||
assert isinstance(values_t, ValuesTransformer)
|
||||
run = GraphRunStream(
|
||||
iter(
|
||||
[
|
||||
_stream_part(
|
||||
"tasks",
|
||||
("agent:abc",),
|
||||
{
|
||||
"id": "t1",
|
||||
"name": "tool",
|
||||
"input": None,
|
||||
"triggers": [],
|
||||
},
|
||||
)
|
||||
]
|
||||
),
|
||||
mux,
|
||||
values_t,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="child finalize boom"):
|
||||
_ = run.output
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_child_finalize_errors_propagate_to_async_run() -> None:
|
||||
mux = StreamMux(
|
||||
factories=[
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
LifecycleTransformer,
|
||||
SubgraphTransformer,
|
||||
_ChildFinalizeExploder,
|
||||
],
|
||||
is_async=True,
|
||||
)
|
||||
values_t = mux.transformer_by_key("values")
|
||||
assert isinstance(values_t, ValuesTransformer)
|
||||
run = AsyncGraphRunStream(
|
||||
_astream_parts(
|
||||
_stream_part(
|
||||
"tasks",
|
||||
("agent:abc",),
|
||||
{
|
||||
"id": "t1",
|
||||
"name": "tool",
|
||||
"input": None,
|
||||
"triggers": [],
|
||||
},
|
||||
)
|
||||
),
|
||||
mux,
|
||||
values_t,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="child afinalize boom"):
|
||||
await run.output()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end real-graph tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
value: str
|
||||
items: Annotated[list[str], operator.add]
|
||||
|
||||
|
||||
def _passthrough(state: _State) -> dict[str, Any]:
|
||||
return {"value": state["value"] + "!", "items": ["x"]}
|
||||
|
||||
|
||||
def _make_two_level_nested() -> Any:
|
||||
"""outer → middle → inner. Three Pregel instances, two nesting levels."""
|
||||
inner_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
inner_b.add_node("inner_node", _passthrough)
|
||||
inner_b.add_edge(START, "inner_node")
|
||||
inner_b.add_edge("inner_node", END)
|
||||
inner = inner_b.compile()
|
||||
|
||||
middle_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
middle_b.add_node("inner", inner)
|
||||
middle_b.add_edge(START, "inner")
|
||||
middle_b.add_edge("inner", END)
|
||||
middle = middle_b.compile()
|
||||
|
||||
outer_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
outer_b.add_node("middle", middle)
|
||||
outer_b.add_edge(START, "middle")
|
||||
outer_b.add_edge("middle", END)
|
||||
return outer_b.compile()
|
||||
|
||||
|
||||
def _item_node(item: str):
|
||||
def node(state: _State) -> dict[str, Any]:
|
||||
return {"items": [item]}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def _make_two_sibling_subgraphs() -> Any:
|
||||
"""outer → one → two, where both nodes are compiled subgraphs."""
|
||||
one_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
one_b.add_node("add_one", _item_node("one"))
|
||||
one_b.add_edge(START, "add_one")
|
||||
one_b.add_edge("add_one", END)
|
||||
one = one_b.compile()
|
||||
|
||||
two_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
two_b.add_node("add_two", _item_node("two"))
|
||||
two_b.add_edge(START, "add_two")
|
||||
two_b.add_edge("add_two", END)
|
||||
two = two_b.compile()
|
||||
|
||||
outer_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
outer_b.add_node("one", one)
|
||||
outer_b.add_node("two", two)
|
||||
outer_b.add_edge(START, "one")
|
||||
outer_b.add_edge("one", "two")
|
||||
outer_b.add_edge("two", END)
|
||||
return outer_b.compile()
|
||||
|
||||
|
||||
def _failing_node(state: _State) -> dict[str, Any]:
|
||||
raise ValueError("child boom")
|
||||
|
||||
|
||||
def _make_failing_nested() -> Any:
|
||||
inner_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
inner_b.add_node("fail", _failing_node)
|
||||
inner_b.add_edge(START, "fail")
|
||||
inner_b.add_edge("fail", END)
|
||||
inner = inner_b.compile()
|
||||
|
||||
outer_b: StateGraph = StateGraph(_State, input_schema=_State)
|
||||
outer_b.add_node("inner", inner)
|
||||
outer_b.add_edge(START, "inner")
|
||||
outer_b.add_edge("inner", END)
|
||||
return outer_b.compile()
|
||||
|
||||
|
||||
def test_stream_v2_real_graph_yields_subgraph_handles() -> None:
|
||||
"""Iterating `run.subgraphs` yields handles for direct-child subgraphs."""
|
||||
graph = _make_two_level_nested()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
handle_paths: list[tuple[str, ...]] = []
|
||||
final_status: dict[tuple[str, ...], str] = {}
|
||||
for handle in run.subgraphs:
|
||||
# Drill into the handle's projections inside the loop body so
|
||||
# the mini-mux is subscribed before the next pump cycle.
|
||||
list(handle.values)
|
||||
handle_paths.append(handle.path)
|
||||
final_status[handle.path] = handle.status
|
||||
|
||||
assert len(handle_paths) == 1
|
||||
assert handle_paths[0][0].startswith("middle:")
|
||||
assert final_status[handle_paths[0]] == "completed"
|
||||
|
||||
|
||||
def test_stream_v2_grandchild_visible_on_child_handle() -> None:
|
||||
"""Drilling into `handle.subgraphs` surfaces nested grandchildren."""
|
||||
graph = _make_two_level_nested()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
grandchild_paths: list[tuple[str, ...]] = []
|
||||
middle_path: tuple[str, ...] | None = None
|
||||
for middle_handle in run.subgraphs:
|
||||
# Subscribe to grandchildren before the next pump cycle.
|
||||
for inner_handle in middle_handle.subgraphs:
|
||||
# Subscribe to inner.values so its mini-mux drains.
|
||||
list(inner_handle.values)
|
||||
grandchild_paths.append(inner_handle.path)
|
||||
middle_path = middle_handle.path
|
||||
|
||||
assert middle_path is not None
|
||||
assert len(grandchild_paths) == 1
|
||||
inner_path = grandchild_paths[0]
|
||||
assert inner_path[1].startswith("inner:")
|
||||
assert inner_path[: len(middle_path)] == middle_path
|
||||
|
||||
|
||||
def test_subgraph_output_stops_at_own_terminal_without_draining_siblings() -> None:
|
||||
"""A handle's `output` must not pump past its terminal event.
|
||||
|
||||
If it over-pumps the root run, the second sibling handle is yielded
|
||||
only after it has already completed, so subscribing to `values`
|
||||
inside the loop body misses its events.
|
||||
"""
|
||||
graph = _make_two_sibling_subgraphs()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
paths: list[tuple[str, ...]] = []
|
||||
second_values: list[dict[str, Any]] = []
|
||||
for handle in run.subgraphs:
|
||||
paths.append(handle.path)
|
||||
if handle.graph_name == "one":
|
||||
assert handle.output is not None
|
||||
assert handle.status == "completed"
|
||||
elif handle.graph_name == "two":
|
||||
second_values = list(handle.values)
|
||||
|
||||
assert [path[0].split(":", 1)[0] for path in paths] == ["one", "two"]
|
||||
assert second_values
|
||||
assert second_values[-1]["items"] == ["one", "two"]
|
||||
|
||||
|
||||
def test_aborted_subgraph_handle_does_not_fail_parent_forwarding() -> None:
|
||||
graph = _make_two_sibling_subgraphs()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
seen: list[str | None] = []
|
||||
for handle in run.subgraphs:
|
||||
seen.append(handle.graph_name)
|
||||
if handle.graph_name == "one":
|
||||
# Subscribe before aborting to ensure forwarding into the
|
||||
# closed mini-mux would have raised without the closed check.
|
||||
iter(handle.values)
|
||||
handle.abort()
|
||||
elif handle.graph_name == "two":
|
||||
assert list(handle.values)
|
||||
|
||||
assert seen == ["one", "two"]
|
||||
|
||||
|
||||
def test_failed_subgraph_output_raises_terminal_error() -> None:
|
||||
graph = _make_failing_nested()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
handle = next(iter(run.subgraphs))
|
||||
with pytest.raises(RuntimeError, match="child boom"):
|
||||
_ = handle.output
|
||||
assert handle.status == "failed"
|
||||
assert handle.error == "child boom"
|
||||
@@ -0,0 +1,792 @@
|
||||
"""End-to-end tests exercising all stream_v2 projections together.
|
||||
|
||||
Each test builds a realistic graph (subgraphs, LLM calls, custom writers,
|
||||
interrupts) and verifies that every projection — values, messages, lifecycle,
|
||||
subgraphs, raw events, output, interleave — produces correct, consistent
|
||||
results through a single stream_v2 / astream_v2 run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
import sys
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.language_models import GenericFakeChatModel
|
||||
from langchain_core.language_models.chat_model_stream import (
|
||||
AsyncChatModelStream,
|
||||
ChatModelStream,
|
||||
)
|
||||
from langchain_core.messages import AIMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import MessagesState, StateGraph
|
||||
from langgraph.stream import StreamChannel, StreamTransformer
|
||||
from langgraph.stream._types import ProtocolEvent
|
||||
from langgraph.types import StreamWriter, interrupt
|
||||
|
||||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="Python 3.11+ is required for async contextvars support",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# State and graph builders
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AgentState(TypedDict):
|
||||
value: str
|
||||
items: Annotated[list[str], operator.add]
|
||||
|
||||
|
||||
def _make_nested_graph():
|
||||
"""Build a two-level graph with pure state transforms.
|
||||
|
||||
Structure:
|
||||
outer:
|
||||
router_node (state transform)
|
||||
inner_graph (compiled subgraph)
|
||||
|
||||
inner_graph:
|
||||
process_node (state transform)
|
||||
"""
|
||||
|
||||
def process_node(state: AgentState) -> dict[str, Any]:
|
||||
return {"value": state["value"] + "_processed", "items": ["processed"]}
|
||||
|
||||
inner_builder: StateGraph = StateGraph(AgentState, input_schema=AgentState)
|
||||
inner_builder.add_node("process_node", process_node)
|
||||
inner_builder.add_edge(START, "process_node")
|
||||
inner_builder.add_edge("process_node", END)
|
||||
inner_graph = inner_builder.compile()
|
||||
|
||||
def router_node(state: AgentState) -> dict[str, Any]:
|
||||
return {"value": state["value"] + "_routed", "items": ["routed"]}
|
||||
|
||||
outer_builder: StateGraph = StateGraph(AgentState, input_schema=AgentState)
|
||||
outer_builder.add_node("router", router_node)
|
||||
outer_builder.add_node("inner", inner_graph)
|
||||
outer_builder.add_edge(START, "router")
|
||||
outer_builder.add_edge("router", "inner")
|
||||
outer_builder.add_edge("inner", END)
|
||||
return outer_builder.compile()
|
||||
|
||||
|
||||
def _make_messages_graph():
|
||||
"""Flat graph with an LLM call for messages projection testing."""
|
||||
model = GenericFakeChatModel(messages=iter(["hello world"]))
|
||||
|
||||
def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": model.invoke(state["messages"])}
|
||||
|
||||
return (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
|
||||
def _make_messages_subgraph():
|
||||
"""Outer graph with a MessagesState subgraph that returns an AIMessage.
|
||||
|
||||
Uses the whole-message fallback path (node returns AIMessage directly)
|
||||
to exercise messages through a subgraph boundary.
|
||||
"""
|
||||
|
||||
def return_message(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": AIMessage(content="from subgraph", id="sub-msg-1")}
|
||||
|
||||
inner = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("return_message", return_message)
|
||||
.add_edge(START, "return_message")
|
||||
.add_edge("return_message", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
class OuterState(TypedDict):
|
||||
messages: Annotated[list[Any], operator.add]
|
||||
done: bool
|
||||
|
||||
def pre_node(state: OuterState) -> dict[str, Any]:
|
||||
return {"done": False}
|
||||
|
||||
return (
|
||||
StateGraph(OuterState)
|
||||
.add_node("pre", pre_node)
|
||||
.add_node("inner", inner)
|
||||
.add_edge(START, "pre")
|
||||
.add_edge("pre", "inner")
|
||||
.add_edge("inner", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
|
||||
def _make_custom_writer_graph():
|
||||
"""Graph where a node emits custom stream events via StreamWriter."""
|
||||
|
||||
def writer_node(state: AgentState, *, writer: StreamWriter) -> dict[str, Any]:
|
||||
writer({"step": "start", "detail": "beginning work"})
|
||||
writer({"step": "middle", "detail": "processing"})
|
||||
writer({"step": "end", "detail": "done"})
|
||||
return {"value": state["value"] + "_custom", "items": ["custom"]}
|
||||
|
||||
builder = StateGraph(AgentState)
|
||||
builder.add_node("writer_node", writer_node)
|
||||
builder.add_edge(START, "writer_node")
|
||||
builder.add_edge("writer_node", END)
|
||||
return builder.compile()
|
||||
|
||||
|
||||
def _make_interrupt_graph():
|
||||
"""Graph that interrupts after the first node."""
|
||||
|
||||
def step_one(state: AgentState) -> dict[str, Any]:
|
||||
return {"value": state["value"] + "_step1", "items": ["step1"]}
|
||||
|
||||
def step_two(state: AgentState) -> dict[str, Any]:
|
||||
answer = interrupt("need approval")
|
||||
return {"value": state["value"] + f"_{answer}", "items": ["step2"]}
|
||||
|
||||
builder = StateGraph(AgentState)
|
||||
builder.add_node("step_one", step_one)
|
||||
builder.add_node("step_two", step_two)
|
||||
builder.add_edge(START, "step_one")
|
||||
builder.add_edge("step_one", "step_two")
|
||||
builder.add_edge("step_two", END)
|
||||
return builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
|
||||
def _make_error_subgraph():
|
||||
"""Graph with a subgraph that raises."""
|
||||
|
||||
def failing_node(state: AgentState) -> dict[str, Any]:
|
||||
raise ValueError("subgraph explosion")
|
||||
|
||||
inner_builder = StateGraph(AgentState)
|
||||
inner_builder.add_node("fail", failing_node)
|
||||
inner_builder.add_edge(START, "fail")
|
||||
inner_builder.add_edge("fail", END)
|
||||
inner = inner_builder.compile()
|
||||
|
||||
outer_builder = StateGraph(AgentState)
|
||||
outer_builder.add_node("inner", inner)
|
||||
outer_builder.add_edge(START, "inner")
|
||||
outer_builder.add_edge("inner", END)
|
||||
return outer_builder.compile()
|
||||
|
||||
|
||||
class _CustomPassthroughTransformer(StreamTransformer):
|
||||
required_stream_modes = ("custom",)
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
class _CounterTransformer(StreamTransformer):
|
||||
"""Custom transformer that counts values events via a StreamChannel."""
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._channel: StreamChannel[int] = StreamChannel("counter")
|
||||
self._count = 0
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"counter": self._channel}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] == "values":
|
||||
self._count += 1
|
||||
self._channel.push(self._count)
|
||||
return True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sync end-to-end: all projections on nested graph
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStreamV2E2ESync:
|
||||
def test_all_projections_nested_graph(self) -> None:
|
||||
"""Run a nested graph through stream_v2 and verify values + lifecycle."""
|
||||
graph = _make_nested_graph()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
values_snapshots: list[dict[str, Any]] = []
|
||||
lifecycle_events: list[dict[str, Any]] = []
|
||||
for name, item in run.interleave("values", "lifecycle"):
|
||||
if name == "values":
|
||||
values_snapshots.append(item)
|
||||
elif name == "lifecycle":
|
||||
lifecycle_events.append(item)
|
||||
|
||||
assert len(values_snapshots) >= 1
|
||||
final = values_snapshots[-1]
|
||||
assert "routed" in final["items"]
|
||||
assert "processed" in final["items"]
|
||||
assert "_routed" in final["value"]
|
||||
assert "_processed" in final["value"]
|
||||
|
||||
assert len(lifecycle_events) >= 2
|
||||
started = [e for e in lifecycle_events if e["event"] == "started"]
|
||||
completed = [e for e in lifecycle_events if e["event"] == "completed"]
|
||||
assert len(started) >= 1
|
||||
assert len(completed) >= 1
|
||||
|
||||
def test_subgraph_handles_with_drill_down(self) -> None:
|
||||
"""Subgraph handles yield and support values drill-down."""
|
||||
graph = _make_nested_graph()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
handles = []
|
||||
for handle in run.subgraphs:
|
||||
child_values = list(handle.values)
|
||||
handles.append(
|
||||
{
|
||||
"path": handle.path,
|
||||
"graph_name": handle.graph_name,
|
||||
"values_count": len(child_values),
|
||||
}
|
||||
)
|
||||
|
||||
assert len(handles) >= 1
|
||||
assert handles[0]["values_count"] >= 1
|
||||
|
||||
output = run.output
|
||||
assert output is not None
|
||||
assert "_routed" in output["value"]
|
||||
assert "_processed" in output["value"]
|
||||
|
||||
def test_raw_events_have_monotonic_seq(self) -> None:
|
||||
"""Raw protocol events have monotonically increasing seq numbers."""
|
||||
graph = _make_nested_graph()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
events = list(run)
|
||||
assert len(events) > 0
|
||||
|
||||
seqs = [e["seq"] for e in events]
|
||||
for i in range(1, len(seqs)):
|
||||
assert seqs[i] > seqs[i - 1], f"seq not monotonic at {i}: {seqs}"
|
||||
|
||||
for event in events:
|
||||
assert event["type"] == "event"
|
||||
assert "method" in event
|
||||
assert isinstance(event["params"]["timestamp"], int)
|
||||
|
||||
def test_output_matches_final_values_snapshot(self) -> None:
|
||||
"""output property returns the same state as the last values snapshot."""
|
||||
run1 = _make_nested_graph().stream_v2({"value": "x", "items": []})
|
||||
snapshots = list(run1.values)
|
||||
final_via_values = snapshots[-1]
|
||||
|
||||
run2 = _make_nested_graph().stream_v2({"value": "x", "items": []})
|
||||
final_via_output = run2.output
|
||||
|
||||
assert final_via_values == final_via_output
|
||||
|
||||
def test_context_manager_and_abort(self) -> None:
|
||||
"""Context manager calls abort, marking the stream exhausted."""
|
||||
graph = _make_nested_graph()
|
||||
with graph.stream_v2({"value": "x", "items": []}) as run:
|
||||
first_val = next(iter(run.values))
|
||||
assert isinstance(first_val, dict)
|
||||
assert run._exhausted is True
|
||||
|
||||
def test_extensions_has_all_native_keys(self) -> None:
|
||||
"""Extensions dict exposes all native projection keys."""
|
||||
graph = _make_nested_graph()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
_ = run.output
|
||||
|
||||
assert "values" in run.extensions
|
||||
assert "messages" in run.extensions
|
||||
assert "lifecycle" in run.extensions
|
||||
assert "subgraphs" in run.extensions
|
||||
assert run.values is run.extensions["values"]
|
||||
assert run.messages is run.extensions["messages"]
|
||||
assert run.lifecycle is run.extensions["lifecycle"]
|
||||
assert run.subgraphs is run.extensions["subgraphs"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sync: messages projection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStreamV2E2EMessages:
|
||||
def test_messages_projection_from_invoke(self) -> None:
|
||||
"""Messages projection captures LLM calls via model.invoke() auto-routing."""
|
||||
graph = _make_messages_graph()
|
||||
run = graph.stream_v2({"messages": "hi"})
|
||||
streams = list(run.messages)
|
||||
|
||||
assert len(streams) >= 1
|
||||
for stream in streams:
|
||||
assert isinstance(stream, ChatModelStream)
|
||||
assert streams[0].output.text == "hello world"
|
||||
|
||||
def test_messages_text_deltas(self) -> None:
|
||||
"""Text deltas from the messages projection concatenate correctly."""
|
||||
model = GenericFakeChatModel(messages=iter(["streamed answer"]))
|
||||
|
||||
def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": model.invoke(state["messages"])}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = graph.stream_v2({"messages": "go"})
|
||||
(stream,) = list(run.messages)
|
||||
assert "".join(stream.text) == "streamed answer"
|
||||
|
||||
def test_messages_from_whole_ai_message(self) -> None:
|
||||
"""Node returning AIMessage directly produces a complete stream."""
|
||||
|
||||
def return_msg(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": AIMessage(content="hardcoded", id="msg-1")}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("return_msg", return_msg)
|
||||
.add_edge(START, "return_msg")
|
||||
.add_edge("return_msg", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = graph.stream_v2({"messages": "hi"})
|
||||
(stream,) = list(run.messages)
|
||||
assert stream.output.text == "hardcoded"
|
||||
assert stream.message_id == "msg-1"
|
||||
|
||||
def test_root_messages_only_shows_root_scope(self) -> None:
|
||||
"""Root messages projection doesn't surface subgraph-scoped messages."""
|
||||
graph = _make_messages_subgraph()
|
||||
run = graph.stream_v2({"messages": ["hi"], "done": False})
|
||||
root_streams = list(run.messages)
|
||||
# The message is emitted inside the subgraph, so the root
|
||||
# messages projection (scoped to root namespace) doesn't see it.
|
||||
assert root_streams == []
|
||||
|
||||
def test_subgraph_handle_messages_drill_down(self) -> None:
|
||||
"""Drilling into subgraph handle's messages surfaces subgraph messages."""
|
||||
graph = _make_messages_subgraph()
|
||||
run = graph.stream_v2({"messages": ["hi"], "done": False})
|
||||
|
||||
found_messages = False
|
||||
for handle in run.subgraphs:
|
||||
child_messages = list(handle.messages)
|
||||
if child_messages:
|
||||
found_messages = True
|
||||
assert isinstance(child_messages[0], ChatModelStream)
|
||||
assert child_messages[0].output.text == "from subgraph"
|
||||
assert found_messages
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sync: custom stream writer + custom transformer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStreamV2E2ECustom:
|
||||
def test_custom_events_with_passthrough_transformer(self) -> None:
|
||||
"""Custom StreamWriter events appear on the main log when a
|
||||
transformer declares the custom mode."""
|
||||
graph = _make_custom_writer_graph()
|
||||
run = graph.stream_v2(
|
||||
{"value": "x", "items": []},
|
||||
transformers=[_CustomPassthroughTransformer],
|
||||
)
|
||||
events = list(run)
|
||||
custom = [e for e in events if e["method"] == "custom"]
|
||||
assert len(custom) == 3
|
||||
steps = [e["params"]["data"]["step"] for e in custom]
|
||||
assert steps == ["start", "middle", "end"]
|
||||
|
||||
def test_custom_events_suppressed_without_transformer(self) -> None:
|
||||
"""Without a custom-mode transformer, custom events don't flow."""
|
||||
graph = _make_custom_writer_graph()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
events = list(run)
|
||||
custom = [e for e in events if e["method"] == "custom"]
|
||||
assert custom == []
|
||||
|
||||
def test_custom_transformer_with_stream_channel(self) -> None:
|
||||
"""A custom transformer with a StreamChannel produces extension data."""
|
||||
graph = _make_nested_graph()
|
||||
run = graph.stream_v2(
|
||||
{"value": "x", "items": []},
|
||||
transformers=[_CounterTransformer],
|
||||
)
|
||||
|
||||
assert "counter" in run.extensions
|
||||
counter_iter = iter(run.extensions["counter"])
|
||||
_ = run.output
|
||||
counts = list(counter_iter)
|
||||
assert len(counts) >= 1
|
||||
assert all(isinstance(c, int) for c in counts)
|
||||
assert counts == sorted(counts)
|
||||
|
||||
def test_custom_channel_events_on_main_log(self) -> None:
|
||||
"""StreamChannel auto-forward injects custom:<name> events into the main log."""
|
||||
graph = _make_nested_graph()
|
||||
run = graph.stream_v2(
|
||||
{"value": "x", "items": []},
|
||||
transformers=[_CounterTransformer],
|
||||
)
|
||||
events = list(run)
|
||||
counter_events = [e for e in events if e["method"] == "custom:counter"]
|
||||
assert len(counter_events) >= 1
|
||||
assert all(isinstance(e["params"]["data"], int) for e in counter_events)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sync: interrupt handling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStreamV2E2EInterrupt:
|
||||
def test_interrupt_sets_flags_and_surfaces_interrupts(self) -> None:
|
||||
"""Interrupted run has correct flags and interrupt payloads."""
|
||||
graph = _make_interrupt_graph()
|
||||
config: dict[str, Any] = {"configurable": {"thread_id": "int-1"}}
|
||||
run = graph.stream_v2({"value": "x", "items": []}, config)
|
||||
|
||||
output = run.output
|
||||
assert output is not None
|
||||
assert run.interrupted is True
|
||||
assert len(run.interrupts) > 0
|
||||
assert output["items"] == ["step1"]
|
||||
assert "_step1" in output["value"]
|
||||
|
||||
def test_interrupt_values_snapshot_has_partial_state(self) -> None:
|
||||
"""Values snapshots captured before the interrupt reflect partial state."""
|
||||
graph = _make_interrupt_graph()
|
||||
config: dict[str, Any] = {"configurable": {"thread_id": "int-2"}}
|
||||
run = graph.stream_v2({"value": "x", "items": []}, config)
|
||||
|
||||
snapshots = list(run.values)
|
||||
assert len(snapshots) >= 1
|
||||
last = snapshots[-1]
|
||||
assert "step1" in last["items"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sync: error propagation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStreamV2E2EErrors:
|
||||
def test_subgraph_error_propagates_through_output(self) -> None:
|
||||
"""Error in a subgraph propagates through output."""
|
||||
graph = _make_error_subgraph()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
with pytest.raises(ValueError, match="subgraph explosion"):
|
||||
_ = run.output
|
||||
|
||||
def test_subgraph_error_propagates_through_raw_events(self) -> None:
|
||||
graph = _make_error_subgraph()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
with pytest.raises(ValueError, match="subgraph explosion"):
|
||||
list(run)
|
||||
|
||||
def test_error_subgraph_handle_status(self) -> None:
|
||||
"""Subgraph handle surfaces the error status."""
|
||||
graph = _make_error_subgraph()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
handle = next(iter(run.subgraphs))
|
||||
with pytest.raises(RuntimeError, match="subgraph explosion"):
|
||||
_ = handle.output
|
||||
assert handle.status == "failed"
|
||||
assert handle.error == "subgraph explosion"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async end-to-end
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@NEEDS_CONTEXTVARS
|
||||
class TestStreamV2E2EAsync:
|
||||
async def test_all_projections_async(self) -> None:
|
||||
"""Async run exercises values projection."""
|
||||
graph = _make_nested_graph()
|
||||
run = await graph.astream_v2({"value": "x", "items": []})
|
||||
|
||||
values_snapshots = [s async for s in run.values]
|
||||
assert len(values_snapshots) >= 1
|
||||
final = values_snapshots[-1]
|
||||
assert "_routed" in final["value"]
|
||||
assert "_processed" in final["value"]
|
||||
|
||||
async def test_async_output(self) -> None:
|
||||
"""Async output returns the final state."""
|
||||
graph = _make_nested_graph()
|
||||
run = await graph.astream_v2({"value": "x", "items": []})
|
||||
output = await run.output()
|
||||
assert output is not None
|
||||
assert output["value"] == "x_routed_processed"
|
||||
assert "routed" in output["items"]
|
||||
assert "processed" in output["items"]
|
||||
|
||||
async def test_async_raw_events(self) -> None:
|
||||
"""Async raw event iteration yields well-formed ProtocolEvents."""
|
||||
graph = _make_nested_graph()
|
||||
run = await graph.astream_v2({"value": "x", "items": []})
|
||||
events = [e async for e in run]
|
||||
assert len(events) > 0
|
||||
seqs = [e["seq"] for e in events]
|
||||
for i in range(1, len(seqs)):
|
||||
assert seqs[i] > seqs[i - 1]
|
||||
|
||||
async def test_async_messages_projection(self) -> None:
|
||||
"""Async messages projection captures LLM streams."""
|
||||
model = GenericFakeChatModel(messages=iter(["async answer"]))
|
||||
|
||||
async def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": await model.ainvoke(state["messages"])}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = await graph.astream_v2({"messages": "hi"})
|
||||
streams = [s async for s in run.messages]
|
||||
assert len(streams) >= 1
|
||||
for s in streams:
|
||||
assert isinstance(s, AsyncChatModelStream)
|
||||
assert (await streams[0].output).text == "async answer"
|
||||
|
||||
async def test_async_interrupt(self) -> None:
|
||||
"""Async interrupted run has correct flags."""
|
||||
graph = _make_interrupt_graph()
|
||||
config: dict[str, Any] = {"configurable": {"thread_id": "async-int-1"}}
|
||||
run = await graph.astream_v2({"value": "x", "items": []}, config)
|
||||
|
||||
output = await run.output()
|
||||
assert output is not None
|
||||
assert await run.interrupted() is True
|
||||
assert len(await run.interrupts()) > 0
|
||||
|
||||
async def test_async_error_propagation(self) -> None:
|
||||
"""Async error from subgraph propagates through output."""
|
||||
graph = _make_error_subgraph()
|
||||
run = await graph.astream_v2({"value": "x", "items": []})
|
||||
with pytest.raises(ValueError, match="subgraph explosion"):
|
||||
await run.output()
|
||||
|
||||
async def test_async_context_manager(self) -> None:
|
||||
"""Async context manager calls abort on exit."""
|
||||
graph = _make_nested_graph()
|
||||
run = await graph.astream_v2({"value": "x", "items": []})
|
||||
async with run:
|
||||
_ = await anext(aiter(run.values))
|
||||
assert run._exhausted is True
|
||||
|
||||
async def test_async_extensions_present(self) -> None:
|
||||
"""Async run has all native extensions."""
|
||||
graph = _make_nested_graph()
|
||||
run = await graph.astream_v2({"value": "x", "items": []})
|
||||
_ = await run.output()
|
||||
assert "values" in run.extensions
|
||||
assert "messages" in run.extensions
|
||||
assert "lifecycle" in run.extensions
|
||||
assert "subgraphs" in run.extensions
|
||||
|
||||
async def test_async_custom_transformer(self) -> None:
|
||||
"""Async custom transformer with StreamChannel works."""
|
||||
graph = _make_nested_graph()
|
||||
run = await graph.astream_v2(
|
||||
{"value": "x", "items": []},
|
||||
transformers=[_CounterTransformer],
|
||||
)
|
||||
assert "counter" in run.extensions
|
||||
counter_cursor = aiter(run.extensions["counter"])
|
||||
_ = await run.output()
|
||||
counts = [c async for c in counter_cursor]
|
||||
assert len(counts) >= 1
|
||||
assert counts == sorted(counts)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sync: combined projections stress test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStreamV2E2ECombined:
|
||||
def test_interleave_all_native_projections(self) -> None:
|
||||
"""Interleave values + messages + lifecycle without deadlock."""
|
||||
graph = _make_nested_graph()
|
||||
run = graph.stream_v2({"value": "x", "items": []})
|
||||
|
||||
seen_names: set[str] = set()
|
||||
for name, _item in run.interleave("values", "messages", "lifecycle"):
|
||||
seen_names.add(name)
|
||||
|
||||
assert "values" in seen_names
|
||||
assert "lifecycle" in seen_names
|
||||
|
||||
def test_multiple_custom_transformers(self) -> None:
|
||||
"""Multiple custom transformers can coexist."""
|
||||
|
||||
class TagTransformer(StreamTransformer):
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._channel: StreamChannel[str] = StreamChannel("tags")
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"tags": self._channel}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] == "values":
|
||||
self._channel.push(
|
||||
f"tag:{event['params']['data'].get('value', '')}"
|
||||
)
|
||||
return True
|
||||
|
||||
graph = _make_nested_graph()
|
||||
run = graph.stream_v2(
|
||||
{"value": "x", "items": []},
|
||||
transformers=[_CounterTransformer, TagTransformer],
|
||||
)
|
||||
|
||||
assert "counter" in run.extensions
|
||||
assert "tags" in run.extensions
|
||||
|
||||
counter_iter = iter(run.extensions["counter"])
|
||||
tags_iter = iter(run.extensions["tags"])
|
||||
_ = run.output
|
||||
counts = list(counter_iter)
|
||||
tags = list(tags_iter)
|
||||
|
||||
assert len(counts) >= 1
|
||||
assert len(tags) >= 1
|
||||
assert all(t.startswith("tag:") for t in tags)
|
||||
|
||||
def test_two_sibling_subgraphs_both_discoverable(self) -> None:
|
||||
"""Two sequential subgraph invocations produce two handles."""
|
||||
|
||||
class _S(TypedDict):
|
||||
items: Annotated[list[str], operator.add]
|
||||
|
||||
def _item(name: str):
|
||||
def node(state: _S) -> dict[str, Any]:
|
||||
return {"items": [name]}
|
||||
|
||||
return node
|
||||
|
||||
inner_a = (
|
||||
StateGraph(_S)
|
||||
.add_node("add_a", _item("a"))
|
||||
.add_edge(START, "add_a")
|
||||
.add_edge("add_a", END)
|
||||
.compile()
|
||||
)
|
||||
inner_b = (
|
||||
StateGraph(_S)
|
||||
.add_node("add_b", _item("b"))
|
||||
.add_edge(START, "add_b")
|
||||
.add_edge("add_b", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
outer = (
|
||||
StateGraph(_S)
|
||||
.add_node("sub_a", inner_a)
|
||||
.add_node("sub_b", inner_b)
|
||||
.add_edge(START, "sub_a")
|
||||
.add_edge("sub_a", "sub_b")
|
||||
.add_edge("sub_b", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = outer.stream_v2({"items": []})
|
||||
handles = []
|
||||
for handle in run.subgraphs:
|
||||
list(handle.values)
|
||||
handles.append(handle)
|
||||
|
||||
assert len(handles) == 2
|
||||
names = [h.graph_name for h in handles]
|
||||
assert "sub_a" in names
|
||||
assert "sub_b" in names
|
||||
assert all(h.status == "completed" for h in handles)
|
||||
|
||||
output = run.output
|
||||
assert output is not None
|
||||
assert set(output["items"]) == {"a", "b"}
|
||||
|
||||
def test_lifecycle_matches_subgraph_handles(self) -> None:
|
||||
"""Lifecycle events and subgraph handles agree on discovered subgraphs."""
|
||||
run1 = _make_nested_graph().stream_v2({"value": "x", "items": []})
|
||||
handle_paths: list[tuple[str, ...]] = []
|
||||
for handle in run1.subgraphs:
|
||||
list(handle.values)
|
||||
handle_paths.append(handle.path)
|
||||
|
||||
run2 = _make_nested_graph().stream_v2({"value": "x", "items": []})
|
||||
lifecycle = list(run2.lifecycle)
|
||||
|
||||
started_ns = [
|
||||
tuple(e["namespace"]) for e in lifecycle if e["event"] == "started"
|
||||
]
|
||||
# Handle paths use format "graph_name:call_id", lifecycle namespaces
|
||||
# use the same format. Both should have the same graph_name prefix.
|
||||
handle_prefixes = {p[0].split(":")[0] for p in handle_paths}
|
||||
lifecycle_prefixes = {ns[0].split(":")[0] for ns in started_ns}
|
||||
assert handle_prefixes == lifecycle_prefixes
|
||||
|
||||
def test_values_plus_messages_plus_custom(self) -> None:
|
||||
"""Values, messages, and a custom transformer all produce data in one run."""
|
||||
model = GenericFakeChatModel(messages=iter(["combined test"]))
|
||||
|
||||
def call_model(state: MessagesState) -> dict[str, Any]:
|
||||
return {"messages": model.invoke(state["messages"])}
|
||||
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("call_model", call_model)
|
||||
.add_edge(START, "call_model")
|
||||
.add_edge("call_model", END)
|
||||
.compile()
|
||||
)
|
||||
|
||||
run = graph.stream_v2(
|
||||
{"messages": "hi"},
|
||||
transformers=[_CounterTransformer],
|
||||
)
|
||||
|
||||
counter_iter = iter(run.extensions["counter"])
|
||||
values_iter = iter(run.values)
|
||||
messages_iter = iter(run.messages)
|
||||
|
||||
values = list(values_iter)
|
||||
messages = list(messages_iter)
|
||||
counts = list(counter_iter)
|
||||
|
||||
assert len(values) >= 1
|
||||
assert len(messages) >= 1
|
||||
assert len(counts) >= 1
|
||||
assert messages[0].output.text == "combined test"
|
||||
@@ -0,0 +1,290 @@
|
||||
"""Tests for StreamToolCallHandler and ToolRuntime.emit_output_delta.
|
||||
|
||||
These tests exercise the langgraph-core piece in isolation — the prebuilt
|
||||
`ToolCallTransformer` has its own test file. Here we feed real graphs
|
||||
through `Pregel.stream(stream_mode=["tools", ...])` and inspect the raw
|
||||
`(ns, mode, payload)` tuples on the `tools` channel.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.tools import tool
|
||||
from langgraph.prebuilt import ToolNode, ToolRuntime
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.pregel._tools import _tool_call_writer
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
|
||||
|
||||
def _caller_sync(tool_name: str, tool_args: dict[str, Any], tc_id: str = "tc1"):
|
||||
def caller(state: _State) -> dict:
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[{"name": tool_name, "args": tool_args, "id": tc_id}],
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
return caller
|
||||
|
||||
|
||||
def _caller_async(tool_name: str, tool_args: dict[str, Any], tc_id: str = "tc1"):
|
||||
async def caller(state: _State) -> dict:
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[{"name": tool_name, "args": tool_args, "id": tc_id}],
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
return caller
|
||||
|
||||
|
||||
def _build_graph(caller, tools) -> Any:
|
||||
sg = StateGraph(_State)
|
||||
sg.add_node("caller", caller)
|
||||
sg.add_node("tools", ToolNode(tools))
|
||||
sg.add_edge(START, "caller")
|
||||
sg.add_edge("caller", "tools")
|
||||
sg.add_edge("tools", END)
|
||||
return sg.compile()
|
||||
|
||||
|
||||
def _tool_events(stream) -> list[tuple[tuple[str, ...], dict]]:
|
||||
"""Collect `(ns, payload)` for every `tools`-mode chunk."""
|
||||
out: list[tuple[tuple[str, ...], dict]] = []
|
||||
for ns, mode, payload in stream:
|
||||
if mode == "tools":
|
||||
out.append((tuple(ns), payload))
|
||||
return out
|
||||
|
||||
|
||||
class TestSyncGraphSyncTool:
|
||||
def test_started_finished_cycle(self) -> None:
|
||||
@tool
|
||||
def echo(text: str) -> str:
|
||||
"""echo."""
|
||||
return f"echoed:{text}"
|
||||
|
||||
graph = _build_graph(_caller_sync("echo", {"text": "hi"}), [echo])
|
||||
events = _tool_events(
|
||||
graph.stream(
|
||||
{"messages": []},
|
||||
stream_mode=["tools"],
|
||||
subgraphs=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert [p["event"] for _, p in events] == [
|
||||
"tool-started",
|
||||
"tool-finished",
|
||||
]
|
||||
assert events[0][1]["tool_call_id"] == "tc1"
|
||||
assert events[0][1]["tool_name"] == "echo"
|
||||
assert events[0][1]["input"] == {"text": "hi"}
|
||||
# ToolNode wraps the return in a ToolMessage.
|
||||
assert events[1][1]["tool_call_id"] == "tc1"
|
||||
|
||||
def test_emit_output_delta_produces_delta_events(self) -> None:
|
||||
@tool
|
||||
def streaming_echo(text: str, runtime: ToolRuntime) -> str:
|
||||
"""stream chunks."""
|
||||
for chunk in ("a", "b", "c"):
|
||||
runtime.emit_output_delta(chunk)
|
||||
return text
|
||||
|
||||
graph = _build_graph(
|
||||
_caller_sync("streaming_echo", {"text": "x"}), [streaming_echo]
|
||||
)
|
||||
events = _tool_events(
|
||||
graph.stream(
|
||||
{"messages": []},
|
||||
stream_mode=["tools"],
|
||||
subgraphs=True,
|
||||
)
|
||||
)
|
||||
|
||||
deltas = [p["delta"] for _, p in events if p["event"] == "tool-output-delta"]
|
||||
assert deltas == ["a", "b", "c"]
|
||||
# The deltas must be bracketed by started and finished.
|
||||
ordered = [p["event"] for _, p in events]
|
||||
assert ordered[0] == "tool-started"
|
||||
assert ordered[-1] == "tool-finished"
|
||||
|
||||
def test_tool_error_event(self) -> None:
|
||||
@tool
|
||||
def boom() -> str:
|
||||
"""raises."""
|
||||
raise ValueError("nope")
|
||||
|
||||
graph = _build_graph(_caller_sync("boom", {}), [boom])
|
||||
events: list[tuple[tuple[str, ...], dict]] = []
|
||||
with pytest.raises(ValueError, match="nope"):
|
||||
for ns, mode, payload in graph.stream(
|
||||
{"messages": []},
|
||||
stream_mode=["tools"],
|
||||
subgraphs=True,
|
||||
):
|
||||
if mode == "tools":
|
||||
events.append((tuple(ns), payload))
|
||||
|
||||
kinds = [p["event"] for _, p in events]
|
||||
assert kinds == ["tool-started", "tool-error"]
|
||||
assert events[1][1]["message"] == "nope"
|
||||
|
||||
def test_writer_unset_outside_tool(self) -> None:
|
||||
# Outside any tool body the ContextVar that ToolRuntime reads
|
||||
# is unset — emitting from there would be a no-op.
|
||||
assert _tool_call_writer.get() is None
|
||||
|
||||
def test_no_events_without_tools_mode(self) -> None:
|
||||
@tool
|
||||
def echo(text: str) -> str:
|
||||
"""echo."""
|
||||
return text
|
||||
|
||||
graph = _build_graph(_caller_sync("echo", {"text": "hi"}), [echo])
|
||||
# No "tools" in stream_mode — handler is not attached and zero
|
||||
# `tools`-method events fire.
|
||||
chunks = list(
|
||||
graph.stream(
|
||||
{"messages": []},
|
||||
stream_mode=["values"],
|
||||
subgraphs=True,
|
||||
)
|
||||
)
|
||||
assert all(
|
||||
not (isinstance(c, tuple) and len(c) == 3 and c[1] == "tools")
|
||||
for c in chunks
|
||||
)
|
||||
|
||||
|
||||
class TestAsyncGraphAsyncTool:
|
||||
@pytest.mark.anyio
|
||||
async def test_async_tool_produces_events(self) -> None:
|
||||
@tool
|
||||
async def aecho(text: str, runtime: ToolRuntime) -> str:
|
||||
"""async echo."""
|
||||
runtime.emit_output_delta(text)
|
||||
return f"got:{text}"
|
||||
|
||||
graph = _build_graph(_caller_async("aecho", {"text": "hi"}), [aecho])
|
||||
events: list[tuple[tuple[str, ...], dict]] = []
|
||||
async for ns, mode, payload in graph.astream(
|
||||
{"messages": []},
|
||||
stream_mode=["tools"],
|
||||
subgraphs=True,
|
||||
):
|
||||
if mode == "tools":
|
||||
events.append((tuple(ns), payload))
|
||||
|
||||
kinds = [p["event"] for _, p in events]
|
||||
assert kinds == ["tool-started", "tool-output-delta", "tool-finished"]
|
||||
assert events[1][1]["delta"] == "hi"
|
||||
|
||||
|
||||
class TestConcurrentToolCalls:
|
||||
def test_parallel_tool_calls_do_not_bleed(self) -> None:
|
||||
@tool
|
||||
def streamer(marker: str, runtime: ToolRuntime) -> str:
|
||||
"""emits marker twice."""
|
||||
runtime.emit_output_delta(f"{marker}-1")
|
||||
runtime.emit_output_delta(f"{marker}-2")
|
||||
return marker
|
||||
|
||||
def caller(state: _State) -> dict:
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{"name": "streamer", "args": {"marker": "A"}, "id": "a"},
|
||||
{"name": "streamer", "args": {"marker": "B"}, "id": "b"},
|
||||
],
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
graph = _build_graph(caller, [streamer])
|
||||
events = _tool_events(
|
||||
graph.stream(
|
||||
{"messages": []},
|
||||
stream_mode=["tools"],
|
||||
subgraphs=True,
|
||||
)
|
||||
)
|
||||
|
||||
# Group deltas by tool_call_id.
|
||||
by_id: dict[str, list[str]] = {}
|
||||
for _, p in events:
|
||||
if p["event"] == "tool-output-delta":
|
||||
by_id.setdefault(p["tool_call_id"], []).append(p["delta"])
|
||||
assert by_id["a"] == ["A-1", "A-2"]
|
||||
assert by_id["b"] == ["B-1", "B-2"]
|
||||
|
||||
|
||||
class TestSubgraphNamespacePropagation:
|
||||
def test_tool_inside_subgraph_emits_with_subgraph_ns(self) -> None:
|
||||
@tool
|
||||
def inner_tool(text: str) -> str:
|
||||
"""inner tool."""
|
||||
return text
|
||||
|
||||
def sub_caller(state: _State) -> dict:
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "inner_tool",
|
||||
"args": {"text": "x"},
|
||||
"id": "tc1",
|
||||
}
|
||||
],
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
inner = StateGraph(_State)
|
||||
inner.add_node("sub_caller", sub_caller)
|
||||
inner.add_node("sub_tools", ToolNode([inner_tool]))
|
||||
inner.add_edge(START, "sub_caller")
|
||||
inner.add_edge("sub_caller", "sub_tools")
|
||||
inner.add_edge("sub_tools", END)
|
||||
inner_graph = inner.compile()
|
||||
|
||||
outer = StateGraph(_State)
|
||||
outer.add_node("sub", inner_graph)
|
||||
outer.add_edge(START, "sub")
|
||||
outer.add_edge("sub", END)
|
||||
graph = outer.compile()
|
||||
|
||||
events = _tool_events(
|
||||
graph.stream(
|
||||
{"messages": []},
|
||||
stream_mode=["tools"],
|
||||
subgraphs=True,
|
||||
)
|
||||
)
|
||||
|
||||
# All `tools` events should carry a non-empty namespace rooted
|
||||
# at the `sub` node.
|
||||
assert events, "expected at least one tools event"
|
||||
for ns, _ in events:
|
||||
assert ns # non-empty
|
||||
assert ns[0].startswith("sub:")
|
||||
Generated
+17
-4
@@ -1348,10 +1348,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.1"
|
||||
version = "1.3.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
{ name = "langchain-protocol" },
|
||||
{ name = "langsmith" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pydantic" },
|
||||
@@ -1360,9 +1361,21 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f4/fe/abeae8d0d2899e191d67c6c7f065f7e52a953f30b21ef327fa49084e4af9/langchain_core-1.3.1.tar.gz", hash = "sha256:41b384055799f93f34520df6bf7b80e2e5e23153cdfd46874251c6c9916ea030", size = 862403, upload-time = "2026-04-23T18:54:01.857Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a1/c2/8493be505921857988db068b7c027f28a9b1587b4425c6a32b1221c9c9fe/langchain_core-1.3.1-py3-none-any.whl", hash = "sha256:8b13d19d3bed3f4768df12c7f6932d2ada715f3ac9fd020c63d28c693968269e", size = 515879, upload-time = "2026-04-23T18:53:59.94Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langchain-protocol"
|
||||
version = "0.0.12"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/5c/51/1157009b6f94e6e58be58fa8b620187d657909a8b36a6bf5b0c52a2711f6/langchain_protocol-0.0.12.tar.gz", hash = "sha256:5e14c434290a705c9510fdb1a83ecf7561a5e6e0dfd053930ade80dba069269f", size = 6408, upload-time = "2026-04-25T01:05:01.489Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/95/82/3431e3061c917439589fa88a6b23c9bc0e154cba0f05d2e895a68c76ff74/langchain_protocol-0.0.12-py3-none-any.whl", hash = "sha256:402b61f42d4139692528cf37226c367bb6efc8ff8165b29380accb0abfece7b2", size = 6639, upload-time = "2026-04-25T01:05:00.487Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1439,7 +1452,7 @@ test = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = ">=1.3.2,<2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
|
||||
{ name = "langgraph-sdk", editable = "../sdk-py" },
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""langgraph.prebuilt exposes a higher-level API for creating and executing agents and tools."""
|
||||
|
||||
from langgraph.prebuilt._tool_call_transformer import ToolCallTransformer
|
||||
from langgraph.prebuilt.chat_agent_executor import create_react_agent
|
||||
from langgraph.prebuilt.tool_node import (
|
||||
InjectedState,
|
||||
@@ -13,6 +14,7 @@ from langgraph.prebuilt.tool_validator import ValidationNode
|
||||
__all__ = [
|
||||
"create_react_agent",
|
||||
"ToolNode",
|
||||
"ToolCallTransformer",
|
||||
"tools_condition",
|
||||
"ValidationNode",
|
||||
"InjectedState",
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
"""In-process handle for a single tool call's streaming execution.
|
||||
|
||||
Mirrors the shape of `ChatModelStream` from langchain-core but simpler —
|
||||
a tool has one output channel, no content-block multiplexing. Populated
|
||||
by `ToolCallTransformer` as `tool-started` / `tool-output-delta` /
|
||||
`tool-finished` / `tool-error` events flow in on the `tools` channel.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
|
||||
|
||||
class ToolCallStream:
|
||||
"""Scoped view of a single tool call's lifecycle.
|
||||
|
||||
Yielded on `run.tool_calls` once per `tool-started` event. Fields
|
||||
are populated as events arrive:
|
||||
|
||||
- `tool_call_id`, `tool_name`, `input`: stable from the start event.
|
||||
- `output_deltas`: an `EventLog` of delta chunks. Iterate (sync or
|
||||
async) to consume partial output in arrival order.
|
||||
- `output`: terminal payload from `tool-finished`, or `None` if the
|
||||
call failed or is still in flight.
|
||||
- `error`: terminal error string from `tool-error`, or `None` if the
|
||||
call succeeded or is still in flight.
|
||||
- `completed`: True once a terminal event (`tool-finished` or
|
||||
`tool-error`) has been observed.
|
||||
|
||||
`ToolCallStream` is not meant to be constructed by end users — it's
|
||||
produced by `ToolCallTransformer` as events flow through the mux.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tool_call_id: str,
|
||||
tool_name: str,
|
||||
input: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Initialize a fresh handle for a tool call.
|
||||
|
||||
Args:
|
||||
tool_call_id: The `tool_call_id` from the AIMessage.
|
||||
tool_name: The tool's name.
|
||||
input: The tool's input arguments (as reported by
|
||||
`on_tool_start`), or `None` if none were captured.
|
||||
"""
|
||||
self.tool_call_id = tool_call_id
|
||||
self.tool_name = tool_name
|
||||
self.input = input
|
||||
self._output_deltas: EventLog[Any] = EventLog()
|
||||
self.output: Any = None
|
||||
self.error: str | None = None
|
||||
self.completed = False
|
||||
|
||||
@property
|
||||
def output_deltas(self) -> EventLog[Any]:
|
||||
"""The EventLog of streamed `tool-output-delta` payloads.
|
||||
|
||||
Iterate (sync or async depending on how the run was started)
|
||||
to consume partial output in arrival order. The log closes when
|
||||
the tool finishes or errors.
|
||||
"""
|
||||
return self._output_deltas
|
||||
|
||||
def _bind(self, *, is_async: bool) -> None:
|
||||
"""Bind the deltas log to sync or async iteration.
|
||||
|
||||
Called by `ToolCallTransformer` when constructing this handle so
|
||||
the log matches the enclosing mux's mode.
|
||||
"""
|
||||
self._output_deltas._bind(is_async=is_async)
|
||||
|
||||
def _push_delta(self, delta: Any) -> None:
|
||||
self._output_deltas.push(delta)
|
||||
|
||||
def _finish(self, output: Any) -> None:
|
||||
self.output = output
|
||||
self.completed = True
|
||||
self._output_deltas.close()
|
||||
|
||||
def _fail(self, message: str) -> None:
|
||||
self.error = message
|
||||
self.completed = True
|
||||
self._output_deltas.close()
|
||||
|
||||
def __iter__(self) -> Iterator[Any]:
|
||||
"""Iterate delta chunks synchronously.
|
||||
|
||||
Equivalent to `iter(self.output_deltas)`. Raises `TypeError` if
|
||||
the underlying log is bound to async mode.
|
||||
"""
|
||||
return iter(self._output_deltas)
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[Any]:
|
||||
"""Iterate delta chunks asynchronously.
|
||||
|
||||
Equivalent to `aiter(self.output_deltas)`. Raises `TypeError`
|
||||
if the underlying log is bound to sync mode.
|
||||
"""
|
||||
return self._output_deltas.__aiter__()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
status = (
|
||||
"completed"
|
||||
if self.completed and self.error is None
|
||||
else "failed"
|
||||
if self.completed
|
||||
else "running"
|
||||
)
|
||||
return (
|
||||
f"ToolCallStream(tool_call_id={self.tool_call_id!r}, "
|
||||
f"tool_name={self.tool_name!r}, status={status})"
|
||||
)
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Transformer that projects `tools` channel events into `ToolCallStream`s."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
|
||||
from langgraph.prebuilt._tool_call_stream import ToolCallStream
|
||||
|
||||
|
||||
class ToolCallTransformer(StreamTransformer):
|
||||
"""Project `tools` channel events into `ToolCallStream` handles.
|
||||
|
||||
Each `tool-started` event spawns a `ToolCallStream`, pushed onto
|
||||
`run.tool_calls`. Subsequent `tool-output-delta` events append to
|
||||
that stream's deltas log; `tool-finished` and `tool-error` close it.
|
||||
|
||||
Native transformer — the `tool_calls` projection is exposed as a
|
||||
direct attribute on the run stream.
|
||||
|
||||
`EventLog[ToolCallStream]` is used (not `StreamChannel`) because the
|
||||
live handles are not serializable and should not be auto-forwarded
|
||||
onto the main event log. Wire consumers subscribe to the `tools`
|
||||
channel instead, where the raw protocol events flow through
|
||||
untouched by this transformer (`process` returns `True`).
|
||||
|
||||
Registered explicitly by users at compile time via
|
||||
`builder.compile(transformers=[ToolCallTransformer])` — not a
|
||||
default built-in, so the `tools` channel is user-opt-in.
|
||||
"""
|
||||
|
||||
_native = True
|
||||
required_stream_modes = ("tools",)
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._log: EventLog[ToolCallStream] = EventLog()
|
||||
self._active: dict[str, ToolCallStream] = {}
|
||||
self._is_async = False
|
||||
self._pump_fn: Callable[[], bool] | None = None
|
||||
self._apump_fn: Callable[[], Awaitable[bool]] | None = None
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"tool_calls": self._log}
|
||||
|
||||
def _bind_pump(self, fn: Callable[[], bool]) -> None:
|
||||
"""Wire the sync pull callback onto this transformer.
|
||||
|
||||
Called by `StreamMux.bind_pump`. Stored so each new
|
||||
`ToolCallStream` created by `process` can wire its deltas log
|
||||
for pump-driven iteration.
|
||||
"""
|
||||
self._pump_fn = fn
|
||||
self._is_async = False
|
||||
|
||||
def _bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:
|
||||
"""Async counterpart to `_bind_pump`."""
|
||||
self._apump_fn = fn
|
||||
self._is_async = True
|
||||
|
||||
def _new_stream(
|
||||
self,
|
||||
tool_call_id: str,
|
||||
tool_name: str,
|
||||
tool_input: dict[str, Any] | None,
|
||||
) -> ToolCallStream:
|
||||
stream = ToolCallStream(tool_call_id, tool_name, tool_input)
|
||||
stream._bind(is_async=self._is_async)
|
||||
if self._apump_fn is not None:
|
||||
stream._output_deltas._arequest_more = self._apump_fn
|
||||
if self._pump_fn is not None:
|
||||
stream._output_deltas._request_more = self._pump_fn
|
||||
return stream
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
# Namespace filtering is handled by the mux via `scope_exact`.
|
||||
if event["method"] != "tools":
|
||||
return True
|
||||
|
||||
data = event["params"]["data"]
|
||||
tool_call_id = data.get("tool_call_id")
|
||||
if tool_call_id is None:
|
||||
return True
|
||||
event_type = data.get("event")
|
||||
|
||||
stream: ToolCallStream | None
|
||||
if event_type == "tool-started":
|
||||
stream = self._new_stream(
|
||||
tool_call_id,
|
||||
data.get("tool_name", ""),
|
||||
data.get("input"),
|
||||
)
|
||||
self._active[tool_call_id] = stream
|
||||
self._log.push(stream)
|
||||
elif event_type == "tool-output-delta":
|
||||
stream = self._active.get(tool_call_id)
|
||||
if stream is not None:
|
||||
stream._push_delta(data.get("delta"))
|
||||
elif event_type == "tool-finished":
|
||||
stream = self._active.pop(tool_call_id, None)
|
||||
if stream is not None:
|
||||
stream._finish(data.get("output"))
|
||||
elif event_type == "tool-error":
|
||||
stream = self._active.pop(tool_call_id, None)
|
||||
if stream is not None:
|
||||
stream._fail(data.get("message", ""))
|
||||
|
||||
# Pass-through — wire consumers subscribe to the `tools` channel
|
||||
# directly and reconstruct handles client-side.
|
||||
return True
|
||||
|
||||
def finalize(self) -> None:
|
||||
"""Close any still-active tool streams left open at run end."""
|
||||
for stream in self._active.values():
|
||||
if not stream.completed:
|
||||
stream._finish(None)
|
||||
self._active.clear()
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Fail any still-active tool streams when the run errors."""
|
||||
message = str(err)
|
||||
for stream in self._active.values():
|
||||
if not stream.completed:
|
||||
stream._fail(message)
|
||||
self._active.clear()
|
||||
@@ -86,6 +86,7 @@ from langgraph._internal._constants import CONF, CONFIG_KEY_READ
|
||||
from langgraph._internal._runnable import RunnableCallable
|
||||
from langgraph.errors import GraphBubbleUp
|
||||
from langgraph.graph.message import REMOVE_ALL_MESSAGES
|
||||
from langgraph.pregel._tools import _tool_call_writer
|
||||
from langgraph.runtime import ExecutionInfo, ServerInfo # noqa: TC002
|
||||
from langgraph.store.base import BaseStore # noqa: TC002
|
||||
from langgraph.types import Command, Send, StreamWriter
|
||||
@@ -1728,6 +1729,26 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
|
||||
execution_info: ExecutionInfo | None = None
|
||||
server_info: ServerInfo | None = None
|
||||
|
||||
def emit_output_delta(self, delta: Any) -> None:
|
||||
"""Stream a partial output chunk on the `tools` stream channel.
|
||||
|
||||
Reads the per-tool-call writer that `StreamToolCallHandler`
|
||||
installs on a ContextVar at `on_tool_start` and forwards `delta`
|
||||
through it. Silent no-op when the graph was not run with
|
||||
`"tools"` in `stream_mode` (no writer is set), so tool authors
|
||||
can leave `emit_output_delta` calls in place without gating
|
||||
them on stream mode.
|
||||
|
||||
Args:
|
||||
delta: Partial output chunk. Any JSON-serializable value;
|
||||
surfaced as-is on the `tools` channel's
|
||||
`tool-output-delta` payload under `"delta"`.
|
||||
"""
|
||||
writer = _tool_call_writer.get()
|
||||
if writer is None:
|
||||
return
|
||||
writer(delta)
|
||||
|
||||
|
||||
class InjectedState(InjectedToolArg):
|
||||
"""Annotation for injecting graph state into tool arguments.
|
||||
|
||||
@@ -0,0 +1,307 @@
|
||||
"""Tests for ToolCallTransformer and the ToolCallStream projection."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.tools import tool
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import ProtocolEvent
|
||||
from langgraph.stream.transformers import (
|
||||
MessagesTransformer,
|
||||
ValuesTransformer,
|
||||
)
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.prebuilt import (
|
||||
ToolCallTransformer,
|
||||
ToolNode,
|
||||
ToolRuntime,
|
||||
)
|
||||
from langgraph.prebuilt._tool_call_stream import ToolCallStream
|
||||
|
||||
TS = int(time.time() * 1000)
|
||||
|
||||
|
||||
def _tool_event(
|
||||
event: str,
|
||||
tool_call_id: str,
|
||||
*,
|
||||
tool_name: str = "",
|
||||
input: dict[str, Any] | None = None,
|
||||
delta: Any = None,
|
||||
output: Any = None,
|
||||
message: str = "",
|
||||
namespace: list[str] | None = None,
|
||||
) -> ProtocolEvent:
|
||||
data: dict[str, Any] = {"event": event, "tool_call_id": tool_call_id}
|
||||
if event == "tool-started":
|
||||
data["tool_name"] = tool_name
|
||||
if input is not None:
|
||||
data["input"] = input
|
||||
elif event == "tool-output-delta":
|
||||
data["delta"] = delta
|
||||
elif event == "tool-finished":
|
||||
data["output"] = output
|
||||
elif event == "tool-error":
|
||||
data["message"] = message
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "tools",
|
||||
"params": {
|
||||
"namespace": namespace or [],
|
||||
"timestamp": TS,
|
||||
"data": data,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _subscribe(log: EventLog) -> None:
|
||||
log._subscribed = True
|
||||
|
||||
|
||||
def _mux() -> tuple[StreamMux, ToolCallTransformer]:
|
||||
transformer = ToolCallTransformer()
|
||||
mux = StreamMux(
|
||||
[
|
||||
ValuesTransformer(),
|
||||
MessagesTransformer(),
|
||||
transformer,
|
||||
],
|
||||
is_async=False,
|
||||
)
|
||||
_subscribe(transformer._log)
|
||||
return mux, transformer
|
||||
|
||||
|
||||
class TestToolCallTransformerUnit:
|
||||
def test_required_stream_modes_declares_tools(self) -> None:
|
||||
assert ToolCallTransformer.required_stream_modes == ("tools",)
|
||||
|
||||
def test_tool_started_yields_handle(self) -> None:
|
||||
mux, transformer = _mux()
|
||||
mux.push(
|
||||
_tool_event(
|
||||
"tool-started",
|
||||
"tc1",
|
||||
tool_name="echo",
|
||||
input={"text": "hi"},
|
||||
)
|
||||
)
|
||||
handles = list(transformer._log._items)
|
||||
assert len(handles) == 1
|
||||
h = handles[0]
|
||||
assert isinstance(h, ToolCallStream)
|
||||
assert h.tool_call_id == "tc1"
|
||||
assert h.tool_name == "echo"
|
||||
assert h.input == {"text": "hi"}
|
||||
assert h.completed is False
|
||||
|
||||
def test_delta_accumulates_on_active_stream(self) -> None:
|
||||
mux, transformer = _mux()
|
||||
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
||||
_subscribe(transformer._active["tc1"]._output_deltas)
|
||||
mux.push(_tool_event("tool-output-delta", "tc1", delta="a"))
|
||||
mux.push(_tool_event("tool-output-delta", "tc1", delta="b"))
|
||||
stream = transformer._active["tc1"]
|
||||
assert list(stream._output_deltas._items) == ["a", "b"]
|
||||
|
||||
def test_finish_closes_stream(self) -> None:
|
||||
mux, transformer = _mux()
|
||||
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
||||
stream = transformer._active["tc1"]
|
||||
mux.push(_tool_event("tool-finished", "tc1", output="done"))
|
||||
assert stream.completed is True
|
||||
assert stream.output == "done"
|
||||
assert stream.error is None
|
||||
assert "tc1" not in transformer._active
|
||||
|
||||
def test_error_closes_stream(self) -> None:
|
||||
mux, transformer = _mux()
|
||||
mux.push(_tool_event("tool-started", "tc1", tool_name="boom"))
|
||||
stream = transformer._active["tc1"]
|
||||
mux.push(_tool_event("tool-error", "tc1", message="nope"))
|
||||
assert stream.completed is True
|
||||
assert stream.output is None
|
||||
assert stream.error == "nope"
|
||||
assert "tc1" not in transformer._active
|
||||
|
||||
def test_concurrent_tool_calls_do_not_bleed(self) -> None:
|
||||
mux, transformer = _mux()
|
||||
mux.push(_tool_event("tool-started", "a", tool_name="t"))
|
||||
mux.push(_tool_event("tool-started", "b", tool_name="t"))
|
||||
for tc in ("a", "b"):
|
||||
_subscribe(transformer._active[tc]._output_deltas)
|
||||
mux.push(_tool_event("tool-output-delta", "a", delta="A1"))
|
||||
mux.push(_tool_event("tool-output-delta", "b", delta="B1"))
|
||||
mux.push(_tool_event("tool-output-delta", "a", delta="A2"))
|
||||
assert list(transformer._active["a"]._output_deltas._items) == ["A1", "A2"]
|
||||
assert list(transformer._active["b"]._output_deltas._items) == ["B1"]
|
||||
|
||||
def test_tools_event_passes_through_main_log(self) -> None:
|
||||
mux, transformer = _mux()
|
||||
_subscribe(mux._events)
|
||||
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
||||
kept = [e for e in mux._events._items if e["method"] == "tools"]
|
||||
assert len(kept) == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end tests with a real graph
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
|
||||
|
||||
def _build_graph(caller, tools):
|
||||
sg = StateGraph(_State)
|
||||
sg.add_node("caller", caller)
|
||||
sg.add_node("tools", ToolNode(tools))
|
||||
sg.add_edge(START, "caller")
|
||||
sg.add_edge("caller", "tools")
|
||||
sg.add_edge("tools", END)
|
||||
return sg.compile()
|
||||
|
||||
|
||||
class TestToolCallTransformerEndToEnd:
|
||||
def test_sync_streaming_tool_populates_tool_calls(self) -> None:
|
||||
@tool
|
||||
def streamer(text: str, runtime: ToolRuntime) -> str:
|
||||
"""streams chunks."""
|
||||
for chunk in ("one", "two"):
|
||||
runtime.emit_output_delta(chunk)
|
||||
return text
|
||||
|
||||
def caller(state: _State) -> dict:
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{"name": "streamer", "args": {"text": "x"}, "id": "tc1"}
|
||||
],
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
graph = _build_graph(caller, [streamer])
|
||||
run = graph.stream_v2({"messages": []}, transformers=[ToolCallTransformer])
|
||||
|
||||
tool_calls: list[ToolCallStream] = []
|
||||
for tc in run.tool_calls:
|
||||
tool_calls.append(tc)
|
||||
deltas = list(tc.output_deltas)
|
||||
assert deltas == ["one", "two"]
|
||||
assert len(tool_calls) == 1
|
||||
tc = tool_calls[0]
|
||||
assert tc.tool_call_id == "tc1"
|
||||
assert tc.tool_name == "streamer"
|
||||
assert tc.completed is True
|
||||
assert tc.error is None
|
||||
|
||||
def test_stream_modes_union_includes_tools(self) -> None:
|
||||
@tool
|
||||
def echo(text: str) -> str:
|
||||
"""echo."""
|
||||
return text
|
||||
|
||||
def caller(state: _State) -> dict:
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{"name": "echo", "args": {"text": "x"}, "id": "tc1"}
|
||||
],
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
graph = _build_graph(caller, [echo])
|
||||
# Without ToolCallTransformer, no tool_calls projection is
|
||||
# exposed and no `tools` events flow through (required_stream_modes
|
||||
# omits it).
|
||||
run_no_tc = graph.stream_v2({"messages": []})
|
||||
assert "tool_calls" not in run_no_tc._mux.extensions # type: ignore[attr-defined]
|
||||
|
||||
# With ToolCallTransformer, the projection is present.
|
||||
run = graph.stream_v2({"messages": []}, transformers=[ToolCallTransformer])
|
||||
assert "tool_calls" in run._mux.extensions # type: ignore[attr-defined]
|
||||
# Drain so the run closes cleanly.
|
||||
list(run.tool_calls)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_streaming_tool_populates_tool_calls(self) -> None:
|
||||
@tool
|
||||
async def astreamer(text: str, runtime: ToolRuntime) -> str:
|
||||
"""async streams."""
|
||||
runtime.emit_output_delta(text)
|
||||
runtime.emit_output_delta(text + "!")
|
||||
return text
|
||||
|
||||
async def caller(state: _State) -> dict:
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{"name": "astreamer", "args": {"text": "hi"}, "id": "tc1"}
|
||||
],
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
graph = _build_graph(caller, [astreamer])
|
||||
run = await graph.astream_v2(
|
||||
{"messages": []}, transformers=[ToolCallTransformer]
|
||||
)
|
||||
|
||||
collected: list[ToolCallStream] = []
|
||||
async for tc in run.tool_calls:
|
||||
collected.append(tc)
|
||||
deltas = [d async for d in tc.output_deltas]
|
||||
assert deltas == ["hi", "hi!"]
|
||||
assert len(collected) == 1
|
||||
assert collected[0].completed is True
|
||||
assert collected[0].error is None
|
||||
|
||||
def test_tool_error_populates_error_field(self) -> None:
|
||||
@tool
|
||||
def boom() -> str:
|
||||
"""raises."""
|
||||
raise ValueError("nope")
|
||||
|
||||
def caller(state: _State) -> dict:
|
||||
return {
|
||||
"messages": [
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[{"name": "boom", "args": {}, "id": "tc1"}],
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
graph = _build_graph(caller, [boom])
|
||||
run = graph.stream_v2({"messages": []}, transformers=[ToolCallTransformer])
|
||||
|
||||
collected: list[ToolCallStream] = []
|
||||
with pytest.raises(ValueError, match="nope"):
|
||||
for tc in run.tool_calls:
|
||||
collected.append(tc)
|
||||
# Drain deltas so the error field is populated before we
|
||||
# inspect it below.
|
||||
list(tc.output_deltas)
|
||||
|
||||
assert len(collected) == 1
|
||||
assert collected[0].error == "nope"
|
||||
assert collected[0].output is None
|
||||
assert collected[0].completed is True
|
||||
Generated
+17
-4
@@ -249,10 +249,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.1"
|
||||
version = "1.3.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
{ name = "langchain-protocol" },
|
||||
{ name = "langsmith" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pydantic" },
|
||||
@@ -261,9 +262,21 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f4/fe/abeae8d0d2899e191d67c6c7f065f7e52a953f30b21ef327fa49084e4af9/langchain_core-1.3.1.tar.gz", hash = "sha256:41b384055799f93f34520df6bf7b80e2e5e23153cdfd46874251c6c9916ea030", size = 862403, upload-time = "2026-04-23T18:54:01.857Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a1/c2/8493be505921857988db068b7c027f28a9b1587b4425c6a32b1221c9c9fe/langchain_core-1.3.1-py3-none-any.whl", hash = "sha256:8b13d19d3bed3f4768df12c7f6932d2ada715f3ac9fd020c63d28c693968269e", size = 515879, upload-time = "2026-04-23T18:53:59.94Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langchain-protocol"
|
||||
version = "0.0.12"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/5c/51/1157009b6f94e6e58be58fa8b620187d657909a8b36a6bf5b0c52a2711f6/langchain_protocol-0.0.12.tar.gz", hash = "sha256:5e14c434290a705c9510fdb1a83ecf7561a5e6e0dfd053930ade80dba069269f", size = 6408, upload-time = "2026-04-25T01:05:01.489Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/95/82/3431e3061c917439589fa88a6b23c9bc0e154cba0f05d2e895a68c76ff74/langchain_protocol-0.0.12-py3-none-any.whl", hash = "sha256:402b61f42d4139692528cf37226c367bb6efc8ff8165b29380accb0abfece7b2", size = 6639, upload-time = "2026-04-25T01:05:00.487Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -281,7 +294,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = ">=1.3.2,<2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "." },
|
||||
{ name = "langgraph-sdk", editable = "../sdk-py" },
|
||||
|
||||
Generated
+17
-4
@@ -266,10 +266,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "1.3.1"
|
||||
version = "1.3.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jsonpatch" },
|
||||
{ name = "langchain-protocol" },
|
||||
{ name = "langsmith" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pydantic" },
|
||||
@@ -278,9 +279,21 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "uuid-utils" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f4/fe/abeae8d0d2899e191d67c6c7f065f7e52a953f30b21ef327fa49084e4af9/langchain_core-1.3.1.tar.gz", hash = "sha256:41b384055799f93f34520df6bf7b80e2e5e23153cdfd46874251c6c9916ea030", size = 862403, upload-time = "2026-04-23T18:54:01.857Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a1/c2/8493be505921857988db068b7c027f28a9b1587b4425c6a32b1221c9c9fe/langchain_core-1.3.1-py3-none-any.whl", hash = "sha256:8b13d19d3bed3f4768df12c7f6932d2ada715f3ac9fd020c63d28c693968269e", size = 515879, upload-time = "2026-04-23T18:53:59.94Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "langchain-protocol"
|
||||
version = "0.0.12"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/5c/51/1157009b6f94e6e58be58fa8b620187d657909a8b36a6bf5b0c52a2711f6/langchain_protocol-0.0.12.tar.gz", hash = "sha256:5e14c434290a705c9510fdb1a83ecf7561a5e6e0dfd053930ade80dba069269f", size = 6408, upload-time = "2026-04-25T01:05:01.489Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/95/82/3431e3061c917439589fa88a6b23c9bc0e154cba0f05d2e895a68c76ff74/langchain_protocol-0.0.12-py3-none-any.whl", hash = "sha256:402b61f42d4139692528cf37226c367bb6efc8ff8165b29380accb0abfece7b2", size = 6639, upload-time = "2026-04-25T01:05:00.487Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -298,7 +311,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "langchain-core", specifier = ">=1.3.0,<2" },
|
||||
{ name = "langchain-core", specifier = ">=1.3.2,<2" },
|
||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
|
||||
{ name = "langgraph-sdk", editable = "." },
|
||||
|
||||
Reference in New Issue
Block a user