diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 2397c1978..21d4d593b 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -99,6 +99,7 @@ from langgraph.types import ( All, CachePolicy, Checkpointer, + Command, Interrupt, Send, StateSnapshot, @@ -2342,7 +2343,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou def stream( self, - input: InputT, + input: InputT | Command | None, config: RunnableConfig | None = None, *, stream_mode: StreamMode | Sequence[StreamMode] | None = None, @@ -2564,7 +2565,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou async def astream( self, - input: InputT, + input: InputT | Command | None, config: RunnableConfig | None = None, *, stream_mode: StreamMode | Sequence[StreamMode] | None = None, @@ -2808,7 +2809,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou def invoke( self, - input: InputT, + input: InputT | Command | None, config: RunnableConfig | None = None, *, stream_mode: StreamMode = "values", @@ -2883,7 +2884,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou async def ainvoke( self, - input: InputT, + input: InputT | Command | None, config: RunnableConfig | None = None, *, stream_mode: StreamMode = "values", diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index e55be3cbc..dce0ef870 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -9,6 +9,7 @@ from langchain_core.runnables.graph import Graph as DrawableGraph from typing_extensions import Self from langgraph.pregel.types import All, StateSnapshot, StateUpdate, StreamMode +from langgraph.types import Command from langgraph.typing import InputT, OutputT, StateT @@ -98,7 +99,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], AB @abstractmethod def stream( self, - input: InputT, + input: InputT | Command | None, config: RunnableConfig | None = None, *, stream_mode: StreamMode | list[StreamMode] | None = None, @@ -110,7 +111,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], AB @abstractmethod def astream( self, - input: InputT, + input: InputT | Command | None, config: RunnableConfig | None = None, *, stream_mode: StreamMode | list[StreamMode] | None = None, @@ -122,7 +123,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], AB @abstractmethod def invoke( self, - input: InputT, + input: InputT | Command | None, config: RunnableConfig | None = None, *, interrupt_before: All | Sequence[str] | None = None, @@ -132,7 +133,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, InputT, OutputT], AB @abstractmethod async def ainvoke( self, - input: InputT, + input: InputT | Command | None, config: RunnableConfig | None = None, *, interrupt_before: All | Sequence[str] | None = None,