From 79562f3f3775b479436711a6714f6eb23c4bce5d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 9 Dec 2024 13:54:42 -0800 Subject: [PATCH 1/8] lib: Add support for invoke(Command(goto=)) --- libs/langgraph/langgraph/graph/state.py | 1 + libs/langgraph/langgraph/pregel/io.py | 12 ++++++++---- 2 files changed, 9 insertions(+), 4 deletions(-) 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(): From 5f869b9e752fab37371db78e3de88e65574c005a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 9 Dec 2024 14:10:57 -0800 Subject: [PATCH 2/8] Update test --- libs/langgraph/tests/test_pregel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 20aaec74c..00f588e58 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", ) ] }, From a9b94f93eec54ab27009a4ff3b1ad25ebf45c214 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 9 Dec 2024 14:12:17 -0800 Subject: [PATCH 3/8] Update again --- libs/langgraph/tests/test_pregel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 00f588e58..d8161a439 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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-0000000000100" ), HumanMessage(content="what is weather in la"), ], From df5d08f689cc3c3ad13191255ff737bf0aad73d8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 9 Dec 2024 14:14:57 -0800 Subject: [PATCH 4/8] Fix --- libs/langgraph/tests/test_pregel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d8161a439..6e610365d 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8255,7 +8255,7 @@ def test_root_graph( ), AIMessage(content="answer", id="ai2"), AIMessage( - content="an extra message", id="00000000-0000-4000-8000-0000000000100" + content="an extra message", id="00000000-0000-4000-8000-000000000100" ), HumanMessage(content="what is weather in la"), ], From dd778f8ed6523e6edc1e065dcc28c617bafa3f4e Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Tue, 10 Dec 2024 11:51:54 -0500 Subject: [PATCH 5/8] qxqx --- libs/langgraph/tests/test_pregel.py | 39 +++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 6e610365d..276dcd141 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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: str + + def node1(state: State): + return { + "foo": state["foo"] + "|node-1", + } + + def node2(state: State): + return { + "foo": state["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-2"} From f9cdfd3ac4fbdb9f477699d676b1fa5d5eb3ccbf Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Tue, 10 Dec 2024 11:54:43 -0500 Subject: [PATCH 6/8] x --- libs/langgraph/tests/test_pregel.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 276dcd141..6c2a20f63 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -14940,9 +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) def test_multistep_plan(request: pytest.FixtureRequest, checkpointer_name: str): From 7cabc0a3dc4a88be4633c1988647b30fed7ee423 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Tue, 10 Dec 2024 11:54:59 -0500 Subject: [PATCH 7/8] reformat --- libs/langgraph/tests/test_pregel.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 6c2a20f63..46dc876f7 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -14943,6 +14943,7 @@ def test_command_with_static_breakpoints( result = graph.invoke(Command(resume="node1"), config) assert result == {"foo": "abc|node-1|node-2"} + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multistep_plan(request: pytest.FixtureRequest, checkpointer_name: str): from langchain_core.messages import AnyMessage From a7ac9ffd4e6a04af1707e8f4d746b53632e19159 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 10 Dec 2024 11:31:41 -0800 Subject: [PATCH 8/8] Update test --- libs/langgraph/tests/test_pregel.py | 8 ++--- libs/langgraph/tests/test_pregel_async.py | 36 +++++++++++++++++++++++ 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 46dc876f7..48c15e133 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -15016,16 +15016,16 @@ def test_command_goto_with_static_breakpoints( class State(TypedDict): """The graph state.""" - foo: str + foo: Annotated[str, operator.add] def node1(state: State): return { - "foo": state["foo"] + "|node-1", + "foo": "|node-1", } def node2(state: State): return { - "foo": state["foo"] + "|node-2", + "foo": "|node-2", } builder = StateGraph(State) @@ -15041,4 +15041,4 @@ def test_command_goto_with_static_breakpoints( # 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-2"} + 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"}