diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 0998b21c6..affca6d7c 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -2454,9 +2454,9 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v1"], **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[StreamPart[OutputT, StateT]]: ... + ) -> Iterator[dict[str, Any] | Any]: ... @overload def stream( @@ -2473,9 +2473,9 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, - stream_version: Literal["v1"] = ..., + stream_version: Literal["v2"] = ..., **kwargs: Unpack[DeprecatedKwargs], - ) -> Iterator[dict[str, Any] | Any]: ... + ) -> Iterator[StreamPart[OutputT, StateT]]: ... def stream( self, @@ -2491,7 +2491,7 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, - stream_version: Literal["v1", "v2"] = "v1", + stream_version: Literal["v1", "v2"] = "v2", **kwargs: Unpack[DeprecatedKwargs], ) -> Iterator[dict[str, Any] | Any]: """Stream graph steps for a single input. @@ -2785,9 +2785,9 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v1"], **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ... + ) -> AsyncIterator[dict[str, Any] | Any]: ... @overload def astream( @@ -2804,9 +2804,9 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, - stream_version: Literal["v1"] = ..., + stream_version: Literal["v2"] = ..., **kwargs: Unpack[DeprecatedKwargs], - ) -> AsyncIterator[dict[str, Any] | Any]: ... + ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ... async def astream( self, @@ -2822,7 +2822,7 @@ class Pregel( durability: Durability | None = None, subgraphs: bool = False, debug: bool | None = None, - stream_version: Literal["v1", "v2"] = "v1", + stream_version: Literal["v1", "v2"] = "v2", **kwargs: Unpack[DeprecatedKwargs], ) -> AsyncIterator[dict[str, Any] | Any]: """Asynchronously stream graph steps for a single input. @@ -3162,6 +3162,23 @@ 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: 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"], + **kwargs: Any, + ) -> dict[str, Any] | Any: ... + @overload def invoke( self, @@ -3175,7 +3192,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v2"] = ..., **kwargs: Any, ) -> GraphOutput[OutputT]: ... @@ -3192,11 +3209,10 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v2"] = ..., **kwargs: Any, ) -> list[StreamPart[OutputT, StateT]]: ... - @overload def invoke( self, input: InputT | Command | None, @@ -3209,23 +3225,7 @@ 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", + stream_version: Literal["v1", "v2"] = "v2", **kwargs: Any, ) -> dict[str, Any] | Any: """Run the graph with a single input and config. @@ -3332,6 +3332,23 @@ class Pregel( else: return chunks + @overload + 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"], + **kwargs: Any, + ) -> dict[str, Any] | Any: ... + @overload async def ainvoke( self, @@ -3345,7 +3362,7 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v2"] = ..., **kwargs: Any, ) -> GraphOutput[OutputT]: ... @@ -3362,11 +3379,10 @@ class Pregel( interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, durability: Durability | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v2"] = ..., **kwargs: Any, ) -> list[StreamPart[OutputT, StateT]]: ... - @overload async def ainvoke( self, input: InputT | Command | None, @@ -3379,23 +3395,7 @@ 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", + stream_version: Literal["v1", "v2"] = "v2", **kwargs: Any, ) -> dict[str, Any] | Any: """Asynchronously run the graph with a single input and config. @@ -3556,7 +3556,7 @@ def _output( stream_subgraphs: bool, getter: Callable[[], tuple[tuple[str, ...], str, Any]], empty_exc: type[Exception], - stream_version: Literal["v1", "v2"] = "v1", + stream_version: Literal["v1", "v2"] = "v2", output_mapper: Callable[[Any], Any] | None = None, state_mapper: Callable[[Any], Any] | None = None, ) -> Iterator: diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index bd147d944..f508614e1 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -116,12 +116,41 @@ 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["v2"], + stream_version: Literal["v1"], + ) -> Iterator[dict[str, Any] | Any]: ... + + @overload + @abstractmethod + def stream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode | list[StreamMode] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + subgraphs: bool = False, + stream_version: Literal["v2"] = ..., + ) -> Iterator[StreamPart[OutputT, StateT]]: ... + + @abstractmethod + def stream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode | list[StreamMode] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + subgraphs: bool = False, + stream_version: Literal["v1", "v2"] = "v2", ) -> Iterator[StreamPart[OutputT, StateT]]: ... @overload @abstractmethod - def stream( + def astream( self, input: InputT | Command | None, config: RunnableConfig | None = None, @@ -131,22 +160,8 @@ 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"] = ..., - ) -> Iterator[dict[str, Any] | Any]: ... - - @abstractmethod - def stream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: StreamMode | list[StreamMode] | None = None, - 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]: ... + stream_version: Literal["v1"], + ) -> AsyncIterator[dict[str, Any] | Any]: ... @overload @abstractmethod @@ -160,37 +175,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["v2"], + stream_version: Literal["v2"] = ..., + ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ... + + @abstractmethod + def astream( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + stream_mode: StreamMode | list[StreamMode] | None = None, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + subgraphs: bool = False, + stream_version: Literal["v1", "v2"] = "v2", ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ... @overload @abstractmethod - def astream( + def invoke( self, input: InputT | Command | None, config: RunnableConfig | None = None, *, context: ContextT | None = None, - stream_mode: StreamMode | list[StreamMode] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - subgraphs: bool = False, - stream_version: Literal["v1"] = ..., - ) -> AsyncIterator[dict[str, Any] | Any]: ... - - @abstractmethod - def astream( - self, - input: InputT | Command | None, - config: RunnableConfig | None = None, - *, - context: ContextT | None = None, - stream_mode: StreamMode | list[StreamMode] | None = None, - 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]: ... + stream_version: Literal["v1"], + ) -> dict[str, Any] | Any: ... @overload @abstractmethod @@ -202,10 +215,9 @@ 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["v2"], + stream_version: Literal["v2"] = ..., ) -> GraphOutput[OutputT]: ... - @overload @abstractmethod def invoke( self, @@ -215,32 +227,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"] = ..., - ) -> dict[str, Any] | Any: ... - - @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", "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"], + stream_version: Literal["v1", "v2"] = "v2", ) -> GraphOutput[OutputT]: ... @overload @@ -253,9 +240,22 @@ 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"] = ..., + stream_version: Literal["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"] = ..., + ) -> GraphOutput[OutputT]: ... + @abstractmethod async def ainvoke( self, @@ -265,8 +265,8 @@ 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: ... + stream_version: Literal["v1", "v2"] = "v2", + ) -> GraphOutput[OutputT]: ... StreamChunk = tuple[tuple[str, ...], str, Any] diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 54d119485..676030cc7 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -697,9 +697,9 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v1"], **kwargs: Any, - ) -> Iterator[StreamPart]: ... + ) -> Iterator[dict[str, Any] | Any]: ... @overload def stream( @@ -713,9 +713,9 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v1"] = ..., + stream_version: Literal["v2"] = ..., **kwargs: Any, - ) -> Iterator[dict[str, Any] | Any]: ... + ) -> Iterator[StreamPart]: ... def stream( self, @@ -728,7 +728,7 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v1", "v2"] = "v1", + stream_version: Literal["v1", "v2"] = "v2", **kwargs: Any, ) -> Iterator[dict[str, Any] | Any]: """Create a run and stream the results. @@ -842,9 +842,9 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v1"], **kwargs: Any, - ) -> AsyncIterator[StreamPart]: ... + ) -> AsyncIterator[dict[str, Any] | Any]: ... @overload def astream( @@ -858,9 +858,9 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v1"] = ..., + stream_version: Literal["v2"] = ..., **kwargs: Any, - ) -> AsyncIterator[dict[str, Any] | Any]: ... + ) -> AsyncIterator[StreamPart]: ... async def astream( self, @@ -873,7 +873,7 @@ class RemoteGraph(PregelProtocol): subgraphs: bool = False, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v1", "v2"] = "v1", + stream_version: Literal["v1", "v2"] = "v2", **kwargs: Any, ) -> AsyncIterator[dict[str, Any] | Any]: """Create a run and stream the results. @@ -1001,9 +1001,9 @@ class RemoteGraph(PregelProtocol): interrupt_after: All | Sequence[str] | None = None, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v1"], **kwargs: Any, - ) -> GraphOutput[dict[str, Any]]: ... + ) -> dict[str, Any] | Any: ... @overload def invoke( @@ -1015,9 +1015,9 @@ class RemoteGraph(PregelProtocol): interrupt_after: All | Sequence[str] | None = None, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v1"] = ..., + stream_version: Literal["v2"] = ..., **kwargs: Any, - ) -> dict[str, Any] | Any: ... + ) -> GraphOutput[dict[str, Any]]: ... def invoke( self, @@ -1028,7 +1028,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", + stream_version: Literal["v1", "v2"] = "v2", **kwargs: Any, ) -> dict[str, Any] | Any: """Create a run, wait until it finishes and return the final state. @@ -1079,9 +1079,9 @@ class RemoteGraph(PregelProtocol): interrupt_after: All | Sequence[str] | None = None, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v2"], + stream_version: Literal["v1"], **kwargs: Any, - ) -> GraphOutput[dict[str, Any]]: ... + ) -> dict[str, Any] | Any: ... @overload async def ainvoke( @@ -1093,9 +1093,9 @@ class RemoteGraph(PregelProtocol): interrupt_after: All | Sequence[str] | None = None, headers: dict[str, str] | None = None, params: QueryParamTypes | None = None, - stream_version: Literal["v1"] = ..., + stream_version: Literal["v2"] = ..., **kwargs: Any, - ) -> dict[str, Any] | Any: ... + ) -> GraphOutput[dict[str, Any]]: ... async def ainvoke( self, @@ -1106,7 +1106,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", + stream_version: Literal["v1", "v2"] = "v2", **kwargs: Any, ) -> dict[str, Any] | Any: """Create a run, wait until it finishes and return the final state.