From 9023ad99318016b5fe7fd94dbbfc0e86252a073e Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Fri, 27 Feb 2026 09:37:09 -0500 Subject: [PATCH] feat(langgraph): backwards compat type safe streaming (#6931) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Type-safe stream parts for v2 streaming ### Review recommendations * Don't fear the diff! It's not so bad, I swear! The PR description below gives a nice overview of changes * Check out my explicit comments below, those should help to orient you to the important changes :) * On a first pass, ignore the test files! Just check out the new types, overloads, and minor logical changes (diff behavior based on the flag) ## Summary Adds a `stream_version="v2"` option that emits typed `{"type", "ns", "data"}` dicts instead of raw tuples/SSE events. Each stream mode gets its own `TypedDict` with a `Literal` type field, enabling full type narrowing on `part["type"]`. This is **opt in**, so it's **non-breaking**!! ```python async for part in graph.astream( inputs, stream_mode=["messages", "custom"], stream_version="v2" ): if part["type"] == "messages": msg, metadata = part["data"] # tuple[AnyMessage, dict] ✅ print(msg.content) elif part["type"] == "custom": part["data"] # Any ✅ ``` Before v2 you'd get `tuple[str, Any]` with no way to narrow `data` based on mode. ### What changed **`langgraph` (core):** New `StreamPart` discriminated union + per-mode TypedDicts in `types.py`. Stream-emit code in `pregel/` refactored to use the new types. `RemoteGraph` gains a `stream_version` param. **`sdk-py`:** Client-side v2 wrapper that converts SSE events into typed dicts. No server API changes — v2 is purely a client-side rewrite of the stream format. ### Stream part types #### `langgraph` (core) | `type` | `data` | |---|---| | `"values"` | `dict[str, Any]` — full state after each step | | `"updates"` | `dict[str, Any]` — node name → output | | `"messages"` | `tuple[AnyMessage, dict]` — message + metadata | | `"custom"` | `Any` — whatever was passed to `StreamWriter` | | `"tasks"` | `TaskPayload \| TaskResultPayload` | | `"checkpoints"` | `CheckpointPayload` | | `"debug"` | `DebugPayload` | #### `sdk-py` (additional types from SSE events) | `type` | `data` | Description | |---|---|---| | `"messages/partial"` | `list[dict]` | Partial message chunks | | `"messages/complete"` | `list[dict]` | Complete messages | | `"messages/metadata"` | `dict` | Message metadata | | `"metadata"` | `RunMetadataPayload` | Run-level metadata (`run_id`, etc.) | All parts share the shape `{"type": Literal[...], "ns": list[str], "data": ...}`. ## Release plan * Release as a part of langgraph 1.1 * Before release, I'd like to do more experimentation with support for pydantic + dataclasses and/or input/output runtime validation, as that would help resolve the lack of typing for `values` mode. ### Notes - Docs need a mass update to cover v2 streaming usage and the new types - Changing the default `stream_version` to `"v2"` in a future release would be breaking but backwards compatible (users can pin `"v1"` to keep current behavior) - I called out specifically relevant parts of the code in comments on the PR :) --- libs/langgraph/langgraph/pregel/debug.py | 45 +- libs/langgraph/langgraph/pregel/main.py | 496 +++++--------- libs/langgraph/langgraph/pregel/protocol.py | 322 +++------ libs/langgraph/langgraph/pregel/remote.py | 359 +++------- libs/langgraph/langgraph/types.py | 192 +++++- libs/langgraph/tests/test_stream_v2.py | 622 ++++++++++++++++++ libs/sdk-py/langgraph_sdk/_async/runs.py | 94 ++- .../sdk-py/langgraph_sdk/_shared/utilities.py | 13 + libs/sdk-py/langgraph_sdk/_sync/runs.py | 95 ++- libs/sdk-py/langgraph_sdk/schema.py | 200 ++++++ libs/sdk-py/tests/test_client_stream.py | 209 +++++- 11 files changed, 1753 insertions(+), 894 deletions(-) create mode 100644 libs/langgraph/tests/test_stream_v2.py 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 2dc26b9fc..4c1dfbb91 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -20,7 +20,6 @@ from dataclasses import is_dataclass from functools import partial from inspect import isclass from typing import ( - TYPE_CHECKING, Any, Generic, Literal, @@ -51,9 +50,6 @@ from langgraph.store.base import BaseStore from pydantic import BaseModel, TypeAdapter from typing_extensions import Self, Unpack, deprecated, is_typeddict -if TYPE_CHECKING: - from langchain_core.messages import AnyMessage - from langgraph._internal._config import ( ensure_config, merge_configs, @@ -148,6 +144,7 @@ from langgraph.types import ( StateSnapshot, StateUpdate, StreamMode, + StreamPart, ensure_valid_checkpointer, ) from langgraph.typing import ContextT, InputT, OutputT, StateT @@ -2417,16 +2414,17 @@ class Pregel( config: RunnableConfig | None = None, *, context: ContextT | None = None, - stream_mode: Literal["values"], + 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: Literal[False] = False, + subgraphs: bool = False, debug: bool | None = None, + stream_version: Literal["v2"], **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[OutputT]: ... + ) -> Iterator[StreamPart]: ... @overload def stream( @@ -2435,160 +2433,17 @@ class Pregel( config: RunnableConfig | None = None, *, context: ContextT | None = None, - stream_mode: Literal["updates"], + 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: Literal[False] = False, + subgraphs: bool = False, debug: bool | None = None, + stream_version: Literal["v1"] = ..., **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[dict[str, Any]]: ... - - @overload - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["messages"], - 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: Literal[False] = False, - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[tuple[AnyMessage, dict[str, Any]]]: ... - - @overload - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["custom"], - 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: Literal[False] = False, - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[Any]: ... - - @overload - def stream( - 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, - subgraphs: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[tuple[tuple[str, ...], OutputT]]: ... - - @overload - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["updates"], - 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: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[tuple[tuple[str, ...], dict[str, Any]]]: ... - - @overload - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["messages"], - 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: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ... - - @overload - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["custom"], - 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: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[tuple[tuple[str, ...], Any]]: ... - - @overload - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: list[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, - subgraphs: Literal[False] = False, - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[tuple[str, Any]]: ... - - @overload - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: list[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, - subgraphs: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[tuple[tuple[str, ...], str, Any]]: ... + ) -> Iterator[dict[str, Any] | Any]: ... def stream( self, @@ -2604,6 +2459,7 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, + stream_version: Literal["v1", "v2"] = "v1", **kwargs: Unpack[DeprecatedKwargs], ) -> Iterator[dict[str, Any] | Any]: """Stream graph steps for a single input. @@ -2837,7 +2693,12 @@ 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, + stream_version, ) loop.after_tick() # wait for checkpoint @@ -2845,7 +2706,12 @@ 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, + stream_version, ) # handle exit if loop.status == "out_of_steps": @@ -2871,16 +2737,17 @@ class Pregel( config: RunnableConfig | None = None, *, context: ContextT | None = None, - stream_mode: Literal["values"], + 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: Literal[False] = False, + subgraphs: bool = False, debug: bool | None = None, + stream_version: Literal["v2"], **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[OutputT]: ... + ) -> AsyncIterator[StreamPart]: ... @overload def astream( @@ -2889,160 +2756,17 @@ class Pregel( config: RunnableConfig | None = None, *, context: ContextT | None = None, - stream_mode: Literal["updates"], + 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: Literal[False] = False, + subgraphs: bool = False, debug: bool | None = None, + stream_version: Literal["v1"] = ..., **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[dict[str, Any]]: ... - - @overload - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["messages"], - 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: Literal[False] = False, - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[tuple[AnyMessage, dict[str, Any]]]: ... - - @overload - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["custom"], - 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: Literal[False] = False, - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[Any]: ... - - @overload - def astream( - 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, - subgraphs: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[tuple[tuple[str, ...], OutputT]]: ... - - @overload - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["updates"], - 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: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[tuple[tuple[str, ...], dict[str, Any]]]: ... - - @overload - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["messages"], - 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: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ... - - @overload - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["custom"], - 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: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[tuple[tuple[str, ...], Any]]: ... - - @overload - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: list[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, - subgraphs: Literal[False] = False, - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[tuple[str, Any]]: ... - - @overload - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: list[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, - subgraphs: Literal[True], - debug: bool | None = None, - **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[tuple[tuple[str, ...], str, Any]]: ... + ) -> AsyncIterator[dict[str, Any] | Any]: ... async def astream( self, @@ -3058,6 +2782,7 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, + stream_version: Literal["v1", "v2"] = "v1", **kwargs: Unpack[DeprecatedKwargs], ) -> AsyncIterator[dict[str, Any] | Any]: """Asynchronously stream graph steps for a single input. @@ -3350,6 +3075,7 @@ class Pregel( subgraphs, stream.get_nowait, asyncio.QueueEmpty, + stream_version, ): yield o loop.after_tick() @@ -3368,6 +3094,7 @@ class Pregel( subgraphs, stream.get_nowait, asyncio.QueueEmpty, + stream_version, ): yield o # handle exit @@ -3387,6 +3114,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, + stream_version: Literal["v2"], + **kwargs: Any, + ) -> dict[str, Any]: ... + + @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, + stream_version: Literal["v2"], + **kwargs: Any, + ) -> list[StreamPart]: ... + + @overload def invoke( self, input: InputT | Command | None, @@ -3399,6 +3161,23 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + stream_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, + stream_version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: """Run the graph with a single input and config. @@ -3422,6 +3201,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. + stream_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: @@ -3439,30 +3221,33 @@ class Pregel( config, context=context, stream_mode=( - ["updates", "values"] # type: ignore[arg-type] - if stream_mode == "values" - else 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, + stream_version=stream_version, # type: ignore[arg-type] **kwargs, ): if stream_mode == "values": - if len(chunk) == 2: - mode, payload = cast(tuple[StreamMode, Any], chunk) + if stream_version == "v2": + mode = chunk["type"] + payload = chunk["data"] else: - _, mode, payload = cast( - tuple[tuple[str, ...], StreamMode, Any], chunk - ) + 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) + interrupts.extend(ints) # type: ignore[arg-type] elif mode == "values": latest = payload else: @@ -3479,6 +3264,41 @@ 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, + stream_version: Literal["v2"], + **kwargs: Any, + ) -> dict[str, Any]: ... + + @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, + stream_version: Literal["v2"], + **kwargs: Any, + ) -> list[StreamPart]: ... + + @overload async def ainvoke( self, input: InputT | Command | None, @@ -3491,6 +3311,23 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, + stream_version: Literal["v1"] = ..., + **kwargs: Any, + ) -> dict[str, Any] | Any: ... + + 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, + stream_version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: """Asynchronously run the graph with a single input and config. @@ -3514,6 +3351,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. + stream_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: @@ -3531,30 +3371,33 @@ class Pregel( config, context=context, stream_mode=( - ["updates", "values"] # type: ignore[arg-type] - if stream_mode == "values" - else 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, + stream_version=stream_version, # type: ignore[arg-type] **kwargs, ): if stream_mode == "values": - if len(chunk) == 2: - mode, payload = cast(tuple[StreamMode, Any], chunk) + if stream_version == "v2": + mode = chunk["type"] + payload = chunk["data"] else: - _, mode, payload = cast( - tuple[tuple[str, ...], StreamMode, Any], chunk - ) + 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) + interrupts.extend(ints) # type: ignore[arg-type] elif mode == "values": latest = payload else: @@ -3625,6 +3468,7 @@ def _output( stream_subgraphs: bool, getter: Callable[[], tuple[tuple[str, ...], str, Any]], empty_exc: type[Exception], + stream_version: Literal["v1", "v2"] = "v1", ) -> Iterator: while True: try: @@ -3652,7 +3496,9 @@ def _output( ) ) if mode in stream_mode: - if stream_subgraphs and isinstance(stream_mode, list): + if stream_version == "v2": + 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) diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index 44ace247b..fb39c76c2 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -2,16 +2,20 @@ from __future__ import annotations from abc import abstractmethod from collections.abc import AsyncIterator, Callable, Iterator, Sequence -from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload +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 - -if TYPE_CHECKING: - from langchain_core.messages import AnyMessage +from langgraph.types import ( + All, + Command, + StateSnapshot, + StateUpdate, + StreamMode, + StreamPart, +) from langgraph.typing import ContextT, InputT, OutputT, StateT __all__ = ("PregelProtocol", "StreamProtocol") @@ -107,11 +111,12 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou config: RunnableConfig | None = None, *, context: ContextT | None = None, - stream_mode: Literal["values"], + stream_mode: StreamMode | list[StreamMode] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> Iterator[OutputT]: ... + subgraphs: bool = False, + stream_version: Literal["v2"], + ) -> Iterator[StreamPart]: ... @overload @abstractmethod @@ -121,123 +126,12 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou config: RunnableConfig | None = None, *, context: ContextT | None = None, - stream_mode: Literal["updates"], + stream_mode: StreamMode | list[StreamMode] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> Iterator[dict[str, Any]]: ... - - @overload - @abstractmethod - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["messages"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> Iterator[tuple[AnyMessage, dict[str, Any]]]: ... - - @overload - @abstractmethod - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["custom"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> Iterator[Any]: ... - - @overload - @abstractmethod - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["values"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> Iterator[tuple[tuple[str, ...], OutputT]]: ... - - @overload - @abstractmethod - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["updates"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> Iterator[tuple[tuple[str, ...], dict[str, Any]]]: ... - - @overload - @abstractmethod - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["messages"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> Iterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ... - - @overload - @abstractmethod - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["custom"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> Iterator[tuple[tuple[str, ...], Any]]: ... - - @overload - @abstractmethod - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: list[StreamMode], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> Iterator[tuple[str, Any]]: ... - - @overload - @abstractmethod - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: list[StreamMode], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> Iterator[tuple[tuple[str, ...], str, Any]]: ... + subgraphs: bool = False, + stream_version: Literal["v1"] = ..., + ) -> Iterator[dict[str, Any] | Any]: ... @abstractmethod def stream( @@ -250,6 +144,7 @@ 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, + stream_version: Literal["v1", "v2"] = "v1", ) -> Iterator[dict[str, Any] | Any]: ... @overload @@ -260,11 +155,12 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou config: RunnableConfig | None = None, *, context: ContextT | None = None, - stream_mode: Literal["values"], + stream_mode: StreamMode | list[StreamMode] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> AsyncIterator[OutputT]: ... + subgraphs: bool = False, + stream_version: Literal["v2"], + ) -> AsyncIterator[StreamPart]: ... @overload @abstractmethod @@ -274,123 +170,12 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou config: RunnableConfig | None = None, *, context: ContextT | None = None, - stream_mode: Literal["updates"], + stream_mode: StreamMode | list[StreamMode] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> AsyncIterator[dict[str, Any]]: ... - - @overload - @abstractmethod - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["messages"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> AsyncIterator[tuple[AnyMessage, dict[str, Any]]]: ... - - @overload - @abstractmethod - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["custom"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> AsyncIterator[Any]: ... - - @overload - @abstractmethod - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["values"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> AsyncIterator[tuple[tuple[str, ...], OutputT]]: ... - - @overload - @abstractmethod - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["updates"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> AsyncIterator[tuple[tuple[str, ...], dict[str, Any]]]: ... - - @overload - @abstractmethod - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["messages"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> AsyncIterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ... - - @overload - @abstractmethod - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: Literal["custom"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> AsyncIterator[tuple[tuple[str, ...], Any]]: ... - - @overload - @abstractmethod - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: list[StreamMode], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - ) -> AsyncIterator[tuple[str, Any]]: ... - - @overload - @abstractmethod - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: list[StreamMode], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - ) -> AsyncIterator[tuple[tuple[str, ...], str, Any]]: ... + subgraphs: bool = False, + stream_version: Literal["v1"] = ..., + ) -> AsyncIterator[dict[str, Any] | Any]: ... @abstractmethod def astream( @@ -403,8 +188,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, + stream_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, + stream_version: Literal["v2"], + ) -> dict[str, 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, + stream_version: Literal["v1"] = ..., + ) -> dict[str, Any] | Any: ... + @abstractmethod def invoke( self, @@ -414,6 +226,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, + stream_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, + stream_version: Literal["v2"], + ) -> dict[str, 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, + stream_version: Literal["v1"] = ..., ) -> dict[str, Any] | Any: ... @abstractmethod @@ -425,6 +264,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, + stream_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 6c11d676e..691c9646d 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -4,7 +4,6 @@ import logging from collections.abc import AsyncIterator, Iterator, Sequence from dataclasses import asdict from typing import ( - TYPE_CHECKING, Any, Literal, cast, @@ -43,9 +42,6 @@ from langgraph_sdk.schema import ( ) from typing_extensions import Self -if TYPE_CHECKING: - from langchain_core.messages import AnyMessage - from langgraph._internal._config import merge_configs from langgraph._internal._constants import ( CONF, @@ -66,6 +62,7 @@ from langgraph.types import ( PregelTask, StateSnapshot, StreamMode, + StreamPart, ) logger = logging.getLogger(__name__) @@ -693,14 +690,15 @@ class RemoteGraph(PregelProtocol): input: dict[str, Any] | Any, config: RunnableConfig | None = None, *, - stream_mode: Literal["values"], + stream_mode: StreamMode | list[StreamMode] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, + subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + stream_version: Literal["v2"], **kwargs: Any, - ) -> Iterator[dict[str, Any]]: ... + ) -> Iterator[StreamPart]: ... @overload def stream( @@ -708,134 +706,15 @@ class RemoteGraph(PregelProtocol): input: dict[str, Any] | Any, config: RunnableConfig | None = None, *, - stream_mode: Literal["updates"], + stream_mode: StreamMode | list[StreamMode] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, + subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + stream_version: Literal["v1"] = ..., **kwargs: Any, - ) -> Iterator[dict[str, Any]]: ... - - @overload - def stream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["messages"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> Iterator[tuple[AnyMessage, dict[str, Any]]]: ... - - @overload - def stream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["custom"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> Iterator[Any]: ... - - @overload - def stream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["values"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> Iterator[tuple[tuple[str, ...], dict[str, Any]]]: ... - - @overload - def stream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["updates"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> Iterator[tuple[tuple[str, ...], dict[str, Any]]]: ... - - @overload - def stream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["messages"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> Iterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ... - - @overload - def stream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["custom"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> Iterator[tuple[tuple[str, ...], Any]]: ... - - @overload - def stream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: list[StreamMode], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> Iterator[tuple[str, Any]]: ... - - @overload - def stream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: list[StreamMode], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> Iterator[tuple[tuple[str, ...], str, Any]]: ... + ) -> Iterator[dict[str, Any] | Any]: ... def stream( self, @@ -848,6 +727,7 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + stream_version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> Iterator[dict[str, Any] | Any]: """Create a run and stream the results. @@ -929,10 +809,12 @@ 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 stream_version == "v2": + yield {"type": mode, "ns": ns, "data": chunk.data} + elif subgraphs: if NS_SEP in chunk.event: mode, ns_ = chunk.event.split(NS_SEP, 1) ns = tuple(ns_.split(NS_SEP)) @@ -953,14 +835,15 @@ class RemoteGraph(PregelProtocol): input: dict[str, Any] | Any, config: RunnableConfig | None = None, *, - stream_mode: Literal["values"], + stream_mode: StreamMode | list[StreamMode] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, + subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + stream_version: Literal["v2"], **kwargs: Any, - ) -> AsyncIterator[dict[str, Any]]: ... + ) -> AsyncIterator[StreamPart]: ... @overload def astream( @@ -968,134 +851,15 @@ class RemoteGraph(PregelProtocol): input: dict[str, Any] | Any, config: RunnableConfig | None = None, *, - stream_mode: Literal["updates"], + stream_mode: StreamMode | list[StreamMode] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, + subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + stream_version: Literal["v1"] = ..., **kwargs: Any, - ) -> AsyncIterator[dict[str, Any]]: ... - - @overload - def astream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["messages"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> AsyncIterator[tuple[AnyMessage, dict[str, Any]]]: ... - - @overload - def astream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["custom"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> AsyncIterator[Any]: ... - - @overload - def astream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["values"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> AsyncIterator[tuple[tuple[str, ...], dict[str, Any]]]: ... - - @overload - def astream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["updates"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> AsyncIterator[tuple[tuple[str, ...], dict[str, Any]]]: ... - - @overload - def astream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["messages"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> AsyncIterator[tuple[tuple[str, ...], tuple[AnyMessage, dict[str, Any]]]]: ... - - @overload - def astream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: Literal["custom"], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> AsyncIterator[tuple[tuple[str, ...], Any]]: ... - - @overload - def astream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: list[StreamMode], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[False] = False, - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> AsyncIterator[tuple[str, Any]]: ... - - @overload - def astream( - self, - input: dict[str, Any] | Any, - config: RunnableConfig | None = None, - *, - stream_mode: list[StreamMode], - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - subgraphs: Literal[True], - headers: dict[str, str] | None = None, - params: QueryParamTypes | None = None, - **kwargs: Any, - ) -> AsyncIterator[tuple[tuple[str, ...], str, Any]]: ... + ) -> AsyncIterator[dict[str, Any] | Any]: ... async def astream( self, @@ -1108,6 +872,7 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + stream_version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> AsyncIterator[dict[str, Any] | Any]: """Create a run and stream the results. @@ -1189,10 +954,12 @@ 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 stream_version == "v2": + yield {"type": mode, "ns": ns, "data": chunk.data} + elif subgraphs: if NS_SEP in chunk.event: mode, ns_ = chunk.event.split(NS_SEP, 1) ns = tuple(ns_.split(NS_SEP)) @@ -1223,6 +990,7 @@ class RemoteGraph(PregelProtocol): ) -> AsyncIterator[dict[str, Any]]: raise NotImplementedError + @overload def invoke( self, input: dict[str, Any] | Any, @@ -1232,6 +1000,34 @@ class RemoteGraph(PregelProtocol): interrupt_after: All | Sequence[str] | None = None, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + stream_version: Literal["v2"], + **kwargs: Any, + ) -> 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, + stream_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, + stream_version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: """Create a run, wait until it finishes and return the final state. @@ -1242,12 +1038,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. + stream_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, @@ -1255,15 +1053,46 @@ class RemoteGraph(PregelProtocol): headers=headers, stream_mode="values", params=params, + stream_version=stream_version, **kwargs, ): pass try: + if stream_version == "v2": + return chunk["data"] 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, + stream_version: Literal["v2"], + **kwargs: Any, + ) -> 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, + stream_version: Literal["v1"] = ..., + **kwargs: Any, + ) -> dict[str, Any] | Any: ... + async def ainvoke( self, input: dict[str, Any] | Any, @@ -1273,6 +1102,7 @@ class RemoteGraph(PregelProtocol): interrupt_after: All | Sequence[str] | None = None, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, + stream_version: Literal["v1", "v2"] = "v1", **kwargs: Any, ) -> dict[str, Any] | Any: """Create a run, wait until it finishes and return the final state. @@ -1283,12 +1113,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. + stream_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, @@ -1296,10 +1128,13 @@ class RemoteGraph(PregelProtocol): headers=headers, stream_mode="values", params=params, + stream_version=stream_version, **kwargs, ): pass try: + if stream_version == "v2": + return chunk["data"] 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..1ed24feaa 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -16,9 +16,10 @@ 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, TypedDict, Unpack, deprecated from xxhash import xxh3_128_hexdigest from langgraph._internal._cache import default_cache_key @@ -44,6 +45,19 @@ __all__ = ( "Checkpointer", "StreamMode", "StreamWriter", + "StreamPart", + "ValuesStreamPart", + "UpdatesStreamPart", + "MessagesStreamPart", + "CustomStreamPart", + "CheckpointStreamPart", + "TasksStreamPart", + "DebugStreamPart", + "TaskPayload", + "TaskResultPayload", + "CheckpointTask", + "CheckpointPayload", + "DebugPayload", "RetryPolicy", "CachePolicy", "Interrupt", @@ -113,6 +127,182 @@ 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 + name: str + input: Any + triggers: list[str] + + +class TaskResultPayload(TypedDict): + """Payload for a task result event.""" + + id: str + name: str + error: str | None + interrupts: list[dict] + result: dict[str, Any] + + +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 + name: str + error: NotRequired[str] + result: NotRequired[Any] + interrupts: NotRequired[list[dict]] + state: StateSnapshot | RunnableConfig | None + + +class CheckpointPayload(TypedDict): + """Payload for a checkpoint event.""" + + config: RunnableConfig | None + metadata: CheckpointMetadata + values: dict[str, Any] + next: list[str] + parent_config: RunnableConfig | None + tasks: list[CheckpointTask] + + +class _DebugCheckpointPayload(TypedDict): + step: int + timestamp: str + type: Literal["checkpoint"] + payload: CheckpointPayload + + +class _DebugTaskPayload(TypedDict): + step: int + timestamp: str + type: Literal["task"] + payload: TaskPayload + + +class _DebugTaskResultPayload(TypedDict): + step: int + timestamp: str + type: Literal["task_result"] + payload: TaskResultPayload + + +DebugPayload = _DebugCheckpointPayload | _DebugTaskPayload | _DebugTaskResultPayload +"""Wrapper payload for debug events. Discriminate on `type`.""" + + +class ValuesStreamPart(TypedDict): + """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: dict[str, Any] + + +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): + """Stream part emitted for `stream_mode="checkpoints"`.""" + + type: Literal["checkpoints"] + ns: tuple[str, ...] + data: CheckpointPayload + + +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): + """Stream part emitted for `stream_mode="debug"`.""" + + type: Literal["debug"] + ns: tuple[str, ...] + data: DebugPayload + + +StreamPart = ( + ValuesStreamPart + | UpdatesStreamPart + | MessagesStreamPart + | CustomStreamPart + | CheckpointStreamPart + | TasksStreamPart + | DebugStreamPart +) +"""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"] # dict[str, Any] — full state + elif part["type"] == "messages": + part["data"] # tuple[BaseMessage, dict] — (message, metadata) + elif part["type"] == "custom": + part["data"] # Any — user-defined +``` +""" + _DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True} diff --git a/libs/langgraph/tests/test_stream_v2.py b/libs/langgraph/tests/test_stream_v2.py new file mode 100644 index 000000000..07b1974f1 --- /dev/null +++ b/libs/langgraph/tests/test_stream_v2.py @@ -0,0 +1,622 @@ +"""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 typing import Annotated, Any + +import pytest +from langchain_core.messages import AIMessage, BaseMessage +from langgraph.checkpoint.memory import InMemorySaver +from typing_extensions import TypedDict, assert_type + +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, + MessagesStreamPart, + StreamPart, + StreamWriter, + TaskPayload, + TaskResultPayload, + TasksStreamPart, + UpdatesStreamPart, + ValuesStreamPart, +) +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) -> 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) + + +# --- 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", stream_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", stream_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", stream_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", stream_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"], + stream_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, + stream_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", + stream_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", + stream_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", + stream_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, + stream_version="v2", + ) + ) + chunks_with_sub = list( + graph.stream( + _SIMPLE_INPUT, + stream_mode="updates", + subgraphs=True, + stream_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, stream_version="v2") + assert isinstance(result, dict) + assert result["value"] == "x_a_b" + assert result["items"] == ["a", "b"] + + def test_updates_mode(self) -> None: + graph = _make_simple_graph().compile() + result = graph.invoke(_SIMPLE_INPUT, stream_mode="updates", stream_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, stream_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", stream_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", stream_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", stream_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", stream_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"], + stream_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, + stream_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", + stream_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", + stream_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", + stream_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, stream_version="v2") + assert isinstance(result, dict) + assert result["value"] == "x_a_b" + assert result["items"] == ["a", "b"] + + @pytest.mark.anyio + async def test_updates_mode(self) -> None: + graph = _make_simple_graph().compile() + result = await graph.ainvoke( + _SIMPLE_INPUT, stream_mode="updates", stream_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, stream_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 narrowing compile-time checks --- +# These assert_type calls verify that mypy narrows the union correctly. + + +def _check_type_narrowing(part: StreamPart) -> None: + """Compile-time type narrowing checks — never called at runtime.""" + 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"] == "messages": + assert_type(part, MessagesStreamPart) + elif part["type"] == "custom": + assert_type(part, CustomStreamPart) + elif part["type"] == "checkpoints": + assert_type(part, CheckpointStreamPart) + 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) + 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..d98edb618 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, + stream_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, + stream_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, + stream_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]: ... + stream_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]: + stream_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 + stream_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 stream_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..0f2e8bdcf 100644 --- a/libs/sdk-py/langgraph_sdk/_shared/utilities.py +++ b/libs/sdk-py/langgraph_sdk/_shared/utilities.py @@ -107,6 +107,19 @@ 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 [] + return {"type": event_type, "ns": ns, "data": data} + + 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..f931792c1 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, + stream_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, + stream_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, + stream_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]: ... + stream_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]: + stream_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 - + stream_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 stream_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..dc9e3c945 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -588,6 +588,206 @@ 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 + name: str + input: Any + triggers: list[str] + + +class TaskResultPayload(TypedDict): + """Payload for a task result event.""" + + id: str + name: str + error: str | None + interrupts: list[dict[str, Any]] + result: dict[str, Any] + + +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 + name: str + error: NotRequired[str] + result: NotRequired[Any] + interrupts: NotRequired[list[dict[str, Any]]] + state: dict[str, Any] | None + + +class CheckpointPayload(TypedDict): + """Payload for a checkpoint event.""" + + config: dict[str, Any] | None + metadata: dict[str, Any] + values: dict[str, Any] + next: list[str] + parent_config: dict[str, Any] | None + tasks: list[CheckpointTaskPayload] + + +class _DebugCheckpointPayload(TypedDict): + step: int + timestamp: str + type: Literal["checkpoint"] + payload: CheckpointPayload + + +class _DebugTaskPayload(TypedDict): + step: int + timestamp: str + type: Literal["task"] + payload: TaskPayload + + +class _DebugTaskResultPayload(TypedDict): + step: int + timestamp: str + type: Literal["task_result"] + payload: TaskResultPayload + + +DebugPayload = _DebugCheckpointPayload | _DebugTaskPayload | _DebugTaskResultPayload +"""Wrapper payload for debug events. Discriminate on ``type``.""" + + +class RunMetadataPayload(TypedDict): + """Payload for the ``metadata`` control event.""" + + run_id: str + + +# --- v2 stream part TypedDicts --- + + +class ValuesStreamPart(TypedDict): + """Stream part emitted for ``stream_mode="values"``.""" + + type: Literal["values"] + ns: list[str] + data: dict[str, Any] + + +class UpdatesStreamPart(TypedDict): + """Stream part emitted for ``stream_mode="updates"``.""" + + type: Literal["updates"] + ns: list[str] + data: dict[str, Any] + + +class MessagesPartialStreamPart(TypedDict): + """Stream part emitted for partial message chunks (``messages/partial``).""" + + type: Literal["messages/partial"] + ns: list[str] + data: list[dict[str, Any]] + + +class MessagesCompleteStreamPart(TypedDict): + """Stream part emitted for complete messages (``messages/complete``).""" + + type: Literal["messages/complete"] + ns: list[str] + data: list[dict[str, Any]] + + +class MessagesMetadataStreamPart(TypedDict): + """Stream part emitted for message metadata (``messages/metadata``).""" + + type: Literal["messages/metadata"] + ns: list[str] + data: dict[str, Any] + + +class MessagesTupleStreamPart(TypedDict): + """Stream part emitted for ``stream_mode="messages"`` (raw message+metadata pair).""" + + type: Literal["messages"] + ns: list[str] + data: list[dict[str, Any]] + + +class CustomStreamPart(TypedDict): + """Stream part emitted for ``stream_mode="custom"``.""" + + type: Literal["custom"] + ns: list[str] + data: Any + + +class CheckpointsStreamPart(TypedDict): + """Stream part emitted for ``stream_mode="checkpoints"``.""" + + type: Literal["checkpoints"] + ns: list[str] + data: CheckpointPayload + + +class TasksStreamPart(TypedDict): + """Stream part emitted for ``stream_mode="tasks"``.""" + + type: Literal["tasks"] + ns: list[str] + data: TaskPayload | TaskResultPayload + + +class DebugStreamPart(TypedDict): + """Stream part emitted for ``stream_mode="debug"``.""" + + type: Literal["debug"] + ns: list[str] + data: DebugPayload + + +class MetadataStreamPart(TypedDict): + """Control event with ``run_id`` and other run metadata.""" + + type: Literal["metadata"] + ns: list[str] + data: RunMetadataPayload + + +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..d6295218b 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,138 @@ 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"}]}, + } + + +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"}, + } + + +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", + } + + +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"}, + } + + +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"}], + } + + +# --- 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"}} + assert parts[1] == { + "type": "values", + "ns": [], + "data": {"messages": [{"role": "user", "content": "hi"}]}, + } + assert parts[2] == { + "type": "updates", + "ns": ["sub:abc"], + "data": {"node": {"out": 1}}, + } - 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"}} + assert parts[1] == {"type": "values", "ns": [], "data": {"state": "full"}} + + +# --- 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)