diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index f29c0b92a..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) @@ -407,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) 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