Merge pull request #39 from langchain-ai/nc/jan16/recursion-error

Add recursion check

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)

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 == {

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