diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 37e8125f9..2dc26b9fc 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/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, diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index c9bf6e5ff..44ace247b 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/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, diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 2535d966a..6c11d676e 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/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,