diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index c416d5f6a..7a5614f91 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -559,6 +559,7 @@ class StateGraph(Graph): for key, node in self.nodes.items(): compiled.attach_node(key, node) + compiled.attach_branch(START, SELF, CONTROL_BRANCH, with_reader=False) for key, node in self.nodes.items(): compiled.attach_branch(key, SELF, CONTROL_BRANCH, with_reader=False) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index f2df972d8..d54f8b8b2 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -14,6 +14,8 @@ from langgraph.constants import ( PUSH, RESUME, RETURN, + SELF, + START, TAG_HIDDEN, TASKS, ) @@ -79,12 +81,14 @@ def map_command( else: sends = [cmd.goto] for send in sends: - if not isinstance(send, Send): + if isinstance(send, Send): + yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send) + elif isinstance(send, str): + yield (NULL_TASK_ID, f"branch:{START}:{SELF}:{send}", START) + else: raise TypeError( - f"In Command.goto, expected Send, got {type(send).__name__}" + f"In Command.goto, expected Send/str, got {type(send).__name__}" ) - yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send) - # TODO handle goto str for state graph if cmd.resume: if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume): for tid, resume in cmd.resume.items(): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 20aaec74c..48c15e133 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7602,7 +7602,7 @@ def test_root_graph( content="result for query", name="search_api", tool_call_id="tool_call123", - id="00000000-0000-4000-8000-000000000033", + id="00000000-0000-4000-8000-000000000037", ) ] }, @@ -7625,7 +7625,7 @@ def test_root_graph( content="result for another", name="search_api", tool_call_id="tool_call456", - id="00000000-0000-4000-8000-000000000041", + id="00000000-0000-4000-8000-000000000045", ) ] }, @@ -8235,7 +8235,7 @@ def test_root_graph( "__root__": [ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000070", + id="00000000-0000-4000-8000-000000000078", ), AIMessage( content="", @@ -8255,7 +8255,7 @@ def test_root_graph( ), AIMessage(content="answer", id="ai2"), AIMessage( - content="an extra message", id="00000000-0000-4000-8000-000000000092" + content="an extra message", id="00000000-0000-4000-8000-000000000100" ), HumanMessage(content="what is weather in la"), ], @@ -14940,8 +14940,8 @@ def test_command_with_static_breakpoints( # Start the graph and interrupt at the first node graph.invoke({"foo": "abc"}, config) - result = graph.invoke(Command(update={"foo": "def"}), config) - assert result == {"foo": "def|node-1|node-2"} + result = graph.invoke(Command(resume="node1"), config) + assert result == {"foo": "abc|node-1|node-2"} @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) @@ -15003,3 +15003,42 @@ def test_multistep_plan(request: pytest.FixtureRequest, checkpointer_name: str): ], "plan": [], } + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_command_goto_with_static_breakpoints( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + """Use Command goto with static breakpoints.""" + + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + """The graph state.""" + + foo: Annotated[str, operator.add] + + def node1(state: State): + return { + "foo": "|node-1", + } + + def node2(state: State): + return { + "foo": "|node-2", + } + + builder = StateGraph(State) + builder.add_node("node1", node1) + builder.add_node("node2", node2) + builder.add_edge(START, "node1") + builder.add_edge("node1", "node2") + + graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) + + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + + # Start the graph and interrupt at the first node + graph.invoke({"foo": "abc"}, config) + result = graph.invoke(Command(goto=["node2"]), config) + assert result == {"foo": "abc|node-1|node-2|node-2"} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 5cc7f3312..cde54ae7c 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -13295,3 +13295,39 @@ async def test_multistep_plan(checkpointer_name: str): ], "plan": [], } + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> None: + """Use Command goto with static breakpoints.""" + + class State(TypedDict): + """The graph state.""" + + foo: Annotated[str, operator.add] + + def node1(state: State): + return { + "foo": "|node-1", + } + + def node2(state: State): + return { + "foo": "|node-2", + } + + builder = StateGraph(State) + builder.add_node("node1", node1) + builder.add_node("node2", node2) + builder.add_edge(START, "node1") + builder.add_edge("node1", "node2") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) + + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + + # Start the graph and interrupt at the first node + await graph.ainvoke({"foo": "abc"}, config) + result = await graph.ainvoke(Command(goto=["node2"]), config) + assert result == {"foo": "abc|node-1|node-2|node-2"}