From 8154114e11dc580ca58176bc1c922377f9f41477 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 16 Jan 2024 16:13:03 -0800 Subject: [PATCH 1/2] Add recursion check --- langgraph/pregel/__init__.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index f29c0b92a..6a1346480 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -346,6 +346,11 @@ class Pregel( # interrupt if any channel written to is in interrupt list if any(chan for chan, _ in pending_writes if chan in self.interrupt): break + else: + raise RuntimeError( + f"Recursion limit of {config['recursion_limit']} reached" + "without hitting a stop condition." + ) # save end of run checkpoint if ( @@ -483,6 +488,11 @@ class Pregel( # interrupt if any channel written to is in interrupt list if any(chan for chan, _ in pending_writes if chan in self.interrupt): break + else: + raise RuntimeError( + f"Recursion limit of {config['recursion_limit']} reached" + "without hitting a stop condition." + ) # save end of run checkpoint if ( From b535b5caa04618ba5795179099a5b092a4a9b97b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 16 Jan 2024 16:23:41 -0800 Subject: [PATCH 2/2] Add tests --- langgraph/pregel/__init__.py | 30 ++++++++++++++++++------------ tests/test_pregel.py | 5 ++++- tests/test_pregel_async.py | 5 ++++- 3 files changed, 26 insertions(+), 14 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 6a1346480..b557b9c5d 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -72,6 +72,10 @@ WriteValue = Union[ ] +class GraphRecursionError(RecursionError): + pass + + def _coerce_write_value(value: WriteValue) -> Runnable[Input, Output]: if not isinstance(value, Runnable) and not callable(value): return coerce_to_runnable(lambda _: value) @@ -275,12 +279,18 @@ class Pregel( # channel updates from step N are only visible in step N+1 # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps - for step in range(config["recursion_limit"]): + for step in range(config["recursion_limit"] + 1): next_tasks = _prepare_next_tasks(checkpoint, processes, channels) # if no more tasks, we're done if not next_tasks: break + elif step == config["recursion_limit"]: + raise GraphRecursionError( + f"Recursion limit of {config['recursion_limit']} reached" + "without hitting a stop condition. You can increase the limit" + "by setting the `recursion_limit` config key." + ) if self.debug: print_step_start(step, next_tasks) @@ -346,11 +356,6 @@ class Pregel( # interrupt if any channel written to is in interrupt list if any(chan for chan, _ in pending_writes if chan in self.interrupt): break - else: - raise RuntimeError( - f"Recursion limit of {config['recursion_limit']} reached" - "without hitting a stop condition." - ) # save end of run checkpoint if ( @@ -412,12 +417,18 @@ class Pregel( # channel updates from step N are only visible in step N+1, # channels are guaranteed to be immutable for the duration of the step, # channel updates being applied only at the transition between steps - for step in range(config["recursion_limit"]): + for step in range(config["recursion_limit"] + 1): next_tasks = _prepare_next_tasks(checkpoint, processes, channels) # if no more tasks, we're done if not next_tasks: break + elif step == config["recursion_limit"]: + raise GraphRecursionError( + f"Recursion limit of {config['recursion_limit']} reached" + "without hitting a stop condition. You can increase the limit" + "by setting the `recursion_limit` config key." + ) if self.debug: print_step_start(step, next_tasks) @@ -488,11 +499,6 @@ class Pregel( # interrupt if any channel written to is in interrupt list if any(chan for chan, _ in pending_writes if chan in self.interrupt): break - else: - raise RuntimeError( - f"Recursion limit of {config['recursion_limit']} reached" - "without hitting a stop condition." - ) # save end of run checkpoint if ( diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 961cd370c..fd4a986b5 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -16,7 +16,7 @@ from langgraph.channels.topic import Topic from langgraph.checkpoint.memory import MemorySaver from langgraph.graph import END, Graph from langgraph.graph.state import StateGraph -from langgraph.pregel import Channel, Pregel +from langgraph.pregel import Channel, GraphRecursionError, Pregel from langgraph.pregel.reserved import ReservedChannels @@ -159,6 +159,9 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert app.invoke(2, input_keys="inbox") == 3 + with pytest.raises(GraphRecursionError): + app.invoke(2, {"recursion_limit": 1}) + for step, values in enumerate(app.stream(2), start=1): if step == 1: assert values == { diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 012992ecb..6799bdd73 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -23,7 +23,7 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.memory import MemorySaver from langgraph.graph import END, Graph, StateGraph -from langgraph.pregel import Channel, Pregel +from langgraph.pregel import Channel, GraphRecursionError, Pregel from langgraph.pregel.reserved import ReservedChannels @@ -168,6 +168,9 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert await app.ainvoke(2, input_keys="inbox") == 3 + with pytest.raises(GraphRecursionError): + await app.ainvoke(2, {"recursion_limit": 1}) + step = 0 async for values in app.astream(2): step += 1