mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-09 11:17:53 +02:00
Compare commits
42
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
623eec565e | ||
|
|
8e2f23e051 | ||
|
|
cdec5bf336 | ||
|
|
f30a11055b | ||
|
|
1bf0f1b7f8 | ||
|
|
133082af71 | ||
|
|
de72fb4b4b | ||
|
|
221ab0c9ad | ||
|
|
2982512adb | ||
|
|
4956512692 | ||
|
|
40055e92cc | ||
|
|
66e6c27155 | ||
|
|
f4a5d56535 | ||
|
|
742488d5a0 | ||
|
|
b319426725 | ||
|
|
e01037d90c | ||
|
|
bebcd20815 | ||
|
|
2ad30132a3 | ||
|
|
7715239e3b | ||
|
|
ad0146a4de | ||
|
|
b6a196fac6 | ||
|
|
910240a930 | ||
|
|
ab1d6980b5 | ||
|
|
7e5df56688 | ||
|
|
0f2f66fc8f | ||
|
|
acaa767542 | ||
|
|
5f24a0356a | ||
|
|
adda5f0341 | ||
|
|
6fcca359df | ||
|
|
28ce32edc7 | ||
|
|
119847f80f | ||
|
|
f43743c3e7 | ||
|
|
dbded7a59e | ||
|
|
986c1cc2e3 | ||
|
|
28cf5ed78d | ||
|
|
ca5d9a6bd7 | ||
|
|
ae3c823499 | ||
|
|
b72b5fefd0 | ||
|
|
5b1f86facc | ||
|
|
cf966419d5 | ||
|
|
8f03bf9f15 | ||
|
|
0076da9008 |
@@ -66,6 +66,9 @@ CONFIG_KEY_RUNTIME = sys.intern("__pregel_runtime")
|
|||||||
# holds a `Runtime` instance with context, store, stream writer, etc.
|
# holds a `Runtime` instance with context, store, stream writer, etc.
|
||||||
CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map")
|
CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map")
|
||||||
# holds a mapping of task ns -> resume value for resuming tasks
|
# 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 GraphStreamer only.
|
||||||
|
|
||||||
# --- Other constants ---
|
# --- Other constants ---
|
||||||
PUSH = sys.intern("__pregel_push")
|
PUSH = sys.intern("__pregel_push")
|
||||||
@@ -107,6 +110,7 @@ RESERVED = {
|
|||||||
CONFIG_KEY_CHECKPOINT_ID,
|
CONFIG_KEY_CHECKPOINT_ID,
|
||||||
CONFIG_KEY_CHECKPOINT_NS,
|
CONFIG_KEY_CHECKPOINT_NS,
|
||||||
CONFIG_KEY_RESUME_MAP,
|
CONFIG_KEY_RESUME_MAP,
|
||||||
|
CONFIG_KEY_STREAM_MESSAGES_V2,
|
||||||
# other constants
|
# other constants
|
||||||
PUSH,
|
PUSH,
|
||||||
PULL,
|
PULL,
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import sys
|
import sys
|
||||||
|
from collections.abc import Callable
|
||||||
|
from contextvars import ContextVar
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from langchain_core.runnables import RunnableConfig
|
from langchain_core.runnables import RunnableConfig
|
||||||
@@ -9,6 +11,18 @@ from langgraph.store.base import BaseStore
|
|||||||
from langgraph._internal._constants import CONF, CONFIG_KEY_RUNTIME
|
from langgraph._internal._constants import CONF, CONFIG_KEY_RUNTIME
|
||||||
from langgraph.types import StreamWriter
|
from langgraph.types import StreamWriter
|
||||||
|
|
||||||
|
_tool_call_writer: ContextVar[Callable[[Any], None] | 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.
|
||||||
|
Defined here (rather than alongside the handler in `pregel/_tools.py`)
|
||||||
|
so `emit_tool_output_delta` can import it without triggering the
|
||||||
|
pregel import chain — user tool code does
|
||||||
|
`from langgraph.config import emit_tool_output_delta` at import time.
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
def _no_op_stream_writer(c: Any) -> None:
|
def _no_op_stream_writer(c: Any) -> None:
|
||||||
pass
|
pass
|
||||||
@@ -194,3 +208,30 @@ def get_stream_writer() -> StreamWriter:
|
|||||||
"""
|
"""
|
||||||
runtime = get_config()[CONF][CONFIG_KEY_RUNTIME]
|
runtime = get_config()[CONF][CONFIG_KEY_RUNTIME]
|
||||||
return runtime.stream_writer
|
return runtime.stream_writer
|
||||||
|
|
||||||
|
|
||||||
|
def emit_tool_output_delta(delta: Any) -> None:
|
||||||
|
"""Emit a `tool-output-delta` event onto the `tools` stream mode.
|
||||||
|
|
||||||
|
Must be called from inside a tool's execution scope (sync or async).
|
||||||
|
While a tool is running, `StreamToolCallHandler.on_tool_start` sets a
|
||||||
|
writer closure on a ContextVar keyed to that call's `tool_call_id`
|
||||||
|
and namespace; this helper reads the ContextVar and forwards `delta`
|
||||||
|
through it.
|
||||||
|
|
||||||
|
When called outside any tool call, or when the graph was not
|
||||||
|
streamed with `"tools"` in `stream_mode`, this is a silent no-op —
|
||||||
|
tool authors can leave `emit_tool_output_delta` calls in place
|
||||||
|
without gating them on stream mode.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
delta: The partial output chunk to stream. Shape is up to the
|
||||||
|
caller — strings are the common case, but any JSON-
|
||||||
|
serializable value is accepted and 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)
|
||||||
|
|||||||
@@ -1045,6 +1045,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
interrupt_after: All | list[str] | None = None,
|
interrupt_after: All | list[str] | None = None,
|
||||||
debug: bool = False,
|
debug: bool = False,
|
||||||
name: str | None = None,
|
name: str | None = None,
|
||||||
|
transformers: Sequence[Callable[..., Any]] | None = None,
|
||||||
) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]:
|
) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]:
|
||||||
"""Compiles the `StateGraph` into a `CompiledStateGraph` object.
|
"""Compiles the `StateGraph` into a `CompiledStateGraph` object.
|
||||||
|
|
||||||
@@ -1077,6 +1078,11 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
interrupt_after: An optional list of node names to interrupt after.
|
interrupt_after: An optional list of node names to interrupt after.
|
||||||
debug: A flag indicating whether to enable debug mode.
|
debug: A flag indicating whether to enable debug mode.
|
||||||
name: The name to use for the compiled graph.
|
name: The name to use for the compiled graph.
|
||||||
|
transformers: Optional sequence of zero-arg factories returning
|
||||||
|
`StreamTransformer` instances. Registered on the compiled
|
||||||
|
graph and instantiated per-run whenever `stream_v2` /
|
||||||
|
`astream_v2` is called. Appended after the built-in
|
||||||
|
`ValuesTransformer` and `MessagesTransformer`.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
CompiledStateGraph: The compiled `StateGraph`.
|
CompiledStateGraph: The compiled `StateGraph`.
|
||||||
@@ -1159,6 +1165,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
store=store,
|
store=store,
|
||||||
cache=cache,
|
cache=cache,
|
||||||
name=name or "LangGraph",
|
name=name or "LangGraph",
|
||||||
|
stream_transformers=transformers,
|
||||||
)
|
)
|
||||||
compiled._serde_allowlist = serde_allowlist
|
compiled._serde_allowlist = serde_allowlist
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,274 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from collections.abc import AsyncIterator, Callable, Iterator
|
||||||
|
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.errors import GraphInterrupt
|
||||||
|
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")
|
||||||
|
|
||||||
|
|
||||||
|
_LANGGRAPH_SENTINEL_NODES = frozenset({"__start__", "__end__"})
|
||||||
|
|
||||||
|
|
||||||
|
def _is_nested_pregel_start(
|
||||||
|
name: str | None,
|
||||||
|
metadata: dict[str, Any] | None,
|
||||||
|
parent_run_id: UUID | None,
|
||||||
|
task_run_ids: set[UUID],
|
||||||
|
) -> bool:
|
||||||
|
"""Recognize a nested `Pregel` invocation from its `on_chain_start` metadata.
|
||||||
|
|
||||||
|
When a compiled graph is added as a node, pregel fires two
|
||||||
|
`on_chain_start` callbacks at that task: first for the node chain
|
||||||
|
(whose `name` matches `metadata["langgraph_node"]`) and second for
|
||||||
|
the inner `Pregel` chain (whose `name` is the graph's `name`, not
|
||||||
|
the node name). Both share the same `langgraph_checkpoint_ns`.
|
||||||
|
|
||||||
|
Primary signal: a `langgraph_checkpoint_ns` is set AND `name`
|
||||||
|
differs from the owning task's `langgraph_node`. This covers the
|
||||||
|
common case where the compiled subgraph's name differs from the
|
||||||
|
node name it was registered under.
|
||||||
|
|
||||||
|
Fallback for name collisions (subgraph compiled with
|
||||||
|
`name == node_name`): the inner `Pregel` start's `parent_run_id`
|
||||||
|
is the run_id of the node chain's start event, which the handler
|
||||||
|
records in `task_run_ids` on the first start. Matching
|
||||||
|
`parent_run_id` to that set identifies the second start as the
|
||||||
|
nested `Pregel` even when names coincide.
|
||||||
|
|
||||||
|
Regular node chains are skipped; the root `Pregel` (which has no
|
||||||
|
`langgraph_node` metadata) isn't observed by this handler because
|
||||||
|
the root's start fires before the handler is attached.
|
||||||
|
|
||||||
|
Metadata-based detection is used because `on_chain_start`'s
|
||||||
|
`serialized` argument is `None` for compiled graphs in this
|
||||||
|
version of langchain-core, so class-based detection via
|
||||||
|
`serialized["id"]` isn't available.
|
||||||
|
|
||||||
|
Sentinel nodes (`__start__` / `__end__`) are excluded: conditional
|
||||||
|
edges from `START` fire an `on_chain_start` with `lg_node=__start__`
|
||||||
|
and the router function's name as `name`, which would otherwise
|
||||||
|
match the discriminator without representing an actual nested
|
||||||
|
`Pregel`.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: The `name` kwarg from `on_chain_start`.
|
||||||
|
metadata: The `metadata` kwarg from `on_chain_start`.
|
||||||
|
parent_run_id: The `parent_run_id` kwarg from `on_chain_start`.
|
||||||
|
task_run_ids: The set of run_ids the handler has already seen
|
||||||
|
as node-chain starts (i.e. `name == langgraph_node`).
|
||||||
|
"""
|
||||||
|
if not metadata:
|
||||||
|
return False
|
||||||
|
if not metadata.get("langgraph_checkpoint_ns"):
|
||||||
|
return False
|
||||||
|
lg_node = metadata.get("langgraph_node")
|
||||||
|
if lg_node is None or lg_node in _LANGGRAPH_SENTINEL_NODES:
|
||||||
|
return False
|
||||||
|
if name != lg_node:
|
||||||
|
return True
|
||||||
|
# Name collision fallback: the inner Pregel's parent_run_id is
|
||||||
|
# the node chain's run_id, which we recorded when that node
|
||||||
|
# chain's start fired.
|
||||||
|
return parent_run_id is not None and parent_run_id in task_run_ids
|
||||||
|
|
||||||
|
|
||||||
|
class StreamLifecycleHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||||
|
"""Callback handler that emits subgraph lifecycle events on the stream.
|
||||||
|
|
||||||
|
Pushes `LifecycleData`-shaped payloads onto the pregel stream under
|
||||||
|
the `"lifecycle"` mode, keyed by the subgraph's namespace tuple.
|
||||||
|
Drives the `started` → `running` → `completed` / `failed` /
|
||||||
|
`interrupted` state machine.
|
||||||
|
|
||||||
|
The handler is attached to `run_manager.inheritable_handlers` inside
|
||||||
|
a `Pregel.stream` / `astream` call, so it sees callbacks for every
|
||||||
|
descendant chain (nodes, nested `Pregel` subgraphs) but *not* for
|
||||||
|
the root `Pregel` whose start event has already fired. The root's
|
||||||
|
`started` event is emitted eagerly at construction; its terminal
|
||||||
|
state is emitted by `SubgraphTransformer.finalize` / `fail`.
|
||||||
|
|
||||||
|
`run_inline = True` keeps event ordering deterministic.
|
||||||
|
"""
|
||||||
|
|
||||||
|
run_inline = True
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
stream: Callable[[StreamChunk], None],
|
||||||
|
*,
|
||||||
|
root_graph_name: str | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Initialize the handler and emit the root graph's `started` event.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
stream: Callable that accepts a `StreamChunk` tuple
|
||||||
|
`(namespace, mode, payload)` and enqueues it.
|
||||||
|
root_graph_name: The root `Pregel` instance's `name`, emitted
|
||||||
|
with the root's `started` lifecycle payload.
|
||||||
|
"""
|
||||||
|
self.stream = stream
|
||||||
|
# Namespaces awaiting the started→running transition.
|
||||||
|
self._pending_running: set[tuple[str, ...]] = set()
|
||||||
|
# run_id → subgraph namespace; populated only for Pregel chains.
|
||||||
|
self._run_to_ns: dict[UUID, tuple[str, ...]] = {}
|
||||||
|
# run_ids of node-chain starts (name == langgraph_node); used
|
||||||
|
# as the parent_run_id fallback when a subgraph's name equals
|
||||||
|
# its node name. Cleared as each chain ends.
|
||||||
|
self._task_run_ids: set[UUID] = set()
|
||||||
|
|
||||||
|
root_payload: dict[str, Any] = {"event": "started"}
|
||||||
|
if root_graph_name is not None:
|
||||||
|
root_payload["graph_name"] = root_graph_name
|
||||||
|
self.stream(((), "lifecycle", root_payload))
|
||||||
|
self._pending_running.add(())
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _subgraph_ns_from_metadata(metadata: dict[str, Any] | None) -> tuple[str, ...]:
|
||||||
|
"""Return the running subgraph's own namespace from task metadata.
|
||||||
|
|
||||||
|
For a nested `Pregel` invoked as a node, `langgraph_checkpoint_ns`
|
||||||
|
ends at the node segment (no inner task appended yet), so
|
||||||
|
splitting on `NS_SEP` gives the subgraph's own namespace.
|
||||||
|
"""
|
||||||
|
if not metadata:
|
||||||
|
return ()
|
||||||
|
nskey = metadata.get("langgraph_checkpoint_ns")
|
||||||
|
if not nskey:
|
||||||
|
return ()
|
||||||
|
return tuple(cast(str, nskey).split(NS_SEP))
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _containing_ns_from_metadata(
|
||||||
|
metadata: dict[str, Any] | None,
|
||||||
|
) -> tuple[str, ...]:
|
||||||
|
"""Return the namespace of the subgraph that contains this task.
|
||||||
|
|
||||||
|
For an inner task with `langgraph_checkpoint_ns`
|
||||||
|
`"seg_a|seg_b"`, the containing subgraph is `("seg_a",)`.
|
||||||
|
"""
|
||||||
|
if not metadata:
|
||||||
|
return ()
|
||||||
|
nskey = metadata.get("langgraph_checkpoint_ns")
|
||||||
|
if not nskey:
|
||||||
|
return ()
|
||||||
|
return tuple(cast(str, nskey).split(NS_SEP))[:-1]
|
||||||
|
|
||||||
|
def _emit(self, ns: tuple[str, ...], payload: dict[str, Any]) -> None:
|
||||||
|
self.stream((ns, "lifecycle", payload))
|
||||||
|
|
||||||
|
def tap_output_aiter(
|
||||||
|
self, run_id: UUID, output: AsyncIterator[T]
|
||||||
|
) -> AsyncIterator[T]:
|
||||||
|
"""Pass-through — required by the `_StreamingCallbackHandler` protocol.
|
||||||
|
|
||||||
|
Returns the iterator unchanged. A missing implementation lets
|
||||||
|
langchain's default `Protocol` body return `None`, which breaks
|
||||||
|
the `_consume_aiter` code path in `_runnable.py:900`.
|
||||||
|
"""
|
||||||
|
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
|
||||||
|
|
||||||
|
def _fire_running_if_pending(self, ns: tuple[str, ...]) -> None:
|
||||||
|
if ns in self._pending_running:
|
||||||
|
self._pending_running.discard(ns)
|
||||||
|
self._emit(ns, {"event": "running"})
|
||||||
|
|
||||||
|
def on_chain_start(
|
||||||
|
self,
|
||||||
|
serialized: dict[str, Any],
|
||||||
|
inputs: dict[str, Any],
|
||||||
|
*,
|
||||||
|
run_id: UUID,
|
||||||
|
parent_run_id: UUID | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
metadata: dict[str, Any] | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
# Any descendant activity transitions the containing subgraph to running.
|
||||||
|
containing = self._containing_ns_from_metadata(metadata)
|
||||||
|
self._fire_running_if_pending(containing)
|
||||||
|
|
||||||
|
name = cast(str | None, kwargs.get("name"))
|
||||||
|
lg_node = (metadata or {}).get("langgraph_node")
|
||||||
|
|
||||||
|
# Record node-chain starts so the name-collision fallback in
|
||||||
|
# `_is_nested_pregel_start` can match the inner Pregel's
|
||||||
|
# parent_run_id to them.
|
||||||
|
if (
|
||||||
|
lg_node is not None
|
||||||
|
and lg_node not in _LANGGRAPH_SENTINEL_NODES
|
||||||
|
and name == lg_node
|
||||||
|
):
|
||||||
|
self._task_run_ids.add(run_id)
|
||||||
|
|
||||||
|
if not _is_nested_pregel_start(
|
||||||
|
name, metadata, parent_run_id, self._task_run_ids
|
||||||
|
):
|
||||||
|
return
|
||||||
|
|
||||||
|
ns = self._subgraph_ns_from_metadata(metadata)
|
||||||
|
if not ns:
|
||||||
|
return
|
||||||
|
|
||||||
|
self._run_to_ns[run_id] = ns
|
||||||
|
payload: dict[str, Any] = {"event": "started"}
|
||||||
|
if name:
|
||||||
|
payload["graph_name"] = name
|
||||||
|
# `cause` is intentionally not populated here: pregel does not know
|
||||||
|
# what on the parent namespace triggered this subgraph. Product-
|
||||||
|
# specific stream transformers populate `cause` before events
|
||||||
|
# reach the wire. See LifecycleCause in the protocol definition.
|
||||||
|
self._emit(ns, payload)
|
||||||
|
self._pending_running.add(ns)
|
||||||
|
|
||||||
|
def on_chain_end(
|
||||||
|
self,
|
||||||
|
response: Any,
|
||||||
|
*,
|
||||||
|
run_id: UUID,
|
||||||
|
parent_run_id: UUID | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
self._task_run_ids.discard(run_id)
|
||||||
|
ns = self._run_to_ns.pop(run_id, None)
|
||||||
|
if ns is None:
|
||||||
|
return
|
||||||
|
# Ensure started→running fired even for empty subgraphs.
|
||||||
|
if ns in self._pending_running:
|
||||||
|
self._pending_running.discard(ns)
|
||||||
|
self._emit(ns, {"event": "running"})
|
||||||
|
self._emit(ns, {"event": "completed"})
|
||||||
|
|
||||||
|
def on_chain_error(
|
||||||
|
self,
|
||||||
|
error: BaseException,
|
||||||
|
*,
|
||||||
|
run_id: UUID,
|
||||||
|
parent_run_id: UUID | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
self._task_run_ids.discard(run_id)
|
||||||
|
ns = self._run_to_ns.pop(run_id, None)
|
||||||
|
if ns is None:
|
||||||
|
return
|
||||||
|
self._pending_running.discard(ns)
|
||||||
|
if isinstance(error, GraphInterrupt):
|
||||||
|
self._emit(ns, {"event": "interrupted"})
|
||||||
|
else:
|
||||||
|
self._emit(ns, {"event": "failed", "error": str(error)})
|
||||||
@@ -24,6 +24,11 @@ try:
|
|||||||
except ImportError:
|
except ImportError:
|
||||||
_StreamingCallbackHandler = object # type: ignore
|
_StreamingCallbackHandler = object # type: ignore
|
||||||
|
|
||||||
|
try:
|
||||||
|
from langchain_core.tracers._streaming import _V2StreamingCallbackHandler
|
||||||
|
except ImportError:
|
||||||
|
_V2StreamingCallbackHandler = object # type: ignore
|
||||||
|
|
||||||
T = TypeVar("T")
|
T = TypeVar("T")
|
||||||
Meta = tuple[tuple[str, ...], dict[str, Any]]
|
Meta = tuple[tuple[str, ...], dict[str, Any]]
|
||||||
|
|
||||||
@@ -256,3 +261,185 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
|||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> Any:
|
) -> Any:
|
||||||
self.metadata.pop(run_id, None)
|
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
|
||||||
|
`GraphStreamer` 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_chat_model_start(
|
||||||
|
self,
|
||||||
|
serialized: dict[str, Any],
|
||||||
|
messages: list[list[BaseMessage]],
|
||||||
|
*,
|
||||||
|
run_id: UUID,
|
||||||
|
parent_run_id: UUID | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
metadata: dict[str, Any] | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
"""Record metadata with the FULL checkpoint namespace for v2.
|
||||||
|
|
||||||
|
v1's ``on_chat_model_start`` (inherited) slices the ns tuple
|
||||||
|
with ``[:-1]`` to re-position chat model tokens onto the
|
||||||
|
*containing pregel's* namespace — historically convenient for
|
||||||
|
consumers of ``stream_mode="messages"`` who want "where did
|
||||||
|
this node produce its output" rather than the chat-model's
|
||||||
|
own task ns.
|
||||||
|
|
||||||
|
For the protocol-v2 wire shape that is wrong: the client
|
||||||
|
subscribes the root feed at ``namespaces=[[]]`` with
|
||||||
|
``depth=1``, and any message emitted at depth ``>=1`` that
|
||||||
|
still carries the containing node's ns must appear at the
|
||||||
|
*full* path from root so that depth filtering cleanly isolates
|
||||||
|
subgraph chatter from the root conversation. JS's
|
||||||
|
``StreamMessagesHandlerV2`` already does
|
||||||
|
``metadata.langgraph_checkpoint_ns.split("|")`` (no slice); this
|
||||||
|
override brings the Python v2 handler to the same shape.
|
||||||
|
|
||||||
|
Without this, a chat model invoked inside a nested subgraph
|
||||||
|
(e.g. ``research -> researcher``, ``research`` being a root
|
||||||
|
node that ``.ainvoke()``s a ``researcher`` subgraph) emits at
|
||||||
|
``["research:<task>"]`` — a single level deep — which slips
|
||||||
|
through the root-feed depth-1 filter and pollutes the main
|
||||||
|
conversation with subgraph tokens. With this override we emit
|
||||||
|
at ``["research:<task>", "researcher:<task>"]`` so the client
|
||||||
|
routes those tokens to the subgraph card instead.
|
||||||
|
"""
|
||||||
|
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
|
||||||
|
task_checkpoint_ns = cast(str, metadata["langgraph_checkpoint_ns"])
|
||||||
|
# Keep the trailing ``:<task_id>`` segment (unlike the v1
|
||||||
|
# handler which strips it via ``[:-1]``). The client's
|
||||||
|
# lifecycle events land on the same ns, so message deltas
|
||||||
|
# now correlate 1:1 with a ``lifecycle: started`` event —
|
||||||
|
# ``useMessages(stream, subgraph)`` picks them up without
|
||||||
|
# needing to collapse sibling namespaces.
|
||||||
|
ns = tuple(task_checkpoint_ns.split(NS_SEP))
|
||||||
|
if not self.subgraphs and len(ns) > 1 and ns != self.parent_ns:
|
||||||
|
return
|
||||||
|
stream_metadata = dict(metadata)
|
||||||
|
# Preserve the v1-shaped ``langgraph_checkpoint_ns`` (task
|
||||||
|
# id stripped, trailing ``NS_END`` retained) so downstream
|
||||||
|
# consumers reading checkpoint metadata off a streamed
|
||||||
|
# message see the same shape they did pre-v2. Only the ns
|
||||||
|
# tuple emitted on the wire changes.
|
||||||
|
checkpoint_ns = (
|
||||||
|
f"{task_checkpoint_ns.rsplit(NS_END, 1)[0]}{NS_END}"
|
||||||
|
if NS_END in task_checkpoint_ns
|
||||||
|
else task_checkpoint_ns
|
||||||
|
)
|
||||||
|
stream_metadata["langgraph_checkpoint_ns"] = checkpoint_ns
|
||||||
|
stream_metadata["checkpoint_ns"] = checkpoint_ns
|
||||||
|
if tags:
|
||||||
|
if filtered_tags := [t for t in tags if not t.startswith("seq:step")]:
|
||||||
|
stream_metadata["tags"] = filtered_tags
|
||||||
|
self.metadata[run_id] = (ns, stream_metadata)
|
||||||
|
|
||||||
|
def on_chain_start(
|
||||||
|
self,
|
||||||
|
serialized: dict[str, Any],
|
||||||
|
inputs: dict[str, Any],
|
||||||
|
*,
|
||||||
|
run_id: UUID,
|
||||||
|
parent_run_id: UUID | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
metadata: dict[str, Any] | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
"""Record chain (node) metadata with the FULL checkpoint ns.
|
||||||
|
|
||||||
|
Mirror of :meth:`on_chat_model_start` for the node-start path,
|
||||||
|
so messages returned by ``on_chain_end`` (``Command`` updates
|
||||||
|
and plain state dict outputs) land at the same full-path ns as
|
||||||
|
any chat-model deltas from within that node. See the
|
||||||
|
:meth:`on_chat_model_start` docstring for why the v1 ``[:-1]``
|
||||||
|
slice is dropped here.
|
||||||
|
"""
|
||||||
|
if (
|
||||||
|
metadata
|
||||||
|
and kwargs.get("name") == metadata.get("langgraph_node")
|
||||||
|
and (not tags or TAG_HIDDEN not in tags)
|
||||||
|
):
|
||||||
|
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))
|
||||||
|
if not self.subgraphs and len(ns) > 1:
|
||||||
|
return
|
||||||
|
self.metadata[run_id] = (ns, metadata)
|
||||||
|
for value in _state_values(inputs):
|
||||||
|
if isinstance(value, BaseMessage):
|
||||||
|
if value.id is not None:
|
||||||
|
self.seen.add(value.id)
|
||||||
|
elif isinstance(value, Sequence) and not isinstance(value, str):
|
||||||
|
for item in value:
|
||||||
|
if isinstance(item, BaseMessage):
|
||||||
|
if item.id is not None:
|
||||||
|
self.seen.add(item.id)
|
||||||
|
|
||||||
|
def on_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 on_stream_event(
|
||||||
|
self,
|
||||||
|
event: dict[str, Any],
|
||||||
|
*,
|
||||||
|
run_id: UUID,
|
||||||
|
parent_run_id: UUID | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
"""Forward a protocol event from `stream_v2` as a messages stream part.
|
||||||
|
|
||||||
|
Fires once per `MessagesData` event (`message-start`, per-block
|
||||||
|
`content-block-*`, `message-finish`). The transformer layer
|
||||||
|
correlates events back to a single `ChatModelStream` via
|
||||||
|
`metadata["run_id"]` — attached here so the v1
|
||||||
|
`stream_mode="messages"` output (which emits
|
||||||
|
`(AIMessageChunk, metadata)` via `on_llm_new_token`) keeps its
|
||||||
|
original metadata shape.
|
||||||
|
|
||||||
|
Lives on the v2 handler rather than the v1 base: content-block
|
||||||
|
events are a v2-only concept, and forwarding them only when the
|
||||||
|
v2 handler is attached keeps the message channel's shape
|
||||||
|
predictable for v1 callers.
|
||||||
|
"""
|
||||||
|
if meta := self.metadata.get(run_id):
|
||||||
|
# Record message_id on message-start so on_chain_end's
|
||||||
|
# dedupe skips the finalized AIMessage the node returns
|
||||||
|
# (otherwise the messages projection double-counts: once
|
||||||
|
# from streaming, once from the chain output).
|
||||||
|
if event.get("event") == "message-start":
|
||||||
|
msg_id = event.get("message_id")
|
||||||
|
if msg_id:
|
||||||
|
self.seen.add(msg_id)
|
||||||
|
v2_meta = {**meta[1], "run_id": str(run_id)}
|
||||||
|
self.stream((meta[0], "messages", (event, v2_meta)))
|
||||||
|
|||||||
@@ -0,0 +1,223 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from collections.abc import AsyncIterator, Callable, Iterator
|
||||||
|
from contextvars import 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.config import _tool_call_writer
|
||||||
|
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."""
|
||||||
|
|
||||||
|
|
||||||
|
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`. The
|
||||||
|
`emit_tool_output_delta` helper in `langgraph.config` 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]) -> None:
|
||||||
|
"""Initialize the handler.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
stream: Callable that accepts a `StreamChunk` tuple
|
||||||
|
`(namespace, mode, payload)` and enqueues it.
|
||||||
|
"""
|
||||||
|
self.stream = stream
|
||||||
|
# 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]]
|
||||||
|
] = {}
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _containing_ns_from_metadata(
|
||||||
|
metadata: dict[str, Any] | None,
|
||||||
|
) -> tuple[str, ...]:
|
||||||
|
"""Return the namespace of the subgraph that contains this tool call.
|
||||||
|
|
||||||
|
`langgraph_checkpoint_ns` on a tool's callback metadata ends with
|
||||||
|
the `node_name:task_id` segment of the node that invoked the
|
||||||
|
tool. Dropping that segment gives the subgraph's own namespace,
|
||||||
|
which matches what other `tools` / `lifecycle` / `messages`
|
||||||
|
emitters use.
|
||||||
|
"""
|
||||||
|
if not metadata:
|
||||||
|
return ()
|
||||||
|
nskey = metadata.get("langgraph_checkpoint_ns")
|
||||||
|
if not nskey:
|
||||||
|
return ()
|
||||||
|
return tuple(cast(str, nskey).split(NS_SEP))[:-1]
|
||||||
|
|
||||||
|
def _start(
|
||||||
|
self,
|
||||||
|
serialized: dict[str, Any] | None,
|
||||||
|
input_str: str,
|
||||||
|
*,
|
||||||
|
run_id: UUID,
|
||||||
|
metadata: dict[str, Any] | None,
|
||||||
|
inputs: dict[str, Any] | None,
|
||||||
|
kwargs: dict[str, Any],
|
||||||
|
) -> None:
|
||||||
|
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 ""
|
||||||
|
)
|
||||||
|
ns = self._containing_ns_from_metadata(metadata)
|
||||||
|
|
||||||
|
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,
|
||||||
|
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_RUNTIME,
|
||||||
CONFIG_KEY_SEND,
|
CONFIG_KEY_SEND,
|
||||||
CONFIG_KEY_STREAM,
|
CONFIG_KEY_STREAM,
|
||||||
|
CONFIG_KEY_STREAM_MESSAGES_V2,
|
||||||
CONFIG_KEY_TASK_ID,
|
CONFIG_KEY_TASK_ID,
|
||||||
CONFIG_KEY_THREAD_ID,
|
CONFIG_KEY_THREAD_ID,
|
||||||
ERROR,
|
ERROR,
|
||||||
@@ -129,14 +130,19 @@ from langgraph.pregel._checkpoint import (
|
|||||||
)
|
)
|
||||||
from langgraph.pregel._draw import draw_graph
|
from langgraph.pregel._draw import draw_graph
|
||||||
from langgraph.pregel._io import map_input, read_channels
|
from langgraph.pregel._io import map_input, read_channels
|
||||||
|
from langgraph.pregel._lifecycle import StreamLifecycleHandler
|
||||||
from langgraph.pregel._loop import (
|
from langgraph.pregel._loop import (
|
||||||
AsyncPregelLoop,
|
AsyncPregelLoop,
|
||||||
SyncPregelLoop,
|
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._read import DEFAULT_BOUND, PregelNode
|
||||||
from langgraph.pregel._retry import RetryPolicy
|
from langgraph.pregel._retry import RetryPolicy
|
||||||
from langgraph.pregel._runner import PregelRunner
|
from langgraph.pregel._runner import PregelRunner
|
||||||
|
from langgraph.pregel._tools import StreamToolCallHandler
|
||||||
from langgraph.pregel._utils import get_new_channel_versions
|
from langgraph.pregel._utils import get_new_channel_versions
|
||||||
from langgraph.pregel._validate import validate_graph, validate_keys
|
from langgraph.pregel._validate import validate_graph, validate_keys
|
||||||
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
|
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
|
||||||
@@ -340,6 +346,65 @@ 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 are registered (or none declares
|
||||||
|
a given mode), the graph does not stream events for that mode.
|
||||||
|
"""
|
||||||
|
modes: set[str] = set()
|
||||||
|
for transformer in mux._transformers:
|
||||||
|
modes.update(transformer.required_stream_modes)
|
||||||
|
return cast("list[StreamMode]", list(modes))
|
||||||
|
|
||||||
|
|
||||||
|
def _build_stream_factories(
|
||||||
|
compile_time: Sequence[Callable[..., Any]],
|
||||||
|
call_site: Sequence[Any] | None,
|
||||||
|
) -> list[Callable[..., Any]]:
|
||||||
|
"""Assemble the factory list handed to `StreamMux(factories=...)`.
|
||||||
|
|
||||||
|
Prepends the built-in `ValuesTransformer`, `MessagesTransformer`,
|
||||||
|
and `SubgraphTransformer` factories, then appends the graph's
|
||||||
|
compile-time `stream_transformers` followed by any call-site
|
||||||
|
additions. Factories flow down into subgraph mini-muxes, so
|
||||||
|
per-scope instances propagate automatically.
|
||||||
|
"""
|
||||||
|
from langgraph.stream.transformers import (
|
||||||
|
MessagesTransformer,
|
||||||
|
SubgraphTransformer,
|
||||||
|
ToolLifecycleTransformer,
|
||||||
|
ValuesTransformer,
|
||||||
|
)
|
||||||
|
|
||||||
|
builtins: list[Callable[..., Any]] = [
|
||||||
|
ValuesTransformer,
|
||||||
|
ToolLifecycleTransformer,
|
||||||
|
MessagesTransformer,
|
||||||
|
SubgraphTransformer,
|
||||||
|
]
|
||||||
|
return [*builtins, *compile_time, *(call_site or ())]
|
||||||
|
|
||||||
|
|
||||||
|
def _merge_v2_messages_flag(
|
||||||
|
config: RunnableConfig | None,
|
||||||
|
) -> RunnableConfig:
|
||||||
|
"""Return a config with the v2 messages flag set in `configurable`.
|
||||||
|
|
||||||
|
Signals to pregel that `stream_mode="messages"` should attach
|
||||||
|
`StreamMessagesHandlerV2` for this call so invoke-time model runs
|
||||||
|
route through the v2 event generator and their protocol events
|
||||||
|
reach the messages channel.
|
||||||
|
"""
|
||||||
|
merged: RunnableConfig = dict(config or {}) # type: ignore[assignment]
|
||||||
|
configurable = dict(merged.get(CONF) or {})
|
||||||
|
configurable[CONFIG_KEY_STREAM_MESSAGES_V2] = True
|
||||||
|
merged[CONF] = configurable
|
||||||
|
return merged
|
||||||
|
|
||||||
|
|
||||||
class Pregel(
|
class Pregel(
|
||||||
PregelProtocol[StateT, ContextT, InputT, OutputT],
|
PregelProtocol[StateT, ContextT, InputT, OutputT],
|
||||||
Generic[StateT, ContextT, InputT, OutputT],
|
Generic[StateT, ContextT, InputT, OutputT],
|
||||||
@@ -671,6 +736,7 @@ class Pregel(
|
|||||||
config: RunnableConfig | None = None,
|
config: RunnableConfig | None = None,
|
||||||
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
|
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
|
||||||
name: str = "LangGraph",
|
name: str = "LangGraph",
|
||||||
|
stream_transformers: Sequence[Callable[..., Any]] | None = None,
|
||||||
**deprecated_kwargs: Unpack[DeprecatedKwargs],
|
**deprecated_kwargs: Unpack[DeprecatedKwargs],
|
||||||
) -> None:
|
) -> None:
|
||||||
if (
|
if (
|
||||||
@@ -717,6 +783,9 @@ class Pregel(
|
|||||||
self.config = config
|
self.config = config
|
||||||
self.trigger_to_nodes = trigger_to_nodes or {}
|
self.trigger_to_nodes = trigger_to_nodes or {}
|
||||||
self.name = name
|
self.name = name
|
||||||
|
self._stream_transformers: tuple[Callable[..., Any], ...] = tuple(
|
||||||
|
stream_transformers or ()
|
||||||
|
)
|
||||||
self._serde_allowlist: set[tuple[str, ...]] | None = None
|
self._serde_allowlist: set[tuple[str, ...]] | None = None
|
||||||
if auto_validate:
|
if auto_validate:
|
||||||
self.validate()
|
self.validate()
|
||||||
@@ -807,6 +876,15 @@ class Pregel(
|
|||||||
|
|
||||||
def copy(self, update: dict[str, Any] | None = None) -> Self:
|
def copy(self, update: dict[str, Any] | None = None) -> Self:
|
||||||
attrs = {k: v for k, v in self.__dict__.items() if k != "__orig_class__"}
|
attrs = {k: v for k, v in self.__dict__.items() if k != "__orig_class__"}
|
||||||
|
# ``__init__`` accepts ``stream_transformers`` (public parameter) but
|
||||||
|
# the attribute is stored as ``_stream_transformers`` (private). Map
|
||||||
|
# the private key back onto the public kwarg so compile-time
|
||||||
|
# transformers survive ``copy()`` / ``with_config()``. Without this,
|
||||||
|
# ``_stream_transformers`` gets captured by ``**deprecated_kwargs``
|
||||||
|
# and the resulting instance silently has an empty transformer
|
||||||
|
# pipeline.
|
||||||
|
if "_stream_transformers" in attrs:
|
||||||
|
attrs["stream_transformers"] = attrs.pop("_stream_transformers")
|
||||||
attrs.update(update or {})
|
attrs.update(update or {})
|
||||||
return self.__class__(**attrs)
|
return self.__class__(**attrs)
|
||||||
|
|
||||||
@@ -2626,14 +2704,34 @@ class Pregel(
|
|||||||
# set up messages stream mode
|
# set up messages stream mode
|
||||||
if "messages" in stream_modes:
|
if "messages" in stream_modes:
|
||||||
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||||
|
messages_handler_cls = (
|
||||||
|
StreamMessagesHandlerV2
|
||||||
|
if config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||||||
|
else StreamMessagesHandler
|
||||||
|
)
|
||||||
run_manager.inheritable_handlers.append(
|
run_manager.inheritable_handlers.append(
|
||||||
StreamMessagesHandler(
|
messages_handler_cls(
|
||||||
stream.put,
|
stream.put,
|
||||||
subgraphs,
|
subgraphs,
|
||||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# set up lifecycle stream mode
|
||||||
|
if "lifecycle" in stream_modes:
|
||||||
|
run_manager.inheritable_handlers.append(
|
||||||
|
StreamLifecycleHandler(
|
||||||
|
stream.put,
|
||||||
|
root_graph_name=self.name,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# set up tools stream mode
|
||||||
|
if "tools" in stream_modes:
|
||||||
|
run_manager.inheritable_handlers.append(
|
||||||
|
StreamToolCallHandler(stream.put)
|
||||||
|
)
|
||||||
|
|
||||||
# set up custom stream mode
|
# set up custom stream mode
|
||||||
if "custom" in stream_modes:
|
if "custom" in stream_modes:
|
||||||
|
|
||||||
@@ -3003,14 +3101,34 @@ class Pregel(
|
|||||||
if "messages" in stream_modes:
|
if "messages" in stream_modes:
|
||||||
# namespace can be None in a root level graph?
|
# namespace can be None in a root level graph?
|
||||||
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||||
|
messages_handler_cls = (
|
||||||
|
StreamMessagesHandlerV2
|
||||||
|
if config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||||||
|
else StreamMessagesHandler
|
||||||
|
)
|
||||||
run_manager.inheritable_handlers.append(
|
run_manager.inheritable_handlers.append(
|
||||||
StreamMessagesHandler(
|
messages_handler_cls(
|
||||||
stream_put,
|
stream_put,
|
||||||
subgraphs,
|
subgraphs,
|
||||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# set up lifecycle stream mode
|
||||||
|
if "lifecycle" in stream_modes:
|
||||||
|
run_manager.inheritable_handlers.append(
|
||||||
|
StreamLifecycleHandler(
|
||||||
|
stream_put,
|
||||||
|
root_graph_name=self.name,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# set up tools stream mode
|
||||||
|
if "tools" in stream_modes:
|
||||||
|
run_manager.inheritable_handlers.append(
|
||||||
|
StreamToolCallHandler(stream_put)
|
||||||
|
)
|
||||||
|
|
||||||
# set up custom stream mode
|
# set up custom stream mode
|
||||||
def stream_writer(c: Any) -> None:
|
def stream_writer(c: Any) -> None:
|
||||||
aioloop.call_soon_threadsafe(
|
aioloop.call_soon_threadsafe(
|
||||||
@@ -3237,6 +3355,106 @@ class Pregel(
|
|||||||
await asyncio.shield(run_manager.on_chain_error(e))
|
await asyncio.shield(run_manager.on_chain_error(e))
|
||||||
raise
|
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[Any] | None = None,
|
||||||
|
stream_modes: Sequence[StreamMode] | None = None,
|
||||||
|
output_keys: str | Sequence[str] | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> 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.
|
||||||
|
|
||||||
|
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 instances appended after
|
||||||
|
compile-time `stream_transformers`.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A `GraphRunStream` the caller iterates to drive the run.
|
||||||
|
"""
|
||||||
|
from langgraph.stream._mux import StreamMux
|
||||||
|
from langgraph.stream.run_stream import GraphRunStream
|
||||||
|
|
||||||
|
factories = _build_stream_factories(self._stream_transformers, transformers)
|
||||||
|
mux = StreamMux(factories=factories, is_async=False)
|
||||||
|
requested_stream_modes = set(_collect_stream_modes(mux))
|
||||||
|
requested_stream_modes.update(stream_modes or ())
|
||||||
|
graph_iter = iter(
|
||||||
|
self.stream(
|
||||||
|
input,
|
||||||
|
_merge_v2_messages_flag(config),
|
||||||
|
stream_mode=list(requested_stream_modes),
|
||||||
|
subgraphs=True,
|
||||||
|
version="v2",
|
||||||
|
output_keys=output_keys,
|
||||||
|
interrupt_before=interrupt_before,
|
||||||
|
interrupt_after=interrupt_after,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return GraphRunStream(graph_iter, mux)
|
||||||
|
|
||||||
|
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[Any] | None = None,
|
||||||
|
stream_modes: Sequence[StreamMode] | None = None,
|
||||||
|
output_keys: str | Sequence[str] | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> 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.
|
||||||
|
|
||||||
|
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 instances appended after
|
||||||
|
compile-time `stream_transformers`.
|
||||||
|
"""
|
||||||
|
from langgraph.stream._mux import StreamMux
|
||||||
|
from langgraph.stream.run_stream import AsyncGraphRunStream
|
||||||
|
|
||||||
|
factories = _build_stream_factories(self._stream_transformers, transformers)
|
||||||
|
mux = StreamMux(factories=factories, is_async=True)
|
||||||
|
requested_stream_modes = set(_collect_stream_modes(mux))
|
||||||
|
requested_stream_modes.update(stream_modes or ())
|
||||||
|
graph_aiter = self.astream(
|
||||||
|
input,
|
||||||
|
_merge_v2_messages_flag(config),
|
||||||
|
stream_mode=list(requested_stream_modes),
|
||||||
|
subgraphs=True,
|
||||||
|
version="v2",
|
||||||
|
output_keys=output_keys,
|
||||||
|
interrupt_before=interrupt_before,
|
||||||
|
interrupt_after=interrupt_after,
|
||||||
|
**kwargs,
|
||||||
|
).__aiter__()
|
||||||
|
return AsyncGraphRunStream(graph_aiter, mux)
|
||||||
|
|
||||||
@overload
|
@overload
|
||||||
def invoke(
|
def invoke(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -650,22 +650,46 @@ class RemoteGraph(PregelProtocol):
|
|||||||
"""
|
"""
|
||||||
updated_stream_modes: list[StreamModeSDK] = []
|
updated_stream_modes: list[StreamModeSDK] = []
|
||||||
req_single = True
|
req_single = True
|
||||||
|
# `"lifecycle"` is emitted locally by the `StreamLifecycleHandler`
|
||||||
|
# attached inside `Pregel.stream` / `astream`. The remote graph
|
||||||
|
# API has no corresponding mode, so requests for it against a
|
||||||
|
# `RemoteGraph` are silently stripped here and a warning is
|
||||||
|
# logged so the caller isn't left wondering why no lifecycle
|
||||||
|
# events arrive.
|
||||||
|
dropped_lifecycle = False
|
||||||
# coerce to list, or add default stream mode
|
# coerce to list, or add default stream mode
|
||||||
if stream_mode:
|
if stream_mode:
|
||||||
if isinstance(stream_mode, str):
|
if isinstance(stream_mode, str):
|
||||||
updated_stream_modes.append(stream_mode)
|
if stream_mode != "lifecycle":
|
||||||
|
updated_stream_modes.append(cast(StreamModeSDK, stream_mode))
|
||||||
|
else:
|
||||||
|
dropped_lifecycle = True
|
||||||
else:
|
else:
|
||||||
req_single = False
|
req_single = False
|
||||||
updated_stream_modes.extend(stream_mode)
|
for m in stream_mode:
|
||||||
|
if m == "lifecycle":
|
||||||
|
dropped_lifecycle = True
|
||||||
|
else:
|
||||||
|
updated_stream_modes.append(cast(StreamModeSDK, m))
|
||||||
else:
|
else:
|
||||||
updated_stream_modes.append(default)
|
updated_stream_modes.append(default) # type: ignore[arg-type]
|
||||||
requested_stream_modes = updated_stream_modes.copy()
|
requested_stream_modes = updated_stream_modes.copy()
|
||||||
# add any from parent graph
|
# add any from parent graph
|
||||||
stream: StreamProtocol | None = (
|
stream: StreamProtocol | None = (
|
||||||
(config or {}).get(CONF, {}).get(CONFIG_KEY_STREAM)
|
(config or {}).get(CONF, {}).get(CONFIG_KEY_STREAM)
|
||||||
)
|
)
|
||||||
if stream:
|
if stream:
|
||||||
updated_stream_modes.extend(stream.modes)
|
for m in stream.modes:
|
||||||
|
if m == "lifecycle":
|
||||||
|
dropped_lifecycle = True
|
||||||
|
else:
|
||||||
|
updated_stream_modes.append(cast(StreamModeSDK, m))
|
||||||
|
if dropped_lifecycle:
|
||||||
|
logger.warning(
|
||||||
|
"Stream mode 'lifecycle' is not supported by RemoteGraph "
|
||||||
|
"and was stripped from the request; no lifecycle events "
|
||||||
|
"will be emitted for this remote run."
|
||||||
|
)
|
||||||
# map "messages" to "messages-tuple"
|
# map "messages" to "messages-tuple"
|
||||||
if "messages" in updated_stream_modes:
|
if "messages" in updated_stream_modes:
|
||||||
updated_stream_modes.remove("messages")
|
updated_stream_modes.remove("messages")
|
||||||
@@ -996,6 +1020,99 @@ class RemoteGraph(PregelProtocol):
|
|||||||
else:
|
else:
|
||||||
yield chunk
|
yield chunk
|
||||||
|
|
||||||
|
def stream_v2(
|
||||||
|
self,
|
||||||
|
input: dict[str, Any] | Command | None,
|
||||||
|
config: RunnableConfig | None = None,
|
||||||
|
*,
|
||||||
|
context: Context | None = None,
|
||||||
|
interrupt_before: All | Sequence[str] | None = None,
|
||||||
|
interrupt_after: All | Sequence[str] | None = None,
|
||||||
|
transformers: Sequence[Any] | None = None,
|
||||||
|
stream_modes: Sequence[StreamMode] | None = None,
|
||||||
|
headers: dict[str, str] | None = None,
|
||||||
|
params: QueryParamTypes | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
"""Start a sync v2 remote run driven by transformer projections."""
|
||||||
|
from langgraph.pregel.main import (
|
||||||
|
_build_stream_factories,
|
||||||
|
_collect_stream_modes,
|
||||||
|
_merge_v2_messages_flag,
|
||||||
|
)
|
||||||
|
from langgraph.stream._convert import convert_to_protocol_event
|
||||||
|
from langgraph.stream._mux import StreamMux
|
||||||
|
from langgraph.stream.run_stream import RemoteGraphRunStream
|
||||||
|
|
||||||
|
factories = _build_stream_factories((), transformers)
|
||||||
|
mux = StreamMux(factories=factories, is_async=False)
|
||||||
|
requested_stream_modes = set(_collect_stream_modes(mux))
|
||||||
|
requested_stream_modes.update(stream_modes or ())
|
||||||
|
remote_iter = (
|
||||||
|
convert_to_protocol_event(part)
|
||||||
|
for part in self.stream(
|
||||||
|
input,
|
||||||
|
_merge_v2_messages_flag(config),
|
||||||
|
context=context,
|
||||||
|
stream_mode=list(requested_stream_modes),
|
||||||
|
interrupt_before=interrupt_before,
|
||||||
|
interrupt_after=interrupt_after,
|
||||||
|
subgraphs=True,
|
||||||
|
headers=headers,
|
||||||
|
params=params,
|
||||||
|
version="v2",
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return RemoteGraphRunStream(iter(remote_iter), mux)
|
||||||
|
|
||||||
|
async def astream_v2(
|
||||||
|
self,
|
||||||
|
input: dict[str, Any] | Command | None,
|
||||||
|
config: RunnableConfig | None = None,
|
||||||
|
*,
|
||||||
|
context: Context | None = None,
|
||||||
|
interrupt_before: All | Sequence[str] | None = None,
|
||||||
|
interrupt_after: All | Sequence[str] | None = None,
|
||||||
|
transformers: Sequence[Any] | None = None,
|
||||||
|
stream_modes: Sequence[StreamMode] | None = None,
|
||||||
|
headers: dict[str, str] | None = None,
|
||||||
|
params: QueryParamTypes | None = None,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
"""Async counterpart to ``stream_v2`` for remote graphs."""
|
||||||
|
from langgraph.pregel.main import (
|
||||||
|
_build_stream_factories,
|
||||||
|
_collect_stream_modes,
|
||||||
|
_merge_v2_messages_flag,
|
||||||
|
)
|
||||||
|
from langgraph.stream._convert import convert_to_protocol_event
|
||||||
|
from langgraph.stream._mux import StreamMux
|
||||||
|
from langgraph.stream.run_stream import AsyncRemoteGraphRunStream
|
||||||
|
|
||||||
|
factories = _build_stream_factories((), transformers)
|
||||||
|
mux = StreamMux(factories=factories, is_async=True)
|
||||||
|
requested_stream_modes = set(_collect_stream_modes(mux))
|
||||||
|
requested_stream_modes.update(stream_modes or ())
|
||||||
|
|
||||||
|
async def remote_events() -> AsyncIterator[Any]:
|
||||||
|
async for part in self.astream(
|
||||||
|
input,
|
||||||
|
_merge_v2_messages_flag(config),
|
||||||
|
context=context,
|
||||||
|
stream_mode=list(requested_stream_modes),
|
||||||
|
interrupt_before=interrupt_before,
|
||||||
|
interrupt_after=interrupt_after,
|
||||||
|
subgraphs=True,
|
||||||
|
headers=headers,
|
||||||
|
params=params,
|
||||||
|
version="v2",
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
yield convert_to_protocol_event(part)
|
||||||
|
|
||||||
|
return AsyncRemoteGraphRunStream(remote_events().__aiter__(), mux)
|
||||||
|
|
||||||
async def astream_events(
|
async def astream_events(
|
||||||
self,
|
self,
|
||||||
input: Any,
|
input: Any,
|
||||||
|
|||||||
@@ -0,0 +1,27 @@
|
|||||||
|
"""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,
|
||||||
|
AsyncRemoteGraphRunStream,
|
||||||
|
GraphRunStream,
|
||||||
|
RemoteGraphRunStream,
|
||||||
|
)
|
||||||
|
from langgraph.stream.stream_channel import StreamChannel
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"AsyncGraphRunStream",
|
||||||
|
"AsyncRemoteGraphRunStream",
|
||||||
|
"EventLog",
|
||||||
|
"GraphRunStream",
|
||||||
|
"RemoteGraphRunStream",
|
||||||
|
"ProtocolEvent",
|
||||||
|
"StreamChannel",
|
||||||
|
"StreamTransformer",
|
||||||
|
]
|
||||||
@@ -0,0 +1,71 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import time
|
||||||
|
from typing import Any, cast
|
||||||
|
|
||||||
|
from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams
|
||||||
|
from langgraph.types import StreamPart
|
||||||
|
|
||||||
|
|
||||||
|
def _is_v2_messages_payload(data: Any) -> bool:
|
||||||
|
return isinstance(data, dict) and isinstance(data.get("event"), str)
|
||||||
|
|
||||||
|
|
||||||
|
def _normalize_messages_data(data: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
"""Normalize Python Core message fields to the protocol wire shape."""
|
||||||
|
normalized = {**data}
|
||||||
|
if (
|
||||||
|
normalized["event"] == "message-start"
|
||||||
|
and "id" not in normalized
|
||||||
|
and isinstance(normalized.get("message_id"), str)
|
||||||
|
):
|
||||||
|
normalized["id"] = normalized["message_id"]
|
||||||
|
if (
|
||||||
|
normalized["event"]
|
||||||
|
in ("content-block-start", "content-block-delta", "content-block-finish")
|
||||||
|
and "content" not in normalized
|
||||||
|
and isinstance(normalized.get("content_block"), dict)
|
||||||
|
):
|
||||||
|
normalized["content"] = normalized["content_block"]
|
||||||
|
normalized.pop("message_id", None)
|
||||||
|
normalized.pop("content_block", None)
|
||||||
|
return normalized
|
||||||
|
|
||||||
|
|
||||||
|
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)
|
||||||
|
data = part_dict["data"]
|
||||||
|
params: _ProtocolEventParams = {
|
||||||
|
"namespace": list(part_dict["ns"]),
|
||||||
|
"timestamp": int(time.time() * 1000),
|
||||||
|
"data": data,
|
||||||
|
}
|
||||||
|
if (
|
||||||
|
part_dict["type"] == "messages"
|
||||||
|
and isinstance(data, tuple)
|
||||||
|
and len(data) == 2
|
||||||
|
and _is_v2_messages_payload(data[0])
|
||||||
|
and isinstance(data[1], dict)
|
||||||
|
):
|
||||||
|
payload, metadata = data
|
||||||
|
params["data"] = _normalize_messages_data(payload)
|
||||||
|
if isinstance(metadata.get("langgraph_node"), str):
|
||||||
|
params["node"] = metadata["langgraph_node"]
|
||||||
|
if isinstance(metadata.get("run_id"), str):
|
||||||
|
params["run_id"] = metadata["run_id"]
|
||||||
|
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,497 @@
|
|||||||
|
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` (root or mini-mux) with the mux's scope
|
||||||
|
— typically a subgraph's namespace or `()` for the root. Standard
|
||||||
|
transformer classes (`ValuesTransformer`, `MessagesTransformer`,
|
||||||
|
`SubgraphTransformer`) 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, ...] = (),
|
||||||
|
) -> None:
|
||||||
|
"""Initialize the mux and register transformers in order.
|
||||||
|
|
||||||
|
Callers pass either `transformers` (pre-built instances) or
|
||||||
|
`factories` (callables producing fresh instances per mux). A
|
||||||
|
factory list is preferred — mini-muxes built by `make_child()`
|
||||||
|
inherit the factory list, so transformers propagate naturally
|
||||||
|
into every subgraph's scope. `transformers` is kept for
|
||||||
|
back-compat tests that exercise the mux directly.
|
||||||
|
|
||||||
|
Each transformer's `init()` is called once during registration,
|
||||||
|
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. Mutually
|
||||||
|
exclusive with `factories`.
|
||||||
|
is_async: True for async dispatch (`apush` / `aclose` /
|
||||||
|
`afail`), False for the sync path.
|
||||||
|
factories: Zero-or-one-argument callables producing
|
||||||
|
transformers. Called with this mux's `scope`.
|
||||||
|
scope: The namespace the mux operates within. The root mux
|
||||||
|
is `()`; mini-muxes for subgraphs use the subgraph's
|
||||||
|
namespace tuple.
|
||||||
|
|
||||||
|
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, or if
|
||||||
|
both `transformers` and `factories` are supplied.
|
||||||
|
"""
|
||||||
|
if transformers is not None and factories is not None:
|
||||||
|
raise ValueError("Pass either `transformers` or `factories`, not both.")
|
||||||
|
|
||||||
|
self._is_async = is_async
|
||||||
|
self._factories: list[TransformerFactory] = list(factories or ())
|
||||||
|
self.scope: tuple[str, ...] = scope
|
||||||
|
self._pump_fn: Callable[[], bool] | None = None
|
||||||
|
self._apump_fn: Callable[[], Awaitable[bool]] | None = None
|
||||||
|
|
||||||
|
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] = {}
|
||||||
|
|
||||||
|
if factories is not None:
|
||||||
|
for factory in factories:
|
||||||
|
self._register(factory(scope))
|
||||||
|
else:
|
||||||
|
for transformer in transformers or ():
|
||||||
|
self._register(transformer)
|
||||||
|
|
||||||
|
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 binding (so cursors on its projection
|
||||||
|
logs drive the root pump) and carries the same factory list
|
||||||
|
forward to any grandchild subgraphs.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If the mux was not built from a factory list
|
||||||
|
(i.e., constructed with `transformers=`). Mini-muxes
|
||||||
|
require factories so each scope gets its own fresh
|
||||||
|
transformer instances.
|
||||||
|
"""
|
||||||
|
if not self._factories:
|
||||||
|
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,
|
||||||
|
)
|
||||||
|
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 bind_pump(self, fn: Callable[[], bool]) -> None:
|
||||||
|
"""Wire the sync pull callback onto every EventLog in the mux.
|
||||||
|
|
||||||
|
Also propagates to transformers that expose `_bind_pump` so
|
||||||
|
nested handles (e.g., `ChatModelStream` instances produced by
|
||||||
|
`MessagesTransformer`) can drive the graph pump from their
|
||||||
|
projection 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 _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}"
|
||||||
|
)
|
||||||
|
self._transformers.append(transformer)
|
||||||
|
self._bind_and_wire(projection)
|
||||||
|
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 getattr(transformer, "_native", False):
|
||||||
|
self.native_keys.update(projection.keys())
|
||||||
|
on_register = getattr(transformer, "_on_register", None)
|
||||||
|
if on_register is not None:
|
||||||
|
on_register(self)
|
||||||
|
|
||||||
|
def transformer_by_key(self, key: str) -> StreamTransformer | None:
|
||||||
|
"""Return the transformer that owns the projection at `key`, if any."""
|
||||||
|
return self._transformer_by_key.get(key)
|
||||||
|
|
||||||
|
def emit(self, event: ProtocolEvent) -> None:
|
||||||
|
"""Append a protocol event directly to the main log.
|
||||||
|
|
||||||
|
Built-in transformers use this for protocol repair events that
|
||||||
|
must appear before the source event they are processing. Direct
|
||||||
|
emission intentionally bypasses the transformer pipeline, but
|
||||||
|
still lets this mux remain the only local sequencing authority.
|
||||||
|
"""
|
||||||
|
self._seq += 1
|
||||||
|
event["seq"] = self._seq
|
||||||
|
self._events.push(event)
|
||||||
|
|
||||||
|
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
|
||||||
|
— except when the transformer has `scope_exact = True` (the
|
||||||
|
default) and the event's namespace differs from the mux's
|
||||||
|
`scope`, in which case the transformer is skipped. Transformers
|
||||||
|
that need to see cross-scope events opt out by setting
|
||||||
|
`scope_exact = False` (e.g. `SubgraphTransformer`).
|
||||||
|
|
||||||
|
If any transformer returns False, the event is suppressed from
|
||||||
|
the main log, but transformers that already saw it keep their
|
||||||
|
side effects.
|
||||||
|
|
||||||
|
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 log.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
event: The protocol event to dispatch.
|
||||||
|
"""
|
||||||
|
ns = tuple(event["params"]["namespace"])
|
||||||
|
in_scope = ns == self.scope
|
||||||
|
keep = True
|
||||||
|
for transformer in self._transformers:
|
||||||
|
if transformer.scope_exact and not in_scope:
|
||||||
|
continue
|
||||||
|
if not transformer.process(event):
|
||||||
|
keep = False
|
||||||
|
if keep:
|
||||||
|
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 — except when the transformer
|
||||||
|
has `scope_exact = True` and the event's namespace differs from
|
||||||
|
`self.scope`, in which case it is skipped. 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. 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.
|
||||||
|
"""
|
||||||
|
ns = tuple(event["params"]["namespace"])
|
||||||
|
in_scope = ns == self.scope
|
||||||
|
keep = True
|
||||||
|
for transformer in self._transformers:
|
||||||
|
if transformer.scope_exact and not in_scope:
|
||||||
|
continue
|
||||||
|
if not await transformer.aprocess(event):
|
||||||
|
keep = False
|
||||||
|
if keep:
|
||||||
|
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]) -> None:
|
||||||
|
"""Bind and wire EventLog / StreamChannel instances in a projection."""
|
||||||
|
for value in projection.values():
|
||||||
|
if isinstance(value, StreamChannel):
|
||||||
|
value._bind(is_async=self._is_async)
|
||||||
|
self._channels.append(value)
|
||||||
|
channel_name = value.name
|
||||||
|
|
||||||
|
def _make_forward(name: str) -> Callable[[Any], None]:
|
||||||
|
def _forward(item: Any) -> None:
|
||||||
|
self._forward(name, item)
|
||||||
|
|
||||||
|
return _forward
|
||||||
|
|
||||||
|
value._wire(_make_forward(channel_name))
|
||||||
|
elif isinstance(value, EventLog):
|
||||||
|
value._bind(is_async=self._is_async)
|
||||||
|
self._logs.append(value)
|
||||||
|
|
||||||
|
def _forward(self, channel_name: 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 the main event log but are not passed through
|
||||||
|
transformers' `process()` methods.
|
||||||
|
"""
|
||||||
|
event: ProtocolEvent = {
|
||||||
|
"type": "event",
|
||||||
|
"method": f"custom:{channel_name}",
|
||||||
|
"params": {
|
||||||
|
"namespace": [],
|
||||||
|
"timestamp": int(time.time() * 1000),
|
||||||
|
"data": item,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
self.emit(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
|
||||||
|
node: NotRequired[str]
|
||||||
|
run_id: NotRequired[str]
|
||||||
|
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 StreamMux.
|
||||||
|
Consumers that need a total order across events should use `seq`, not
|
||||||
|
`params.timestamp` (which is wall-clock and not monotonic).
|
||||||
|
"""
|
||||||
|
|
||||||
|
type: Literal["event"]
|
||||||
|
event_id: 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, a subgraph's namespace tuple inside a mini-mux.
|
||||||
|
Set at construction from the mux's scope (each factory is
|
||||||
|
called as `factory(scope)`). Transformers that only care
|
||||||
|
about events at their own namespace compare against
|
||||||
|
`self.scope`; subgraph-aware transformers can treat it as
|
||||||
|
a parent path.
|
||||||
|
scope_exact: If True (the default), the mux only calls
|
||||||
|
`process` / `aprocess` for events whose namespace equals
|
||||||
|
`self.scope` — user transformers get scope-scoped events
|
||||||
|
for free with no boilerplate. Set False for transformers
|
||||||
|
that need to see events across scopes (e.g.
|
||||||
|
`SubgraphTransformer` forwards deeper events into child
|
||||||
|
mini-muxes).
|
||||||
|
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.
|
||||||
|
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 `GraphStreamer` run requests from the
|
||||||
|
graph. Empty tuple means the transformer consumes only
|
||||||
|
synthetic events (or is purely passive).
|
||||||
|
"""
|
||||||
|
|
||||||
|
requires_async: ClassVar[bool] = False
|
||||||
|
scope_exact: ClassVar[bool] = True
|
||||||
|
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, the subgraph's namespace inside a
|
||||||
|
mini-mux. 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 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`).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
transformer: The transformer to inspect.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if the transformer cannot run under sync `stream()`.
|
||||||
|
"""
|
||||||
|
if transformer.requires_async:
|
||||||
|
return True
|
||||||
|
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,471 @@
|
|||||||
|
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 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 BaseRunStream:
|
||||||
|
"""Shared shape for any object that wraps a `StreamMux`.
|
||||||
|
|
||||||
|
Root (`GraphRunStream` / `AsyncGraphRunStream`) and scoped
|
||||||
|
(`SubgraphRunStream`) streams both compose a `StreamMux`. The mux
|
||||||
|
owns the projections — `values`, `messages`, `subgraphs`, and any
|
||||||
|
user-registered keys — all exposed via `extensions`. Native
|
||||||
|
projections (`_native = True`) are also bound as direct attributes
|
||||||
|
(`run.values`, `run.messages`, …) for ergonomics.
|
||||||
|
|
||||||
|
Raw iteration (`for event in run` / `async for event in run`) and
|
||||||
|
the `interleave(...)` helper both live here so every subclass
|
||||||
|
behaves consistently. Subclasses only add pump ownership, scope
|
||||||
|
metadata, or sync/async flavor.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, mux: StreamMux) -> None:
|
||||||
|
self._mux = mux
|
||||||
|
self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
|
||||||
|
for key in mux.native_keys:
|
||||||
|
setattr(self, key, mux.extensions[key])
|
||||||
|
|
||||||
|
@property
|
||||||
|
def _values_transformer(self) -> ValuesTransformer:
|
||||||
|
"""Look up the `ValuesTransformer` backing `output` / `interrupted`.
|
||||||
|
|
||||||
|
Resolved lazily off the mux so subclasses don't have to thread
|
||||||
|
it through their constructors. Raises if no `ValuesTransformer`
|
||||||
|
is registered — `output` / `interrupted` / `interrupts` have
|
||||||
|
nothing to return in that case, so failing loudly is better
|
||||||
|
than returning `None` silently.
|
||||||
|
"""
|
||||||
|
from langgraph.stream.transformers import ValuesTransformer
|
||||||
|
|
||||||
|
vt = self._mux.transformer_by_key("values")
|
||||||
|
if not isinstance(vt, ValuesTransformer):
|
||||||
|
raise RuntimeError(
|
||||||
|
"No ValuesTransformer is registered on this mux — "
|
||||||
|
"`output`, `interrupted`, and `interrupts` require one. "
|
||||||
|
"Add it to your GraphStreamer subclass's "
|
||||||
|
"`builtin_factories` or pass it via `transformers=`."
|
||||||
|
)
|
||||||
|
return vt
|
||||||
|
|
||||||
|
def __iter__(self) -> Iterator[ProtocolEvent]:
|
||||||
|
"""Sync iteration of protocol events on this mux's main log.
|
||||||
|
|
||||||
|
Raises at the EventLog level if the mux is async-bound.
|
||||||
|
"""
|
||||||
|
return iter(self._mux._events)
|
||||||
|
|
||||||
|
def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
|
||||||
|
"""Async iteration of protocol events on this mux's main log.
|
||||||
|
|
||||||
|
Raises at the EventLog level if the mux is sync-bound.
|
||||||
|
"""
|
||||||
|
return self._mux._events.__aiter__()
|
||||||
|
|
||||||
|
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 GraphRunStream(BaseRunStream):
|
||||||
|
"""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.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
graph_iter: Iterator[Any],
|
||||||
|
mux: StreamMux,
|
||||||
|
) -> None:
|
||||||
|
"""Initialize the run stream.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
graph_iter: Pull-based iterator over the graph's stream.
|
||||||
|
mux: The StreamMux owning projections and the main log.
|
||||||
|
Must have a `ValuesTransformer` registered under the
|
||||||
|
`"values"` key — `output` / `interrupted` / `interrupts`
|
||||||
|
read from it lazily.
|
||||||
|
"""
|
||||||
|
super().__init__(mux)
|
||||||
|
self._graph_iter = graph_iter
|
||||||
|
self._exhausted = False
|
||||||
|
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.
|
||||||
|
"""
|
||||||
|
if self._exhausted:
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
part = next(self._graph_iter)
|
||||||
|
except StopIteration:
|
||||||
|
self._mux.close()
|
||||||
|
self._exhausted = True
|
||||||
|
return False
|
||||||
|
except Exception as e:
|
||||||
|
self._mux.fail(e)
|
||||||
|
self._exhausted = True
|
||||||
|
return False
|
||||||
|
self._mux.push(convert_to_protocol_event(part))
|
||||||
|
return True
|
||||||
|
|
||||||
|
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)
|
||||||
|
vt = self._values_transformer
|
||||||
|
if vt.error is not None:
|
||||||
|
raise vt.error
|
||||||
|
return vt._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)
|
||||||
|
vt = self._values_transformer
|
||||||
|
if vt.error is not None:
|
||||||
|
raise vt.error
|
||||||
|
return vt._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)
|
||||||
|
vt = self._values_transformer
|
||||||
|
if vt.error is not None:
|
||||||
|
raise vt.error
|
||||||
|
return vt._interrupts
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncGraphRunStream(BaseRunStream):
|
||||||
|
"""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],
|
||||||
|
mux: StreamMux,
|
||||||
|
) -> None:
|
||||||
|
"""Initialize the async run stream.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
graph_aiter: Async iterator over the graph's stream.
|
||||||
|
mux: The StreamMux owning projections and the main log.
|
||||||
|
Must have a `ValuesTransformer` registered under the
|
||||||
|
`"values"` key — `output` / `interrupted` / `interrupts`
|
||||||
|
read from it lazily.
|
||||||
|
"""
|
||||||
|
super().__init__(mux)
|
||||||
|
self._graph_aiter = graph_aiter
|
||||||
|
self._exhausted = False
|
||||||
|
self._pump_cond = asyncio.Condition()
|
||||||
|
self._pumping = False
|
||||||
|
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:
|
||||||
|
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__()
|
||||||
|
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
|
||||||
|
await self._mux.apush(convert_to_protocol_event(part))
|
||||||
|
return True
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
class RemoteGraphRunStream(GraphRunStream):
|
||||||
|
"""Sync run stream fed by already-normalized remote protocol events."""
|
||||||
|
|
||||||
|
def __init__(self, events: Iterator[ProtocolEvent], mux: StreamMux) -> None:
|
||||||
|
super().__init__(events, mux)
|
||||||
|
|
||||||
|
def _pump_next(self) -> bool:
|
||||||
|
"""Pull one remote protocol event and push it through the mux."""
|
||||||
|
if self._exhausted:
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
event = next(self._graph_iter)
|
||||||
|
except StopIteration:
|
||||||
|
self._mux.close()
|
||||||
|
self._exhausted = True
|
||||||
|
return False
|
||||||
|
except Exception as e:
|
||||||
|
self._mux.fail(e)
|
||||||
|
self._exhausted = True
|
||||||
|
return False
|
||||||
|
self._mux.push(event)
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
class AsyncRemoteGraphRunStream(AsyncGraphRunStream):
|
||||||
|
"""Async run stream fed by already-normalized remote protocol events."""
|
||||||
|
|
||||||
|
def __init__(self, events: AsyncIterator[ProtocolEvent], mux: StreamMux) -> None:
|
||||||
|
super().__init__(events, mux)
|
||||||
|
|
||||||
|
async def _apump_next(self) -> bool:
|
||||||
|
"""Pull one remote protocol event and push it through the mux."""
|
||||||
|
async with self._pump_cond:
|
||||||
|
if self._exhausted:
|
||||||
|
return False
|
||||||
|
if self._pumping:
|
||||||
|
await self._pump_cond.wait()
|
||||||
|
return not self._exhausted
|
||||||
|
self._pumping = True
|
||||||
|
|
||||||
|
try:
|
||||||
|
try:
|
||||||
|
event = await self._graph_aiter.__anext__()
|
||||||
|
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
|
||||||
|
await self._mux.apush(event)
|
||||||
|
return True
|
||||||
|
finally:
|
||||||
|
async with self._pump_cond:
|
||||||
|
self._pumping = False
|
||||||
|
self._pump_cond.notify_all()
|
||||||
@@ -0,0 +1,109 @@
|
|||||||
|
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 (`custom:<name>` on the wire).
|
||||||
|
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,750 @@
|
|||||||
|
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 (
|
||||||
|
CheckpointRef,
|
||||||
|
LifecycleCause,
|
||||||
|
LifecycleData,
|
||||||
|
MessagesData,
|
||||||
|
)
|
||||||
|
|
||||||
|
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 BaseRunStream
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import Awaitable, Callable
|
||||||
|
|
||||||
|
from langgraph.stream._mux import StreamMux
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
SubgraphStatus = Literal["started", "running", "completed", "failed", "interrupted"]
|
||||||
|
_TERMINAL_STATUSES: frozenset[SubgraphStatus] = frozenset(
|
||||||
|
{"completed", "failed", "interrupted"}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _is_record(value: Any) -> bool:
|
||||||
|
return isinstance(value, dict)
|
||||||
|
|
||||||
|
|
||||||
|
def _to_chat_model_stream_event(event: MessagesData) -> MessagesData:
|
||||||
|
"""Convert wire-shaped message fields to ChatModelStream's internal shape."""
|
||||||
|
event_type = event.get("event")
|
||||||
|
converted: dict[str, Any] = dict(event)
|
||||||
|
if (
|
||||||
|
event_type == "message-start"
|
||||||
|
and "message_id" not in converted
|
||||||
|
and isinstance(converted.get("id"), str)
|
||||||
|
):
|
||||||
|
converted["message_id"] = converted["id"]
|
||||||
|
if (
|
||||||
|
event_type in ("content-block-start", "content-block-delta", "content-block-finish")
|
||||||
|
and "content_block" not in converted
|
||||||
|
and isinstance(converted.get("content"), dict)
|
||||||
|
):
|
||||||
|
converted["content_block"] = converted["content"]
|
||||||
|
return cast("MessagesData", converted)
|
||||||
|
|
||||||
|
|
||||||
|
def _message_event_id(event: MessagesData) -> str | None:
|
||||||
|
raw_id = event.get("id") or event.get("message_id")
|
||||||
|
return str(raw_id) if raw_id is not None else None
|
||||||
|
|
||||||
|
|
||||||
|
def _content_block_start_skeleton(content: Any) -> dict[str, Any] | None:
|
||||||
|
"""Return a minimal content-block-start payload for a delta/finish block."""
|
||||||
|
if not _is_record(content) or not isinstance(content.get("type"), str):
|
||||||
|
return None
|
||||||
|
|
||||||
|
block_type = content["type"]
|
||||||
|
skeleton: dict[str, Any] = {"type": block_type}
|
||||||
|
if block_type == "text":
|
||||||
|
skeleton["text"] = ""
|
||||||
|
elif block_type == "reasoning":
|
||||||
|
skeleton["reasoning"] = ""
|
||||||
|
elif block_type in ("tool_call", "tool_call_chunk"):
|
||||||
|
skeleton["type"] = "tool_call_chunk"
|
||||||
|
if isinstance(content.get("id"), str):
|
||||||
|
skeleton["id"] = content["id"]
|
||||||
|
if isinstance(content.get("name"), str):
|
||||||
|
skeleton["name"] = content["name"]
|
||||||
|
skeleton["args"] = ""
|
||||||
|
elif block_type in ("server_tool_call", "server_tool_call_chunk"):
|
||||||
|
skeleton["type"] = "server_tool_call_chunk"
|
||||||
|
if isinstance(content.get("id"), str):
|
||||||
|
skeleton["id"] = content["id"]
|
||||||
|
if isinstance(content.get("name"), str):
|
||||||
|
skeleton["name"] = content["name"]
|
||||||
|
skeleton["args"] = ""
|
||||||
|
return skeleton
|
||||||
|
|
||||||
|
|
||||||
|
def _copy_event(
|
||||||
|
source: ProtocolEvent,
|
||||||
|
*,
|
||||||
|
method: str,
|
||||||
|
namespace: list[str],
|
||||||
|
data: Any,
|
||||||
|
) -> ProtocolEvent:
|
||||||
|
params = {**source["params"], "namespace": namespace, "data": data}
|
||||||
|
return {"type": "event", "method": method, "params": params}
|
||||||
|
|
||||||
|
|
||||||
|
def _message_repair_key(event: ProtocolEvent, run_id: str) -> str:
|
||||||
|
namespace_key = "\x1f".join(event["params"]["namespace"])
|
||||||
|
return f"{namespace_key}\x1e{run_id}"
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_tool_calls_from_values(data: Any) -> dict[str, dict[str, Any]]:
|
||||||
|
if not _is_record(data):
|
||||||
|
return {}
|
||||||
|
messages = data.get("messages")
|
||||||
|
if not isinstance(messages, list):
|
||||||
|
return {}
|
||||||
|
known: dict[str, dict[str, Any]] = {}
|
||||||
|
for message in messages:
|
||||||
|
if not _is_record(message):
|
||||||
|
continue
|
||||||
|
tool_calls = message.get("tool_calls")
|
||||||
|
if not isinstance(tool_calls, list):
|
||||||
|
continue
|
||||||
|
for tool_call in tool_calls:
|
||||||
|
if not _is_record(tool_call):
|
||||||
|
continue
|
||||||
|
tool_call_id = tool_call.get("id")
|
||||||
|
if not isinstance(tool_call_id, str):
|
||||||
|
continue
|
||||||
|
name = tool_call.get("name")
|
||||||
|
args = tool_call.get("args")
|
||||||
|
known[tool_call_id] = {
|
||||||
|
"tool_name": name if isinstance(name, str) else "",
|
||||||
|
"input": args if _is_record(args) else {},
|
||||||
|
}
|
||||||
|
return known
|
||||||
|
|
||||||
|
|
||||||
|
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`).
|
||||||
|
|
||||||
|
`scope` (inherited from `StreamTransformer`) is the namespace the
|
||||||
|
transformer captures values for. `()` matches the root graph;
|
||||||
|
subgraph mini-muxes pass their subgraph's namespace, so each
|
||||||
|
instance sees only its own level.
|
||||||
|
"""
|
||||||
|
|
||||||
|
_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] = []
|
||||||
|
|
||||||
|
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:
|
||||||
|
# Namespace filtering is handled by the mux via `scope_exact`.
|
||||||
|
if event["method"] != "values":
|
||||||
|
return True
|
||||||
|
params = event["params"]
|
||||||
|
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 ToolLifecycleTransformer(StreamTransformer):
|
||||||
|
"""Repair tool-start events needed for deterministic subagent discovery.
|
||||||
|
|
||||||
|
Some subagent frameworks expose a tool-caused subgraph lifecycle before
|
||||||
|
a LangChain tool callback has emitted the matching `tool-started`
|
||||||
|
frame. Core can infer the missing start from the latest values snapshot
|
||||||
|
(`messages[*].tool_calls`) and emit it before the lifecycle event leaves
|
||||||
|
the mux, keeping remote clients from guessing from values snapshots.
|
||||||
|
"""
|
||||||
|
|
||||||
|
scope_exact = False
|
||||||
|
required_stream_modes = ("values", "tools", "lifecycle")
|
||||||
|
|
||||||
|
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||||
|
super().__init__(scope)
|
||||||
|
self._known_tool_calls: dict[str, dict[str, Any]] = {}
|
||||||
|
self._emitted_tool_starts: set[str] = set()
|
||||||
|
self._mux: StreamMux | None = None
|
||||||
|
|
||||||
|
def init(self) -> dict[str, Any]:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
def _on_register(self, mux: StreamMux) -> None:
|
||||||
|
self._mux = mux
|
||||||
|
|
||||||
|
def process(self, event: ProtocolEvent) -> bool:
|
||||||
|
method = event["method"]
|
||||||
|
data = event["params"]["data"]
|
||||||
|
if method == "values":
|
||||||
|
self._known_tool_calls.update(_extract_tool_calls_from_values(data))
|
||||||
|
return True
|
||||||
|
if method == "tools" and _is_record(data):
|
||||||
|
if (
|
||||||
|
data.get("event") == "tool-started"
|
||||||
|
and isinstance(data.get("tool_call_id"), str)
|
||||||
|
):
|
||||||
|
tool_call_id = cast("str", data["tool_call_id"])
|
||||||
|
if tool_call_id in self._emitted_tool_starts:
|
||||||
|
return False
|
||||||
|
self._emitted_tool_starts.add(tool_call_id)
|
||||||
|
return True
|
||||||
|
if method == "lifecycle":
|
||||||
|
self._emit_missing_tool_started(event)
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _emit_missing_tool_started(self, event: ProtocolEvent) -> None:
|
||||||
|
if self._mux is None:
|
||||||
|
return
|
||||||
|
data = event["params"]["data"]
|
||||||
|
if not _is_record(data) or data.get("event") != "started":
|
||||||
|
return
|
||||||
|
cause = data.get("cause")
|
||||||
|
if not _is_record(cause) or cause.get("type") != "toolCall":
|
||||||
|
return
|
||||||
|
tool_call_id = cause.get("tool_call_id")
|
||||||
|
if not isinstance(tool_call_id, str):
|
||||||
|
return
|
||||||
|
if tool_call_id in self._emitted_tool_starts:
|
||||||
|
return
|
||||||
|
known = self._known_tool_calls.get(tool_call_id)
|
||||||
|
if known is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
self._emitted_tool_starts.add(tool_call_id)
|
||||||
|
namespace = event["params"]["namespace"]
|
||||||
|
self._mux.emit(
|
||||||
|
_copy_event(
|
||||||
|
event,
|
||||||
|
method="tools",
|
||||||
|
namespace=namespace[:-1],
|
||||||
|
data={
|
||||||
|
"event": "tool-started",
|
||||||
|
"tool_call_id": tool_call_id,
|
||||||
|
"tool_name": known["tool_name"],
|
||||||
|
"input": known["input"],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
|
`scope` (inherited from `StreamTransformer`) is the namespace the
|
||||||
|
transformer captures messages for. `()` matches the root graph;
|
||||||
|
subgraph mini-muxes pass their subgraph's namespace, so each
|
||||||
|
instance sees only its own level.
|
||||||
|
|
||||||
|
Native transformer — the `messages` projection is exposed as a
|
||||||
|
direct attribute on the run stream.
|
||||||
|
|
||||||
|
`scope_exact = False`: matches events at the transformer's own
|
||||||
|
namespace **or** exactly one segment deeper (the chat-model /
|
||||||
|
node's own task ns). Mirrors JS's root-feed filter
|
||||||
|
(`namespaces=[[]], depth=1`) — root accepts depth-0 events plus
|
||||||
|
its own nodes' depth-1 tokens; subgraph mini-muxes accept their
|
||||||
|
own scope plus their internal nodes' tokens. Events deeper than
|
||||||
|
scope + 1 are dropped (the enclosing `SubgraphTransformer` has
|
||||||
|
already forwarded them to the matching child mini-mux).
|
||||||
|
"""
|
||||||
|
|
||||||
|
_native = True
|
||||||
|
scope_exact = False
|
||||||
|
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._started_blocks: dict[str, set[int]] = {}
|
||||||
|
self._mux: StreamMux | None = None
|
||||||
|
self._pump_fn: Callable[[], bool] | None = None
|
||||||
|
self._apump_fn: Callable[[], Awaitable[bool]] | None = None
|
||||||
|
|
||||||
|
def init(self) -> dict[str, Any]:
|
||||||
|
return {"messages": self._log}
|
||||||
|
|
||||||
|
def _on_register(self, mux: StreamMux) -> None:
|
||||||
|
self._mux = mux
|
||||||
|
|
||||||
|
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"]
|
||||||
|
# Accept events at our scope or exactly one segment deeper
|
||||||
|
# (the chat-model / node's own task ns). Deeper events belong
|
||||||
|
# to a subgraph and are routed by `SubgraphTransformer`.
|
||||||
|
ns = tuple(params["namespace"])
|
||||||
|
depth = len(self.scope)
|
||||||
|
if ns[:depth] != self.scope:
|
||||||
|
return True
|
||||||
|
|
||||||
|
raw_data = params["data"]
|
||||||
|
metadata: dict[str, Any] = {}
|
||||||
|
if isinstance(raw_data, tuple) and len(raw_data) == 2:
|
||||||
|
payload, raw_metadata = raw_data
|
||||||
|
metadata = raw_metadata if isinstance(raw_metadata, dict) else {}
|
||||||
|
else:
|
||||||
|
payload = raw_data
|
||||||
|
node = params.get("node")
|
||||||
|
if not isinstance(node, str):
|
||||||
|
node = metadata.get("langgraph_node")
|
||||||
|
if not isinstance(node, str):
|
||||||
|
node = None
|
||||||
|
raw_run_id = params.get("run_id", metadata.get("run_id"))
|
||||||
|
run_id = str(raw_run_id) if raw_run_id is not None else ""
|
||||||
|
|
||||||
|
if isinstance(payload, dict) and "event" in payload:
|
||||||
|
self._repair_content_block_lifecycle(
|
||||||
|
event, cast("MessagesData", payload), run_id=run_id
|
||||||
|
)
|
||||||
|
if len(ns) > depth + 1:
|
||||||
|
return True
|
||||||
|
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:
|
||||||
|
stream_event = _to_chat_model_stream_event(event)
|
||||||
|
event_type = event.get("event")
|
||||||
|
if event_type == "message-start":
|
||||||
|
message_id = _message_event_id(event)
|
||||||
|
stream = self._make_stream(
|
||||||
|
namespace=list(self.scope),
|
||||||
|
node=node,
|
||||||
|
message_id=message_id,
|
||||||
|
)
|
||||||
|
self._by_run[run_id or message_id or ""] = stream
|
||||||
|
self._log.push(stream)
|
||||||
|
stream.dispatch(stream_event)
|
||||||
|
elif run_id in self._by_run:
|
||||||
|
stream = self._by_run[run_id]
|
||||||
|
stream.dispatch(stream_event)
|
||||||
|
if event_type == "message-finish":
|
||||||
|
del self._by_run[run_id]
|
||||||
|
|
||||||
|
def _repair_content_block_lifecycle(
|
||||||
|
self,
|
||||||
|
source: ProtocolEvent,
|
||||||
|
event: MessagesData,
|
||||||
|
*,
|
||||||
|
run_id: str,
|
||||||
|
) -> None:
|
||||||
|
if self._mux is None:
|
||||||
|
return
|
||||||
|
event_type = event.get("event")
|
||||||
|
key = _message_repair_key(source, run_id)
|
||||||
|
if event_type == "message-start":
|
||||||
|
self._started_blocks[key] = set()
|
||||||
|
return
|
||||||
|
if event_type == "content-block-start":
|
||||||
|
index = event.get("index")
|
||||||
|
if isinstance(index, int):
|
||||||
|
self._started_blocks.setdefault(key, set()).add(index)
|
||||||
|
return
|
||||||
|
if event_type in ("content-block-delta", "content-block-finish"):
|
||||||
|
index = event.get("index")
|
||||||
|
if not isinstance(index, int):
|
||||||
|
return
|
||||||
|
started = self._started_blocks.setdefault(key, set())
|
||||||
|
if index in started:
|
||||||
|
return
|
||||||
|
skeleton = _content_block_start_skeleton(event.get("content"))
|
||||||
|
if skeleton is None:
|
||||||
|
return
|
||||||
|
started.add(index)
|
||||||
|
self._mux.emit(
|
||||||
|
_copy_event(
|
||||||
|
source,
|
||||||
|
method="messages",
|
||||||
|
namespace=list(source["params"]["namespace"]),
|
||||||
|
data={
|
||||||
|
"event": "content-block-start",
|
||||||
|
"index": index,
|
||||||
|
"content": skeleton,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
elif event_type == "message-finish":
|
||||||
|
self._started_blocks.pop(key, None)
|
||||||
|
|
||||||
|
def _route_whole_message(self, message: BaseMessage, *, node: str | None) -> None:
|
||||||
|
stream = self._make_stream(
|
||||||
|
namespace=list(self.scope),
|
||||||
|
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()
|
||||||
|
self._started_blocks.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()
|
||||||
|
self._started_blocks.clear()
|
||||||
|
|
||||||
|
|
||||||
|
class SubgraphRunStream(BaseRunStream):
|
||||||
|
"""Scoped view of a single nested subgraph execution.
|
||||||
|
|
||||||
|
Yielded on `run.subgraphs` (or `parent.subgraphs` for grandchildren)
|
||||||
|
when a nested `Pregel` spawns. Wraps a mini-`StreamMux` built with
|
||||||
|
the same transformer factories as the root mux, so `.values`,
|
||||||
|
`.messages`, `.subgraphs` are populated by the standard
|
||||||
|
transformers scoped to this handle's namespace — no duplicated
|
||||||
|
routing logic. The mini-mux borrows the root's pump via
|
||||||
|
`make_child`'s pump inheritance, so any cursor on a subagent
|
||||||
|
projection drives the whole run forward.
|
||||||
|
|
||||||
|
Lifecycle fields update in place as events arrive:
|
||||||
|
|
||||||
|
- `path`: the namespace tuple — stable for the life of the handle.
|
||||||
|
- `graph_name` / `cause`: set once from the `started` payload.
|
||||||
|
`cause` is populated by product-specific stream transformers
|
||||||
|
(see `LifecycleCause` in the protocol definition); pregel itself
|
||||||
|
emits no `cause`, so it may be `None` for subgraphs not covered
|
||||||
|
by a product transformer.
|
||||||
|
- `status`: advances `started` → `running` → `completed` /
|
||||||
|
`failed` / `interrupted`.
|
||||||
|
- `error` / `checkpoint`: set on the terminal event when present.
|
||||||
|
|
||||||
|
`.output` is a snapshot of the latest values seen at this
|
||||||
|
namespace — it doesn't drive the pump (unlike root's
|
||||||
|
`GraphRunStream.output`), because advancing a subgraph to
|
||||||
|
completion is only meaningful as part of advancing the whole run.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
path: tuple[str, ...],
|
||||||
|
mux: StreamMux,
|
||||||
|
*,
|
||||||
|
graph_name: str | None = None,
|
||||||
|
cause: LifecycleCause | None = None,
|
||||||
|
) -> None:
|
||||||
|
super().__init__(mux)
|
||||||
|
self.path: tuple[str, ...] = path
|
||||||
|
self.graph_name: str | None = graph_name
|
||||||
|
self.cause: LifecycleCause | None = cause
|
||||||
|
self.status: SubgraphStatus = "started"
|
||||||
|
self.error: str | None = None
|
||||||
|
self.checkpoint: CheckpointRef | None = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def output(self) -> dict[str, Any] | None:
|
||||||
|
"""Latest values snapshot at this namespace, or `None`.
|
||||||
|
|
||||||
|
Snapshot-only — iterating other projections or the root's
|
||||||
|
`.output` is what drives the pump.
|
||||||
|
"""
|
||||||
|
values_t = self._mux.transformer_by_key("values")
|
||||||
|
if isinstance(values_t, ValuesTransformer):
|
||||||
|
return values_t._latest
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class SubgraphTransformer(StreamTransformer):
|
||||||
|
"""Discover subgraphs and route events into per-subgraph mini-muxes.
|
||||||
|
|
||||||
|
Thin state-machine + dispatcher. At its own `scope` (inherited
|
||||||
|
from `StreamTransformer`, determined by the enclosing mux), it
|
||||||
|
watches for `lifecycle` events at exactly one level deeper to
|
||||||
|
discover direct children. Each discovered child gets its own
|
||||||
|
`SubgraphRunStream` backed by a mini-`StreamMux` — built via
|
||||||
|
`parent_mux.make_child(path)`, so the same factory list produces
|
||||||
|
fresh transformer instances at the child's scope.
|
||||||
|
|
||||||
|
Every incoming event that falls under one of the direct children
|
||||||
|
(ns starts with a child's `path`) is forwarded into that child's
|
||||||
|
mini-mux via `push`. The standard transformers in that mini-mux
|
||||||
|
(`ValuesTransformer`, `MessagesTransformer`, and another
|
||||||
|
`SubgraphTransformer` for grandchildren) handle the rest. No
|
||||||
|
duplicated routing or assembly logic.
|
||||||
|
|
||||||
|
Lifecycle state for each handle (running / completed / failed /
|
||||||
|
interrupted) is updated in place as events fire. On terminal
|
||||||
|
events, the handle's mini-mux is closed so any subscribed cursors
|
||||||
|
unblock. `finalize` / `fail` handle dangling handles left mid-run.
|
||||||
|
|
||||||
|
Native transformer — `subgraphs` exposes the direct-children log.
|
||||||
|
|
||||||
|
`scope_exact = False`: this transformer sees events at any
|
||||||
|
namespace, because it forwards out-of-scope events to the matching
|
||||||
|
direct-child mini-mux.
|
||||||
|
"""
|
||||||
|
|
||||||
|
_native = True
|
||||||
|
scope_exact = False
|
||||||
|
required_stream_modes = ("lifecycle",)
|
||||||
|
|
||||||
|
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||||
|
super().__init__(scope)
|
||||||
|
self._root_log: EventLog[SubgraphRunStream] = EventLog()
|
||||||
|
# Direct children only (namespace = scope + one segment).
|
||||||
|
self._by_ns: dict[tuple[str, ...], SubgraphRunStream] = {}
|
||||||
|
self._mux: StreamMux | None = None
|
||||||
|
|
||||||
|
def init(self) -> dict[str, Any]:
|
||||||
|
return {"subgraphs": self._root_log}
|
||||||
|
|
||||||
|
def _on_register(self, mux: StreamMux) -> None:
|
||||||
|
"""Capture the enclosing mux so we can build child mini-muxes."""
|
||||||
|
self._mux = mux
|
||||||
|
|
||||||
|
def process(self, event: ProtocolEvent) -> bool:
|
||||||
|
ns = tuple(event["params"]["namespace"])
|
||||||
|
method = event["method"]
|
||||||
|
depth = len(self.scope)
|
||||||
|
|
||||||
|
# 1. On `started` for a direct child (ns depth = mine + 1 and
|
||||||
|
# ns prefix matches mine), register the handle.
|
||||||
|
if method == "lifecycle" and len(ns) == depth + 1 and ns[:-1] == self.scope:
|
||||||
|
data = cast(LifecycleData, event["params"]["data"])
|
||||||
|
if data.get("event") == "started":
|
||||||
|
self._on_started(ns, data)
|
||||||
|
|
||||||
|
# 2. Forward the event to the matching direct-child mini-mux
|
||||||
|
# before the status-change step below so that terminal events
|
||||||
|
# reach the child's log and grandchild transformers *before*
|
||||||
|
# the child's mini-mux is closed. Prefix-match: ns must start
|
||||||
|
# with some child's path.
|
||||||
|
direct_child_ns = ns[: depth + 1] if len(ns) > depth else None
|
||||||
|
if direct_child_ns is not None and direct_child_ns in self._by_ns:
|
||||||
|
self._by_ns[direct_child_ns]._mux.push(event)
|
||||||
|
|
||||||
|
# 3. Status change for a direct child (ns = child's path, method
|
||||||
|
# = lifecycle). Update handle fields, close mini-mux on
|
||||||
|
# terminal.
|
||||||
|
if (
|
||||||
|
method == "lifecycle"
|
||||||
|
and ns in self._by_ns
|
||||||
|
and len(ns) == depth + 1
|
||||||
|
and ns[:-1] == self.scope
|
||||||
|
):
|
||||||
|
data = cast(LifecycleData, event["params"]["data"])
|
||||||
|
event_type = data.get("event")
|
||||||
|
if event_type in ("running", "completed", "failed", "interrupted"):
|
||||||
|
self._on_status_change(ns, event_type, data)
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _on_started(self, ns: tuple[str, ...], data: LifecycleData) -> None:
|
||||||
|
if ns in self._by_ns:
|
||||||
|
# Duplicate started — ignore.
|
||||||
|
return
|
||||||
|
# `_on_register` is called by the mux during registration, which
|
||||||
|
# happens before any event can be dispatched — so this should
|
||||||
|
# always be set by the time we process an event.
|
||||||
|
assert self._mux is not None, (
|
||||||
|
"SubgraphTransformer processed an event before _on_register; "
|
||||||
|
"transformer registration ordering is broken."
|
||||||
|
)
|
||||||
|
child_mux = self._mux.make_child(ns)
|
||||||
|
handle = SubgraphRunStream(
|
||||||
|
path=ns,
|
||||||
|
mux=child_mux,
|
||||||
|
graph_name=data.get("graph_name"),
|
||||||
|
cause=data.get("cause"),
|
||||||
|
)
|
||||||
|
self._by_ns[ns] = handle
|
||||||
|
self._root_log.push(handle)
|
||||||
|
|
||||||
|
def _on_status_change(
|
||||||
|
self,
|
||||||
|
ns: tuple[str, ...],
|
||||||
|
event_type: SubgraphStatus,
|
||||||
|
data: LifecycleData,
|
||||||
|
) -> None:
|
||||||
|
handle = self._by_ns[ns]
|
||||||
|
handle.status = event_type
|
||||||
|
err = data.get("error")
|
||||||
|
if err is not None:
|
||||||
|
handle.error = err
|
||||||
|
checkpoint = data.get("checkpoint")
|
||||||
|
if checkpoint is not None:
|
||||||
|
handle.checkpoint = checkpoint
|
||||||
|
if event_type in _TERMINAL_STATUSES:
|
||||||
|
self._close_handle_mux(handle)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _close_handle_mux(handle: SubgraphRunStream) -> None:
|
||||||
|
# Idempotent close — mux.close() runs finalize on its transformers
|
||||||
|
# (which cascades through grandchildren) and closes projection logs.
|
||||||
|
if not handle._mux._events._closed:
|
||||||
|
try:
|
||||||
|
handle._mux.close()
|
||||||
|
except Exception:
|
||||||
|
logger.warning(
|
||||||
|
"Error closing subgraph mini-mux at %s; subscribers "
|
||||||
|
"may not see a clean close.",
|
||||||
|
handle.path,
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
def finalize(self) -> None:
|
||||||
|
"""Transition any still-open direct children to `completed`."""
|
||||||
|
for handle in self._by_ns.values():
|
||||||
|
if handle.status not in _TERMINAL_STATUSES:
|
||||||
|
handle.status = "completed"
|
||||||
|
self._close_handle_mux(handle)
|
||||||
|
|
||||||
|
def fail(self, err: BaseException) -> None:
|
||||||
|
"""Transition any still-open direct children to `failed` / `interrupted`."""
|
||||||
|
is_interrupt = isinstance(err, GraphInterrupt)
|
||||||
|
terminal: SubgraphStatus = "interrupted" if is_interrupt else "failed"
|
||||||
|
error_str = None if is_interrupt else str(err)
|
||||||
|
for handle in self._by_ns.values():
|
||||||
|
if handle.status not in _TERMINAL_STATUSES:
|
||||||
|
handle.status = terminal
|
||||||
|
if error_str is not None and handle.error is None:
|
||||||
|
handle.error = error_str
|
||||||
|
if 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,
|
||||||
|
)
|
||||||
@@ -116,7 +116,15 @@ def ensure_valid_checkpointer(checkpointer: Checkpointer) -> Checkpointer:
|
|||||||
|
|
||||||
|
|
||||||
StreamMode = Literal[
|
StreamMode = Literal[
|
||||||
"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"
|
"values",
|
||||||
|
"updates",
|
||||||
|
"checkpoints",
|
||||||
|
"tasks",
|
||||||
|
"debug",
|
||||||
|
"messages",
|
||||||
|
"custom",
|
||||||
|
"lifecycle",
|
||||||
|
"tools",
|
||||||
]
|
]
|
||||||
"""How the stream method should emit outputs.
|
"""How the stream method should emit outputs.
|
||||||
|
|
||||||
@@ -129,6 +137,8 @@ StreamMode = Literal[
|
|||||||
- `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by `get_state()`.
|
- `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by `get_state()`.
|
||||||
- `"tasks"`: Emit events when tasks start and finish, including their results and errors.
|
- `"tasks"`: Emit events when tasks start and finish, including their results and errors.
|
||||||
- `"debug"`: Emit `"checkpoints"` and `"tasks"` events for debugging purposes.
|
- `"debug"`: Emit `"checkpoints"` and `"tasks"` events for debugging purposes.
|
||||||
|
- `"lifecycle"`: Emit subgraph lifecycle events (`started`, `running`, `completed`, `failed`, `interrupted`) with payloads matching `LifecycleData`.
|
||||||
|
- `"tools"`: Emit tool-call lifecycle events (`tool-started`, `tool-output-delta`, `tool-finished`, `tool-error`) keyed by `tool_call_id`.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
StreamWriter = Callable[[Any], None]
|
StreamWriter = Callable[[Any], None]
|
||||||
|
|||||||
@@ -24,7 +24,7 @@ classifiers = [
|
|||||||
'Programming Language :: Python :: 3.13',
|
'Programming Language :: Python :: 3.13',
|
||||||
]
|
]
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"langchain-core==1.3.0a2",
|
"langchain-core>=1.3.2",
|
||||||
"langgraph-checkpoint>=2.1.0,<5.0.0",
|
"langgraph-checkpoint>=2.1.0,<5.0.0",
|
||||||
"langgraph-sdk>=0.3.0,<0.4.0",
|
"langgraph-sdk>=0.3.0,<0.4.0",
|
||||||
"langgraph-prebuilt>=1.0.9,<1.1.0",
|
"langgraph-prebuilt>=1.0.9,<1.1.0",
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,121 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||||
|
from typing import Any, Literal
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from langchain_core.runnables import RunnableConfig
|
||||||
|
|
||||||
|
from langgraph.pregel.remote import RemoteGraph
|
||||||
|
from langgraph.types import All, StreamMode
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeRemoteGraph(RemoteGraph):
|
||||||
|
def __init__(self) -> None:
|
||||||
|
super().__init__("agent", url="http://unused")
|
||||||
|
self.last_stream_modes: list[StreamMode] | None = None
|
||||||
|
|
||||||
|
def stream(
|
||||||
|
self,
|
||||||
|
input: dict[str, Any] | Any,
|
||||||
|
config: RunnableConfig | None = None,
|
||||||
|
*,
|
||||||
|
stream_mode: StreamMode | list[StreamMode] | None = None,
|
||||||
|
interrupt_before: All | Sequence[str] | None = None,
|
||||||
|
interrupt_after: All | Sequence[str] | None = None,
|
||||||
|
subgraphs: bool = False,
|
||||||
|
version: Literal["v2"],
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Iterator[dict[str, Any]]:
|
||||||
|
assert version == "v2"
|
||||||
|
assert subgraphs is True
|
||||||
|
self.last_stream_modes = (
|
||||||
|
[stream_mode] if isinstance(stream_mode, str) else list(stream_mode or [])
|
||||||
|
)
|
||||||
|
yield {
|
||||||
|
"type": "values",
|
||||||
|
"ns": (),
|
||||||
|
"data": {"value": input["value"] + "A"},
|
||||||
|
"interrupts": (),
|
||||||
|
}
|
||||||
|
yield {
|
||||||
|
"type": "values",
|
||||||
|
"ns": (),
|
||||||
|
"data": {"value": input["value"] + "AB"},
|
||||||
|
"interrupts": (),
|
||||||
|
}
|
||||||
|
|
||||||
|
async def astream(
|
||||||
|
self,
|
||||||
|
input: dict[str, Any] | Any,
|
||||||
|
config: RunnableConfig | None = None,
|
||||||
|
*,
|
||||||
|
stream_mode: StreamMode | list[StreamMode] | None = None,
|
||||||
|
interrupt_before: All | Sequence[str] | None = None,
|
||||||
|
interrupt_after: All | Sequence[str] | None = None,
|
||||||
|
subgraphs: bool = False,
|
||||||
|
version: Literal["v2"],
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> AsyncIterator[dict[str, Any]]:
|
||||||
|
assert version == "v2"
|
||||||
|
assert subgraphs is True
|
||||||
|
self.last_stream_modes = (
|
||||||
|
[stream_mode] if isinstance(stream_mode, str) else list(stream_mode or [])
|
||||||
|
)
|
||||||
|
yield {
|
||||||
|
"type": "values",
|
||||||
|
"ns": (),
|
||||||
|
"data": {"value": input["value"] + "A"},
|
||||||
|
"interrupts": (),
|
||||||
|
}
|
||||||
|
yield {
|
||||||
|
"type": "values",
|
||||||
|
"ns": (),
|
||||||
|
"data": {"value": input["value"] + "AB"},
|
||||||
|
"interrupts": (),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_remote_stream_v2_values_and_output() -> None:
|
||||||
|
remote = _FakeRemoteGraph()
|
||||||
|
run = remote.stream_v2({"value": "x"})
|
||||||
|
|
||||||
|
assert list(run.values) == [{"value": "xA"}, {"value": "xAB"}]
|
||||||
|
assert "values" in (remote.last_stream_modes or [])
|
||||||
|
|
||||||
|
|
||||||
|
def test_remote_stream_v2_output_drains_remote_events() -> None:
|
||||||
|
remote = _FakeRemoteGraph()
|
||||||
|
run = remote.stream_v2({"value": "x"})
|
||||||
|
|
||||||
|
assert run.output == {"value": "xAB"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_remote_stream_v2_raw_events() -> None:
|
||||||
|
remote = _FakeRemoteGraph()
|
||||||
|
run = remote.stream_v2({"value": "x"})
|
||||||
|
|
||||||
|
events = list(run)
|
||||||
|
assert [event["method"] for event in events] == ["values", "values"]
|
||||||
|
assert [event["seq"] for event in events] == [1, 2]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_remote_astream_v2_values_and_output() -> None:
|
||||||
|
remote = _FakeRemoteGraph()
|
||||||
|
run = await remote.astream_v2({"value": "x"})
|
||||||
|
|
||||||
|
assert [item async for item in run.values] == [
|
||||||
|
{"value": "xA"},
|
||||||
|
{"value": "xAB"},
|
||||||
|
]
|
||||||
|
assert "values" in (remote.last_stream_modes or [])
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_remote_astream_v2_output_drains_remote_events() -> None:
|
||||||
|
remote = _FakeRemoteGraph()
|
||||||
|
run = await remote.astream_v2({"value": "x"})
|
||||||
|
|
||||||
|
assert await run.output() == {"value": "xAB"}
|
||||||
|
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,628 @@
|
|||||||
|
"""Tests for subgraph lifecycle events and the SubgraphTransformer."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import operator
|
||||||
|
import time
|
||||||
|
from typing import Annotated, Any
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from langgraph.checkpoint.memory import InMemorySaver
|
||||||
|
from typing_extensions import TypedDict
|
||||||
|
|
||||||
|
from langgraph.constants import END, START
|
||||||
|
from langgraph.errors import GraphInterrupt
|
||||||
|
from langgraph.graph import StateGraph
|
||||||
|
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,
|
||||||
|
SubgraphRunStream,
|
||||||
|
SubgraphTransformer,
|
||||||
|
ToolLifecycleTransformer,
|
||||||
|
ValuesTransformer,
|
||||||
|
)
|
||||||
|
from langgraph.types import interrupt
|
||||||
|
|
||||||
|
TS = int(time.time() * 1000)
|
||||||
|
|
||||||
|
|
||||||
|
def _lifecycle(
|
||||||
|
event: str,
|
||||||
|
*,
|
||||||
|
namespace: list[str] | None = None,
|
||||||
|
graph_name: str | None = None,
|
||||||
|
cause: dict[str, Any] | None = None,
|
||||||
|
error: str | None = None,
|
||||||
|
) -> ProtocolEvent:
|
||||||
|
data: dict[str, Any] = {"event": event}
|
||||||
|
if graph_name is not None:
|
||||||
|
data["graph_name"] = graph_name
|
||||||
|
if cause is not None:
|
||||||
|
data["cause"] = cause
|
||||||
|
if error is not None:
|
||||||
|
data["error"] = error
|
||||||
|
return {
|
||||||
|
"type": "event",
|
||||||
|
"method": "lifecycle",
|
||||||
|
"params": {
|
||||||
|
"namespace": namespace or [],
|
||||||
|
"timestamp": TS,
|
||||||
|
"data": data,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _values(payload: dict[str, Any], *, namespace: list[str]) -> ProtocolEvent:
|
||||||
|
return {
|
||||||
|
"type": "event",
|
||||||
|
"method": "values",
|
||||||
|
"params": {
|
||||||
|
"namespace": namespace,
|
||||||
|
"timestamp": TS,
|
||||||
|
"data": payload,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _subscribe(log: EventLog) -> None:
|
||||||
|
"""Flip `_subscribed = True` so pushes retain items for test inspection."""
|
||||||
|
log._subscribed = True
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Unit tests: feed events directly into the transformer
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
_FACTORIES = [
|
||||||
|
ValuesTransformer,
|
||||||
|
ToolLifecycleTransformer,
|
||||||
|
MessagesTransformer,
|
||||||
|
SubgraphTransformer,
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def _handle_values_items(handle: SubgraphRunStream) -> list:
|
||||||
|
return list(handle._mux.extensions["values"]._items) # type: ignore[attr-defined]
|
||||||
|
|
||||||
|
|
||||||
|
def _handle_subgraphs_items(handle: SubgraphRunStream) -> list:
|
||||||
|
return list(handle._mux.extensions["subgraphs"]._items) # type: ignore[attr-defined]
|
||||||
|
|
||||||
|
|
||||||
|
def _pre_subscribe_handle(handle: SubgraphRunStream) -> None:
|
||||||
|
"""Flip `_subscribed` on every EventLog inside the handle's mini-mux.
|
||||||
|
|
||||||
|
The mini-mux is built via `make_child` with the full factory list,
|
||||||
|
so values / messages / subgraphs logs all exist as projections.
|
||||||
|
Tests that feed events directly need them subscribed so pushes
|
||||||
|
retain items in the deque for `_items` inspection.
|
||||||
|
"""
|
||||||
|
for value in handle._mux.extensions.values():
|
||||||
|
if isinstance(value, EventLog):
|
||||||
|
_subscribe(value)
|
||||||
|
|
||||||
|
|
||||||
|
class TestSubgraphTransformerUnit:
|
||||||
|
def _mux(self) -> tuple[StreamMux, SubgraphTransformer]:
|
||||||
|
mux = StreamMux(factories=_FACTORIES, is_async=False)
|
||||||
|
transformer = mux.transformer_by_key("subgraphs")
|
||||||
|
assert isinstance(transformer, SubgraphTransformer)
|
||||||
|
_subscribe(transformer._root_log)
|
||||||
|
return mux, transformer
|
||||||
|
|
||||||
|
def _handle(self, transformer: SubgraphTransformer) -> SubgraphRunStream:
|
||||||
|
"""Return the single root handle after pushing one lifecycle started."""
|
||||||
|
(handle,) = list(transformer._root_log._items)
|
||||||
|
return handle
|
||||||
|
|
||||||
|
def test_root_started_is_ignored(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", graph_name="root"))
|
||||||
|
assert list(transformer._root_log._items) == []
|
||||||
|
assert transformer._by_ns == {}
|
||||||
|
|
||||||
|
def test_child_started_yields_handle(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(
|
||||||
|
_lifecycle(
|
||||||
|
"started",
|
||||||
|
namespace=["task_a:child"],
|
||||||
|
graph_name="child",
|
||||||
|
cause={"type": "toolCall", "tool_call_id": "call_abc"},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
handle = self._handle(transformer)
|
||||||
|
assert handle.path == ("task_a:child",)
|
||||||
|
assert handle.graph_name == "child"
|
||||||
|
assert handle.cause == {"type": "toolCall", "tool_call_id": "call_abc"}
|
||||||
|
assert handle.status == "started"
|
||||||
|
|
||||||
|
def test_tool_started_is_synthesized_before_tool_caused_lifecycle(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
events = iter(mux._events)
|
||||||
|
mux.push(
|
||||||
|
_values(
|
||||||
|
{
|
||||||
|
"messages": [
|
||||||
|
{
|
||||||
|
"tool_calls": [
|
||||||
|
{
|
||||||
|
"id": "call_abc",
|
||||||
|
"name": "task",
|
||||||
|
"args": {"subagent_type": "researcher"},
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
},
|
||||||
|
namespace=[],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
mux.push(
|
||||||
|
_lifecycle(
|
||||||
|
"started",
|
||||||
|
namespace=["task:child"],
|
||||||
|
graph_name="child",
|
||||||
|
cause={"type": "toolCall", "tool_call_id": "call_abc"},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
mux.close()
|
||||||
|
tool_started, lifecycle_started = list(events)[1:3]
|
||||||
|
assert tool_started["method"] == "tools"
|
||||||
|
assert tool_started["params"]["namespace"] == []
|
||||||
|
assert tool_started["params"]["data"] == {
|
||||||
|
"event": "tool-started",
|
||||||
|
"tool_call_id": "call_abc",
|
||||||
|
"tool_name": "task",
|
||||||
|
"input": {"subagent_type": "researcher"},
|
||||||
|
}
|
||||||
|
assert lifecycle_started["method"] == "lifecycle"
|
||||||
|
assert tool_started["seq"] < lifecycle_started["seq"]
|
||||||
|
assert self._handle(transformer).path == ("task:child",)
|
||||||
|
|
||||||
|
def test_core_golden_trace_uses_js_wire_shape_and_ordering(self) -> None:
|
||||||
|
mux, _transformer = self._mux()
|
||||||
|
events = iter(mux._events)
|
||||||
|
mux.push(
|
||||||
|
_values(
|
||||||
|
{
|
||||||
|
"messages": [
|
||||||
|
{
|
||||||
|
"tool_calls": [
|
||||||
|
{
|
||||||
|
"id": "call_abc",
|
||||||
|
"name": "task",
|
||||||
|
"args": {"subagent_type": "researcher"},
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
},
|
||||||
|
namespace=[],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
for data in (
|
||||||
|
{"event": "message-start", "id": "msg-1", "role": "ai"},
|
||||||
|
{
|
||||||
|
"event": "content-block-delta",
|
||||||
|
"index": 0,
|
||||||
|
"content": {"type": "text", "text": "hi"},
|
||||||
|
},
|
||||||
|
):
|
||||||
|
mux.push(
|
||||||
|
{
|
||||||
|
"type": "event",
|
||||||
|
"method": "messages",
|
||||||
|
"params": {
|
||||||
|
"namespace": ["call_model:task-1"],
|
||||||
|
"timestamp": TS,
|
||||||
|
"data": data,
|
||||||
|
"run_id": "run-1",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
mux.push(
|
||||||
|
_lifecycle(
|
||||||
|
"started",
|
||||||
|
namespace=["task:child"],
|
||||||
|
graph_name="child",
|
||||||
|
cause={"type": "toolCall", "tool_call_id": "call_abc"},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
mux.close()
|
||||||
|
|
||||||
|
trace = [
|
||||||
|
(event["method"], event["params"]["namespace"], event["params"]["data"])
|
||||||
|
for event in events
|
||||||
|
]
|
||||||
|
assert trace == [
|
||||||
|
(
|
||||||
|
"values",
|
||||||
|
[],
|
||||||
|
{
|
||||||
|
"messages": [
|
||||||
|
{
|
||||||
|
"tool_calls": [
|
||||||
|
{
|
||||||
|
"id": "call_abc",
|
||||||
|
"name": "task",
|
||||||
|
"args": {"subagent_type": "researcher"},
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
},
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"messages",
|
||||||
|
["call_model:task-1"],
|
||||||
|
{"event": "message-start", "id": "msg-1", "role": "ai"},
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"messages",
|
||||||
|
["call_model:task-1"],
|
||||||
|
{
|
||||||
|
"event": "content-block-start",
|
||||||
|
"index": 0,
|
||||||
|
"content": {"type": "text", "text": ""},
|
||||||
|
},
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"messages",
|
||||||
|
["call_model:task-1"],
|
||||||
|
{
|
||||||
|
"event": "content-block-delta",
|
||||||
|
"index": 0,
|
||||||
|
"content": {"type": "text", "text": "hi"},
|
||||||
|
},
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"tools",
|
||||||
|
[],
|
||||||
|
{
|
||||||
|
"event": "tool-started",
|
||||||
|
"tool_call_id": "call_abc",
|
||||||
|
"tool_name": "task",
|
||||||
|
"input": {"subagent_type": "researcher"},
|
||||||
|
},
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"lifecycle",
|
||||||
|
["task:child"],
|
||||||
|
{
|
||||||
|
"event": "started",
|
||||||
|
"graph_name": "child",
|
||||||
|
"cause": {"type": "toolCall", "tool_call_id": "call_abc"},
|
||||||
|
},
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
def test_status_transitions(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||||
|
mux.push(_lifecycle("running", namespace=["t:c"]))
|
||||||
|
mux.push(_lifecycle("completed", namespace=["t:c"]))
|
||||||
|
|
||||||
|
handle = self._handle(transformer)
|
||||||
|
assert handle.status == "completed"
|
||||||
|
|
||||||
|
def test_grandchild_surfaces_under_child(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:child"], graph_name="child"))
|
||||||
|
|
||||||
|
child = self._handle(transformer)
|
||||||
|
_pre_subscribe_handle(child)
|
||||||
|
|
||||||
|
mux.push(
|
||||||
|
_lifecycle(
|
||||||
|
"started",
|
||||||
|
namespace=["t:child", "u:grand"],
|
||||||
|
graph_name="grand",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
(grand,) = _handle_subgraphs_items(child)
|
||||||
|
assert grand.path == ("t:child", "u:grand")
|
||||||
|
assert grand.graph_name == "grand"
|
||||||
|
|
||||||
|
def test_failed_stores_error(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||||
|
mux.push(_lifecycle("failed", namespace=["t:c"], error="boom"))
|
||||||
|
|
||||||
|
handle = self._handle(transformer)
|
||||||
|
assert handle.status == "failed"
|
||||||
|
assert handle.error == "boom"
|
||||||
|
|
||||||
|
def test_values_routed_into_handle(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||||
|
|
||||||
|
handle = self._handle(transformer)
|
||||||
|
_pre_subscribe_handle(handle)
|
||||||
|
|
||||||
|
mux.push(_values({"value": 1}, namespace=["t:c"]))
|
||||||
|
mux.push(_values({"value": 2}, namespace=["t:c"]))
|
||||||
|
|
||||||
|
assert _handle_values_items(handle) == [{"value": 1}, {"value": 2}]
|
||||||
|
assert handle.output == {"value": 2}
|
||||||
|
|
||||||
|
def test_root_values_not_routed(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||||
|
handle = self._handle(transformer)
|
||||||
|
_pre_subscribe_handle(handle)
|
||||||
|
|
||||||
|
# Values event at root namespace — must not leak into child handle.
|
||||||
|
mux.push(_values({"value": "root"}, namespace=[]))
|
||||||
|
assert _handle_values_items(handle) == []
|
||||||
|
|
||||||
|
def test_finalize_closes_dangling(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||||
|
handle = self._handle(transformer)
|
||||||
|
|
||||||
|
mux.close()
|
||||||
|
assert handle.status == "completed"
|
||||||
|
assert handle._mux.extensions["values"]._closed
|
||||||
|
assert handle._mux.extensions["subgraphs"]._closed
|
||||||
|
|
||||||
|
def test_fail_with_graph_interrupt_marks_interrupted(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||||
|
handle = self._handle(transformer)
|
||||||
|
|
||||||
|
mux.fail(GraphInterrupt())
|
||||||
|
assert handle.status == "interrupted"
|
||||||
|
|
||||||
|
def test_fail_with_generic_error_marks_failed(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||||
|
handle = self._handle(transformer)
|
||||||
|
|
||||||
|
mux.fail(RuntimeError("explode"))
|
||||||
|
assert handle.status == "failed"
|
||||||
|
assert handle.error == "explode"
|
||||||
|
|
||||||
|
def test_duplicate_started_ignored(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||||
|
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="other"))
|
||||||
|
|
||||||
|
handles = list(transformer._root_log._items)
|
||||||
|
assert len(handles) == 1
|
||||||
|
assert handles[0].graph_name == "c"
|
||||||
|
|
||||||
|
def test_non_lifecycle_non_values_passthrough(self) -> None:
|
||||||
|
mux, transformer = self._mux()
|
||||||
|
mux.push(
|
||||||
|
{
|
||||||
|
"type": "event",
|
||||||
|
"method": "messages",
|
||||||
|
"params": {
|
||||||
|
"namespace": ["t:c"],
|
||||||
|
"timestamp": TS,
|
||||||
|
"data": (
|
||||||
|
{"event": "message-start", "message_id": "m1"},
|
||||||
|
{"run_id": "m1"},
|
||||||
|
),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
assert list(transformer._root_log._items) == []
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# End-to-end tests via stream_v2 on real graphs
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
class SimpleState(TypedDict):
|
||||||
|
value: str
|
||||||
|
items: Annotated[list[str], operator.add]
|
||||||
|
|
||||||
|
|
||||||
|
def _build_nested_graph():
|
||||||
|
"""Parent graph with a compiled subgraph node."""
|
||||||
|
|
||||||
|
def inner_node(state: SimpleState) -> dict:
|
||||||
|
return {"value": state["value"] + "X", "items": ["x"]}
|
||||||
|
|
||||||
|
inner_builder = StateGraph(SimpleState)
|
||||||
|
inner_builder.add_node("inner_node", inner_node)
|
||||||
|
inner_builder.add_edge(START, "inner_node")
|
||||||
|
inner_builder.add_edge("inner_node", END)
|
||||||
|
inner = inner_builder.compile()
|
||||||
|
|
||||||
|
def outer_node(state: SimpleState) -> dict:
|
||||||
|
return {"value": state["value"] + "Y", "items": ["y"]}
|
||||||
|
|
||||||
|
outer_builder = StateGraph(SimpleState)
|
||||||
|
outer_builder.add_node("outer_node", outer_node)
|
||||||
|
outer_builder.add_node("sub", inner)
|
||||||
|
outer_builder.add_edge(START, "outer_node")
|
||||||
|
outer_builder.add_edge("outer_node", "sub")
|
||||||
|
outer_builder.add_edge("sub", END)
|
||||||
|
return outer_builder.compile()
|
||||||
|
|
||||||
|
|
||||||
|
class TestSubgraphTransformerEndToEnd:
|
||||||
|
def test_flat_graph_yields_no_subgraphs(self) -> None:
|
||||||
|
builder = StateGraph(SimpleState)
|
||||||
|
builder.add_node("n", lambda s: {"value": s["value"] + "!", "items": ["!"]})
|
||||||
|
builder.add_edge(START, "n")
|
||||||
|
builder.add_edge("n", END)
|
||||||
|
graph = builder.compile()
|
||||||
|
|
||||||
|
run = graph.stream_v2({"value": "", "items": []})
|
||||||
|
|
||||||
|
collected: list[SubgraphRunStream] = []
|
||||||
|
for sub in run.subgraphs:
|
||||||
|
collected.append(sub)
|
||||||
|
assert collected == []
|
||||||
|
# Output still resolves.
|
||||||
|
assert run.output is not None
|
||||||
|
|
||||||
|
def test_nested_graph_yields_one_child(self) -> None:
|
||||||
|
graph = _build_nested_graph()
|
||||||
|
run = graph.stream_v2({"value": "", "items": []})
|
||||||
|
|
||||||
|
collected: list[SubgraphRunStream] = []
|
||||||
|
for sub in run.subgraphs:
|
||||||
|
collected.append(sub)
|
||||||
|
|
||||||
|
assert len(collected) == 1
|
||||||
|
child = collected[0]
|
||||||
|
assert len(child.path) == 1
|
||||||
|
assert child.path[0].startswith("sub:")
|
||||||
|
assert child.status == "completed"
|
||||||
|
|
||||||
|
def test_error_in_subgraph_fails_child(self) -> None:
|
||||||
|
def boom(state: SimpleState) -> dict:
|
||||||
|
raise RuntimeError("subgraph_failed")
|
||||||
|
|
||||||
|
inner_builder = StateGraph(SimpleState)
|
||||||
|
inner_builder.add_node("inner", boom)
|
||||||
|
inner_builder.add_edge(START, "inner")
|
||||||
|
inner_builder.add_edge("inner", END)
|
||||||
|
inner = inner_builder.compile()
|
||||||
|
|
||||||
|
outer_builder = StateGraph(SimpleState)
|
||||||
|
outer_builder.add_node("sub", inner)
|
||||||
|
outer_builder.add_edge(START, "sub")
|
||||||
|
outer_builder.add_edge("sub", END)
|
||||||
|
graph = outer_builder.compile()
|
||||||
|
|
||||||
|
run = graph.stream_v2({"value": "", "items": []})
|
||||||
|
|
||||||
|
collected: list[SubgraphRunStream] = []
|
||||||
|
with pytest.raises(RuntimeError):
|
||||||
|
for sub in run.subgraphs:
|
||||||
|
collected.append(sub)
|
||||||
|
|
||||||
|
assert len(collected) == 1
|
||||||
|
assert collected[0].status == "failed"
|
||||||
|
|
||||||
|
|
||||||
|
class TestSubgraphTransformerAsyncEndToEnd:
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_nested_graph_yields_one_child(self) -> None:
|
||||||
|
async def inner(state: SimpleState) -> dict:
|
||||||
|
return {"value": state["value"] + "X", "items": ["x"]}
|
||||||
|
|
||||||
|
inner_builder = StateGraph(SimpleState)
|
||||||
|
inner_builder.add_node("inner", inner)
|
||||||
|
inner_builder.add_edge(START, "inner")
|
||||||
|
inner_builder.add_edge("inner", END)
|
||||||
|
inner_graph = inner_builder.compile()
|
||||||
|
|
||||||
|
outer_builder = StateGraph(SimpleState)
|
||||||
|
outer_builder.add_node("sub", inner_graph)
|
||||||
|
outer_builder.add_edge(START, "sub")
|
||||||
|
outer_builder.add_edge("sub", END)
|
||||||
|
graph = outer_builder.compile()
|
||||||
|
|
||||||
|
run = await graph.astream_v2({"value": "", "items": []})
|
||||||
|
|
||||||
|
collected: list[SubgraphRunStream] = []
|
||||||
|
async for sub in run.subgraphs:
|
||||||
|
collected.append(sub)
|
||||||
|
|
||||||
|
assert len(collected) == 1
|
||||||
|
child = collected[0]
|
||||||
|
assert child.status == "completed"
|
||||||
|
|
||||||
|
|
||||||
|
class TestSubgraphCause:
|
||||||
|
"""Pregel core emits no `cause`; product transformers populate it."""
|
||||||
|
|
||||||
|
def test_cause_not_populated_by_pregel(self) -> None:
|
||||||
|
graph = _build_nested_graph()
|
||||||
|
run = graph.stream_v2({"value": "", "items": []})
|
||||||
|
|
||||||
|
collected: list[SubgraphRunStream] = list(run.subgraphs)
|
||||||
|
assert len(collected) == 1
|
||||||
|
child = collected[0]
|
||||||
|
|
||||||
|
# The child's single-segment path still encodes `node_name:task_id`
|
||||||
|
# (that's pregel's internal namespace format), but `cause` is now
|
||||||
|
# product-agnostic and must be populated by a stream transformer,
|
||||||
|
# not by pregel itself.
|
||||||
|
assert ":" in child.path[0]
|
||||||
|
node_name, _, task_id = child.path[0].partition(":")
|
||||||
|
assert node_name == "sub"
|
||||||
|
assert task_id # non-empty
|
||||||
|
assert child.cause is None
|
||||||
|
|
||||||
|
|
||||||
|
class TestSubgraphInterrupt:
|
||||||
|
"""Interrupts raised inside a subgraph surface as status=interrupted."""
|
||||||
|
|
||||||
|
def _build_interrupt_subgraph(self):
|
||||||
|
def inner_node(state: SimpleState) -> dict:
|
||||||
|
interrupt("need approval")
|
||||||
|
return {"value": state["value"] + "X", "items": ["x"]}
|
||||||
|
|
||||||
|
inner_builder = StateGraph(SimpleState)
|
||||||
|
inner_builder.add_node("inner_node", inner_node)
|
||||||
|
inner_builder.add_edge(START, "inner_node")
|
||||||
|
inner_builder.add_edge("inner_node", END)
|
||||||
|
inner = inner_builder.compile()
|
||||||
|
|
||||||
|
outer_builder = StateGraph(SimpleState)
|
||||||
|
outer_builder.add_node("sub", inner)
|
||||||
|
outer_builder.add_edge(START, "sub")
|
||||||
|
outer_builder.add_edge("sub", END)
|
||||||
|
return outer_builder.compile(checkpointer=InMemorySaver())
|
||||||
|
|
||||||
|
def test_interrupt_in_subgraph_marks_handle_interrupted(self) -> None:
|
||||||
|
graph = self._build_interrupt_subgraph()
|
||||||
|
run = graph.stream_v2(
|
||||||
|
{"value": "", "items": []},
|
||||||
|
config={"configurable": {"thread_id": "t1"}},
|
||||||
|
)
|
||||||
|
|
||||||
|
collected: list[SubgraphRunStream] = list(run.subgraphs)
|
||||||
|
|
||||||
|
assert run.interrupted is True
|
||||||
|
assert len(collected) == 1
|
||||||
|
assert collected[0].status == "interrupted"
|
||||||
|
|
||||||
|
|
||||||
|
class TestSubgraphNameCollision:
|
||||||
|
"""The subgraph's compiled `name` equaling its node name is detected.
|
||||||
|
|
||||||
|
Primary detector `name != langgraph_node` fails here; the
|
||||||
|
parent_run_id fallback in `_is_nested_pregel_start` is what keeps
|
||||||
|
the subgraph visible.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def test_name_equals_node_name_still_detected(self) -> None:
|
||||||
|
def inner_node(state: SimpleState) -> dict:
|
||||||
|
return {"value": state["value"] + "X", "items": ["x"]}
|
||||||
|
|
||||||
|
inner_builder = StateGraph(SimpleState)
|
||||||
|
inner_builder.add_node("inner_node", inner_node)
|
||||||
|
inner_builder.add_edge(START, "inner_node")
|
||||||
|
inner_builder.add_edge("inner_node", END)
|
||||||
|
# Compile with the same name as the node it will be registered as.
|
||||||
|
inner = inner_builder.compile(name="sub")
|
||||||
|
|
||||||
|
outer_builder = StateGraph(SimpleState)
|
||||||
|
outer_builder.add_node("sub", inner)
|
||||||
|
outer_builder.add_edge(START, "sub")
|
||||||
|
outer_builder.add_edge("sub", END)
|
||||||
|
graph = outer_builder.compile()
|
||||||
|
|
||||||
|
run = graph.stream_v2({"value": "", "items": []})
|
||||||
|
|
||||||
|
collected: list[SubgraphRunStream] = list(run.subgraphs)
|
||||||
|
assert len(collected) == 1
|
||||||
|
child = collected[0]
|
||||||
|
assert child.graph_name == "sub"
|
||||||
|
assert child.status == "completed"
|
||||||
@@ -0,0 +1,290 @@
|
|||||||
|
"""Tests for StreamToolCallHandler and emit_tool_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
|
||||||
|
from typing_extensions import TypedDict
|
||||||
|
|
||||||
|
from langgraph.config import emit_tool_output_delta
|
||||||
|
from langgraph.constants import END, START
|
||||||
|
from langgraph.graph import StateGraph
|
||||||
|
from langgraph.graph.message import add_messages
|
||||||
|
|
||||||
|
|
||||||
|
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_tool_output_delta_produces_delta_events(self) -> None:
|
||||||
|
@tool
|
||||||
|
def streaming_echo(text: str) -> str:
|
||||||
|
"""stream chunks."""
|
||||||
|
for chunk in ("a", "b", "c"):
|
||||||
|
emit_tool_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_emit_outside_tool_is_noop(self) -> None:
|
||||||
|
# Called at import time (outside any tool body) — must not raise.
|
||||||
|
emit_tool_output_delta("ignored")
|
||||||
|
emit_tool_output_delta({"any": "payload"})
|
||||||
|
|
||||||
|
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) -> str:
|
||||||
|
"""async echo."""
|
||||||
|
emit_tool_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) -> str:
|
||||||
|
"""emits marker twice."""
|
||||||
|
emit_tool_output_delta(f"{marker}-1")
|
||||||
|
emit_tool_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]]
|
[[package]]
|
||||||
name = "langchain-core"
|
name = "langchain-core"
|
||||||
version = "1.3.0a2"
|
version = "1.3.2"
|
||||||
source = { registry = "https://pypi.org/simple" }
|
source = { registry = "https://pypi.org/simple" }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
{ name = "jsonpatch" },
|
{ name = "jsonpatch" },
|
||||||
|
{ name = "langchain-protocol" },
|
||||||
{ name = "langsmith" },
|
{ name = "langsmith" },
|
||||||
{ name = "packaging" },
|
{ name = "packaging" },
|
||||||
{ name = "pydantic" },
|
{ name = "pydantic" },
|
||||||
@@ -1360,9 +1361,21 @@ dependencies = [
|
|||||||
{ name = "typing-extensions" },
|
{ name = "typing-extensions" },
|
||||||
{ name = "uuid-utils" },
|
{ name = "uuid-utils" },
|
||||||
]
|
]
|
||||||
sdist = { url = "https://files.pythonhosted.org/packages/af/bc/0bff31fcaff174d86031cc713471a3e85ed4ec8e5cd95ad0217f2aced20e/langchain_core-1.3.0a2.tar.gz", hash = "sha256:52d978c84552b74b9a3f16c1fced84f9e27cc96d7a67c601925ce6cbc4ea3cf9", size = 854580, upload-time = "2026-04-13T14:37:55.745Z" }
|
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 = [
|
wheels = [
|
||||||
{ url = "https://files.pythonhosted.org/packages/0e/14/03c09686602567059f26af29de0c44546a83af2f2aa29925e61040e43ea2/langchain_core-1.3.0a2-py3-none-any.whl", hash = "sha256:9e929a34f0b0c6c1255e395a1de34f8626893ceb4cdae550a22a0bd18c87be54", size = 510233, upload-time = "2026-04-13T14:37:54.277Z" },
|
{ 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]]
|
[[package]]
|
||||||
@@ -1439,7 +1452,7 @@ test = [
|
|||||||
|
|
||||||
[package.metadata]
|
[package.metadata]
|
||||||
requires-dist = [
|
requires-dist = [
|
||||||
{ name = "langchain-core", specifier = "==1.3.0a2" },
|
{ name = "langchain-core", specifier = ">=1.3.2" },
|
||||||
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
|
||||||
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
|
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
|
||||||
{ name = "langgraph-sdk", editable = "../sdk-py" },
|
{ name = "langgraph-sdk", editable = "../sdk-py" },
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
"""langgraph.prebuilt exposes a higher-level API for creating and executing agents and tools."""
|
"""langgraph.prebuilt exposes a higher-level API for creating and executing agents and tools."""
|
||||||
|
|
||||||
|
from langgraph.prebuilt._tool_call_stream import ToolCallStream
|
||||||
|
from langgraph.prebuilt._tool_call_transformer import ToolCallTransformer
|
||||||
from langgraph.prebuilt.chat_agent_executor import create_react_agent
|
from langgraph.prebuilt.chat_agent_executor import create_react_agent
|
||||||
from langgraph.prebuilt.tool_node import (
|
from langgraph.prebuilt.tool_node import (
|
||||||
InjectedState,
|
InjectedState,
|
||||||
@@ -13,6 +15,8 @@ from langgraph.prebuilt.tool_validator import ValidationNode
|
|||||||
__all__ = [
|
__all__ = [
|
||||||
"create_react_agent",
|
"create_react_agent",
|
||||||
"ToolNode",
|
"ToolNode",
|
||||||
|
"ToolCallStream",
|
||||||
|
"ToolCallTransformer",
|
||||||
"tools_condition",
|
"tools_condition",
|
||||||
"ValidationNode",
|
"ValidationNode",
|
||||||
"InjectedState",
|
"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()
|
||||||
@@ -0,0 +1,306 @@
|
|||||||
|
"""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.config import emit_tool_output_delta
|
||||||
|
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,
|
||||||
|
SubgraphTransformer,
|
||||||
|
ValuesTransformer,
|
||||||
|
)
|
||||||
|
from typing_extensions import TypedDict
|
||||||
|
|
||||||
|
from langgraph.prebuilt import ToolCallStream, ToolCallTransformer, ToolNode
|
||||||
|
|
||||||
|
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]:
|
||||||
|
mux = StreamMux(
|
||||||
|
factories=[
|
||||||
|
ValuesTransformer,
|
||||||
|
MessagesTransformer,
|
||||||
|
SubgraphTransformer,
|
||||||
|
ToolCallTransformer,
|
||||||
|
],
|
||||||
|
is_async=False,
|
||||||
|
)
|
||||||
|
transformer = mux.transformer_by_key("tool_calls")
|
||||||
|
assert isinstance(transformer, ToolCallTransformer)
|
||||||
|
_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) -> str:
|
||||||
|
"""streams chunks."""
|
||||||
|
for chunk in ("one", "two"):
|
||||||
|
emit_tool_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) -> str:
|
||||||
|
"""async streams."""
|
||||||
|
emit_tool_output_delta(text)
|
||||||
|
emit_tool_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
|
||||||
@@ -0,0 +1,228 @@
|
|||||||
|
"""Async thread-centric streaming primitives."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import time
|
||||||
|
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||||
|
from typing import Any, cast
|
||||||
|
|
||||||
|
from langgraph_sdk._async.http import HttpClient
|
||||||
|
from langgraph_sdk.protocol import (
|
||||||
|
Channel,
|
||||||
|
Command,
|
||||||
|
CommandMethod,
|
||||||
|
CommandResponse,
|
||||||
|
ErrorResponse,
|
||||||
|
Event,
|
||||||
|
SubscribeParams,
|
||||||
|
)
|
||||||
|
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||||
|
|
||||||
|
|
||||||
|
def _stream_part_to_event(part: StreamPart) -> Event:
|
||||||
|
"""Normalize an SSE ``StreamPart`` into a protocol event envelope."""
|
||||||
|
if isinstance(part.data, dict) and part.data.get("type") == "event":
|
||||||
|
event = cast(Event, part.data)
|
||||||
|
else:
|
||||||
|
event = {
|
||||||
|
"type": "event",
|
||||||
|
"method": part.event,
|
||||||
|
"params": {
|
||||||
|
"namespace": [],
|
||||||
|
"timestamp": int(time.time() * 1000),
|
||||||
|
"data": part.data,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if part.id is not None and "event_id" not in event:
|
||||||
|
event["event_id"] = part.id
|
||||||
|
return event
|
||||||
|
|
||||||
|
|
||||||
|
class EventSubscription:
|
||||||
|
"""Async iterable handle for a filtered event subscription."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
subscription_id: str,
|
||||||
|
params: SubscribeParams,
|
||||||
|
events: AsyncIterator[Event],
|
||||||
|
on_unsubscribe: Any,
|
||||||
|
) -> None:
|
||||||
|
self.subscription_id = subscription_id
|
||||||
|
self.params = params
|
||||||
|
self._events = events
|
||||||
|
self._on_unsubscribe = on_unsubscribe
|
||||||
|
|
||||||
|
def __aiter__(self) -> AsyncIterator[Event]:
|
||||||
|
return self._events
|
||||||
|
|
||||||
|
async def unsubscribe(self) -> None:
|
||||||
|
await self._on_unsubscribe(self.subscription_id)
|
||||||
|
|
||||||
|
|
||||||
|
class ProtocolSseTransport:
|
||||||
|
"""SSE transport for the thread-centric protocol."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
http: HttpClient,
|
||||||
|
thread_id: str,
|
||||||
|
*,
|
||||||
|
headers: Mapping[str, str] | None = None,
|
||||||
|
params: QueryParamTypes | None = None,
|
||||||
|
) -> None:
|
||||||
|
self.http = http
|
||||||
|
self.thread_id = thread_id
|
||||||
|
self.headers = headers
|
||||||
|
self.params = params
|
||||||
|
self.commands_path = f"/v2/threads/{thread_id}/commands"
|
||||||
|
self.events_path = f"/v2/threads/{thread_id}/events"
|
||||||
|
|
||||||
|
async def send(self, command: Command) -> CommandResponse | ErrorResponse | None:
|
||||||
|
return await self.http.post(
|
||||||
|
self.commands_path,
|
||||||
|
json=cast(dict[str, Any], command),
|
||||||
|
headers=self.headers,
|
||||||
|
params=self.params,
|
||||||
|
)
|
||||||
|
|
||||||
|
def open_event_stream(self, params: SubscribeParams) -> AsyncIterator[Event]:
|
||||||
|
async def iterate() -> AsyncIterator[Event]:
|
||||||
|
async for part in self.http.stream(
|
||||||
|
self.events_path,
|
||||||
|
"POST",
|
||||||
|
json=cast(dict[str, Any], params),
|
||||||
|
headers=self.headers,
|
||||||
|
params=self.params,
|
||||||
|
):
|
||||||
|
yield _stream_part_to_event(part)
|
||||||
|
|
||||||
|
return iterate()
|
||||||
|
|
||||||
|
|
||||||
|
class RunModule:
|
||||||
|
"""Run commands exposed by ``ThreadStream.run``."""
|
||||||
|
|
||||||
|
def __init__(self, stream: ThreadStream) -> None:
|
||||||
|
self._stream = stream
|
||||||
|
|
||||||
|
async def input(self, params: Mapping[str, Any]) -> Any:
|
||||||
|
return await self._stream.command("run.input", dict(params))
|
||||||
|
|
||||||
|
|
||||||
|
class InputModule:
|
||||||
|
"""Human-input commands exposed by ``ThreadStream.input``."""
|
||||||
|
|
||||||
|
def __init__(self, stream: ThreadStream) -> None:
|
||||||
|
self._stream = stream
|
||||||
|
|
||||||
|
async def respond(self, params: Mapping[str, Any]) -> Any:
|
||||||
|
return await self._stream.command("input.respond", dict(params))
|
||||||
|
|
||||||
|
async def inject(self, params: Mapping[str, Any]) -> Any:
|
||||||
|
return await self._stream.command("input.inject", dict(params))
|
||||||
|
|
||||||
|
|
||||||
|
class StateModule:
|
||||||
|
"""State commands exposed by ``ThreadStream.state``."""
|
||||||
|
|
||||||
|
def __init__(self, stream: ThreadStream) -> None:
|
||||||
|
self._stream = stream
|
||||||
|
|
||||||
|
async def get(self, params: Mapping[str, Any] | None = None) -> Any:
|
||||||
|
return await self._stream.command("state.get", dict(params or {}))
|
||||||
|
|
||||||
|
async def list_checkpoints(self, params: Mapping[str, Any] | None = None) -> Any:
|
||||||
|
return await self._stream.command("state.listCheckpoints", dict(params or {}))
|
||||||
|
|
||||||
|
async def fork(self, params: Mapping[str, Any]) -> Any:
|
||||||
|
return await self._stream.command("state.fork", dict(params))
|
||||||
|
|
||||||
|
|
||||||
|
class AgentModule:
|
||||||
|
"""Agent commands exposed by ``ThreadStream.agent``."""
|
||||||
|
|
||||||
|
def __init__(self, stream: ThreadStream) -> None:
|
||||||
|
self._stream = stream
|
||||||
|
|
||||||
|
async def get_tree(self, params: Mapping[str, Any] | None = None) -> Any:
|
||||||
|
return await self._stream.command("agent.getTree", dict(params or {}))
|
||||||
|
|
||||||
|
|
||||||
|
class ThreadStream:
|
||||||
|
"""High-level async wrapper around a thread protocol transport."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
transport: ProtocolSseTransport,
|
||||||
|
*,
|
||||||
|
assistant_id: str,
|
||||||
|
starting_command_id: int = 0,
|
||||||
|
) -> None:
|
||||||
|
if not assistant_id:
|
||||||
|
raise ValueError("assistant_id is required")
|
||||||
|
self.transport = transport
|
||||||
|
self.assistant_id = assistant_id
|
||||||
|
self._next_command_id = starting_command_id
|
||||||
|
self._next_subscription_id = 0
|
||||||
|
self.run = RunModule(self)
|
||||||
|
self.input = InputModule(self)
|
||||||
|
self.state = StateModule(self)
|
||||||
|
self.agent = AgentModule(self)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def thread_id(self) -> str:
|
||||||
|
return self.transport.thread_id
|
||||||
|
|
||||||
|
def _command_id(self) -> str:
|
||||||
|
self._next_command_id += 1
|
||||||
|
return str(self._next_command_id)
|
||||||
|
|
||||||
|
def _subscription_id(self) -> str:
|
||||||
|
self._next_subscription_id += 1
|
||||||
|
return f"sub-{self._next_subscription_id}"
|
||||||
|
|
||||||
|
async def command(self, method: CommandMethod, params: dict[str, Any]) -> Any:
|
||||||
|
command: Command = {
|
||||||
|
"id": self._command_id(),
|
||||||
|
"method": method,
|
||||||
|
"params": params,
|
||||||
|
}
|
||||||
|
if method == "run.input":
|
||||||
|
command["params"] = {"assistant_id": self.assistant_id, **params}
|
||||||
|
response = await self.transport.send(command)
|
||||||
|
if response is None:
|
||||||
|
return None
|
||||||
|
if "error" in response:
|
||||||
|
raise RuntimeError(response["error"])
|
||||||
|
return response.get("result")
|
||||||
|
|
||||||
|
async def subscribe(
|
||||||
|
self,
|
||||||
|
channels: Sequence[Channel | str] | SubscribeParams,
|
||||||
|
*,
|
||||||
|
namespaces: Sequence[Sequence[str]] | None = None,
|
||||||
|
depth: int | None = None,
|
||||||
|
) -> EventSubscription:
|
||||||
|
if isinstance(channels, dict):
|
||||||
|
params = SubscribeParams(**channels)
|
||||||
|
else:
|
||||||
|
params = SubscribeParams(channels=list(channels))
|
||||||
|
if namespaces is not None:
|
||||||
|
params["namespaces"] = [list(ns) for ns in namespaces]
|
||||||
|
if depth is not None:
|
||||||
|
params["depth"] = depth
|
||||||
|
|
||||||
|
subscription_id = self._subscription_id()
|
||||||
|
events = self.transport.open_event_stream(params)
|
||||||
|
return EventSubscription(subscription_id, params, events, self._unsubscribe)
|
||||||
|
|
||||||
|
async def _unsubscribe(self, subscription_id: str) -> None:
|
||||||
|
await self.command(
|
||||||
|
"subscription.unsubscribe",
|
||||||
|
{"subscription_id": subscription_id},
|
||||||
|
)
|
||||||
|
|
||||||
|
async def close(self) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
@@ -6,6 +6,7 @@ from collections.abc import AsyncIterator, Mapping, Sequence
|
|||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from langgraph_sdk._async.http import HttpClient
|
from langgraph_sdk._async.http import HttpClient
|
||||||
|
from langgraph_sdk._async.stream import ProtocolSseTransport, ThreadStream
|
||||||
from langgraph_sdk.schema import (
|
from langgraph_sdk.schema import (
|
||||||
Checkpoint,
|
Checkpoint,
|
||||||
Json,
|
Json,
|
||||||
@@ -42,6 +43,28 @@ class ThreadsClient:
|
|||||||
def __init__(self, http: HttpClient) -> None:
|
def __init__(self, http: HttpClient) -> None:
|
||||||
self.http = http
|
self.http = http
|
||||||
|
|
||||||
|
def stream(
|
||||||
|
self,
|
||||||
|
thread_id: str,
|
||||||
|
*,
|
||||||
|
assistant_id: str,
|
||||||
|
headers: Mapping[str, str] | None = None,
|
||||||
|
params: QueryParamTypes | None = None,
|
||||||
|
) -> ThreadStream:
|
||||||
|
"""Create a thread-centric protocol stream.
|
||||||
|
|
||||||
|
This mirrors the JavaScript SDK's ``client.threads.stream(...)`` API
|
||||||
|
and is intentionally separate from ``join_stream()``, which follows
|
||||||
|
the older thread event stream endpoint.
|
||||||
|
"""
|
||||||
|
transport = ProtocolSseTransport(
|
||||||
|
self.http,
|
||||||
|
thread_id,
|
||||||
|
headers=headers,
|
||||||
|
params=params,
|
||||||
|
)
|
||||||
|
return ThreadStream(transport, assistant_id=assistant_id)
|
||||||
|
|
||||||
async def get(
|
async def get(
|
||||||
self,
|
self,
|
||||||
thread_id: str,
|
thread_id: str,
|
||||||
|
|||||||
@@ -0,0 +1,225 @@
|
|||||||
|
"""Synchronous thread-centric streaming primitives."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import time
|
||||||
|
from collections.abc import Iterator, Mapping, Sequence
|
||||||
|
from typing import Any, cast
|
||||||
|
|
||||||
|
from langgraph_sdk._sync.http import SyncHttpClient
|
||||||
|
from langgraph_sdk.protocol import (
|
||||||
|
Channel,
|
||||||
|
Command,
|
||||||
|
CommandMethod,
|
||||||
|
CommandResponse,
|
||||||
|
ErrorResponse,
|
||||||
|
Event,
|
||||||
|
SubscribeParams,
|
||||||
|
)
|
||||||
|
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||||
|
|
||||||
|
|
||||||
|
def _stream_part_to_event(part: StreamPart) -> Event:
|
||||||
|
"""Normalize an SSE ``StreamPart`` into a protocol event envelope."""
|
||||||
|
if isinstance(part.data, dict) and part.data.get("type") == "event":
|
||||||
|
event = cast(Event, part.data)
|
||||||
|
else:
|
||||||
|
event = {
|
||||||
|
"type": "event",
|
||||||
|
"method": part.event,
|
||||||
|
"params": {
|
||||||
|
"namespace": [],
|
||||||
|
"timestamp": int(time.time() * 1000),
|
||||||
|
"data": part.data,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if part.id is not None and "event_id" not in event:
|
||||||
|
event["event_id"] = part.id
|
||||||
|
return event
|
||||||
|
|
||||||
|
|
||||||
|
class SyncEventSubscription:
|
||||||
|
"""Iterator handle for a filtered event subscription."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
subscription_id: str,
|
||||||
|
params: SubscribeParams,
|
||||||
|
events: Iterator[Event],
|
||||||
|
on_unsubscribe: Any,
|
||||||
|
) -> None:
|
||||||
|
self.subscription_id = subscription_id
|
||||||
|
self.params = params
|
||||||
|
self._events = events
|
||||||
|
self._on_unsubscribe = on_unsubscribe
|
||||||
|
|
||||||
|
def __iter__(self) -> Iterator[Event]:
|
||||||
|
return self._events
|
||||||
|
|
||||||
|
def unsubscribe(self) -> None:
|
||||||
|
self._on_unsubscribe(self.subscription_id)
|
||||||
|
|
||||||
|
|
||||||
|
class SyncProtocolSseTransport:
|
||||||
|
"""SSE transport for the thread-centric protocol."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
http: SyncHttpClient,
|
||||||
|
thread_id: str,
|
||||||
|
*,
|
||||||
|
headers: Mapping[str, str] | None = None,
|
||||||
|
params: QueryParamTypes | None = None,
|
||||||
|
) -> None:
|
||||||
|
self.http = http
|
||||||
|
self.thread_id = thread_id
|
||||||
|
self.headers = headers
|
||||||
|
self.params = params
|
||||||
|
self.commands_path = f"/v2/threads/{thread_id}/commands"
|
||||||
|
self.events_path = f"/v2/threads/{thread_id}/events"
|
||||||
|
|
||||||
|
def send(self, command: Command) -> CommandResponse | ErrorResponse | None:
|
||||||
|
return self.http.post(
|
||||||
|
self.commands_path,
|
||||||
|
json=cast(dict[str, Any], command),
|
||||||
|
headers=self.headers,
|
||||||
|
params=self.params,
|
||||||
|
)
|
||||||
|
|
||||||
|
def open_event_stream(self, params: SubscribeParams) -> Iterator[Event]:
|
||||||
|
for part in self.http.stream(
|
||||||
|
self.events_path,
|
||||||
|
"POST",
|
||||||
|
json=cast(dict[str, Any], params),
|
||||||
|
headers=self.headers,
|
||||||
|
params=self.params,
|
||||||
|
):
|
||||||
|
yield _stream_part_to_event(part)
|
||||||
|
|
||||||
|
|
||||||
|
class SyncRunModule:
|
||||||
|
"""Run commands exposed by ``SyncThreadStream.run``."""
|
||||||
|
|
||||||
|
def __init__(self, stream: SyncThreadStream) -> None:
|
||||||
|
self._stream = stream
|
||||||
|
|
||||||
|
def input(self, params: Mapping[str, Any]) -> Any:
|
||||||
|
return self._stream.command("run.input", dict(params))
|
||||||
|
|
||||||
|
|
||||||
|
class SyncInputModule:
|
||||||
|
"""Human-input commands exposed by ``SyncThreadStream.input``."""
|
||||||
|
|
||||||
|
def __init__(self, stream: SyncThreadStream) -> None:
|
||||||
|
self._stream = stream
|
||||||
|
|
||||||
|
def respond(self, params: Mapping[str, Any]) -> Any:
|
||||||
|
return self._stream.command("input.respond", dict(params))
|
||||||
|
|
||||||
|
def inject(self, params: Mapping[str, Any]) -> Any:
|
||||||
|
return self._stream.command("input.inject", dict(params))
|
||||||
|
|
||||||
|
|
||||||
|
class SyncStateModule:
|
||||||
|
"""State commands exposed by ``SyncThreadStream.state``."""
|
||||||
|
|
||||||
|
def __init__(self, stream: SyncThreadStream) -> None:
|
||||||
|
self._stream = stream
|
||||||
|
|
||||||
|
def get(self, params: Mapping[str, Any] | None = None) -> Any:
|
||||||
|
return self._stream.command("state.get", dict(params or {}))
|
||||||
|
|
||||||
|
def list_checkpoints(self, params: Mapping[str, Any] | None = None) -> Any:
|
||||||
|
return self._stream.command("state.listCheckpoints", dict(params or {}))
|
||||||
|
|
||||||
|
def fork(self, params: Mapping[str, Any]) -> Any:
|
||||||
|
return self._stream.command("state.fork", dict(params))
|
||||||
|
|
||||||
|
|
||||||
|
class SyncAgentModule:
|
||||||
|
"""Agent commands exposed by ``SyncThreadStream.agent``."""
|
||||||
|
|
||||||
|
def __init__(self, stream: SyncThreadStream) -> None:
|
||||||
|
self._stream = stream
|
||||||
|
|
||||||
|
def get_tree(self, params: Mapping[str, Any] | None = None) -> Any:
|
||||||
|
return self._stream.command("agent.getTree", dict(params or {}))
|
||||||
|
|
||||||
|
|
||||||
|
class SyncThreadStream:
|
||||||
|
"""High-level sync wrapper around a thread protocol transport."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
transport: SyncProtocolSseTransport,
|
||||||
|
*,
|
||||||
|
assistant_id: str,
|
||||||
|
starting_command_id: int = 0,
|
||||||
|
) -> None:
|
||||||
|
if not assistant_id:
|
||||||
|
raise ValueError("assistant_id is required")
|
||||||
|
self.transport = transport
|
||||||
|
self.assistant_id = assistant_id
|
||||||
|
self._next_command_id = starting_command_id
|
||||||
|
self._next_subscription_id = 0
|
||||||
|
self.run = SyncRunModule(self)
|
||||||
|
self.input = SyncInputModule(self)
|
||||||
|
self.state = SyncStateModule(self)
|
||||||
|
self.agent = SyncAgentModule(self)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def thread_id(self) -> str:
|
||||||
|
return self.transport.thread_id
|
||||||
|
|
||||||
|
def _command_id(self) -> str:
|
||||||
|
self._next_command_id += 1
|
||||||
|
return str(self._next_command_id)
|
||||||
|
|
||||||
|
def _subscription_id(self) -> str:
|
||||||
|
self._next_subscription_id += 1
|
||||||
|
return f"sub-{self._next_subscription_id}"
|
||||||
|
|
||||||
|
def command(self, method: CommandMethod, params: dict[str, Any]) -> Any:
|
||||||
|
command: Command = {
|
||||||
|
"id": self._command_id(),
|
||||||
|
"method": method,
|
||||||
|
"params": params,
|
||||||
|
}
|
||||||
|
if method == "run.input":
|
||||||
|
command["params"] = {"assistant_id": self.assistant_id, **params}
|
||||||
|
response = self.transport.send(command)
|
||||||
|
if response is None:
|
||||||
|
return None
|
||||||
|
if "error" in response:
|
||||||
|
raise RuntimeError(response["error"])
|
||||||
|
return response.get("result")
|
||||||
|
|
||||||
|
def subscribe(
|
||||||
|
self,
|
||||||
|
channels: Sequence[Channel | str] | SubscribeParams,
|
||||||
|
*,
|
||||||
|
namespaces: Sequence[Sequence[str]] | None = None,
|
||||||
|
depth: int | None = None,
|
||||||
|
) -> SyncEventSubscription:
|
||||||
|
if isinstance(channels, dict):
|
||||||
|
params = SubscribeParams(**channels)
|
||||||
|
else:
|
||||||
|
params = SubscribeParams(channels=list(channels))
|
||||||
|
if namespaces is not None:
|
||||||
|
params["namespaces"] = [list(ns) for ns in namespaces]
|
||||||
|
if depth is not None:
|
||||||
|
params["depth"] = depth
|
||||||
|
|
||||||
|
subscription_id = self._subscription_id()
|
||||||
|
events = self.transport.open_event_stream(params)
|
||||||
|
return SyncEventSubscription(subscription_id, params, events, self._unsubscribe)
|
||||||
|
|
||||||
|
def _unsubscribe(self, subscription_id: str) -> None:
|
||||||
|
self.command(
|
||||||
|
"subscription.unsubscribe",
|
||||||
|
{"subscription_id": subscription_id},
|
||||||
|
)
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
@@ -6,6 +6,7 @@ from collections.abc import Iterator, Mapping, Sequence
|
|||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from langgraph_sdk._sync.http import SyncHttpClient
|
from langgraph_sdk._sync.http import SyncHttpClient
|
||||||
|
from langgraph_sdk._sync.stream import SyncProtocolSseTransport, SyncThreadStream
|
||||||
from langgraph_sdk.schema import (
|
from langgraph_sdk.schema import (
|
||||||
Checkpoint,
|
Checkpoint,
|
||||||
Json,
|
Json,
|
||||||
@@ -41,6 +42,28 @@ class SyncThreadsClient:
|
|||||||
def __init__(self, http: SyncHttpClient) -> None:
|
def __init__(self, http: SyncHttpClient) -> None:
|
||||||
self.http = http
|
self.http = http
|
||||||
|
|
||||||
|
def stream(
|
||||||
|
self,
|
||||||
|
thread_id: str,
|
||||||
|
*,
|
||||||
|
assistant_id: str,
|
||||||
|
headers: Mapping[str, str] | None = None,
|
||||||
|
params: QueryParamTypes | None = None,
|
||||||
|
) -> SyncThreadStream:
|
||||||
|
"""Create a thread-centric protocol stream.
|
||||||
|
|
||||||
|
This mirrors the JavaScript SDK's ``client.threads.stream(...)`` API
|
||||||
|
and is intentionally separate from ``join_stream()``, which follows
|
||||||
|
the older thread event stream endpoint.
|
||||||
|
"""
|
||||||
|
transport = SyncProtocolSseTransport(
|
||||||
|
self.http,
|
||||||
|
thread_id,
|
||||||
|
headers=headers,
|
||||||
|
params=params,
|
||||||
|
)
|
||||||
|
return SyncThreadStream(transport, assistant_id=assistant_id)
|
||||||
|
|
||||||
def get(
|
def get(
|
||||||
self,
|
self,
|
||||||
thread_id: str,
|
thread_id: str,
|
||||||
|
|||||||
@@ -17,6 +17,11 @@ from langgraph_sdk._async.client import LangGraphClient, get_client
|
|||||||
from langgraph_sdk._async.cron import CronClient
|
from langgraph_sdk._async.cron import CronClient
|
||||||
from langgraph_sdk._async.http import HttpClient, _adecode_json, _aencode_json
|
from langgraph_sdk._async.http import HttpClient, _adecode_json, _aencode_json
|
||||||
from langgraph_sdk._async.runs import RunsClient
|
from langgraph_sdk._async.runs import RunsClient
|
||||||
|
from langgraph_sdk._async.stream import (
|
||||||
|
EventSubscription,
|
||||||
|
ProtocolSseTransport,
|
||||||
|
ThreadStream,
|
||||||
|
)
|
||||||
from langgraph_sdk._async.store import StoreClient
|
from langgraph_sdk._async.store import StoreClient
|
||||||
from langgraph_sdk._async.threads import ThreadsClient
|
from langgraph_sdk._async.threads import ThreadsClient
|
||||||
from langgraph_sdk._shared.utilities import configure_loopback_transports
|
from langgraph_sdk._shared.utilities import configure_loopback_transports
|
||||||
@@ -27,6 +32,11 @@ from langgraph_sdk._sync.client import SyncLangGraphClient, get_sync_client
|
|||||||
from langgraph_sdk._sync.cron import SyncCronClient
|
from langgraph_sdk._sync.cron import SyncCronClient
|
||||||
from langgraph_sdk._sync.http import SyncHttpClient, _decode_json, _encode_json
|
from langgraph_sdk._sync.http import SyncHttpClient, _decode_json, _encode_json
|
||||||
from langgraph_sdk._sync.runs import SyncRunsClient
|
from langgraph_sdk._sync.runs import SyncRunsClient
|
||||||
|
from langgraph_sdk._sync.stream import (
|
||||||
|
SyncEventSubscription,
|
||||||
|
SyncProtocolSseTransport,
|
||||||
|
SyncThreadStream,
|
||||||
|
)
|
||||||
from langgraph_sdk._sync.store import SyncStoreClient
|
from langgraph_sdk._sync.store import SyncStoreClient
|
||||||
from langgraph_sdk._sync.threads import SyncThreadsClient
|
from langgraph_sdk._sync.threads import SyncThreadsClient
|
||||||
|
|
||||||
@@ -35,15 +45,21 @@ __all__ = [
|
|||||||
"CronClient",
|
"CronClient",
|
||||||
"HttpClient",
|
"HttpClient",
|
||||||
"LangGraphClient",
|
"LangGraphClient",
|
||||||
|
"EventSubscription",
|
||||||
|
"ProtocolSseTransport",
|
||||||
"RunsClient",
|
"RunsClient",
|
||||||
"StoreClient",
|
"StoreClient",
|
||||||
"SyncAssistantsClient",
|
"SyncAssistantsClient",
|
||||||
"SyncCronClient",
|
"SyncCronClient",
|
||||||
"SyncHttpClient",
|
"SyncHttpClient",
|
||||||
"SyncLangGraphClient",
|
"SyncLangGraphClient",
|
||||||
|
"SyncEventSubscription",
|
||||||
|
"SyncProtocolSseTransport",
|
||||||
"SyncRunsClient",
|
"SyncRunsClient",
|
||||||
"SyncStoreClient",
|
"SyncStoreClient",
|
||||||
|
"SyncThreadStream",
|
||||||
"SyncThreadsClient",
|
"SyncThreadsClient",
|
||||||
|
"ThreadStream",
|
||||||
"ThreadsClient",
|
"ThreadsClient",
|
||||||
"_adecode_json",
|
"_adecode_json",
|
||||||
"_aencode_json",
|
"_aencode_json",
|
||||||
|
|||||||
@@ -0,0 +1,105 @@
|
|||||||
|
"""Typed protocol messages for thread-centric remote streaming.
|
||||||
|
|
||||||
|
These shapes mirror the JSON protocol used by the JavaScript SDK's
|
||||||
|
``ThreadStream`` layer. The SDK keeps them lightweight and dependency-free so
|
||||||
|
``langgraph`` can build higher-level projections on top without creating a
|
||||||
|
package cycle.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Any, Literal
|
||||||
|
|
||||||
|
from typing_extensions import NotRequired, TypedDict
|
||||||
|
|
||||||
|
Channel = Literal[
|
||||||
|
"values",
|
||||||
|
"updates",
|
||||||
|
"messages",
|
||||||
|
"tools",
|
||||||
|
"custom",
|
||||||
|
"lifecycle",
|
||||||
|
"input",
|
||||||
|
"debug",
|
||||||
|
"checkpoints",
|
||||||
|
"tasks",
|
||||||
|
]
|
||||||
|
"""Built-in subscribable protocol channels."""
|
||||||
|
|
||||||
|
CommandMethod = Literal[
|
||||||
|
"run.input",
|
||||||
|
"subscription.subscribe",
|
||||||
|
"subscription.unsubscribe",
|
||||||
|
"agent.getTree",
|
||||||
|
"input.respond",
|
||||||
|
"input.inject",
|
||||||
|
"state.get",
|
||||||
|
"state.listCheckpoints",
|
||||||
|
"state.fork",
|
||||||
|
]
|
||||||
|
"""Command methods understood by the thread stream protocol."""
|
||||||
|
|
||||||
|
|
||||||
|
class SubscribeParams(TypedDict, total=False):
|
||||||
|
"""Filter used when subscribing to thread protocol events."""
|
||||||
|
|
||||||
|
channels: list[Channel | str]
|
||||||
|
namespaces: NotRequired[list[list[str]]]
|
||||||
|
depth: NotRequired[int]
|
||||||
|
|
||||||
|
|
||||||
|
class Command(TypedDict):
|
||||||
|
"""Command sent to a thread protocol transport."""
|
||||||
|
|
||||||
|
id: str
|
||||||
|
method: CommandMethod
|
||||||
|
params: dict[str, Any]
|
||||||
|
|
||||||
|
|
||||||
|
class CommandResponse(TypedDict):
|
||||||
|
"""Successful command response."""
|
||||||
|
|
||||||
|
id: str
|
||||||
|
result: Any
|
||||||
|
|
||||||
|
|
||||||
|
class ErrorResponse(TypedDict):
|
||||||
|
"""Error command response."""
|
||||||
|
|
||||||
|
id: str
|
||||||
|
error: dict[str, Any]
|
||||||
|
|
||||||
|
|
||||||
|
class EventParams(TypedDict):
|
||||||
|
"""Protocol event parameters."""
|
||||||
|
|
||||||
|
namespace: list[str]
|
||||||
|
timestamp: int
|
||||||
|
data: Any
|
||||||
|
node: NotRequired[str]
|
||||||
|
run_id: NotRequired[str]
|
||||||
|
interrupts: NotRequired[list[Any]]
|
||||||
|
|
||||||
|
|
||||||
|
class Event(TypedDict):
|
||||||
|
"""Protocol event envelope."""
|
||||||
|
|
||||||
|
type: Literal["event"]
|
||||||
|
method: str
|
||||||
|
params: EventParams
|
||||||
|
event_id: NotRequired[str]
|
||||||
|
seq: NotRequired[int]
|
||||||
|
|
||||||
|
|
||||||
|
class RunInputResult(TypedDict, total=False):
|
||||||
|
"""Result returned by ``run.input`` commands."""
|
||||||
|
|
||||||
|
run_id: str
|
||||||
|
thread_id: str
|
||||||
|
|
||||||
|
|
||||||
|
class SubscribeResult(TypedDict):
|
||||||
|
"""Result returned by ``subscription.subscribe`` commands."""
|
||||||
|
|
||||||
|
subscription_id: str
|
||||||
|
|
||||||
@@ -41,6 +41,7 @@ def _normalize_return_annotation(ann: object) -> str:
|
|||||||
s = str(ann)
|
s = str(ann)
|
||||||
s = re.sub(r"\s+", "", s)
|
s = re.sub(r"\s+", "", s)
|
||||||
s = s.replace("typing.", "").replace("collections.abc.", "")
|
s = s.replace("typing.", "").replace("collections.abc.", "")
|
||||||
|
s = s.replace("SyncThreadStream", "ThreadStream")
|
||||||
s = re.sub(r"AsyncGenerator\[([^,\]]+)(?:,[^\]]*)?\]", r"Iterator[\1]", s)
|
s = re.sub(r"AsyncGenerator\[([^,\]]+)(?:,[^\]]*)?\]", r"Iterator[\1]", s)
|
||||||
s = re.sub(r"Generator\[([^,\]]+)(?:,[^\]]*)?\]", r"Iterator[\1]", s)
|
s = re.sub(r"Generator\[([^,\]]+)(?:,[^\]]*)?\]", r"Iterator[\1]", s)
|
||||||
s = re.sub(r"AsyncIterator\[(.+)\]", r"Iterator[\1]", s)
|
s = re.sub(r"AsyncIterator\[(.+)\]", r"Iterator[\1]", s)
|
||||||
|
|||||||
@@ -10,17 +10,23 @@ from langgraph_sdk import get_sync_client as public_get_sync_client
|
|||||||
from langgraph_sdk.client import (
|
from langgraph_sdk.client import (
|
||||||
AssistantsClient,
|
AssistantsClient,
|
||||||
CronClient,
|
CronClient,
|
||||||
|
EventSubscription,
|
||||||
HttpClient,
|
HttpClient,
|
||||||
LangGraphClient,
|
LangGraphClient,
|
||||||
RunsClient,
|
RunsClient,
|
||||||
StoreClient,
|
StoreClient,
|
||||||
|
ProtocolSseTransport,
|
||||||
SyncAssistantsClient,
|
SyncAssistantsClient,
|
||||||
SyncCronClient,
|
SyncCronClient,
|
||||||
|
SyncEventSubscription,
|
||||||
SyncHttpClient,
|
SyncHttpClient,
|
||||||
SyncLangGraphClient,
|
SyncLangGraphClient,
|
||||||
|
SyncProtocolSseTransport,
|
||||||
SyncRunsClient,
|
SyncRunsClient,
|
||||||
SyncStoreClient,
|
SyncStoreClient,
|
||||||
|
SyncThreadStream,
|
||||||
SyncThreadsClient,
|
SyncThreadsClient,
|
||||||
|
ThreadStream,
|
||||||
ThreadsClient,
|
ThreadsClient,
|
||||||
_adecode_json,
|
_adecode_json,
|
||||||
_aencode_json,
|
_aencode_json,
|
||||||
@@ -52,6 +58,9 @@ def test_client_exports():
|
|||||||
assert RunsClient is not None
|
assert RunsClient is not None
|
||||||
assert CronClient is not None
|
assert CronClient is not None
|
||||||
assert StoreClient is not None
|
assert StoreClient is not None
|
||||||
|
assert EventSubscription is not None
|
||||||
|
assert ProtocolSseTransport is not None
|
||||||
|
assert ThreadStream is not None
|
||||||
|
|
||||||
# Resource client classes - Sync
|
# Resource client classes - Sync
|
||||||
assert SyncAssistantsClient is not None
|
assert SyncAssistantsClient is not None
|
||||||
@@ -59,6 +68,9 @@ def test_client_exports():
|
|||||||
assert SyncRunsClient is not None
|
assert SyncRunsClient is not None
|
||||||
assert SyncCronClient is not None
|
assert SyncCronClient is not None
|
||||||
assert SyncStoreClient is not None
|
assert SyncStoreClient is not None
|
||||||
|
assert SyncEventSubscription is not None
|
||||||
|
assert SyncProtocolSseTransport is not None
|
||||||
|
assert SyncThreadStream is not None
|
||||||
|
|
||||||
# Internal utilities (used by tests)
|
# Internal utilities (used by tests)
|
||||||
assert callable(_aencode_json)
|
assert callable(_aencode_json)
|
||||||
|
|||||||
@@ -9,7 +9,14 @@ import pytest
|
|||||||
from typing_extensions import assert_type
|
from typing_extensions import assert_type
|
||||||
|
|
||||||
from langgraph_sdk._shared.utilities import _sse_to_v2_dict
|
from langgraph_sdk._shared.utilities import _sse_to_v2_dict
|
||||||
from langgraph_sdk.client import HttpClient, SyncHttpClient
|
from langgraph_sdk.client import (
|
||||||
|
HttpClient,
|
||||||
|
SyncHttpClient,
|
||||||
|
SyncThreadStream,
|
||||||
|
SyncThreadsClient,
|
||||||
|
ThreadStream,
|
||||||
|
ThreadsClient,
|
||||||
|
)
|
||||||
from langgraph_sdk.schema import (
|
from langgraph_sdk.schema import (
|
||||||
CheckpointPayload,
|
CheckpointPayload,
|
||||||
CheckpointsStreamPart,
|
CheckpointsStreamPart,
|
||||||
@@ -154,6 +161,107 @@ def test_sync_http_client_stream_flushes_trailing_event():
|
|||||||
assert parts == [StreamPart(event="foo", data={"bar": 1})]
|
assert parts == [StreamPart(event="foo", data={"bar": 1})]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_threads_stream_sends_commands_and_subscribes() -> None:
|
||||||
|
requests: list[tuple[str, str]] = []
|
||||||
|
|
||||||
|
async def handler(request: httpx.Request) -> httpx.Response:
|
||||||
|
body = await request.aread()
|
||||||
|
requests.append((request.method, request.url.path))
|
||||||
|
if request.url.path == "/v2/threads/thread-1/commands":
|
||||||
|
assert request.method == "POST"
|
||||||
|
payload = httpx.Response(200, content=body).json()
|
||||||
|
if payload["method"] == "run.input":
|
||||||
|
assert payload["params"]["assistant_id"] == "agent"
|
||||||
|
return httpx.Response(200, json={"id": payload["id"], "result": "ok"})
|
||||||
|
if payload["method"] == "subscription.unsubscribe":
|
||||||
|
return httpx.Response(200, json={"id": payload["id"], "result": {}})
|
||||||
|
if request.url.path == "/v2/threads/thread-1/events":
|
||||||
|
assert request.method == "POST"
|
||||||
|
assert httpx.Response(200, content=body).json() == {
|
||||||
|
"channels": ["messages"]
|
||||||
|
}
|
||||||
|
return httpx.Response(
|
||||||
|
200,
|
||||||
|
headers={"Content-Type": "text/event-stream"},
|
||||||
|
content=(
|
||||||
|
b"id: evt-1\n"
|
||||||
|
b"event: messages\n"
|
||||||
|
b'data: {"type":"event","method":"messages","params":{"namespace":[],"timestamp":1,"data":{"event":"message-start","id":"m1"}}}\n\n'
|
||||||
|
),
|
||||||
|
)
|
||||||
|
raise AssertionError(f"unexpected request: {request.method} {request.url.path}")
|
||||||
|
|
||||||
|
transport = httpx.MockTransport(handler)
|
||||||
|
async with httpx.AsyncClient(
|
||||||
|
transport=transport, base_url="https://example.com"
|
||||||
|
) as client:
|
||||||
|
threads = ThreadsClient(HttpClient(client))
|
||||||
|
thread = threads.stream("thread-1", assistant_id="agent")
|
||||||
|
assert isinstance(thread, ThreadStream)
|
||||||
|
assert await thread.run.input({"input": {"messages": []}}) == "ok"
|
||||||
|
subscription = await thread.subscribe(["messages"])
|
||||||
|
events = [event async for event in subscription]
|
||||||
|
await subscription.unsubscribe()
|
||||||
|
|
||||||
|
assert events[0]["method"] == "messages"
|
||||||
|
assert events[0]["event_id"] == "evt-1"
|
||||||
|
assert requests == [
|
||||||
|
("POST", "/v2/threads/thread-1/commands"),
|
||||||
|
("POST", "/v2/threads/thread-1/events"),
|
||||||
|
("POST", "/v2/threads/thread-1/commands"),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def test_sync_threads_stream_sends_commands_and_subscribes() -> None:
|
||||||
|
requests: list[tuple[str, str]] = []
|
||||||
|
|
||||||
|
def handler(request: httpx.Request) -> httpx.Response:
|
||||||
|
body = request.read()
|
||||||
|
requests.append((request.method, request.url.path))
|
||||||
|
if request.url.path == "/v2/threads/thread-1/commands":
|
||||||
|
assert request.method == "POST"
|
||||||
|
payload = httpx.Response(200, content=body).json()
|
||||||
|
if payload["method"] == "run.input":
|
||||||
|
assert payload["params"]["assistant_id"] == "agent"
|
||||||
|
return httpx.Response(200, json={"id": payload["id"], "result": "ok"})
|
||||||
|
if payload["method"] == "subscription.unsubscribe":
|
||||||
|
return httpx.Response(200, json={"id": payload["id"], "result": {}})
|
||||||
|
if request.url.path == "/v2/threads/thread-1/events":
|
||||||
|
assert request.method == "POST"
|
||||||
|
assert httpx.Response(200, content=body).json() == {
|
||||||
|
"channels": ["messages"]
|
||||||
|
}
|
||||||
|
return httpx.Response(
|
||||||
|
200,
|
||||||
|
headers={"Content-Type": "text/event-stream"},
|
||||||
|
content=(
|
||||||
|
b"id: evt-1\n"
|
||||||
|
b"event: messages\n"
|
||||||
|
b'data: {"type":"event","method":"messages","params":{"namespace":[],"timestamp":1,"data":{"event":"message-start","id":"m1"}}}\n\n'
|
||||||
|
),
|
||||||
|
)
|
||||||
|
raise AssertionError(f"unexpected request: {request.method} {request.url.path}")
|
||||||
|
|
||||||
|
transport = httpx.MockTransport(handler)
|
||||||
|
with httpx.Client(transport=transport, base_url="https://example.com") as client:
|
||||||
|
threads = SyncThreadsClient(SyncHttpClient(client))
|
||||||
|
thread = threads.stream("thread-1", assistant_id="agent")
|
||||||
|
assert isinstance(thread, SyncThreadStream)
|
||||||
|
assert thread.run.input({"input": {"messages": []}}) == "ok"
|
||||||
|
subscription = thread.subscribe(["messages"])
|
||||||
|
events = list(subscription)
|
||||||
|
subscription.unsubscribe()
|
||||||
|
|
||||||
|
assert events[0]["method"] == "messages"
|
||||||
|
assert events[0]["event_id"] == "evt-1"
|
||||||
|
assert requests == [
|
||||||
|
("POST", "/v2/threads/thread-1/commands"),
|
||||||
|
("POST", "/v2/threads/thread-1/events"),
|
||||||
|
("POST", "/v2/threads/thread-1/commands"),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
def test_sync_http_client_stream_recovers_after_disconnect():
|
def test_sync_http_client_stream_recovers_after_disconnect():
|
||||||
reconnect_path = "/reconnect"
|
reconnect_path = "/reconnect"
|
||||||
first_chunks = [
|
first_chunks = [
|
||||||
|
|||||||
Reference in New Issue
Block a user