From 91aa66f4cf54e1479fb200463befa002f6d50da5 Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Fri, 17 Jan 2025 15:16:01 -0500 Subject: [PATCH] tests: add tests for calling multiple subgraphs in a parent node (#3070) Co-authored-by: Nuno Campos --- libs/langgraph/langgraph/pregel/__init__.py | 6 + libs/langgraph/tests/test_prebuilt.py | 79 ++++++ libs/langgraph/tests/test_pregel.py | 261 ++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 278 ++++++++++++++++++++ 4 files changed, 624 insertions(+) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 877ddcf4e..bac334e76 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1831,6 +1831,12 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after, debug=debug, ) + # set up subgraph checkpointing + if self.checkpointer is True: + ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) + config[CONF][CONFIG_KEY_CHECKPOINT_NS] = NS_SEP.join( + part.split(NS_END)[0] for part in ns.split(NS_SEP) + ) # set up messages stream mode if "messages" in stream_modes: run_manager.inheritable_handlers.append( diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 2a7f57001..99263062a 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -2088,3 +2088,82 @@ def test_inspect_react() -> None: model = FakeToolCallingModel(tool_calls=[]) agent = create_react_agent(model, []) inspect.getclosurevars(agent.nodes["agent"].bound.func) + + +def test_react_with_subgraph_tools() -> None: + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + # Define the subgraphs + def add(state): + return {"result": state["a"] + state["b"]} + + add_subgraph = ( + StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile() + ) + + def multiply(state): + return {"result": state["a"] * state["b"]} + + multiply_subgraph = ( + StateGraph(State, output=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) + + multiply_subgraph.invoke({"a": 2, "b": 3}) + + # Add subgraphs as tools + + def addition(a: int, b: int): + """Add two numbers""" + return add_subgraph.invoke({"a": a, "b": b})["result"] + + def multiplication(a: int, b: int): + """Multiply two numbers""" + return multiply_subgraph.invoke({"a": a, "b": b})["result"] + + model = FakeToolCallingModel( + tool_calls=[ + [ + {"args": {"a": 2, "b": 3}, "id": "1", "name": "addition"}, + {"args": {"a": 2, "b": 3}, "id": "2", "name": "multiplication"}, + ], + [], + ] + ) + checkpointer = MemorySaver() + tool_node = ToolNode([addition, multiplication], handle_tool_errors=False) + agent = create_react_agent(model, tool_node, checkpointer=checkpointer) + result = agent.invoke( + {"messages": [HumanMessage(content="What's 2 + 3 and 2 * 3?")]}, + config={"configurable": {"thread_id": "1"}}, + ) + assert result["messages"] == [ + _AnyIdHumanMessage(content="What's 2 + 3 and 2 * 3?"), + AIMessage( + content="What's 2 + 3 and 2 * 3?", + id="0", + tool_calls=[ + ToolCall(name="addition", args={"a": 2, "b": 3}, id="1"), + ToolCall(name="multiplication", args={"a": 2, "b": 3}, id="2"), + ], + ), + ToolMessage( + content="5", name="addition", tool_call_id="1", id=result["messages"][2].id + ), + ToolMessage( + content="6", + name="multiplication", + tool_call_id="2", + id=result["messages"][3].id, + ), + AIMessage( + content="What's 2 + 3 and 2 * 3?-What's 2 + 3 and 2 * 3?-5-6", id="1" + ), + ] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a30a66cba..5995cd94f 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5797,3 +5797,264 @@ async def test_entrypoint_from_async_generator() -> None: assert previous_return_values == [None] assert list(foo.invoke({"a": "2"}, config)) == ["a", "b"] assert previous_return_values == [None, ["a", "b"]] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_multiple_subgraphs( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + # Define the subgraphs + def add(state): + return {"result": state["a"] + state["b"]} + + add_subgraph = ( + StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile() + ) + + def multiply(state): + return {"result": state["a"] * state["b"]} + + multiply_subgraph = ( + StateGraph(State, output=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) + + # Test calling the same subgraph multiple times + def call_same_subgraph(state): + result = add_subgraph.invoke(state) + another_result = add_subgraph.invoke({"a": result["result"], "b": 10}) + return another_result + + parent_call_same_subgraph = ( + StateGraph(State, output=Output) + .add_node(call_same_subgraph) + .add_edge(START, "call_same_subgraph") + .compile(checkpointer=checkpointer) + ) + config = {"configurable": {"thread_id": "1"}} + assert parent_call_same_subgraph.invoke({"a": 2, "b": 3}, config) == {"result": 15} + + # Test calling multiple subgraphs + class Output(TypedDict): + add_result: int + multiply_result: int + + def call_multiple_subgraphs(state): + add_result = add_subgraph.invoke(state) + multiply_result = multiply_subgraph.invoke(state) + return { + "add_result": add_result["result"], + "multiply_result": multiply_result["result"], + } + + parent_call_multiple_subgraphs = ( + StateGraph(State, output=Output) + .add_node(call_multiple_subgraphs) + .add_edge(START, "call_multiple_subgraphs") + .compile(checkpointer=checkpointer) + ) + config = {"configurable": {"thread_id": "2"}} + assert parent_call_multiple_subgraphs.invoke({"a": 2, "b": 3}, config) == { + "add_result": 5, + "multiply_result": 6, + } + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_multiple_subgraphs_functional( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + # Define addition subgraph + @entrypoint() + def add(inputs): + a, b = inputs + return a + b + + # Define multiplication subgraph using tasks + @task + def multiply_task(a, b): + return a * b + + @entrypoint() + def multiply(inputs): + return multiply_task(*inputs).result() + + # Test calling the same subgraph multiple times + @task + def call_same_subgraph(a, b): + result = add.invoke([a, b]) + another_result = add.invoke([result, 10]) + return another_result + + @entrypoint(checkpointer=checkpointer) + def parent_call_same_subgraph(inputs): + return call_same_subgraph(*inputs).result() + + config = {"configurable": {"thread_id": "1"}} + assert parent_call_same_subgraph.invoke([2, 3], config) == 15 + + # Test calling multiple subgraphs + @task + def call_multiple_subgraphs(a, b): + add_result = add.invoke([a, b]) + multiply_result = multiply.invoke([a, b]) + return [add_result, multiply_result] + + @entrypoint(checkpointer=checkpointer) + def parent_call_multiple_subgraphs(inputs): + return call_multiple_subgraphs(*inputs).result() + + config = {"configurable": {"thread_id": "2"}} + assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_multiple_subgraphs_mixed( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + # Define the subgraphs + def add(state): + return {"result": state["a"] + state["b"]} + + add_subgraph = ( + StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile() + ) + + def multiply(state): + return {"result": state["a"] * state["b"]} + + multiply_subgraph = ( + StateGraph(State, output=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) + + # Test calling the same subgraph multiple times + @task + def call_same_subgraph(a, b): + result = add_subgraph.invoke({"a": a, "b": b})["result"] + another_result = add_subgraph.invoke({"a": result, "b": 10})["result"] + return another_result + + @entrypoint(checkpointer=checkpointer) + def parent_call_same_subgraph(inputs): + return call_same_subgraph(*inputs).result() + + config = {"configurable": {"thread_id": "1"}} + assert parent_call_same_subgraph.invoke([2, 3], config) == 15 + + # Test calling multiple subgraphs + @task + def call_multiple_subgraphs(a, b): + add_result = add_subgraph.invoke({"a": a, "b": b})["result"] + multiply_result = multiply_subgraph.invoke({"a": a, "b": b})["result"] + return [add_result, multiply_result] + + @entrypoint(checkpointer=checkpointer) + def parent_call_multiple_subgraphs(inputs): + return call_multiple_subgraphs(*inputs).result() + + config = {"configurable": {"thread_id": "2"}} + assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_multiple_subgraphs_mixed_checkpointer( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class SubgraphState(TypedDict): + sub_counter: Annotated[int, operator.add] + + def subgraph_node(state): + return {"sub_counter": 2} + + sub_graph_1 = ( + StateGraph(SubgraphState) + .add_node(subgraph_node) + .add_edge(START, "subgraph_node") + .compile(checkpointer=True) + ) + + class OtherSubgraphState(TypedDict): + other_sub_counter: Annotated[int, operator.add] + + def other_subgraph_node(state): + return {"other_sub_counter": 3} + + sub_graph_2 = ( + StateGraph(OtherSubgraphState) + .add_node(other_subgraph_node) + .add_edge(START, "other_subgraph_node") + .compile() + ) + + class ParentState(TypedDict): + parent_counter: int + + def parent_node(state): + result = sub_graph_1.invoke({"sub_counter": state["parent_counter"]}) + other_result = sub_graph_2.invoke({"other_sub_counter": result["sub_counter"]}) + return {"parent_counter": other_result["other_sub_counter"]} + + parent_graph = ( + StateGraph(ParentState) + .add_node(parent_node) + .add_edge(START, "parent_node") + .compile(checkpointer=checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + assert parent_graph.invoke({"parent_counter": 0}, config) == {"parent_counter": 5} + assert parent_graph.invoke({"parent_counter": 0}, config) == {"parent_counter": 7} + config = {"configurable": {"thread_id": "2"}} + assert [ + c + for c in parent_graph.stream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" + ) + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 5}}), + ] + assert [ + c + for c in parent_graph.stream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" + ) + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 7}}), + ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 6767a1510..566ea364c 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6928,3 +6928,281 @@ async def test_async_streaming_with_functional_api() -> None: delta = arrival_times[1] - arrival_times[0] # Delta cannot be less than 10 ms if it is streaming as results are generated. assert delta > time_delay + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_multiple_subgraphs(checkpointer_name: str) -> None: + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + async with awith_checkpointer(checkpointer_name) as checkpointer: + # Define the subgraphs + async def add(state): + return {"result": state["a"] + state["b"]} + + add_subgraph = ( + StateGraph(State, output=Output) + .add_node(add) + .add_edge(START, "add") + .compile() + ) + + async def multiply(state): + return {"result": state["a"] * state["b"]} + + multiply_subgraph = ( + StateGraph(State, output=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) + + # Test calling the same subgraph multiple times + async def call_same_subgraph(state): + result = await add_subgraph.ainvoke(state) + another_result = await add_subgraph.ainvoke( + {"a": result["result"], "b": 10} + ) + return another_result + + parent_call_same_subgraph = ( + StateGraph(State, output=Output) + .add_node(call_same_subgraph) + .add_edge(START, "call_same_subgraph") + .compile(checkpointer=checkpointer) + ) + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == { + "result": 15 + } + + # Test calling multiple subgraphs + class Output(TypedDict): + add_result: int + multiply_result: int + + async def call_multiple_subgraphs(state): + add_result = await add_subgraph.ainvoke(state) + multiply_result = await multiply_subgraph.ainvoke(state) + return { + "add_result": add_result["result"], + "multiply_result": multiply_result["result"], + } + + parent_call_multiple_subgraphs = ( + StateGraph(State, output=Output) + .add_node(call_multiple_subgraphs) + .add_edge(START, "call_multiple_subgraphs") + .compile(checkpointer=checkpointer) + ) + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke( + {"a": 2, "b": 3}, config + ) == { + "add_result": 5, + "multiply_result": 6, + } + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_multiple_subgraphs_functional(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + # Define addition subgraph + @entrypoint() + async def add(inputs): + a, b = inputs + return a + b + + # Define multiplication subgraph using tasks + @task + async def multiply_task(a, b): + return a * b + + @entrypoint() + async def multiply(inputs): + return await multiply_task(*inputs) + + # Test calling the same subgraph multiple times + @task + async def call_same_subgraph(a, b): + result = await add.ainvoke([a, b]) + another_result = await add.ainvoke([result, 10]) + return another_result + + @entrypoint(checkpointer=checkpointer) + async def parent_call_same_subgraph(inputs): + return await call_same_subgraph(*inputs) + + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 + + # Test calling multiple subgraphs + @task + async def call_multiple_subgraphs(a, b): + add_result = await add.ainvoke([a, b]) + multiply_result = await multiply.ainvoke([a, b]) + return [add_result, multiply_result] + + @entrypoint(checkpointer=checkpointer) + async def parent_call_multiple_subgraphs(inputs): + return await call_multiple_subgraphs(*inputs) + + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_multiple_subgraphs_mixed(checkpointer_name: str) -> None: + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + async with awith_checkpointer(checkpointer_name) as checkpointer: + # Define the subgraphs + async def add(state): + return {"result": state["a"] + state["b"]} + + add_subgraph = ( + StateGraph(State, output=Output) + .add_node(add) + .add_edge(START, "add") + .compile() + ) + + async def multiply(state): + return {"result": state["a"] * state["b"]} + + multiply_subgraph = ( + StateGraph(State, output=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) + + # Test calling the same subgraph multiple times + @task + async def call_same_subgraph(a, b): + result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] + another_result = (await add_subgraph.ainvoke({"a": result, "b": 10}))[ + "result" + ] + return another_result + + @entrypoint(checkpointer=checkpointer) + async def parent_call_same_subgraph(inputs): + return await call_same_subgraph(*inputs) + + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 + + # Test calling multiple subgraphs + @task + async def call_multiple_subgraphs(a, b): + add_result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] + multiply_result = (await multiply_subgraph.ainvoke({"a": a, "b": b}))[ + "result" + ] + return [add_result, multiply_result] + + @entrypoint(checkpointer=checkpointer) + async def parent_call_multiple_subgraphs(inputs): + return await call_multiple_subgraphs(*inputs) + + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_multiple_subgraphs_mixed_checkpointer( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + class SubgraphState(TypedDict): + sub_counter: Annotated[int, operator.add] + + async def subgraph_node(state): + return {"sub_counter": 2} + + sub_graph_1 = ( + StateGraph(SubgraphState) + .add_node(subgraph_node) + .add_edge(START, "subgraph_node") + .compile(checkpointer=True) + ) + + class OtherSubgraphState(TypedDict): + other_sub_counter: Annotated[int, operator.add] + + async def other_subgraph_node(state): + return {"other_sub_counter": 3} + + sub_graph_2 = ( + StateGraph(OtherSubgraphState) + .add_node(other_subgraph_node) + .add_edge(START, "other_subgraph_node") + .compile() + ) + + class ParentState(TypedDict): + parent_counter: int + + async def parent_node(state): + result = await sub_graph_1.ainvoke({"sub_counter": state["parent_counter"]}) + other_result = await sub_graph_2.ainvoke( + {"other_sub_counter": result["sub_counter"]} + ) + return {"parent_counter": other_result["other_sub_counter"]} + + parent_graph = ( + StateGraph(ParentState) + .add_node(parent_node) + .add_edge(START, "parent_node") + .compile(checkpointer=checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { + "parent_counter": 5 + } + assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { + "parent_counter": 7 + } + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in parent_graph.astream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" + ) + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 5}}), + ] + assert [ + c + async for c in parent_graph.astream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" + ) + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 7}}), + ]