chore: parametrize StreamPart (#7009)

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(

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

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.
 

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, ...])