POC streaming hints

libs/langgraph/langgraph/pregel/main.py

@@ -20,10 +20,13 @@ from dataclasses import is_dataclass
 from functools import partial
 from inspect import isclass
 from typing import (
+    TYPE_CHECKING,
     Any,
     Generic,
+    Literal,
     cast,
     get_type_hints,
+    overload,
 )
 from uuid import UUID, uuid5
 
@@ -48,6 +51,9 @@ 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,
@@ -2404,6 +2410,186 @@ class Pregel(
             durability,
         )
 
+    @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[False] = False,
+        debug: bool | None = None,
+        **kwargs: Unpack[DeprecatedKwargs],
+    ) -> Iterator[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[False] = False,
+        debug: bool | None = None,
+        **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]]: ...
+
     def stream(
         self,
         input: InputT | Command | None,
@@ -2678,6 +2864,186 @@ class Pregel(
             run_manager.on_chain_error(e)
             raise
 
+    @overload
+    def astream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: 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[False] = False,
+        debug: bool | None = None,
+        **kwargs: Unpack[DeprecatedKwargs],
+    ) -> AsyncIterator[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[False] = False,
+        debug: bool | None = None,
+        **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]]: ...
+
     async def astream(
         self,
         input: InputT | Command | None,
@@ -3068,13 +3434,15 @@ class Pregel(
         chunks: list[dict[str, Any] | Any] = []
         interrupts: list[Interrupt] = []
 
-        for chunk in self.stream(
+        for chunk in self.stream(  # type: ignore[misc]
             input,
             config,
             context=context,
-            stream_mode=["updates", "values"]
-            if stream_mode == "values"
-            else stream_mode,
+            stream_mode=(
+                ["updates", "values"]  # type: ignore[arg-type]
+                if stream_mode == "values"
+                else stream_mode
+            ),
             print_mode=print_mode,
             output_keys=output_keys,
             interrupt_before=interrupt_before,
@@ -3158,13 +3526,15 @@ class Pregel(
         chunks: list[dict[str, Any] | Any] = []
         interrupts: list[Interrupt] = []
 
-        async for chunk in self.astream(
+        async for chunk in self.astream(  # type: ignore[misc]
             input,
             config,
             context=context,
-            stream_mode=["updates", "values"]
-            if stream_mode == "values"
-            else stream_mode,
+            stream_mode=(
+                ["updates", "values"]  # type: ignore[arg-type]
+                if stream_mode == "values"
+                else stream_mode
+            ),
             print_mode=print_mode,
             output_keys=output_keys,
             interrupt_before=interrupt_before,

libs/langgraph/langgraph/pregel/protocol.py

@@ -2,13 +2,16 @@ from __future__ import annotations
 
 from abc import abstractmethod
 from collections.abc import AsyncIterator, Callable, Iterator, Sequence
-from typing import Any, Generic, cast
+from typing import TYPE_CHECKING, 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.typing import ContextT, InputT, OutputT, StateT
 
 __all__ = ("PregelProtocol", "StreamProtocol")
@@ -96,6 +99,146 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
         as_node: str | None = None,
     ) -> RunnableConfig: ...
 
+    @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[False] = False,
+    ) -> Iterator[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[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]]: ...
+
     @abstractmethod
     def stream(
         self,
@@ -109,6 +252,146 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
         subgraphs: bool = False,
     ) -> Iterator[dict[str, Any] | 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[False] = False,
+    ) -> AsyncIterator[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[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]]: ...
+
     @abstractmethod
     def astream(
         self,

libs/langgraph/langgraph/pregel/remote.py

@@ -4,9 +4,11 @@ import logging
 from collections.abc import AsyncIterator, Iterator, Sequence
 from dataclasses import asdict
 from typing import (
+    TYPE_CHECKING,
     Any,
     Literal,
     cast,
+    overload,
 )
 from uuid import UUID
 
@@ -41,6 +43,9 @@ 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,
@@ -682,6 +687,156 @@ class RemoteGraph(PregelProtocol):
             updated_stream_modes.remove("events")
         return (updated_stream_modes, requested_stream_modes, req_single, stream)
 
+    @overload
+    def stream(
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        *,
+        stream_mode: Literal["values"],
+        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[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[False] = False,
+        headers: dict[str, str] | None = None,
+        params: QueryParamTypes | None = None,
+        **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]]: ...
+
     def stream(
         self,
         input: dict[str, Any] | Any,
@@ -792,6 +947,156 @@ class RemoteGraph(PregelProtocol):
             else:
                 yield chunk
 
+    @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[False] = False,
+        headers: dict[str, str] | None = None,
+        params: QueryParamTypes | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[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[False] = False,
+        headers: dict[str, str] | None = None,
+        params: QueryParamTypes | None = None,
+        **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]]: ...
+
     async def astream(
         self,
         input: dict[str, Any] | Any,