From dec15d13dc47e032886816a8cbd16e4e99d6d45a Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Tue, 3 Mar 2026 13:40:16 -0500 Subject: [PATCH] chore: parametrize `StreamPart` (#7009) --- libs/langgraph/langgraph/pregel/main.py | 8 ++++---- libs/langgraph/langgraph/pregel/protocol.py | 4 ++-- libs/langgraph/langgraph/types.py | 16 ++++++++++------ libs/langgraph/tests/test_stream_v2.py | 21 ++++++++++++--------- 4 files changed, 28 insertions(+), 21 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index b3bfb2d9d..0998b21c6 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -2456,7 +2456,7 @@ class Pregel( debug: bool | None = None, stream_version: Literal["v2"], **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[StreamPart]: ... + ) -> Iterator[StreamPart[OutputT, StateT]]: ... @overload def stream( @@ -2787,7 +2787,7 @@ class Pregel( debug: bool | None = None, stream_version: Literal["v2"], **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[StreamPart]: ... + ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ... @overload def astream( @@ -3194,7 +3194,7 @@ class Pregel( durability: Durability | None = None, stream_version: Literal["v2"], **kwargs: Any, - ) -> list[StreamPart]: ... + ) -> list[StreamPart[OutputT, StateT]]: ... @overload def invoke( @@ -3364,7 +3364,7 @@ class Pregel( durability: Durability | None = None, stream_version: Literal["v2"], **kwargs: Any, - ) -> list[StreamPart]: ... + ) -> list[StreamPart[OutputT, StateT]]: ... @overload async def ainvoke( diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index f5880455c..bd147d944 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -117,7 +117,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou interrupt_after: All | Sequence[str] | None = None, subgraphs: bool = False, stream_version: Literal["v2"], - ) -> Iterator[StreamPart]: ... + ) -> Iterator[StreamPart[OutputT, StateT]]: ... @overload @abstractmethod @@ -161,7 +161,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou interrupt_after: All | Sequence[str] | None = None, subgraphs: bool = False, stream_version: Literal["v2"], - ) -> AsyncIterator[StreamPart]: ... + ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ... @overload @abstractmethod diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index d00eb3c61..7da21adc0 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -19,7 +19,7 @@ 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 NotRequired, TypedDict, Unpack, deprecated +from typing_extensions import NotRequired, TypeAliasType, TypedDict, Unpack, deprecated from xxhash import xxh3_128_hexdigest from langgraph._internal._cache import default_cache_key @@ -206,8 +206,10 @@ class _DebugTaskResultPayload(TypedDict): payload: TaskResultPayload -DebugPayload = ( - _DebugCheckpointPayload[StateT] | _DebugTaskPayload | _DebugTaskResultPayload +DebugPayload = TypeAliasType( + "DebugPayload", + _DebugCheckpointPayload[StateT] | _DebugTaskPayload | _DebugTaskResultPayload, + type_params=(StateT,), ) """Wrapper payload for debug events. Discriminate on `type`.""" @@ -288,17 +290,19 @@ class DebugStreamPart(TypedDict, Generic[StateT]): type: Literal["debug"] ns: tuple[str, ...] - data: DebugPayload + data: DebugPayload[StateT] -StreamPart = ( +StreamPart = TypeAliasType( + "StreamPart", ValuesStreamPart[OutputT] | UpdatesStreamPart | MessagesStreamPart | CustomStreamPart | CheckpointStreamPart[StateT] | TasksStreamPart - | DebugStreamPart[StateT] + | DebugStreamPart[StateT], + type_params=(OutputT, StateT), ) """A discriminated union of all v2 stream part types. diff --git a/libs/langgraph/tests/test_stream_v2.py b/libs/langgraph/tests/test_stream_v2.py index 772c00223..ac64f76c0 100644 --- a/libs/langgraph/tests/test_stream_v2.py +++ b/libs/langgraph/tests/test_stream_v2.py @@ -9,7 +9,7 @@ from __future__ import annotations import operator import sys from dataclasses import dataclass -from typing import Annotated, Any +from typing import Annotated, Any, TypeVar import pytest from langchain_core.messages import AIMessage, BaseMessage @@ -114,7 +114,7 @@ def _make_subgraph() -> Any: _STREAM_PART_KEYS = {"type", "ns", "data"} -def _assert_stream_part_shape(part: StreamPart) -> None: +def _assert_stream_part_shape(part: StreamPart[Any, Any]) -> None: """Assert a v2 stream part has the required keys and correct types.""" assert isinstance(part, dict), f"Expected dict, got {type(part)}" assert _STREAM_PART_KEYS <= part.keys(), ( @@ -1141,11 +1141,14 @@ class TestV2ValidationErrors: # These assert_type calls verify that mypy narrows the union correctly. -def _check_type_narrowing(part: StreamPart) -> None: +_OutputT = TypeVar("_OutputT") +_StateT = TypeVar("_StateT") + + +def _check_type_narrowing(part: StreamPart[_OutputT, _StateT]) -> None: """Compile-time type narrowing checks — never called at runtime.""" if part["type"] == "values": - assert_type(part, ValuesStreamPart) - assert_type(part["data"], dict[str, Any]) + assert_type(part, ValuesStreamPart[_OutputT]) elif part["type"] == "updates": assert_type(part, UpdatesStreamPart) assert_type(part["data"], dict[str, Any]) @@ -1154,12 +1157,12 @@ def _check_type_narrowing(part: StreamPart) -> None: elif part["type"] == "custom": assert_type(part, CustomStreamPart) elif part["type"] == "checkpoints": - assert_type(part, CheckpointStreamPart) - assert_type(part["data"], CheckpointPayload) + assert_type(part, CheckpointStreamPart[_StateT]) + assert_type(part["data"], CheckpointPayload[_StateT]) elif part["type"] == "tasks": assert_type(part, TasksStreamPart) assert_type(part["data"], TaskPayload | TaskResultPayload) elif part["type"] == "debug": - assert_type(part, DebugStreamPart) - assert_type(part["data"], DebugPayload) + assert_type(part, DebugStreamPart[_StateT]) + assert_type(part["data"], DebugPayload[_StateT]) assert_type(part["ns"], tuple[str, ...])