From 46fed9d16156ebdad49806bb65579b7d0756ef3f Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Mon, 9 Mar 2026 13:52:36 -0400 Subject: [PATCH] feat: type safe stream/invoke w/ proper output type coercion (#6961) Adding type safe streaming + more robust pydantic + dataclass support on graph outputs * type safe streaming via #6931 * type safe invoke via https://github.com/langchain-ai/langgraph/pull/6963 * actually thread through types via parametrization in https://github.com/langchain-ai/langgraph/pull/7009 * deprecation of backwards compatible accessor via https://github.com/langchain-ai/langgraph/pull/7011 full spec of changes: https://github.com/langchain-ai/langgraph/issues/7008 proving out that required changes are minimal even when we update the default to v2: https://github.com/langchain-ai/langchain/pull/35541 --- libs/langgraph/langgraph/graph/state.py | 29 +- libs/langgraph/langgraph/pregel/debug.py | 45 +- libs/langgraph/langgraph/pregel/main.py | 450 ++++++- libs/langgraph/langgraph/pregel/protocol.py | 128 +- libs/langgraph/langgraph/pregel/remote.py | 171 ++- libs/langgraph/langgraph/types.py | 289 ++++- libs/langgraph/langgraph/warnings.py | 8 + libs/langgraph/tests/test_deprecation.py | 40 +- libs/langgraph/tests/test_stream_v2.py | 1152 +++++++++++++++++ libs/sdk-py/langgraph_sdk/_async/runs.py | 94 +- .../sdk-py/langgraph_sdk/_shared/utilities.py | 18 + libs/sdk-py/langgraph_sdk/_sync/runs.py | 95 +- libs/sdk-py/langgraph_sdk/schema.py | 269 ++++ libs/sdk-py/tests/test_client_stream.py | 249 +++- 14 files changed, 2892 insertions(+), 145 deletions(-) create mode 100644 libs/langgraph/tests/test_stream_v2.py diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index a87c17527..b1c24de2b 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -6,6 +6,7 @@ import typing import warnings from collections import defaultdict from collections.abc import Awaitable, Callable, Hashable, Sequence +from dataclasses import is_dataclass from functools import partial from inspect import isclass, isfunction, ismethod, signature from types import FunctionType @@ -14,6 +15,7 @@ from typing import ( Any, Generic, Literal, + TypeVar, Union, cast, get_args, @@ -1164,6 +1166,20 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): for key, node in self.nodes.items(): compiled.attach_node(key, node) + # Record output/state mappers for v2 stream coercion (pydantic/dataclass only) + compiled._output_mapper = _pick_mapper( + list(output_channels) + if isinstance(output_channels, list) + else [output_channels], + self.output_schema, + ) + compiled._state_mapper = _pick_mapper( + list(stream_channels) + if isinstance(stream_channels, list) + else [stream_channels], + self.state_schema, + ) + for start, end in self.edges: compiled.attach_edge(start, end) @@ -1183,6 +1199,8 @@ class CompiledStateGraph( ): builder: StateGraph[StateT, ContextT, InputT, OutputT] schema_to_mapper: dict[type[Any], Callable[[Any], Any] | None] + _output_mapper: Callable[[Any], Any] | None + _state_mapper: Callable[[Any], Any] | None def __init__( self, @@ -1504,12 +1522,15 @@ def _pick_mapper( ) -> Callable[[Any], Any] | None: if state_keys == ["__root__"]: return None - if isclass(schema) and issubclass(schema, dict): - return None - return partial(_coerce_state, schema) + if isclass(schema) and (issubclass(schema, BaseModel) or is_dataclass(schema)): + return partial(_coerce_state, schema) + return None -def _coerce_state(schema: type[Any], input: dict[str, Any]) -> dict[str, Any]: +_S = TypeVar("_S") + + +def _coerce_state(schema: type[_S], input: dict[str, Any]) -> _S: return schema(**input) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 0de53e120..519a36488 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -7,7 +7,6 @@ from uuid import UUID from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import CheckpointMetadata, PendingWrite -from typing_extensions import TypedDict from langgraph._internal._config import patch_checkpoint_map from langgraph._internal._constants import ( @@ -23,42 +22,14 @@ from langgraph._internal._typing import MISSING from langgraph.channels.base import BaseChannel from langgraph.constants import TAG_HIDDEN from langgraph.pregel._io import read_channels -from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot - -__all__ = ("TaskPayload", "TaskResultPayload", "CheckpointTask", "CheckpointPayload") - - -class TaskPayload(TypedDict): - id: str - name: str - input: Any - triggers: list[str] - - -class TaskResultPayload(TypedDict): - id: str - name: str - error: str | None - interrupts: list[dict] - result: dict[str, Any] - - -class CheckpointTask(TypedDict): - id: str - name: str - error: str | None - interrupts: list[dict] - state: StateSnapshot | RunnableConfig | None - - -class CheckpointPayload(TypedDict): - config: RunnableConfig | None - metadata: CheckpointMetadata - values: dict[str, Any] - next: list[str] - parent_config: RunnableConfig | None - tasks: list[CheckpointTask] - +from langgraph.types import ( + CheckpointPayload, + PregelExecutableTask, + PregelTask, + StateSnapshot, + TaskPayload, + TaskResultPayload, +) TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8") diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index ae6a0014d..d4190a0e6 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -22,8 +22,10 @@ from inspect import isclass from typing import ( Any, Generic, + Literal, cast, get_type_hints, + overload, ) from uuid import UUID, uuid5 @@ -121,7 +123,10 @@ from langgraph.pregel._checkpoint import ( ) from langgraph.pregel._draw import draw_graph from langgraph.pregel._io import map_input, read_channels -from langgraph.pregel._loop import AsyncPregelLoop, SyncPregelLoop +from langgraph.pregel._loop import ( + AsyncPregelLoop, + SyncPregelLoop, +) from langgraph.pregel._messages import StreamMessagesHandler from langgraph.pregel._read import DEFAULT_BOUND, PregelNode from langgraph.pregel._retry import RetryPolicy @@ -138,11 +143,13 @@ from langgraph.types import ( Checkpointer, Command, Durability, + GraphOutput, Interrupt, Send, StateSnapshot, StateUpdate, StreamMode, + StreamPart, ensure_valid_checkpointer, ) from langgraph.typing import ContextT, InputT, OutputT, StateT @@ -993,6 +1000,11 @@ class Pregel( for name, node in self.get_subgraphs(namespace=namespace, recurse=recurse): yield name, node + # Mappers for v2 stream coercion (pydantic/dataclass). + # Set by CompiledStateGraph; None for base Pregel. + _output_mapper: Callable[[Any], Any] | None = None + _state_mapper: Callable[[Any], Any] | None = None + def _migrate_checkpoint(self, checkpoint: Checkpoint) -> None: """Migrate a saved checkpoint to new channel layout.""" if checkpoint["v"] < 4 and checkpoint.get("pending_sends"): @@ -2427,6 +2439,7 @@ class Pregel( durability, ) + @overload def stream( self, input: InputT | Command | None, @@ -2441,6 +2454,44 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, + version: Literal["v2"], + **kwargs: Unpack[DeprecatedKwargs], + ) -> Iterator[StreamPart[OutputT, StateT]]: ... + + @overload + def stream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode | Sequence[StreamMode] | None = None, + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + subgraphs: bool = False, + debug: bool | None = None, + version: Literal["v1"] = ..., + **kwargs: Unpack[DeprecatedKwargs], + ) -> Iterator[dict[str, Any] | Any]: ... + + def stream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode | Sequence[StreamMode] | None = None, + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + subgraphs: bool = False, + debug: bool | None = None, + version: Literal["v1", "v2"] = "v1", **kwargs: Unpack[DeprecatedKwargs], ) -> Iterator[dict[str, Any] | Any]: """Stream graph steps for a single input. @@ -2602,6 +2653,10 @@ class Pregel( runtime = parent_runtime.merge(runtime) config[CONF][CONFIG_KEY_RUNTIME] = runtime + # resolve mappers for v2 stream coercion + _output_mapper = self._output_mapper if version == "v2" else None + _state_mapper = self._state_mapper if version == "v2" else None + with SyncPregelLoop( input, stream=StreamProtocol(stream.put, stream_modes), @@ -2674,7 +2729,14 @@ class Pregel( ): # emit output yield from _output( - stream_mode, print_mode, subgraphs, stream.get, queue.Empty + stream_mode, + print_mode, + subgraphs, + stream.get, + queue.Empty, + version, + _output_mapper, + _state_mapper, ) loop.after_tick() # wait for checkpoint @@ -2682,7 +2744,14 @@ class Pregel( loop._put_checkpoint_fut.result() # emit output yield from _output( - stream_mode, print_mode, subgraphs, stream.get, queue.Empty + stream_mode, + print_mode, + subgraphs, + stream.get, + queue.Empty, + version, + _output_mapper, + _state_mapper, ) # handle exit if loop.status == "out_of_steps": @@ -2701,6 +2770,44 @@ class Pregel( run_manager.on_chain_error(e) raise + @overload + def astream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode | Sequence[StreamMode] | None = None, + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + subgraphs: bool = False, + debug: bool | None = None, + version: Literal["v2"], + **kwargs: Unpack[DeprecatedKwargs], + ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ... + + @overload + def astream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode | Sequence[StreamMode] | None = None, + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + subgraphs: bool = False, + debug: bool | None = None, + version: Literal["v1"] = ..., + **kwargs: Unpack[DeprecatedKwargs], + ) -> AsyncIterator[dict[str, Any] | Any]: ... + async def astream( self, input: InputT | Command | None, @@ -2715,6 +2822,7 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, + version: Literal["v1", "v2"] = "v1", **kwargs: Unpack[DeprecatedKwargs], ) -> AsyncIterator[dict[str, Any] | Any]: """Asynchronously stream graph steps for a single input. @@ -2911,6 +3019,10 @@ class Pregel( runtime = parent_runtime.merge(runtime) config[CONF][CONFIG_KEY_RUNTIME] = runtime + # resolve mappers for v2 stream coercion + _output_mapper = self._output_mapper if version == "v2" else None + _state_mapper = self._state_mapper if version == "v2" else None + async with AsyncPregelLoop( input, stream=StreamProtocol(stream.put_nowait, stream_modes), @@ -3007,6 +3119,9 @@ class Pregel( subgraphs, stream.get_nowait, asyncio.QueueEmpty, + version, + _output_mapper, + _state_mapper, ): yield o loop.after_tick() @@ -3025,6 +3140,9 @@ class Pregel( subgraphs, stream.get_nowait, asyncio.QueueEmpty, + version, + _output_mapper, + _state_mapper, ): yield o # handle exit @@ -3044,6 +3162,41 @@ class Pregel( await asyncio.shield(run_manager.on_chain_error(e)) raise + @overload + def invoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: Literal["values"] = ..., + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + version: Literal["v2"], + **kwargs: Any, + ) -> GraphOutput[OutputT]: ... + + @overload + def invoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode, + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + version: Literal["v2"], + **kwargs: Any, + ) -> list[StreamPart[OutputT, StateT]]: ... + + @overload def invoke( self, input: InputT | Command | None, @@ -3056,6 +3209,23 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + version: Literal["v1"] = ..., + **kwargs: Any, + ) -> dict[str, Any] | Any: ... + + def invoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode = "values", + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: """Run the graph with a single input and config. @@ -3079,6 +3249,9 @@ class Pregel( - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. + version: The streaming format version. `"v1"` (default) returns the + traditional format, `"v2"` returns `StreamPart` typed dicts when + `stream_mode` is not `"values"`. **kwargs: Additional keyword arguments to pass to the graph run. Returns: @@ -3091,39 +3264,64 @@ class Pregel( chunks: list[dict[str, Any] | Any] = [] interrupts: list[Interrupt] = [] - for chunk in self.stream( - input, - config, - context=context, - stream_mode=["updates", "values"] - if stream_mode == "values" - else stream_mode, - print_mode=print_mode, - output_keys=output_keys, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - durability=durability, - **kwargs, - ): - if stream_mode == "values": - if len(chunk) == 2: - mode, payload = cast(tuple[StreamMode, Any], chunk) + if version == "v2": + # v2: values stream parts carry interrupts directly + for chunk in self.stream( + input, + config, + context=context, + stream_mode="values" if stream_mode == "values" else stream_mode, + print_mode=print_mode, + output_keys=output_keys, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + durability=durability, + version=version, + **kwargs, + ): + if stream_mode == "values": + latest = chunk["data"] + if chunk_ints := chunk.get("interrupts", ()): + interrupts.extend(chunk_ints) # type: ignore[arg-type] else: - _, mode, payload = cast( - tuple[tuple[str, ...], StreamMode, Any], chunk - ) - if ( - mode == "updates" - and isinstance(payload, dict) - and (ints := payload.get(INTERRUPT)) is not None - ): - interrupts.extend(ints) - elif mode == "values": - latest = payload - else: - chunks.append(chunk) + chunks.append(chunk) + else: + # v1: collect interrupts from updates stream + for chunk in self.stream( + input, + config, + context=context, + stream_mode=( + ["updates", "values"] if stream_mode == "values" else stream_mode + ), + print_mode=print_mode, + output_keys=output_keys, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + durability=durability, + **kwargs, + ): + if stream_mode == "values": + if len(chunk) == 2: + mode, payload = cast(tuple[StreamMode, Any], chunk) + else: + _, mode, payload = cast( + tuple[tuple[str, ...], StreamMode, Any], chunk + ) + if ( + mode == "updates" + and isinstance(payload, dict) + and (ints := payload.get(INTERRUPT)) is not None + ): + interrupts.extend(ints) + elif mode == "values": + latest = payload + else: + chunks.append(chunk) if stream_mode == "values": + if version == "v2": + return GraphOutput(value=latest, interrupts=tuple(interrupts)) if interrupts: return ( {**latest, INTERRUPT: interrupts} @@ -3134,6 +3332,57 @@ class Pregel( else: return chunks + @overload + async def ainvoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: Literal["values"] = ..., + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + version: Literal["v2"], + **kwargs: Any, + ) -> GraphOutput[OutputT]: ... + + @overload + async def ainvoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode, + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + version: Literal["v2"], + **kwargs: Any, + ) -> list[StreamPart[OutputT, StateT]]: ... + + @overload + async def ainvoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode = "values", + print_mode: StreamMode | Sequence[StreamMode] = (), + output_keys: str | Sequence[str] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + durability: Durability | None = None, + version: Literal["v1"] = ..., + **kwargs: Any, + ) -> dict[str, Any] | Any: ... + async def ainvoke( self, input: InputT | Command | None, @@ -3146,6 +3395,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: """Asynchronously run the graph with a single input and config. @@ -3169,6 +3419,9 @@ class Pregel( - `"sync"`: Changes are persisted synchronously before the next step starts. - `"async"`: Changes are persisted asynchronously while the next step executes. - `"exit"`: Changes are persisted only when the graph exits. + version: The streaming format version. `"v1"` (default) returns the + traditional format, `"v2"` returns `StreamPart` typed dicts when + `stream_mode` is not `"values"`. **kwargs: Additional keyword arguments to pass to the graph run. Returns: @@ -3181,39 +3434,64 @@ class Pregel( chunks: list[dict[str, Any] | Any] = [] interrupts: list[Interrupt] = [] - async for chunk in self.astream( - input, - config, - context=context, - stream_mode=["updates", "values"] - if stream_mode == "values" - else stream_mode, - print_mode=print_mode, - output_keys=output_keys, - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - durability=durability, - **kwargs, - ): - if stream_mode == "values": - if len(chunk) == 2: - mode, payload = cast(tuple[StreamMode, Any], chunk) + if version == "v2": + # v2: values stream parts carry interrupts directly + async for chunk in self.astream( + input, + config, + context=context, + stream_mode="values" if stream_mode == "values" else stream_mode, + print_mode=print_mode, + output_keys=output_keys, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + durability=durability, + version=version, + **kwargs, + ): + if stream_mode == "values": + latest = chunk["data"] + if chunk_ints := chunk.get("interrupts", ()): + interrupts.extend(chunk_ints) # type: ignore[arg-type] else: - _, mode, payload = cast( - tuple[tuple[str, ...], StreamMode, Any], chunk - ) - if ( - mode == "updates" - and isinstance(payload, dict) - and (ints := payload.get(INTERRUPT)) is not None - ): - interrupts.extend(ints) - elif mode == "values": - latest = payload - else: - chunks.append(chunk) + chunks.append(chunk) + else: + # v1: collect interrupts from updates stream + async for chunk in self.astream( + input, + config, + context=context, + stream_mode=( + ["updates", "values"] if stream_mode == "values" else stream_mode + ), + print_mode=print_mode, + output_keys=output_keys, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + durability=durability, + **kwargs, + ): + if stream_mode == "values": + if len(chunk) == 2: + mode, payload = cast(tuple[StreamMode, Any], chunk) + else: + _, mode, payload = cast( + tuple[tuple[str, ...], StreamMode, Any], chunk + ) + if ( + mode == "updates" + and isinstance(payload, dict) + and (ints := payload.get(INTERRUPT)) is not None + ): + interrupts.extend(ints) + elif mode == "values": + latest = payload + else: + chunks.append(chunk) if stream_mode == "values": + if version == "v2": + return GraphOutput(value=latest, interrupts=tuple(interrupts)) if interrupts: return ( {**latest, INTERRUPT: interrupts} @@ -3278,6 +3556,9 @@ def _output( stream_subgraphs: bool, getter: Callable[[], tuple[tuple[str, ...], str, Any]], empty_exc: type[Exception], + version: Literal["v1", "v2"] = "v1", + output_mapper: Callable[[Any], Any] | None = None, + state_mapper: Callable[[Any], Any] | None = None, ) -> Iterator: while True: try: @@ -3305,7 +3586,23 @@ def _output( ) ) if mode in stream_mode: - if stream_subgraphs and isinstance(stream_mode, list): + if version == "v2": + if mode == "values": + # pop __interrupt__ into typed field, coerce data + ints: tuple[Interrupt, ...] = () + if isinstance(payload, dict): + ints = payload.pop(INTERRUPT, ()) + if output_mapper: + payload = output_mapper(payload) + yield {"type": mode, "ns": ns, "data": payload, "interrupts": ints} + elif mode in ("checkpoints", "debug"): + # coerce state values in checkpoint/debug payloads + if state_mapper: + _coerce_checkpoint_values(payload, state_mapper) + yield {"type": mode, "ns": ns, "data": payload} + else: + yield {"type": mode, "ns": ns, "data": payload} + elif stream_subgraphs and isinstance(stream_mode, list): yield (ns, mode, payload) elif isinstance(stream_mode, list): yield (mode, payload) @@ -3315,6 +3612,31 @@ def _output( yield payload +def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> None: + """Coerce `values` dicts inside checkpoint or debug payloads in-place. + + Skips the initial checkpoint (where next contains ``__start__``) because + not all channels are populated yet and coercion would fail. + """ + _START = "__start__" + # debug wrapper: {"type": "checkpoint", "payload": {"values": dict, ...}} + if ( + isinstance(payload, dict) + and payload.get("type") == "checkpoint" + and isinstance(payload.get("payload"), dict) + and isinstance(payload["payload"].get("values"), dict) + and _START not in payload["payload"].get("next", ()) + ): + payload["payload"]["values"] = mapper(payload["payload"]["values"]) + # direct checkpoint payload: {"values": dict, ...} + elif ( + isinstance(payload, dict) + and isinstance(payload.get("values"), dict) + and _START not in payload.get("next", ()) + ): + payload["values"] = mapper(payload["values"]) + + def _coerce_context( context_schema: type[ContextT] | None, context: Any ) -> ContextT | None: diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index c9bf6e5ff..e5e957f50 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -2,13 +2,21 @@ from __future__ import annotations from abc import abstractmethod from collections.abc import AsyncIterator, Callable, Iterator, Sequence -from typing import Any, Generic, cast +from typing import Any, Generic, Literal, cast, overload from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.graph import Graph as DrawableGraph from typing_extensions import Self -from langgraph.types import All, Command, StateSnapshot, StateUpdate, StreamMode +from langgraph.types import ( + All, + Command, + GraphOutput, + StateSnapshot, + StateUpdate, + StreamMode, + StreamPart, +) from langgraph.typing import ContextT, InputT, OutputT, StateT __all__ = ("PregelProtocol", "StreamProtocol") @@ -96,6 +104,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou as_node: str | None = None, ) -> RunnableConfig: ... + @overload @abstractmethod def stream( self, @@ -107,8 +116,68 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, subgraphs: bool = False, + version: Literal["v2"], + ) -> Iterator[StreamPart[OutputT, StateT]]: ... + + @overload + @abstractmethod + def stream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | 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["v1"] = ..., ) -> Iterator[dict[str, Any] | Any]: ... + @abstractmethod + def stream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | 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["v1", "v2"] = "v1", + ) -> Iterator[dict[str, Any] | Any]: ... + + @overload + @abstractmethod + def astream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | 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"], + ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ... + + @overload + @abstractmethod + def astream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | 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["v1"] = ..., + ) -> AsyncIterator[dict[str, Any] | Any]: ... + @abstractmethod def astream( self, @@ -120,8 +189,35 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, subgraphs: bool = False, + version: Literal["v1", "v2"] = "v1", ) -> AsyncIterator[dict[str, Any] | Any]: ... + @overload + @abstractmethod + def invoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + version: Literal["v2"], + ) -> GraphOutput[OutputT]: ... + + @overload + @abstractmethod + def invoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + version: Literal["v1"] = ..., + ) -> dict[str, Any] | Any: ... + @abstractmethod def invoke( self, @@ -131,6 +227,33 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou context: ContextT | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, + version: Literal["v1", "v2"] = "v1", + ) -> dict[str, Any] | Any: ... + + @overload + @abstractmethod + async def ainvoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + version: Literal["v2"], + ) -> GraphOutput[OutputT]: ... + + @overload + @abstractmethod + async def ainvoke( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + version: Literal["v1"] = ..., ) -> dict[str, Any] | Any: ... @abstractmethod @@ -142,6 +265,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou context: ContextT | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, + version: Literal["v1", "v2"] = "v1", ) -> dict[str, Any] | Any: ... diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 2535d966a..7b9602d81 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -7,6 +7,7 @@ from typing import ( Any, Literal, cast, + overload, ) from uuid import UUID @@ -57,10 +58,12 @@ from langgraph.pregel.protocol import PregelProtocol, StreamProtocol from langgraph.types import ( All, Command, + GraphOutput, Interrupt, PregelTask, StateSnapshot, StreamMode, + StreamPart, ) logger = logging.getLogger(__name__) @@ -682,6 +685,7 @@ class RemoteGraph(PregelProtocol): updated_stream_modes.remove("events") return (updated_stream_modes, requested_stream_modes, req_single, stream) + @overload def stream( self, input: dict[str, Any] | Any, @@ -693,6 +697,38 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + version: Literal["v2"], + **kwargs: Any, + ) -> Iterator[StreamPart]: ... + + @overload + 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, + headers: dict[str, str] | None = None, + params: QueryParamTypes | None = None, + version: Literal["v1"] = ..., + **kwargs: Any, + ) -> Iterator[dict[str, Any] | Any]: ... + + 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, + headers: dict[str, str] | None = None, + params: QueryParamTypes | None = None, + version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> Iterator[dict[str, Any] | Any]: """Create a run and stream the results. @@ -774,10 +810,18 @@ class RemoteGraph(PregelProtocol): continue if chunk.event.startswith("messages"): - chunk = chunk._replace(data=tuple(chunk.data)) # type: ignore + chunk = chunk._replace(data=tuple(chunk.data)) # emit chunk - if subgraphs: + if version == "v2": + ints: tuple[Interrupt, ...] = () + if mode == "values" and isinstance(chunk.data, dict): + ints = tuple( + Interrupt(**i) if isinstance(i, dict) else i + for i in chunk.data.pop(INTERRUPT, ()) + ) + yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints} + elif subgraphs: if NS_SEP in chunk.event: mode, ns_ = chunk.event.split(NS_SEP, 1) ns = tuple(ns_.split(NS_SEP)) @@ -792,6 +836,38 @@ class RemoteGraph(PregelProtocol): else: yield chunk + @overload + 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, + headers: dict[str, str] | None = None, + params: QueryParamTypes | None = None, + version: Literal["v2"], + **kwargs: Any, + ) -> AsyncIterator[StreamPart]: ... + + @overload + 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, + headers: dict[str, str] | None = None, + params: QueryParamTypes | None = None, + version: Literal["v1"] = ..., + **kwargs: Any, + ) -> AsyncIterator[dict[str, Any] | Any]: ... + async def astream( self, input: dict[str, Any] | Any, @@ -803,6 +879,7 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> AsyncIterator[dict[str, Any] | Any]: """Create a run and stream the results. @@ -884,10 +961,18 @@ class RemoteGraph(PregelProtocol): continue if chunk.event.startswith("messages"): - chunk = chunk._replace(data=tuple(chunk.data)) # type: ignore + chunk = chunk._replace(data=tuple(chunk.data)) # emit chunk - if subgraphs: + if version == "v2": + ints: tuple[Interrupt, ...] = () + if mode == "values" and isinstance(chunk.data, dict): + ints = tuple( + Interrupt(**i) if isinstance(i, dict) else i + for i in chunk.data.pop(INTERRUPT, ()) + ) + yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints} + elif subgraphs: if NS_SEP in chunk.event: mode, ns_ = chunk.event.split(NS_SEP, 1) ns = tuple(ns_.split(NS_SEP)) @@ -918,6 +1003,7 @@ class RemoteGraph(PregelProtocol): ) -> AsyncIterator[dict[str, Any]]: raise NotImplementedError + @overload def invoke( self, input: dict[str, Any] | Any, @@ -927,6 +1013,34 @@ class RemoteGraph(PregelProtocol): interrupt_after: All | Sequence[str] | None = None, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + version: Literal["v2"], + **kwargs: Any, + ) -> GraphOutput[dict[str, Any]]: ... + + @overload + def invoke( + self, + input: dict[str, Any] | Any, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + headers: dict[str, str] | None = None, + params: QueryParamTypes | None = None, + version: Literal["v1"] = ..., + **kwargs: Any, + ) -> dict[str, Any] | Any: ... + + def invoke( + self, + input: dict[str, Any] | Any, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + headers: dict[str, str] | None = None, + params: QueryParamTypes | None = None, + version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: """Create a run, wait until it finishes and return the final state. @@ -937,12 +1051,14 @@ class RemoteGraph(PregelProtocol): interrupt_before: Interrupt the graph before these nodes. interrupt_after: Interrupt the graph after these nodes. headers: Additional headers to pass to the request. + version: The streaming format version. `"v1"` (default) returns the + traditional format, `"v2"` returns `StreamPart` typed dicts. **kwargs: Additional params to pass to RemoteGraph.stream. Returns: The output of the graph. """ - for chunk in self.stream( + for chunk in self.stream( # type: ignore[misc, call-overload] input, config=config, interrupt_before=interrupt_before, @@ -950,15 +1066,49 @@ class RemoteGraph(PregelProtocol): headers=headers, stream_mode="values", params=params, + version=version, **kwargs, ): pass try: + if version == "v2": + return GraphOutput( + value=chunk["data"], + interrupts=tuple(chunk.get("interrupts", ())), + ) return chunk except UnboundLocalError: logger.warning("No events received from remote graph") return None + @overload + async def ainvoke( + self, + input: dict[str, Any] | Any, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + headers: dict[str, str] | None = None, + params: QueryParamTypes | None = None, + version: Literal["v2"], + **kwargs: Any, + ) -> GraphOutput[dict[str, Any]]: ... + + @overload + async def ainvoke( + self, + input: dict[str, Any] | Any, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + headers: dict[str, str] | None = None, + params: QueryParamTypes | None = None, + version: Literal["v1"] = ..., + **kwargs: Any, + ) -> dict[str, Any] | Any: ... + async def ainvoke( self, input: dict[str, Any] | Any, @@ -968,6 +1118,7 @@ class RemoteGraph(PregelProtocol): interrupt_after: All | Sequence[str] | None = None, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: """Create a run, wait until it finishes and return the final state. @@ -978,12 +1129,14 @@ class RemoteGraph(PregelProtocol): interrupt_before: Interrupt the graph before these nodes. interrupt_after: Interrupt the graph after these nodes. headers: Additional headers to pass to the request. + version: The streaming format version. `"v1"` (default) returns the + traditional format, `"v2"` returns `StreamPart` typed dicts. **kwargs: Additional params to pass to RemoteGraph.astream. Returns: The output of the graph. """ - async for chunk in self.astream( + async for chunk in self.astream( # type: ignore[misc, call-overload] input, config=config, interrupt_before=interrupt_before, @@ -991,10 +1144,16 @@ class RemoteGraph(PregelProtocol): headers=headers, stream_mode="values", params=params, + version=version, **kwargs, ): pass try: + if version == "v2": + return GraphOutput( + value=chunk["data"], + interrupts=tuple(chunk.get("interrupts", ())), + ) return chunk except UnboundLocalError: logger.warning("No events received from remote graph") diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3f95deb97..547513498 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -16,16 +16,25 @@ from typing import ( ) from warnings import warn +from langchain_core.messages import AnyMessage from langchain_core.runnables import Runnable, RunnableConfig from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata -from typing_extensions import Unpack, deprecated +from typing_extensions import NotRequired, TypeAliasType, TypedDict, Unpack, deprecated from xxhash import xxh3_128_hexdigest from langgraph._internal._cache import default_cache_key +from langgraph._internal._constants import INTERRUPT as _INTERRUPT_KEY from langgraph._internal._fields import get_cached_annotated_keys, get_update_as_tuples from langgraph._internal._retry import default_retry_on from langgraph._internal._typing import MISSING, DeprecatedKwargs -from langgraph.warnings import LangGraphDeprecatedSinceV10 +from langgraph.warnings import LangGraphDeprecatedSinceV10, LangGraphDeprecatedSinceV11 + +# Local TypeVars for generic stream TypedDicts. +# We use separate TypeVars here (rather than importing from langgraph.typing) +# because the typing module TypeVars have defaults that cause mypy issues +# when used in standalone type aliases. +StateT = TypeVar("StateT") +OutputT = TypeVar("OutputT") if TYPE_CHECKING: from langgraph.pregel.protocol import PregelProtocol @@ -44,6 +53,19 @@ __all__ = ( "Checkpointer", "StreamMode", "StreamWriter", + "StreamPart", + "ValuesStreamPart", + "UpdatesStreamPart", + "MessagesStreamPart", + "CustomStreamPart", + "CheckpointStreamPart", + "TasksStreamPart", + "DebugStreamPart", + "TaskPayload", + "TaskResultPayload", + "CheckpointTask", + "CheckpointPayload", + "DebugPayload", "RetryPolicy", "CachePolicy", "Interrupt", @@ -56,6 +78,7 @@ __all__ = ( "Durability", "interrupt", "Overwrite", + "GraphOutput", "ensure_valid_checkpointer", ) @@ -113,6 +136,268 @@ StreamWriter = Callable[[Any], None] Always injected into nodes if requested as a keyword argument, but it's a no-op when not using `stream_mode="custom"`.""" + +class TaskPayload(TypedDict): + """Payload for a task start event.""" + + id: str + """Unique identifier for this task.""" + name: str + """Name of the node being executed.""" + input: Any + """Input data passed to the task.""" + triggers: list[str] + """List of triggers that caused this task to be executed (e.g. channel writes).""" + + +class TaskResultPayload(TypedDict): + """Payload for a task result event.""" + + id: str + """Unique identifier for this task.""" + name: str + """Name of the node that was executed.""" + error: str | None + """Error message if the task failed, otherwise `None`.""" + interrupts: list[dict] + """List of interrupts that occurred during task execution.""" + result: dict[str, Any] + """Mapping of channel names to the values written by this task.""" + + +class CheckpointTask(TypedDict): + """A task entry within a `CheckpointPayload`. + + The keys present depend on the task's state: + + - **Error:** `id`, `name`, `error`, `state` + - **Has result:** `id`, `name`, `result`, `interrupts`, `state` + - **Pending:** `id`, `name`, `interrupts`, `state` + """ + + id: str + """Unique identifier for this task.""" + name: str + """Name of the node being executed.""" + error: NotRequired[str] + """Error message, present only if the task failed.""" + result: NotRequired[Any] + """Result of the task, present only if the task completed successfully.""" + interrupts: NotRequired[list[dict]] + """List of interrupts, present when the task has been interrupted or completed.""" + state: StateSnapshot | RunnableConfig | None + """Snapshot of the subgraph state, or a `RunnableConfig` pointing to it. `None` if not a subgraph.""" + + +class CheckpointPayload(TypedDict, Generic[StateT]): + """Payload for a checkpoint event.""" + + config: RunnableConfig | None + """Configuration for this checkpoint, including the `thread_id` and `checkpoint_id`.""" + metadata: CheckpointMetadata + """Metadata associated with this checkpoint (e.g. step number, source, writes).""" + values: StateT + """Current state values at the time of this checkpoint.""" + next: list[str] + """Names of the nodes scheduled to execute next.""" + parent_config: RunnableConfig | None + """Configuration of the parent checkpoint, or `None` if this is the first checkpoint.""" + tasks: list[CheckpointTask] + """List of tasks associated with this checkpoint.""" + + +class _DebugCheckpointPayload(TypedDict, Generic[StateT]): + step: int + """The step number in the graph execution.""" + timestamp: str + """ISO 8601 timestamp of when this event occurred.""" + type: Literal["checkpoint"] + """Event type discriminator, always `"checkpoint"`.""" + payload: CheckpointPayload[StateT] + """The checkpoint payload.""" + + +class _DebugTaskPayload(TypedDict): + step: int + """The step number in the graph execution.""" + timestamp: str + """ISO 8601 timestamp of when this event occurred.""" + type: Literal["task"] + """Event type discriminator, always `"task"`.""" + payload: TaskPayload + """The task start payload.""" + + +class _DebugTaskResultPayload(TypedDict): + step: int + """The step number in the graph execution.""" + timestamp: str + """ISO 8601 timestamp of when this event occurred.""" + type: Literal["task_result"] + """Event type discriminator, always `"task_result"`.""" + payload: TaskResultPayload + """The task result payload.""" + + +DebugPayload = TypeAliasType( + "DebugPayload", + _DebugCheckpointPayload[StateT] | _DebugTaskPayload | _DebugTaskResultPayload, + type_params=(StateT,), +) +"""Wrapper payload for debug events. Discriminate on `type`.""" + + +class ValuesStreamPart(TypedDict, Generic[OutputT]): + """Stream part emitted for `stream_mode="values"`. + + `data` contains the full state after each step, as returned by `read_channels()`. + """ + + type: Literal["values"] + ns: tuple[str, ...] + data: OutputT + interrupts: tuple[Interrupt, ...] + + +class UpdatesStreamPart(TypedDict): + """Stream part emitted for `stream_mode="updates"`. + + `data` maps node names to their outputs. May also contain + `__interrupt__` (tuple of `Interrupt` dicts) and `__metadata__` keys. + """ + + type: Literal["updates"] + ns: tuple[str, ...] + data: dict[str, Any] + + +class MessagesStreamPart(TypedDict): + """Stream part emitted for `stream_mode="messages"`. + + `data` is a 2-tuple of `(message, metadata)` where `message` is a + `BaseMessage` (e.g. `AIMessageChunk`) and `metadata` is a dict containing + keys like `langgraph_step`, `langgraph_node`, `langgraph_triggers`, etc. + """ + + type: Literal["messages"] + ns: tuple[str, ...] + data: tuple[AnyMessage, dict[str, Any]] + + +class CustomStreamPart(TypedDict): + """Stream part emitted for `stream_mode="custom"`. + + `data` is whatever value was passed to `StreamWriter` inside a node. + """ + + type: Literal["custom"] + ns: tuple[str, ...] + data: Any + + +class CheckpointStreamPart(TypedDict, Generic[StateT]): + """Stream part emitted for `stream_mode="checkpoints"`.""" + + type: Literal["checkpoints"] + ns: tuple[str, ...] + data: CheckpointPayload[StateT] + + +class TasksStreamPart(TypedDict): + """Stream part emitted for `stream_mode="tasks"`. + + For task start events, `data` is a `TaskPayload` with `id`, `name`, + `input`, and `triggers` keys. + + For task result events, `data` is a `TaskResultPayload` with `id`, + `name`, `error`, `interrupts`, and `result` keys. + """ + + type: Literal["tasks"] + ns: tuple[str, ...] + data: TaskPayload | TaskResultPayload + + +class DebugStreamPart(TypedDict, Generic[StateT]): + """Stream part emitted for `stream_mode="debug"`.""" + + type: Literal["debug"] + ns: tuple[str, ...] + data: DebugPayload[StateT] + + +StreamPart = TypeAliasType( + "StreamPart", + ValuesStreamPart[OutputT] + | UpdatesStreamPart + | MessagesStreamPart + | CustomStreamPart + | CheckpointStreamPart[StateT] + | TasksStreamPart + | DebugStreamPart[StateT], + type_params=(OutputT, StateT), +) +"""A discriminated union of all v2 stream part types. + +Use `part["type"]` to narrow the type: + +```python +async for part in graph.astream(input, version="v2"): + if part["type"] == "values": + part["data"] # OutputT — full state (pydantic/dataclass/dict) + elif part["type"] == "messages": + part["data"] # tuple[BaseMessage, dict] — (message, metadata) + elif part["type"] == "custom": + part["data"] # Any — user-defined +``` +""" + + +@dataclass(frozen=True) +class GraphOutput(Generic[OutputT]): + """Typed container returned by `invoke()` / `ainvoke()` with `version="v2"`. + + Attributes: + value: The final output of the graph (dict, Pydantic model, dataclass, etc.). + interrupts: Any interrupts that occurred during execution. + """ + + value: OutputT + interrupts: tuple[Interrupt, ...] = () + + def __getitem__(self, key: str) -> Any: + """Backward compat: `result['__interrupt__']` and dict-key access.""" + warn( + "Accessing GraphOutput via `result[key]` is deprecated. " + "Use `result.value` to access the output value directly, " + "or `result.interrupts` for interrupts.", + LangGraphDeprecatedSinceV11, + stacklevel=2, + ) + if key == _INTERRUPT_KEY: + return self.interrupts + if isinstance(self.value, dict): + return self.value[key] + try: + return getattr(self.value, key) + except AttributeError: + raise KeyError(key) + + def __contains__(self, key: object) -> bool: + warn( + "Accessing GraphOutput via `key in result` is deprecated. " + "Use `result.value` to access the output value directly, " + "or `result.interrupts` for interrupts.", + LangGraphDeprecatedSinceV11, + stacklevel=2, + ) + if key == _INTERRUPT_KEY: + return bool(self.interrupts) + if isinstance(self.value, dict): + return key in self.value + return isinstance(key, str) and hasattr(self.value, key) + + _DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True} diff --git a/libs/langgraph/langgraph/warnings.py b/libs/langgraph/langgraph/warnings.py index 638f247e3..aa447fb02 100644 --- a/libs/langgraph/langgraph/warnings.py +++ b/libs/langgraph/langgraph/warnings.py @@ -6,6 +6,7 @@ __all__ = ( "LangGraphDeprecationWarning", "LangGraphDeprecatedSinceV05", "LangGraphDeprecatedSinceV10", + "LangGraphDeprecatedSinceV11", ) @@ -59,3 +60,10 @@ class LangGraphDeprecatedSinceV10(LangGraphDeprecationWarning): def __init__(self, message: str, *args: object) -> None: super().__init__(message, *args, since=(1, 0), expected_removal=(2, 0)) + + +class LangGraphDeprecatedSinceV11(LangGraphDeprecationWarning): + """A specific `LangGraphDeprecationWarning` subclass defining functionality deprecated since LangGraph v1.1.0""" + + def __init__(self, message: str, *args: object) -> None: + super().__init__(message, *args, since=(1, 1), expected_removal=(3, 0)) diff --git a/libs/langgraph/tests/test_deprecation.py b/libs/langgraph/tests/test_deprecation.py index 919a7ca27..54e1e3812 100644 --- a/libs/langgraph/tests/test_deprecation.py +++ b/libs/langgraph/tests/test_deprecation.py @@ -14,8 +14,12 @@ from langgraph.func import entrypoint, task from langgraph.graph import StateGraph from langgraph.graph.message import MessageGraph from langgraph.pregel import NodeBuilder, Pregel -from langgraph.types import Interrupt, RetryPolicy -from langgraph.warnings import LangGraphDeprecatedSinceV05, LangGraphDeprecatedSinceV10 +from langgraph.types import GraphOutput, Interrupt, RetryPolicy +from langgraph.warnings import ( + LangGraphDeprecatedSinceV05, + LangGraphDeprecatedSinceV10, + LangGraphDeprecatedSinceV11, +) class PlainState(TypedDict): ... @@ -197,6 +201,7 @@ def test_deprecated_import() -> None: @pytest.mark.filterwarnings( "ignore:`durability` has no effect when no checkpointer is present" ) +@pytest.mark.filterwarnings("ignore:Accessing GraphOutput via") def test_checkpoint_during_deprecation_state_graph() -> None: class CheckDurability(TypedDict): durability: NotRequired[str] @@ -341,3 +346,34 @@ def test_message_graph_deprecation() -> None: match="MessageGraph is deprecated in LangGraph v1.0.0, to be removed in v2.0.0. Please use StateGraph with a `messages` key instead.", ): MessageGraph() + + +def test_graph_output_getitem_deprecation() -> None: + output = GraphOutput(value={"foo": "bar"}) + + with pytest.warns( + LangGraphDeprecatedSinceV11, + match=r"Accessing GraphOutput via `result\[key\]` is deprecated", + ): + assert output["foo"] == "bar" + + +def test_graph_output_contains_deprecation() -> None: + output = GraphOutput(value={"foo": "bar"}) + + with pytest.warns( + LangGraphDeprecatedSinceV11, + match=r"Accessing GraphOutput via `key in result` is deprecated", + ): + assert "foo" in output + + +def test_graph_output_getitem_interrupt_deprecation() -> None: + interrupts = (Interrupt(value="q", id="abc"),) + output = GraphOutput(value={"foo": "bar"}, interrupts=interrupts) + + with pytest.warns( + LangGraphDeprecatedSinceV11, + match=r"Accessing GraphOutput via `result\[key\]` is deprecated", + ): + assert output["__interrupt__"] == interrupts diff --git a/libs/langgraph/tests/test_stream_v2.py b/libs/langgraph/tests/test_stream_v2.py new file mode 100644 index 000000000..50fb83031 --- /dev/null +++ b/libs/langgraph/tests/test_stream_v2.py @@ -0,0 +1,1152 @@ +"""Tests for v2 streaming format (StreamPart TypedDicts). + +This file is checked by mypy directly — no subprocess workarounds. +Type-narrowing is validated via `assert_type` calls in `_check_type_narrowing`. +""" + +from __future__ import annotations + +import operator +import sys +from dataclasses import dataclass +from typing import Annotated, Any, TypeVar + +import pytest +from langchain_core.messages import AIMessage, BaseMessage +from langgraph.checkpoint.memory import InMemorySaver +from pydantic import BaseModel, ValidationError +from typing_extensions import TypedDict, assert_type + +from langgraph._internal._constants import INTERRUPT +from langgraph.constants import END, START +from langgraph.func import entrypoint +from langgraph.graph import StateGraph +from langgraph.graph.message import MessagesState +from langgraph.types import ( + CheckpointPayload, + CheckpointStreamPart, + CustomStreamPart, + DebugPayload, + DebugStreamPart, + GraphOutput, + Interrupt, + MessagesStreamPart, + StreamPart, + StreamWriter, + TaskPayload, + TaskResultPayload, + TasksStreamPart, + UpdatesStreamPart, + ValuesStreamPart, + interrupt, +) +from tests.fake_chat import FakeChatModel + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + +# --- state and graph builders --- + + +class SimpleState(TypedDict): + value: str + items: Annotated[list[str], operator.add] + + +_SIMPLE_INPUT: SimpleState = {"value": "x", "items": []} +_MSG_INPUT: MessagesState = {"messages": "hi"} + + +def _make_simple_graph() -> StateGraph[SimpleState, None, SimpleState, SimpleState]: + def node_a(state: SimpleState) -> dict[str, Any]: + return {"value": state["value"] + "_a", "items": ["a"]} + + def node_b(state: SimpleState) -> dict[str, Any]: + return {"value": state["value"] + "_b", "items": ["b"]} + + builder = StateGraph(SimpleState, input_schema=SimpleState) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_b", END) + return builder + + +def _make_messages_graph() -> StateGraph[ + MessagesState, None, MessagesState, MessagesState +]: + model = FakeChatModel(messages=[AIMessage(content="hello world")]) + + def call_model(state: MessagesState) -> dict[str, Any]: + return {"messages": model.invoke(state["messages"])} + + builder = StateGraph(MessagesState, input_schema=MessagesState) + builder.add_node("call_model", call_model) + builder.add_edge(START, "call_model") + builder.add_edge("call_model", END) + return builder + + +def _make_custom_graph() -> Any: + @entrypoint() + def graph(inputs: Any, *, writer: StreamWriter) -> Any: + writer("hello") + writer(42) + return inputs + + return graph + + +def _make_subgraph() -> Any: + inner = _make_simple_graph().compile() + outer_builder = StateGraph(SimpleState, input_schema=SimpleState) + outer_builder.add_node("inner", inner) + outer_builder.add_edge(START, "inner") + outer_builder.add_edge("inner", END) + return outer_builder.compile() + + +# --- shared assertion helpers --- + +_STREAM_PART_KEYS = {"type", "ns", "data"} + + +def _assert_stream_part_shape(part: StreamPart[Any, Any]) -> None: + """Assert a v2 stream part has the required keys and correct types.""" + assert isinstance(part, dict), f"Expected dict, got {type(part)}" + assert _STREAM_PART_KEYS <= part.keys(), ( + f"Missing keys: {_STREAM_PART_KEYS - part.keys()}" + ) + assert isinstance(part["type"], str) + assert isinstance(part["ns"], tuple) + for elem in part["ns"]: + assert isinstance(elem, str) + if part["type"] == "values": + assert "interrupts" in part, "values stream part missing 'interrupts' field" + assert isinstance(part["interrupts"], tuple) + + +# --- v1 backwards compatibility --- + + +class TestV1BackwardsCompat: + def test_stream_default_is_v1(self) -> None: + graph = _make_simple_graph().compile() + chunks = list(graph.stream(_SIMPLE_INPUT)) + for chunk in chunks: + assert isinstance(chunk, dict) + + def test_stream_v1_updates_mode(self) -> None: + graph = _make_simple_graph().compile() + chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode="updates")) + assert len(chunks) == 2 + assert "node_a" in chunks[0] + assert "node_b" in chunks[1] + + def test_stream_v1_list_mode(self) -> None: + graph = _make_simple_graph().compile() + chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode=["values", "updates"])) + for chunk in chunks: + assert isinstance(chunk, tuple) and len(chunk) == 2 + mode, _data = chunk + assert mode in ("values", "updates") + + def test_stream_v1_subgraphs(self) -> None: + graph = _make_simple_graph().compile() + chunks = list( + graph.stream(_SIMPLE_INPUT, stream_mode="updates", subgraphs=True) + ) + for chunk in chunks: + assert isinstance(chunk, tuple) and len(chunk) == 2 + ns, _data = chunk + assert isinstance(ns, tuple) + + +# --- v2 sync stream --- + + +class TestV2Stream: + def test_values(self) -> None: + graph = _make_simple_graph().compile() + chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode="values", version="v2")) + assert len(chunks) >= 1 + for c in chunks: + _assert_stream_part_shape(c) + assert c["type"] == "values" + assert c["ns"] == () + assert isinstance(c["data"], dict) + + def test_updates(self) -> None: + graph = _make_simple_graph().compile() + chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode="updates", version="v2")) + assert len(chunks) == 2 + for c in chunks: + _assert_stream_part_shape(c) + assert c["type"] == "updates" + assert c["ns"] == () + assert "node_a" in chunks[0]["data"] + assert "node_b" in chunks[1]["data"] + + def test_messages(self) -> None: + graph = _make_messages_graph().compile() + chunks = list(graph.stream(_MSG_INPUT, stream_mode="messages", version="v2")) + msg_chunks = [c for c in chunks if c["type"] == "messages"] + assert len(msg_chunks) >= 1 + for c in msg_chunks: + _assert_stream_part_shape(c) + assert c["ns"] == () + data = c["data"] + assert isinstance(data, tuple) and len(data) == 2 + message, metadata = data + assert isinstance(message, BaseMessage) + assert isinstance(metadata, dict) + assert "langgraph_node" in metadata + + def test_custom(self) -> None: + graph = _make_custom_graph() + chunks = list(graph.stream({"key": "val"}, stream_mode="custom", version="v2")) + custom = [c for c in chunks if c["type"] == "custom"] + assert len(custom) == 2 + for c in custom: + _assert_stream_part_shape(c) + assert custom[0]["data"] == "hello" + assert custom[1]["data"] == 42 + + def test_multiple_modes(self) -> None: + graph = _make_simple_graph().compile() + chunks = list( + graph.stream( + _SIMPLE_INPUT, + stream_mode=["values", "updates"], + version="v2", + ) + ) + types_seen = {c["type"] for c in chunks} + assert {"values", "updates"} <= types_seen + for c in chunks: + _assert_stream_part_shape(c) + + def test_subgraphs_ns(self) -> None: + outer = _make_subgraph() + chunks = list( + outer.stream( + _SIMPLE_INPUT, + stream_mode="updates", + subgraphs=True, + version="v2", + ) + ) + for c in chunks: + _assert_stream_part_shape(c) + root = [c for c in chunks if c["ns"] == ()] + sub = [c for c in chunks if c["ns"] != ()] + assert len(root) >= 1 + assert len(sub) >= 1 + + def test_checkpoints(self) -> None: + graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) + config: Any = {"configurable": {"thread_id": "test-v2-ckpt"}} + chunks = list( + graph.stream( + _SIMPLE_INPUT, + config, + stream_mode="checkpoints", + version="v2", + ) + ) + ckpt = [c for c in chunks if c["type"] == "checkpoints"] + assert len(ckpt) >= 1 + for c in ckpt: + _assert_stream_part_shape(c) + assert c["ns"] == () + payload = c["data"] + assert {"config", "metadata", "values", "next", "tasks"} <= payload.keys() + + def test_tasks(self) -> None: + graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) + config: Any = {"configurable": {"thread_id": "test-v2-tasks"}} + chunks = list( + graph.stream( + _SIMPLE_INPUT, + config, + stream_mode="tasks", + version="v2", + ) + ) + tasks = [c for c in chunks if c["type"] == "tasks"] + assert len(tasks) >= 2 + for c in tasks: + _assert_stream_part_shape(c) + assert c["ns"] == () + assert "id" in c["data"] and "name" in c["data"] + starts = [c for c in tasks if "triggers" in c["data"]] + results = [c for c in tasks if "result" in c["data"]] + assert len(starts) >= 2 + assert len(results) >= 2 + + def test_debug(self) -> None: + graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) + config: Any = {"configurable": {"thread_id": "test-v2-debug"}} + chunks = list( + graph.stream( + _SIMPLE_INPUT, + config, + stream_mode="debug", + version="v2", + ) + ) + debug = [c for c in chunks if c["type"] == "debug"] + assert len(debug) >= 1 + for c in debug: + _assert_stream_part_shape(c) + assert c["ns"] == () + envelope = c["data"] + assert {"step", "timestamp", "type", "payload"} <= envelope.keys() + assert envelope["type"] in ("checkpoint", "task", "task_result") + + def test_subgraphs_param_does_not_change_format(self) -> None: + """In v2, subgraphs=True/False should not change the output format.""" + graph = _make_simple_graph().compile() + chunks_no_sub = list( + graph.stream( + _SIMPLE_INPUT, + stream_mode="updates", + subgraphs=False, + version="v2", + ) + ) + chunks_with_sub = list( + graph.stream( + _SIMPLE_INPUT, + stream_mode="updates", + subgraphs=True, + version="v2", + ) + ) + for c in chunks_no_sub + chunks_with_sub: + _assert_stream_part_shape(c) + + +# --- v2 sync invoke --- + + +class TestV2Invoke: + def test_values_default(self) -> None: + graph = _make_simple_graph().compile() + result = graph.invoke(_SIMPLE_INPUT, version="v2") + assert isinstance(result, GraphOutput) + assert result.value == {"value": "x_a_b", "items": ["a", "b"]} + assert result.interrupts == () + # backward compat dict access + assert result["value"] == "x_a_b" + assert result["items"] == ["a", "b"] + + def test_invoke_v2_graph_output_with_interrupts(self) -> None: + def my_node(state: SimpleState) -> dict[str, Any]: + answer = interrupt("what is your name?") + return {"value": answer, "items": ["done"]} + + builder: StateGraph = StateGraph(SimpleState) + builder.add_node("my_node", my_node) + builder.add_edge(START, "my_node") + builder.add_edge("my_node", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + config: Any = {"configurable": {"thread_id": "test-invoke-v2-interrupts"}} + result = graph.invoke({"value": "x", "items": []}, config, version="v2") + assert isinstance(result, GraphOutput) + assert len(result.interrupts) > 0 + for intr in result.interrupts: + assert isinstance(intr, Interrupt) + # value should still be the state (not None or empty) + assert isinstance(result.value, dict) + + def test_invoke_v2_graph_output_interrupt_compat(self) -> None: + """result['__interrupt__'] works via __getitem__.""" + + def my_node(state: SimpleState) -> dict[str, Any]: + answer = interrupt("what is your name?") + return {"value": answer, "items": ["done"]} + + builder: StateGraph = StateGraph(SimpleState) + builder.add_node("my_node", my_node) + builder.add_edge(START, "my_node") + builder.add_edge("my_node", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + config: Any = {"configurable": {"thread_id": "test-invoke-v2-compat"}} + result = graph.invoke({"value": "x", "items": []}, config, version="v2") + assert isinstance(result, GraphOutput) + assert INTERRUPT in result + assert result[INTERRUPT] == result.interrupts + assert len(result[INTERRUPT]) > 0 + + def test_invoke_v2_graph_output_no_interrupts(self) -> None: + graph = _make_simple_graph().compile() + result = graph.invoke(_SIMPLE_INPUT, version="v2") + assert isinstance(result, GraphOutput) + assert result.interrupts == () + assert INTERRUPT not in result + + def test_invoke_v2_pydantic_state(self) -> None: + """invoke with v2 and pydantic state returns GraphOutput with pydantic value.""" + + def node_a(state: PydanticState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(PydanticState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + graph = builder.compile() + + result = graph.invoke({"value": "x", "items": []}, version="v2") + assert isinstance(result, GraphOutput) + assert isinstance(result.value, PydanticState) + assert result.value.value == "x_a" + assert result.interrupts == () + + def test_invoke_v2_dataclass_state(self) -> None: + """invoke with v2 and dataclass state returns GraphOutput with dataclass value.""" + + def node_a(state: DataclassState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(DataclassState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + graph = builder.compile() + + result = graph.invoke({"value": "x", "items": []}, version="v2") + assert isinstance(result, GraphOutput) + assert isinstance(result.value, DataclassState) + assert result.value.value == "x_a" + assert result.value.items == ["a"] + assert result.interrupts == () + + def test_invoke_v2_non_values_mode_pydantic(self) -> None: + """invoke with v2 + non-values mode + pydantic state returns list[StreamPart].""" + + def node_a(state: PydanticState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(PydanticState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + graph = builder.compile() + + result = graph.invoke( + {"value": "x", "items": []}, stream_mode="updates", version="v2" + ) + assert isinstance(result, list) + for chunk in result: + _assert_stream_part_shape(chunk) + assert chunk["type"] == "updates" + # updates data should be plain dicts, not coerced to pydantic + assert isinstance(chunk["data"], dict) + + def test_updates_mode(self) -> None: + graph = _make_simple_graph().compile() + result = graph.invoke(_SIMPLE_INPUT, stream_mode="updates", version="v2") + assert isinstance(result, list) and len(result) == 2 + for chunk in result: + _assert_stream_part_shape(chunk) + assert chunk["type"] == "updates" + assert chunk["ns"] == () + assert "node_a" in result[0]["data"] + assert "node_b" in result[1]["data"] + + def test_multiple_modes(self) -> None: + graph = _make_simple_graph().compile() + modes: Any = ["values", "updates"] + result = graph.invoke(_SIMPLE_INPUT, stream_mode=modes, version="v2") + assert isinstance(result, list) + types_seen = {c["type"] for c in result} + assert {"values", "updates"} <= types_seen + for c in result: + _assert_stream_part_shape(c) + + def test_v1_default_unchanged(self) -> None: + graph = _make_simple_graph().compile() + result = graph.invoke(_SIMPLE_INPUT) + assert isinstance(result, dict) + assert result["value"] == "x_a_b" + assert result["items"] == ["a", "b"] + + def test_v1_updates_unchanged(self) -> None: + graph = _make_simple_graph().compile() + result = graph.invoke(_SIMPLE_INPUT, stream_mode="updates") + assert isinstance(result, list) + for chunk in result: + assert "node_a" in chunk or "node_b" in chunk + + +# --- v2 async stream --- + + +class TestV2StreamAsync: + @pytest.mark.anyio + async def test_values(self) -> None: + graph = _make_simple_graph().compile() + chunks = [ + c + async for c in graph.astream( + _SIMPLE_INPUT, stream_mode="values", version="v2" + ) + ] + assert len(chunks) >= 1 + for c in chunks: + _assert_stream_part_shape(c) + assert c["type"] == "values" + assert c["ns"] == () + assert isinstance(c["data"], dict) + + @pytest.mark.anyio + async def test_updates(self) -> None: + graph = _make_simple_graph().compile() + chunks = [ + c + async for c in graph.astream( + _SIMPLE_INPUT, stream_mode="updates", version="v2" + ) + ] + assert len(chunks) == 2 + for c in chunks: + _assert_stream_part_shape(c) + assert c["type"] == "updates" + assert c["ns"] == () + assert "node_a" in chunks[0]["data"] + assert "node_b" in chunks[1]["data"] + + @pytest.mark.anyio + async def test_messages(self) -> None: + graph = _make_messages_graph().compile() + chunks = [ + c + async for c in graph.astream( + _MSG_INPUT, stream_mode="messages", version="v2" + ) + ] + msg_chunks = [c for c in chunks if c["type"] == "messages"] + assert len(msg_chunks) >= 1 + for c in msg_chunks: + _assert_stream_part_shape(c) + assert c["ns"] == () + data = c["data"] + assert isinstance(data, tuple) and len(data) == 2 + message, metadata = data + assert isinstance(message, BaseMessage) + assert isinstance(metadata, dict) + assert "langgraph_node" in metadata + + @NEEDS_CONTEXTVARS + @pytest.mark.anyio + async def test_custom(self) -> None: + graph = _make_custom_graph() + chunks = [ + c + async for c in graph.astream( + {"key": "val"}, stream_mode="custom", version="v2" + ) + ] + custom = [c for c in chunks if c["type"] == "custom"] + assert len(custom) == 2 + for c in custom: + _assert_stream_part_shape(c) + assert custom[0]["data"] == "hello" + assert custom[1]["data"] == 42 + + @pytest.mark.anyio + async def test_multiple_modes(self) -> None: + graph = _make_simple_graph().compile() + chunks = [ + c + async for c in graph.astream( + _SIMPLE_INPUT, + stream_mode=["values", "updates"], + version="v2", + ) + ] + types_seen = {c["type"] for c in chunks} + assert {"values", "updates"} <= types_seen + for c in chunks: + _assert_stream_part_shape(c) + + @pytest.mark.anyio + async def test_subgraphs_ns(self) -> None: + outer = _make_subgraph() + chunks = [ + c + async for c in outer.astream( + _SIMPLE_INPUT, + stream_mode="updates", + subgraphs=True, + version="v2", + ) + ] + for c in chunks: + _assert_stream_part_shape(c) + root = [c for c in chunks if c["ns"] == ()] + sub = [c for c in chunks if c["ns"] != ()] + assert len(root) >= 1 + assert len(sub) >= 1 + + @pytest.mark.anyio + async def test_checkpoints(self) -> None: + graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) + config: Any = {"configurable": {"thread_id": "test-v2-ckpt-async"}} + chunks = [ + c + async for c in graph.astream( + _SIMPLE_INPUT, + config, + stream_mode="checkpoints", + version="v2", + ) + ] + ckpt = [c for c in chunks if c["type"] == "checkpoints"] + assert len(ckpt) >= 1 + for c in ckpt: + _assert_stream_part_shape(c) + assert c["ns"] == () + payload = c["data"] + assert {"config", "metadata", "values", "next", "tasks"} <= payload.keys() + + @pytest.mark.anyio + async def test_tasks(self) -> None: + graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) + config: Any = {"configurable": {"thread_id": "test-v2-tasks-async"}} + chunks = [ + c + async for c in graph.astream( + _SIMPLE_INPUT, + config, + stream_mode="tasks", + version="v2", + ) + ] + tasks = [c for c in chunks if c["type"] == "tasks"] + assert len(tasks) >= 2 + for c in tasks: + _assert_stream_part_shape(c) + assert c["ns"] == () + assert "id" in c["data"] and "name" in c["data"] + starts = [c for c in tasks if "triggers" in c["data"]] + results = [c for c in tasks if "result" in c["data"]] + assert len(starts) >= 2 + assert len(results) >= 2 + + @pytest.mark.anyio + async def test_debug(self) -> None: + graph = _make_simple_graph().compile(checkpointer=InMemorySaver()) + config: Any = {"configurable": {"thread_id": "test-v2-debug-async"}} + chunks = [ + c + async for c in graph.astream( + _SIMPLE_INPUT, + config, + stream_mode="debug", + version="v2", + ) + ] + debug = [c for c in chunks if c["type"] == "debug"] + assert len(debug) >= 1 + for c in debug: + _assert_stream_part_shape(c) + assert c["ns"] == () + envelope = c["data"] + assert {"step", "timestamp", "type", "payload"} <= envelope.keys() + assert envelope["type"] in ("checkpoint", "task", "task_result") + + +# --- v2 async invoke --- + + +class TestV2InvokeAsync: + @pytest.mark.anyio + async def test_values_default(self) -> None: + graph = _make_simple_graph().compile() + result = await graph.ainvoke(_SIMPLE_INPUT, version="v2") + assert isinstance(result, GraphOutput) + assert result.value == {"value": "x_a_b", "items": ["a", "b"]} + assert result.interrupts == () + # backward compat dict access + assert result["value"] == "x_a_b" + assert result["items"] == ["a", "b"] + + @NEEDS_CONTEXTVARS + @pytest.mark.anyio + async def test_ainvoke_v2_graph_output_with_interrupts(self) -> None: + def my_node(state: SimpleState) -> dict[str, Any]: + answer = interrupt("what is your name?") + return {"value": answer, "items": ["done"]} + + builder: StateGraph = StateGraph(SimpleState) + builder.add_node("my_node", my_node) + builder.add_edge(START, "my_node") + builder.add_edge("my_node", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + config: Any = {"configurable": {"thread_id": "test-ainvoke-v2-interrupts"}} + result = await graph.ainvoke({"value": "x", "items": []}, config, version="v2") + assert isinstance(result, GraphOutput) + assert len(result.interrupts) > 0 + for intr in result.interrupts: + assert isinstance(intr, Interrupt) + assert isinstance(result.value, dict) + + @pytest.mark.anyio + async def test_ainvoke_v2_pydantic_state(self) -> None: + """ainvoke with v2 and pydantic state returns GraphOutput with pydantic value.""" + + def node_a(state: PydanticState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(PydanticState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + graph = builder.compile() + + result = await graph.ainvoke({"value": "x", "items": []}, version="v2") + assert isinstance(result, GraphOutput) + assert isinstance(result.value, PydanticState) + assert result.value.value == "x_a" + assert result.value.items == ["a"] + assert result.interrupts == () + + @pytest.mark.anyio + async def test_ainvoke_v2_dataclass_state(self) -> None: + """ainvoke with v2 and dataclass state returns GraphOutput with dataclass value.""" + + def node_a(state: DataclassState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(DataclassState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + graph = builder.compile() + + result = await graph.ainvoke({"value": "x", "items": []}, version="v2") + assert isinstance(result, GraphOutput) + assert isinstance(result.value, DataclassState) + assert result.value.value == "x_a" + assert result.value.items == ["a"] + assert result.interrupts == () + + @pytest.mark.anyio + async def test_ainvoke_v2_graph_output_no_interrupts(self) -> None: + graph = _make_simple_graph().compile() + result = await graph.ainvoke(_SIMPLE_INPUT, version="v2") + assert isinstance(result, GraphOutput) + assert result.interrupts == () + assert INTERRUPT not in result + + @pytest.mark.anyio + async def test_updates_mode(self) -> None: + graph = _make_simple_graph().compile() + result = await graph.ainvoke(_SIMPLE_INPUT, stream_mode="updates", version="v2") + assert isinstance(result, list) and len(result) == 2 + for chunk in result: + _assert_stream_part_shape(chunk) + assert chunk["type"] == "updates" + assert chunk["ns"] == () + assert "node_a" in result[0]["data"] + assert "node_b" in result[1]["data"] + + @pytest.mark.anyio + async def test_multiple_modes(self) -> None: + graph = _make_simple_graph().compile() + modes: Any = ["values", "updates"] + result = await graph.ainvoke(_SIMPLE_INPUT, stream_mode=modes, version="v2") + assert isinstance(result, list) + types_seen = {c["type"] for c in result} + assert {"values", "updates"} <= types_seen + for c in result: + _assert_stream_part_shape(c) + + +# --- type-safe streaming: coercion + interrupt separation --- + + +class PydanticState(BaseModel): + value: str + items: Annotated[list[str], operator.add] + + +@dataclass +class DataclassState: + value: str + items: Annotated[list[str], operator.add] + + +class TestV2TypeSafeStreaming: + """Test that v2 streaming coerces values to pydantic/dataclass instances + and separates interrupts into a dedicated field.""" + + def test_values_pydantic_state(self) -> None: + """v2 values + pydantic state -> data is pydantic model instance.""" + + def node_a(state: PydanticState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(PydanticState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + graph = builder.compile() + + chunks = list( + graph.stream( + {"value": "x", "items": []}, + stream_mode="values", + version="v2", + ) + ) + assert len(chunks) >= 1 + for c in chunks: + _assert_stream_part_shape(c) + assert c["type"] == "values" + assert isinstance(c["data"], PydanticState), ( + f"Expected PydanticState, got {type(c['data'])}" + ) + assert c["interrupts"] == () + + def test_values_dataclass_state(self) -> None: + """v2 values + dataclass state -> data is dataclass instance.""" + + def node_a(state: DataclassState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(DataclassState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + graph = builder.compile() + + chunks = list( + graph.stream( + {"value": "x", "items": []}, + stream_mode="values", + version="v2", + ) + ) + assert len(chunks) >= 1 + for c in chunks: + _assert_stream_part_shape(c) + assert c["type"] == "values" + assert isinstance(c["data"], DataclassState), ( + f"Expected DataclassState, got {type(c['data'])}" + ) + + def test_values_typeddict_state(self) -> None: + """v2 values + TypedDict state -> data stays plain dict (no coercion).""" + graph = _make_simple_graph().compile() + chunks = list(graph.stream(_SIMPLE_INPUT, stream_mode="values", version="v2")) + assert len(chunks) >= 1 + for c in chunks: + _assert_stream_part_shape(c) + assert c["type"] == "values" + # TypedDict state should remain a plain dict + assert isinstance(c["data"], dict) + assert type(c["data"]) is dict + + def test_values_interrupt_v2(self) -> None: + """v2 values + interrupt -> interrupts in typed field, not in data.""" + + def my_node(state: SimpleState) -> dict[str, Any]: + answer = interrupt("what is your name?") + return {"value": answer, "items": ["done"]} + + builder: StateGraph = StateGraph(SimpleState) + builder.add_node("my_node", my_node) + builder.add_edge(START, "my_node") + builder.add_edge("my_node", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + config: Any = {"configurable": {"thread_id": "test-v2-interrupt"}} + chunks = list( + graph.stream( + {"value": "x", "items": []}, + config, + stream_mode="values", + version="v2", + ) + ) + # should have at least one values chunk with interrupts + interrupt_chunks = [c for c in chunks if c.get("interrupts", ())] + assert len(interrupt_chunks) >= 1, f"Expected interrupt chunks, got {chunks}" + for c in interrupt_chunks: + assert c["type"] == "values" + assert isinstance(c["interrupts"], tuple) + assert len(c["interrupts"]) > 0 + for intr in c["interrupts"]: + assert isinstance(intr, Interrupt) + # __interrupt__ should NOT be in data + if isinstance(c["data"], dict): + assert INTERRUPT not in c["data"] + + def test_values_interrupt_v1_compat(self) -> None: + """v1 values + interrupt -> __interrupt__ still in dict (v1 compat).""" + + def my_node(state: SimpleState) -> dict[str, Any]: + answer = interrupt("what is your name?") + return {"value": answer, "items": ["done"]} + + builder: StateGraph = StateGraph(SimpleState) + builder.add_node("my_node", my_node) + builder.add_edge(START, "my_node") + builder.add_edge("my_node", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + config: Any = {"configurable": {"thread_id": "test-v1-interrupt-compat"}} + chunks = list( + graph.stream( + {"value": "x", "items": []}, + config, + stream_mode="values", + ) + ) + # v1 format: should have __interrupt__ in dict + interrupt_chunks = [c for c in chunks if isinstance(c, dict) and INTERRUPT in c] + assert len(interrupt_chunks) >= 1, ( + f"Expected v1 interrupt chunks with {INTERRUPT}, got {chunks}" + ) + + def test_checkpoints_pydantic_state(self) -> None: + """v2 checkpoints + pydantic state -> values is pydantic model instance + (at least for checkpoints emitted after all channels are populated).""" + + def node_a(state: PydanticState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(PydanticState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + config: Any = {"configurable": {"thread_id": "test-v2-ckpt-pydantic"}} + chunks = list( + graph.stream( + {"value": "x", "items": []}, + config, + stream_mode="checkpoints", + version="v2", + ) + ) + ckpt_chunks = [c for c in chunks if c["type"] == "checkpoints"] + assert len(ckpt_chunks) >= 1 + # At least one checkpoint (after first node runs) should have coerced values + coerced_ckpts = [ + c for c in ckpt_chunks if isinstance(c["data"]["values"], PydanticState) + ] + assert len(coerced_ckpts) >= 1, ( + f"Expected at least one checkpoint with PydanticState values, got types: " + f"{[type(c['data']['values']) for c in ckpt_chunks]}" + ) + + def test_debug_pydantic_state(self) -> None: + """v2 debug + pydantic state -> inner checkpoint payload has coerced values + (at least for checkpoints emitted after all channels are populated).""" + + def node_a(state: PydanticState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(PydanticState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + config: Any = {"configurable": {"thread_id": "test-v2-debug-pydantic"}} + chunks = list( + graph.stream( + {"value": "x", "items": []}, + config, + stream_mode="debug", + version="v2", + ) + ) + debug_chunks = [c for c in chunks if c["type"] == "debug"] + checkpoint_debug = [ + c for c in debug_chunks if c["data"]["type"] == "checkpoint" + ] + assert len(checkpoint_debug) >= 1 + # At least one debug checkpoint should have coerced values + coerced_debug = [ + c + for c in checkpoint_debug + if isinstance(c["data"]["payload"]["values"], PydanticState) + ] + assert len(coerced_debug) >= 1, ( + f"Expected at least one debug checkpoint with PydanticState values, got types: " + f"{[type(c['data']['payload']['values']) for c in checkpoint_debug]}" + ) + + def test_values_pydantic_interrupt(self) -> None: + """v2 values + pydantic state + interrupt -> data is model, interrupts separated.""" + + def my_node(state: PydanticState) -> dict[str, Any]: + answer = interrupt("what is your name?") + return {"value": answer, "items": ["done"]} + + builder: StateGraph = StateGraph(PydanticState) + builder.add_node("my_node", my_node) + builder.add_edge(START, "my_node") + builder.add_edge("my_node", END) + graph = builder.compile(checkpointer=InMemorySaver()) + + config: Any = {"configurable": {"thread_id": "test-v2-pydantic-interrupt"}} + chunks = list( + graph.stream( + {"value": "x", "items": []}, + config, + stream_mode="values", + version="v2", + ) + ) + interrupt_chunks = [c for c in chunks if c.get("interrupts", ())] + assert len(interrupt_chunks) >= 1 + for c in interrupt_chunks: + assert isinstance(c["data"], PydanticState), ( + f"Expected PydanticState, got {type(c['data'])}" + ) + assert isinstance(c["interrupts"], tuple) + assert len(c["interrupts"]) > 0 + + def test_subgraph_different_pydantic_schema(self) -> None: + """Subgraph with different pydantic schema -> subgraph data coerced with subgraph's schema.""" + + class InnerState(BaseModel): + value: str + + class OuterState(BaseModel): + value: str + + def inner_node(state: InnerState) -> dict[str, Any]: + return {"value": state.value + "_inner"} + + def outer_node(state: OuterState) -> dict[str, Any]: + return {"value": state.value + "_outer"} + + inner_builder: StateGraph = StateGraph(InnerState) + inner_builder.add_node("inner_node", inner_node) + inner_builder.add_edge(START, "inner_node") + inner_builder.add_edge("inner_node", END) + inner_graph = inner_builder.compile() + + outer_builder: StateGraph = StateGraph(OuterState) + outer_builder.add_node("outer_node", outer_node) + outer_builder.add_node("inner", inner_graph) + outer_builder.add_edge(START, "outer_node") + outer_builder.add_edge("outer_node", "inner") + outer_builder.add_edge("inner", END) + outer = outer_builder.compile() + + chunks = list( + outer.stream( + {"value": "x"}, + stream_mode="values", + subgraphs=True, + version="v2", + ) + ) + # Root-level values should be OuterState instances + root_values = [c for c in chunks if c["type"] == "values" and c["ns"] == ()] + assert len(root_values) >= 1 + for c in root_values: + assert isinstance(c["data"], OuterState), ( + f"Expected OuterState, got {type(c['data'])}" + ) + # Subgraph values are streamed from the subgraph's own stream() + # which runs with default version="v1", so no coercion + sub_values = [c for c in chunks if c["type"] == "values" and c["ns"] != ()] + assert len(sub_values) >= 1 + + +# --- v2 validation errors --- + + +def _make_pydantic_graph() -> Any: + """Build a simple graph with PydanticState for validation error tests.""" + + def node_a(state: PydanticState) -> dict[str, Any]: + return {"value": state.value + "_a", "items": ["a"]} + + builder: StateGraph = StateGraph(PydanticState) + builder.add_node("node_a", node_a) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", END) + return builder.compile() + + +class TestV2ValidationErrors: + """Validation errors propagate for pydantic state in both v1 and v2. + + Uses `value=[1, 2, 3]` which channels accept (LastValue stores anything) + but pydantic rejects (list is not coercible to str even in lax mode). + """ + + _INVALID_INPUT: dict[str, Any] = {"value": [1, 2, 3], "items": []} + + def test_stream_v2_pydantic_validation_error(self) -> None: + """Invalid input to stream with v2 + pydantic state raises ValidationError.""" + graph = _make_pydantic_graph() + with pytest.raises(ValidationError): + list( + graph.stream( + self._INVALID_INPUT, + stream_mode="values", + version="v2", + ) + ) + + def test_invoke_v2_pydantic_validation_error(self) -> None: + """Invalid input to invoke with v2 + pydantic state raises ValidationError.""" + graph = _make_pydantic_graph() + with pytest.raises(ValidationError): + graph.invoke(self._INVALID_INPUT, version="v2") + + def test_invoke_v1_pydantic_validation_error(self) -> None: + """Regression: invalid input to invoke without version raises ValidationError.""" + graph = _make_pydantic_graph() + with pytest.raises(ValidationError): + graph.invoke(self._INVALID_INPUT) + + +# --- type narrowing compile-time checks --- +# These assert_type calls verify that mypy narrows the union correctly. + + +_OutputT = TypeVar("_OutputT") +_StateT = TypeVar("_StateT") + + +def _check_type_narrowing(part: StreamPart[_OutputT, _StateT]) -> None: + """Compile-time type narrowing checks — never called at runtime.""" + if part["type"] == "values": + assert_type(part, ValuesStreamPart[_OutputT]) + elif part["type"] == "updates": + assert_type(part, UpdatesStreamPart) + assert_type(part["data"], dict[str, Any]) + elif part["type"] == "messages": + assert_type(part, MessagesStreamPart) + elif part["type"] == "custom": + assert_type(part, CustomStreamPart) + elif part["type"] == "checkpoints": + assert_type(part, CheckpointStreamPart[_StateT]) + assert_type(part["data"], CheckpointPayload[_StateT]) + elif part["type"] == "tasks": + assert_type(part, TasksStreamPart) + assert_type(part["data"], TaskPayload | TaskResultPayload) + elif part["type"] == "debug": + assert_type(part, DebugStreamPart[_StateT]) + assert_type(part["data"], DebugPayload[_StateT]) + assert_type(part["ns"], tuple[str, ...]) diff --git a/libs/sdk-py/langgraph_sdk/_async/runs.py b/libs/sdk-py/langgraph_sdk/_async/runs.py index 028c2e71a..c6d1fa162 100644 --- a/libs/sdk-py/langgraph_sdk/_async/runs.py +++ b/libs/sdk-py/langgraph_sdk/_async/runs.py @@ -5,12 +5,15 @@ from __future__ import annotations import builtins import warnings from collections.abc import AsyncIterator, Callable, Mapping, Sequence -from typing import Any, overload +from typing import Any, Literal, overload import httpx from langgraph_sdk._async.http import HttpClient -from langgraph_sdk._shared.utilities import _get_run_metadata_from_response +from langgraph_sdk._shared.utilities import ( + _get_run_metadata_from_response, + _sse_to_v2_dict, +) from langgraph_sdk.schema import ( All, BulkCancelRunsStatus, @@ -33,9 +36,21 @@ from langgraph_sdk.schema import ( RunStatus, StreamMode, StreamPart, + StreamPartV2, + StreamVersion, ) +async def _wrap_stream_v2( + raw: AsyncIterator[StreamPart], +) -> AsyncIterator[StreamPartV2]: + """Wrap a raw SSE stream, converting each event to a v2 dict.""" + async for part in raw: + v2 = _sse_to_v2_dict(part.event, part.data) + if v2 is not None: + yield v2 + + class RunsClient: """Client for managing runs in LangGraph. @@ -81,6 +96,66 @@ class RunsClient: headers: Mapping[str, str] | None = None, params: QueryParamTypes | None = None, on_run_created: Callable[[RunCreateMetadata], None] | None = None, + version: Literal["v1"] = "v1", + ) -> AsyncIterator[StreamPart]: ... + + @overload + def stream( + self, + thread_id: str, + assistant_id: str, + *, + input: Input | None = None, + command: Command | None = None, + stream_mode: StreamMode | Sequence[StreamMode] = "values", + stream_subgraphs: bool = False, + stream_resumable: bool = False, + metadata: Mapping[str, Any] | None = None, + config: Config | None = None, + context: Context | None = None, + checkpoint: Checkpoint | None = None, + checkpoint_id: str | None = None, + checkpoint_during: bool | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + feedback_keys: Sequence[str] | None = None, + on_disconnect: DisconnectMode | None = None, + webhook: str | None = None, + multitask_strategy: MultitaskStrategy | None = None, + if_not_exists: IfNotExists | None = None, + after_seconds: int | None = None, + headers: Mapping[str, str] | None = None, + params: QueryParamTypes | None = None, + on_run_created: Callable[[RunCreateMetadata], None] | None = None, + version: Literal["v2"], + ) -> AsyncIterator[StreamPartV2]: ... + + @overload + def stream( + self, + thread_id: None, + assistant_id: str, + *, + input: Input | None = None, + command: Command | None = None, + stream_mode: StreamMode | Sequence[StreamMode] = "values", + stream_subgraphs: bool = False, + stream_resumable: bool = False, + metadata: Mapping[str, Any] | None = None, + config: Config | None = None, + checkpoint_during: bool | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + feedback_keys: Sequence[str] | None = None, + on_disconnect: DisconnectMode | None = None, + on_completion: OnCompletionBehavior | None = None, + if_not_exists: IfNotExists | None = None, + webhook: str | None = None, + after_seconds: int | None = None, + headers: Mapping[str, str] | None = None, + params: QueryParamTypes | None = None, + on_run_created: Callable[[RunCreateMetadata], None] | None = None, + version: Literal["v1"] = "v1", ) -> AsyncIterator[StreamPart]: ... @overload @@ -108,7 +183,8 @@ class RunsClient: headers: Mapping[str, str] | None = None, params: QueryParamTypes | None = None, on_run_created: Callable[[RunCreateMetadata], None] | None = None, - ) -> AsyncIterator[StreamPart]: ... + version: Literal["v2"], + ) -> AsyncIterator[StreamPartV2]: ... def stream( self, @@ -139,7 +215,8 @@ class RunsClient: params: QueryParamTypes | None = None, on_run_created: Callable[[RunCreateMetadata], None] | None = None, durability: Durability | None = None, - ) -> AsyncIterator[StreamPart]: + version: StreamVersion = "v1", + ) -> AsyncIterator[StreamPart | StreamPartV2]: """Create a run and stream the results. Args: @@ -180,6 +257,8 @@ class RunsClient: "async" means checkpoints are persisted async while next graph step executes, replaces checkpoint_during=True "sync" means checkpoints are persisted sync after graph step executes, replaces checkpoint_during=False "exit" means checkpoints are only persisted when the run exits, does not save intermediate steps + version: Stream format version. "v1" (default) returns raw SSE StreamPart + NamedTuples. "v2" returns typed dicts with `type`, `ns`, and `data` keys. Returns: Asynchronous iterator of stream results. @@ -222,7 +301,7 @@ class RunsClient: stacklevel=2, ) - payload = { + payload: dict[str, Any] = { "input": input, "command": ( {k: v for k, v in command.items() if v is not None} if command else None @@ -259,7 +338,7 @@ class RunsClient: if on_run_created and (metadata := _get_run_metadata_from_response(res)): on_run_created(metadata) - return self.http.stream( + raw = self.http.stream( endpoint, "POST", json={k: v for k, v in payload.items() if v is not None}, @@ -267,6 +346,9 @@ class RunsClient: headers=headers, on_response=on_response if on_run_created else None, ) + if version == "v2": + return _wrap_stream_v2(raw) + return raw @overload async def create( diff --git a/libs/sdk-py/langgraph_sdk/_shared/utilities.py b/libs/sdk-py/langgraph_sdk/_shared/utilities.py index 02f28fd9a..54d55580c 100644 --- a/libs/sdk-py/langgraph_sdk/_shared/utilities.py +++ b/libs/sdk-py/langgraph_sdk/_shared/utilities.py @@ -107,6 +107,24 @@ def _get_run_metadata_from_response( return None +def _sse_to_v2_dict(event: str, data: Any) -> dict[str, Any] | None: + """Convert an SSE event+data pair into a v2 stream part dict. + + Returns None for ``end`` events (signals end of stream). + """ + if event == "end": + return None + parts = event.split("|") + event_type = parts[0] + ns = parts[1:] if len(parts) > 1 else [] + result: dict[str, Any] = {"type": event_type, "ns": ns, "data": data} + if event_type == "values" and isinstance(data, dict): + result["interrupts"] = data.pop("__interrupt__", []) + else: + result["interrupts"] = [] + return result + + def _provided_vals(d: Mapping[str, Any]) -> dict[str, Any]: return {k: v for k, v in d.items() if v is not None} diff --git a/libs/sdk-py/langgraph_sdk/_sync/runs.py b/libs/sdk-py/langgraph_sdk/_sync/runs.py index febd017c0..1b52b8113 100644 --- a/libs/sdk-py/langgraph_sdk/_sync/runs.py +++ b/libs/sdk-py/langgraph_sdk/_sync/runs.py @@ -5,11 +5,14 @@ from __future__ import annotations import builtins import warnings from collections.abc import Callable, Iterator, Mapping, Sequence -from typing import Any, overload +from typing import Any, Literal, overload import httpx -from langgraph_sdk._shared.utilities import _get_run_metadata_from_response +from langgraph_sdk._shared.utilities import ( + _get_run_metadata_from_response, + _sse_to_v2_dict, +) from langgraph_sdk._sync.http import SyncHttpClient from langgraph_sdk.schema import ( All, @@ -33,9 +36,21 @@ from langgraph_sdk.schema import ( RunStatus, StreamMode, StreamPart, + StreamPartV2, + StreamVersion, ) +def _wrap_stream_v2_sync( + raw: Iterator[StreamPart], +) -> Iterator[StreamPartV2]: + """Wrap a raw SSE stream, converting each event to a v2 dict.""" + for part in raw: + v2 = _sse_to_v2_dict(part.event, part.data) + if v2 is not None: + yield v2 + + class SyncRunsClient: """Synchronous client for managing runs in LangGraph. @@ -80,6 +95,66 @@ class SyncRunsClient: headers: Mapping[str, str] | None = None, params: QueryParamTypes | None = None, on_run_created: Callable[[RunCreateMetadata], None] | None = None, + version: Literal["v1"] = "v1", + ) -> Iterator[StreamPart]: ... + + @overload + def stream( + self, + thread_id: str, + assistant_id: str, + *, + input: Input | None = None, + command: Command | None = None, + stream_mode: StreamMode | Sequence[StreamMode] = "values", + stream_subgraphs: bool = False, + metadata: Mapping[str, Any] | None = None, + config: Config | None = None, + context: Context | None = None, + checkpoint: Checkpoint | None = None, + checkpoint_id: str | None = None, + checkpoint_during: bool | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + feedback_keys: Sequence[str] | None = None, + on_disconnect: DisconnectMode | None = None, + webhook: str | None = None, + multitask_strategy: MultitaskStrategy | None = None, + if_not_exists: IfNotExists | None = None, + after_seconds: int | None = None, + headers: Mapping[str, str] | None = None, + params: QueryParamTypes | None = None, + on_run_created: Callable[[RunCreateMetadata], None] | None = None, + version: Literal["v2"], + ) -> Iterator[StreamPartV2]: ... + + @overload + def stream( + self, + thread_id: None, + assistant_id: str, + *, + input: Input | None = None, + command: Command | None = None, + stream_mode: StreamMode | Sequence[StreamMode] = "values", + stream_subgraphs: bool = False, + stream_resumable: bool = False, + metadata: Mapping[str, Any] | None = None, + config: Config | None = None, + context: Context | None = None, + checkpoint_during: bool | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + feedback_keys: Sequence[str] | None = None, + on_disconnect: DisconnectMode | None = None, + on_completion: OnCompletionBehavior | None = None, + if_not_exists: IfNotExists | None = None, + webhook: str | None = None, + after_seconds: int | None = None, + headers: Mapping[str, str] | None = None, + params: QueryParamTypes | None = None, + on_run_created: Callable[[RunCreateMetadata], None] | None = None, + version: Literal["v1"] = "v1", ) -> Iterator[StreamPart]: ... @overload @@ -108,7 +183,8 @@ class SyncRunsClient: headers: Mapping[str, str] | None = None, params: QueryParamTypes | None = None, on_run_created: Callable[[RunCreateMetadata], None] | None = None, - ) -> Iterator[StreamPart]: ... + version: Literal["v2"], + ) -> Iterator[StreamPartV2]: ... def stream( self, @@ -139,7 +215,8 @@ class SyncRunsClient: params: QueryParamTypes | None = None, on_run_created: Callable[[RunCreateMetadata], None] | None = None, durability: Durability | None = None, - ) -> Iterator[StreamPart]: + version: StreamVersion = "v1", + ) -> Iterator[StreamPart | StreamPartV2]: """Create a run and stream the results. Args: @@ -179,7 +256,8 @@ class SyncRunsClient: "async" means checkpoints are persisted async while next graph step executes, replaces checkpoint_during=True "sync" means checkpoints are persisted sync after graph step executes, replaces checkpoint_during=False "exit" means checkpoints are only persisted when the run exits, does not save intermediate steps - + version: Stream format version. "v1" (default) returns raw SSE StreamPart + NamedTuples. "v2" returns typed dicts with `type`, `ns`, and `data` keys. Returns: Iterator of stream results. @@ -218,7 +296,7 @@ class SyncRunsClient: DeprecationWarning, stacklevel=2, ) - payload = { + payload: dict[str, Any] = { "input": input, "command": ( {k: v for k, v in command.items() if v is not None} if command else None @@ -255,7 +333,7 @@ class SyncRunsClient: if on_run_created and (metadata := _get_run_metadata_from_response(res)): on_run_created(metadata) - return self.http.stream( + raw = self.http.stream( endpoint, "POST", json={k: v for k, v in payload.items() if v is not None}, @@ -263,6 +341,9 @@ class SyncRunsClient: headers=headers, on_response=on_response if on_run_created else None, ) + if version == "v2": + return _wrap_stream_v2_sync(raw) + return raw @overload def create( diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py index 163f13bb7..8f9a53149 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -588,6 +588,275 @@ class StreamPart(NamedTuple): """The ID of the event.""" +StreamVersion = Literal["v1", "v2"] +"""Stream format version. + +- `"v1"`: Traditional format — raw SSE `StreamPart` NamedTuples. +- `"v2"`: Each event is a typed dict with `type`, `ns`, and `data` keys. +""" + + +# --- Typed payload dicts (JSON-deserialized from the server) --- + + +class TaskPayload(TypedDict): + """Payload for a task start event.""" + + id: str + """Unique identifier for this task.""" + name: str + """Name of the node being executed.""" + input: Any + """Input data passed to the task.""" + triggers: list[str] + """List of triggers that caused this task to be executed (e.g. channel writes).""" + + +class TaskResultPayload(TypedDict): + """Payload for a task result event.""" + + id: str + """Unique identifier for this task.""" + name: str + """Name of the node that was executed.""" + error: str | None + """Error message if the task failed, otherwise `None`.""" + interrupts: list[dict[str, Any]] + """List of interrupts that occurred during task execution.""" + result: dict[str, Any] + """Mapping of channel names to the values written by this task.""" + + +class CheckpointTaskPayload(TypedDict): + """A task entry within a `CheckpointPayload`. + + The keys present depend on the task's state: + + - **Error:** `id`, `name`, `error`, `state` + - **Has result:** `id`, `name`, `result`, `interrupts`, `state` + - **Pending:** `id`, `name`, `interrupts`, `state` + """ + + id: str + """Unique identifier for this task.""" + name: str + """Name of the node being executed.""" + error: NotRequired[str] + """Error message, present only if the task failed.""" + result: NotRequired[Any] + """Result of the task, present only if the task completed successfully.""" + interrupts: NotRequired[list[dict[str, Any]]] + """List of interrupts, present when the task has been interrupted or completed.""" + state: dict[str, Any] | None + """Snapshot of the subgraph state. `None` if not a subgraph.""" + + +class CheckpointPayload(TypedDict): + """Payload for a checkpoint event.""" + + config: dict[str, Any] | None + """Configuration for this checkpoint, including the `thread_id` and `checkpoint_id`.""" + metadata: dict[str, Any] + """Metadata associated with this checkpoint (e.g. step number, source, writes).""" + values: dict[str, Any] + """Current state values at the time of this checkpoint.""" + next: list[str] + """Names of the nodes scheduled to execute next.""" + parent_config: dict[str, Any] | None + """Configuration of the parent checkpoint, or `None` if this is the first checkpoint.""" + tasks: list[CheckpointTaskPayload] + """List of tasks associated with this checkpoint.""" + + +class _DebugCheckpointPayload(TypedDict): + step: int + """The step number in the graph execution.""" + timestamp: str + """ISO 8601 timestamp of when this event occurred.""" + type: Literal["checkpoint"] + """Event type discriminator, always `"checkpoint"`.""" + payload: CheckpointPayload + """The checkpoint payload.""" + + +class _DebugTaskPayload(TypedDict): + step: int + """The step number in the graph execution.""" + timestamp: str + """ISO 8601 timestamp of when this event occurred.""" + type: Literal["task"] + """Event type discriminator, always `"task"`.""" + payload: TaskPayload + """The task start payload.""" + + +class _DebugTaskResultPayload(TypedDict): + step: int + """The step number in the graph execution.""" + timestamp: str + """ISO 8601 timestamp of when this event occurred.""" + type: Literal["task_result"] + """Event type discriminator, always `"task_result"`.""" + payload: TaskResultPayload + """The task result payload.""" + + +DebugPayload = _DebugCheckpointPayload | _DebugTaskPayload | _DebugTaskResultPayload +"""Wrapper payload for debug events. Discriminate on `type`.""" + + +class RunMetadataPayload(TypedDict): + """Payload for the `metadata` control event.""" + + run_id: str + """The unique identifier of the run.""" + + +# --- v2 stream part TypedDicts --- + + +class ValuesStreamPart(TypedDict): + """Stream part emitted for `stream_mode="values"`.""" + + type: Literal["values"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: dict[str, Any] + """Full state values after the step.""" + interrupts: list[dict[str, Any]] + """List of interrupts that occurred during this step.""" + + +class UpdatesStreamPart(TypedDict): + """Stream part emitted for `stream_mode="updates"`.""" + + type: Literal["updates"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: dict[str, Any] + """Mapping of node names to their outputs.""" + + +class MessagesPartialStreamPart(TypedDict): + """Stream part emitted for partial message chunks (`messages/partial`).""" + + type: Literal["messages/partial"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: list[dict[str, Any]] + """List of partial message chunk dicts.""" + + +class MessagesCompleteStreamPart(TypedDict): + """Stream part emitted for complete messages (`messages/complete`).""" + + type: Literal["messages/complete"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: list[dict[str, Any]] + """List of complete message dicts.""" + + +class MessagesMetadataStreamPart(TypedDict): + """Stream part emitted for message metadata (`messages/metadata`).""" + + type: Literal["messages/metadata"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: dict[str, Any] + """Metadata dict for the message (e.g. `langgraph_step`, `langgraph_node`).""" + + +class MessagesTupleStreamPart(TypedDict): + """Stream part emitted for `stream_mode="messages"` (raw message+metadata pair).""" + + type: Literal["messages"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: list[dict[str, Any]] + """Two-element list of `[message_dict, metadata_dict]`.""" + + +class CustomStreamPart(TypedDict): + """Stream part emitted for `stream_mode="custom"`.""" + + type: Literal["custom"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: Any + """User-defined data passed to `StreamWriter` inside a node.""" + + +class CheckpointsStreamPart(TypedDict): + """Stream part emitted for `stream_mode="checkpoints"`.""" + + type: Literal["checkpoints"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: CheckpointPayload + """The checkpoint payload.""" + + +class TasksStreamPart(TypedDict): + """Stream part emitted for `stream_mode="tasks"`.""" + + type: Literal["tasks"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: TaskPayload | TaskResultPayload + """Task start or task result payload.""" + + +class DebugStreamPart(TypedDict): + """Stream part emitted for `stream_mode="debug"`.""" + + type: Literal["debug"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path of the emitting node (empty for root graph).""" + data: DebugPayload + """The debug event payload.""" + + +class MetadataStreamPart(TypedDict): + """Control event with `run_id` and other run metadata.""" + + type: Literal["metadata"] + """Stream part type discriminator.""" + ns: list[str] + """Namespace path (empty for root graph).""" + data: RunMetadataPayload + """The run metadata payload.""" + + +StreamPartV2 = ( + ValuesStreamPart + | UpdatesStreamPart + | MessagesPartialStreamPart + | MessagesCompleteStreamPart + | MessagesMetadataStreamPart + | MessagesTupleStreamPart + | CustomStreamPart + | CheckpointsStreamPart + | TasksStreamPart + | DebugStreamPart + | MetadataStreamPart +) +"""Discriminated union of all v2 stream part types. + +Use `part["type"]` to narrow the type. +""" + + class Send(TypedDict): """Represents a message to be sent to a specific node in the graph. diff --git a/libs/sdk-py/tests/test_client_stream.py b/libs/sdk-py/tests/test_client_stream.py index b2da417f1..9ac455f9f 100644 --- a/libs/sdk-py/tests/test_client_stream.py +++ b/libs/sdk-py/tests/test_client_stream.py @@ -2,18 +2,39 @@ from __future__ import annotations from collections.abc import Iterator, Sequence from pathlib import Path +from typing import Any import httpx import pytest +from typing_extensions import assert_type +from langgraph_sdk._shared.utilities import _sse_to_v2_dict from langgraph_sdk.client import HttpClient, SyncHttpClient -from langgraph_sdk.schema import StreamPart +from langgraph_sdk.schema import ( + CheckpointPayload, + CheckpointsStreamPart, + CustomStreamPart, + DebugPayload, + DebugStreamPart, + MetadataStreamPart, + RunMetadataPayload, + StreamPart, + StreamPartV2, + TaskPayload, + TaskResultPayload, + TasksStreamPart, + UpdatesStreamPart, + ValuesStreamPart, +) from langgraph_sdk.sse import BytesLike, BytesLineDecoder, SSEDecoder with open(Path(__file__).parent / "fixtures" / "response.txt", "rb") as f: RESPONSE_PAYLOAD = f.read() +# --- test helpers --- + + class AsyncListByteStream(httpx.AsyncByteStream): def __init__(self, chunks: Sequence[bytes], exc: Exception | None = None) -> None: self._chunks = list(chunks) @@ -50,6 +71,24 @@ def iter_lines_raw(payload: list[bytes]) -> Iterator[BytesLike]: yield from decoder.flush() +_V2_REQUIRED_KEYS = {"type", "ns", "data"} + + +def _assert_v2_shape(part: Any) -> None: + """Assert a v2 stream part has the required keys and types.""" + assert isinstance(part, dict), f"Expected dict, got {type(part)}" + assert part.keys() >= _V2_REQUIRED_KEYS, ( + f"Missing keys: {_V2_REQUIRED_KEYS - part.keys()}" + ) + assert isinstance(part["type"], str) + assert isinstance(part["ns"], list) + for elem in part["ns"]: + assert isinstance(elem, str) + + +# --- SSE parsing --- + + def test_stream_sse(): for groups in ( [RESPONSE_PAYLOAD], @@ -69,6 +108,9 @@ def test_stream_sse(): assert len(parts) == 79 +# --- HTTP client streaming --- + + @pytest.mark.asyncio async def test_http_client_stream_flushes_trailing_event(): payload = b'event: foo\ndata: {"bar": 1}\n' @@ -92,6 +134,26 @@ async def test_http_client_stream_flushes_trailing_event(): assert parts == [StreamPart(event="foo", data={"bar": 1})] +def test_sync_http_client_stream_flushes_trailing_event(): + payload = b'event: foo\ndata: {"bar": 1}\n' + + def handler(request: httpx.Request) -> httpx.Response: + assert request.headers["accept"] == "text/event-stream" + assert request.headers["cache-control"] == "no-store" + return httpx.Response( + 200, + headers={"Content-Type": "text/event-stream"}, + content=payload, + ) + + transport = httpx.MockTransport(handler) + with httpx.Client(transport=transport, base_url="https://example.com") as client: + http_client = SyncHttpClient(client) + parts = list(http_client.stream("/stream", "GET")) + + assert parts == [StreamPart(event="foo", data={"bar": 1})] + + def test_sync_http_client_stream_recovers_after_disconnect(): reconnect_path = "/reconnect" first_chunks = [ @@ -228,21 +290,178 @@ async def test_http_client_stream_recovers_after_disconnect(): ] -def test_sync_http_client_stream_flushes_trailing_event(): - payload = b'event: foo\ndata: {"bar": 1}\n' +# --- _sse_to_v2_dict conversion --- - def handler(request: httpx.Request) -> httpx.Response: - assert request.headers["accept"] == "text/event-stream" - assert request.headers["cache-control"] == "no-store" - return httpx.Response( - 200, - headers={"Content-Type": "text/event-stream"}, - content=payload, + +def test_sse_to_v2_dict_basic() -> None: + result = _sse_to_v2_dict("values", {"messages": [{"role": "user"}]}) + assert result is not None + _assert_v2_shape(result) + assert result == { + "type": "values", + "ns": [], + "data": {"messages": [{"role": "user"}]}, + "interrupts": [], + } + + +def test_sse_to_v2_dict_with_namespace() -> None: + result = _sse_to_v2_dict("updates|sub:abc", {"key": "val"}) + assert result is not None + _assert_v2_shape(result) + assert result == { + "type": "updates", + "ns": ["sub:abc"], + "data": {"key": "val"}, + "interrupts": [], + } + + +def test_sse_to_v2_dict_with_multiple_ns() -> None: + result = _sse_to_v2_dict("custom|parent|child:123", "hello") + assert result is not None + _assert_v2_shape(result) + assert result == { + "type": "custom", + "ns": ["parent", "child:123"], + "data": "hello", + "interrupts": [], + } + + +def test_sse_to_v2_dict_end_event() -> None: + assert _sse_to_v2_dict("end", None) is None + + +def test_sse_to_v2_dict_metadata_event() -> None: + result = _sse_to_v2_dict("metadata", {"run_id": "abc-123"}) + assert result is not None + _assert_v2_shape(result) + assert result == { + "type": "metadata", + "ns": [], + "data": {"run_id": "abc-123"}, + "interrupts": [], + } + + +def test_sse_to_v2_dict_messages_partial() -> None: + result = _sse_to_v2_dict("messages/partial", [{"type": "ai", "content": "hi"}]) + assert result is not None + _assert_v2_shape(result) + assert result == { + "type": "messages/partial", + "ns": [], + "data": [{"type": "ai", "content": "hi"}], + "interrupts": [], + } + + +def test_sse_to_v2_dict_values_with_interrupts() -> None: + data = { + "messages": [{"role": "user"}], + "__interrupt__": [{"value": "confirm?", "resumable": True}], + } + result = _sse_to_v2_dict("values", data) + assert result is not None + _assert_v2_shape(result) + assert result == { + "type": "values", + "ns": [], + "data": {"messages": [{"role": "user"}]}, + "interrupts": [{"value": "confirm?", "resumable": True}], + } + # __interrupt__ should be popped from data + assert "__interrupt__" not in result["data"] + + +# --- client-side v2 stream wrapping --- + + +@pytest.mark.asyncio +async def test_async_stream_v2_client_side_conversion() -> None: + from langgraph_sdk._async.runs import _wrap_stream_v2 + + async def mock_stream() -> Any: + yield StreamPart(event="metadata", data={"run_id": "r1"}) + yield StreamPart( + event="values", data={"messages": [{"role": "user", "content": "hi"}]} ) + yield StreamPart(event="updates|sub:abc", data={"node": {"out": 1}}) + yield StreamPart(event="end", data=None) # type: ignore[arg-type] - transport = httpx.MockTransport(handler) - with httpx.Client(transport=transport, base_url="https://example.com") as client: - http_client = SyncHttpClient(client) - parts = list(http_client.stream("/stream", "GET")) + parts: list[StreamPartV2] = [part async for part in _wrap_stream_v2(mock_stream())] + assert len(parts) == 3 + for part in parts: + _assert_v2_shape(part) + assert parts[0] == { + "type": "metadata", + "ns": [], + "data": {"run_id": "r1"}, + "interrupts": [], + } + assert parts[1] == { + "type": "values", + "ns": [], + "data": {"messages": [{"role": "user", "content": "hi"}]}, + "interrupts": [], + } + assert parts[2] == { + "type": "updates", + "ns": ["sub:abc"], + "data": {"node": {"out": 1}}, + "interrupts": [], + } - assert parts == [StreamPart(event="foo", data={"bar": 1})] + +def test_sync_stream_v2_client_side_conversion() -> None: + from langgraph_sdk._sync.runs import _wrap_stream_v2_sync + + def mock_stream() -> Any: + yield StreamPart(event="metadata", data={"run_id": "r1"}) + yield StreamPart(event="values", data={"state": "full"}) + yield StreamPart(event="end", data=None) # type: ignore[arg-type] + + parts: list[StreamPartV2] = list(_wrap_stream_v2_sync(mock_stream())) + assert len(parts) == 2 + for part in parts: + _assert_v2_shape(part) + assert parts[0] == { + "type": "metadata", + "ns": [], + "data": {"run_id": "r1"}, + "interrupts": [], + } + assert parts[1] == { + "type": "values", + "ns": [], + "data": {"state": "full"}, + "interrupts": [], + } + + +# --- type narrowing compile-time checks --- + + +def _check_v2_type_narrowing(part: StreamPartV2) -> None: + """Compile-time type narrowing checks — validates mypy narrows the union.""" + if part["type"] == "values": + assert_type(part, ValuesStreamPart) + assert_type(part["data"], dict[str, Any]) + elif part["type"] == "updates": + assert_type(part, UpdatesStreamPart) + assert_type(part["data"], dict[str, Any]) + elif part["type"] == "custom": + assert_type(part, CustomStreamPart) + elif part["type"] == "checkpoints": + assert_type(part, CheckpointsStreamPart) + assert_type(part["data"], CheckpointPayload) + elif part["type"] == "tasks": + assert_type(part, TasksStreamPart) + assert_type(part["data"], TaskPayload | TaskResultPayload) + elif part["type"] == "debug": + assert_type(part, DebugStreamPart) + assert_type(part["data"], DebugPayload) + elif part["type"] == "metadata": + assert_type(part, MetadataStreamPart) + assert_type(part["data"], RunMetadataPayload)