diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index d3b95d9e2..c5b0cd958 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -14,7 +14,6 @@ from typing import ( Optional, Sequence, Type, - TypeVar, Union, cast, get_args, @@ -50,15 +49,13 @@ from langgraph.managed.base import ( from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import All, Checkpointer, Command, RetryPolicy +from langgraph.types import All, Checkpointer, Command, N, RetryPolicy from langgraph.utils.fields import get_field_default from langgraph.utils.pydantic import create_model from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable logger = logging.getLogger(__name__) -N = TypeVar("N") - def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None: if isinstance(schema, type): @@ -81,7 +78,7 @@ def _get_node_name(node: RunnableLike) -> str: raise TypeError(f"Unsupported node type: {type(node)}") -class GraphCommand(Command, Generic[N]): +class GraphCommand(Generic[N], Command[N]): """One or more commands to update a StateGraph's state and go to, or send messages to nodes.""" __slots__ = ("goto",) @@ -90,9 +87,9 @@ class GraphCommand(Command, Generic[N]): self, *, update: Optional[dict[str, Any]] = None, - goto: Union[str, Sequence[str]] = (), send: Union[Send, Sequence[Send]] = (), resume: Optional[Union[Any, dict[str, Any]]] = None, + goto: Union[str, Sequence[str]] = (), ) -> None: super().__init__(update=update, send=send, resume=resume) self.goto = goto @@ -390,7 +387,7 @@ class StateGraph(Graph): input = input_hint if ( (rtn := hints.get("return")) - and get_origin(rtn) is GraphCommand + and get_origin(rtn) in (Command, GraphCommand) and (rargs := get_args(rtn)) and get_origin(rargs[0]) is Literal and (vals := get_args(rargs[0])) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index d867d60be..b910616a2 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -4,11 +4,14 @@ from typing import ( TYPE_CHECKING, Any, Callable, + Generic, + Hashable, Literal, NamedTuple, Optional, Sequence, Type, + TypeVar, Union, cast, ) @@ -221,7 +224,10 @@ class Send: ) -class Command: +N = TypeVar("N", bound=Hashable) + + +class Command(Generic[N]): """One or more commands to update the graph's state and send messages to nodes.""" __slots__ = ("update", "send", "resume") diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 4268a310a..1cfdd9b76 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -5108,6 +5108,81 @@ ''' # --- +# name: test_send_react_interrupt_control[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([
__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[sqlite] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_simple_multi_edge ''' graph TD; diff --git a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr index 3d4021111..46916c7a4 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr @@ -1302,6 +1302,81 @@ +---------+ ''' # --- +# name: test_send_react_interrupt_control[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_aio_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_aio_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[sqlite_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_weather_subgraph[duckdb_aio] ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d801d9c98..81d45f4ee 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2691,7 +2691,7 @@ def test_send_react_interrupt( @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_react_interrupt_control( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion ) -> None: from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage @@ -2721,6 +2721,7 @@ def test_send_react_interrupt_control( builder.add_node(foo) builder.add_edge(START, "agent") graph = builder.compile() + assert graph.get_graph().draw_mermaid() == snapshot assert graph.invoke({"messages": [HumanMessage("hello")]}) == { "messages": [ diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 022076fd9..df85ed650 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2973,7 +2973,9 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_react_interrupt_control(checkpointer_name: str) -> None: +async def test_send_react_interrupt_control( + checkpointer_name: str, snapshot: SnapshotAssertion +) -> None: from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage ai_message = AIMessage( @@ -2982,7 +2984,7 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], ) - async def agent(state) -> GraphCommand[Literal["foo"]]: + async def agent(state) -> Command[Literal["foo"]]: return GraphCommand( update={"messages": ai_message}, send=[Send(call["name"], call) for call in ai_message.tool_calls], @@ -3000,6 +3002,7 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: builder.add_node(foo) builder.add_edge(START, "agent") graph = builder.compile() + assert graph.get_graph().draw_mermaid() == snapshot assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == { "messages": [