diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index e08f0eff4..84de26c05 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -506,6 +506,7 @@ def prepare_single_task( CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, CONFIG_KEY_SCRATCHPAD: _scratchpad( + config, pending_writes, task_id, ), @@ -615,6 +616,7 @@ def prepare_single_task( CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, CONFIG_KEY_SCRATCHPAD: _scratchpad( + config, pending_writes, task_id, ), @@ -740,6 +742,7 @@ def prepare_single_task( CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, CONFIG_KEY_SCRATCHPAD: _scratchpad( + config, pending_writes, task_id, ), @@ -761,15 +764,23 @@ def prepare_single_task( def _scratchpad( + config: RunnableConfig, pending_writes: list[PendingWrite], task_id: str, ) -> PregelScratchpad: + # None cannot be used as a resume value, because it would be difficult to + # distinguish from missing when used over http null_resume_write = next( (w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None ) + parent_scratchpad: Optional[PregelScratchpad] = config[CONF].get( + CONFIG_KEY_SCRATCHPAD + ) def get_null_resume(consume: bool = False) -> Any: if null_resume_write is None: + if parent_scratchpad is not None: + return parent_scratchpad.get_null_resume(consume) return None if consume: try: diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 6ef009058..7ab3431b7 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -587,14 +587,6 @@ class PregelLoop(LoopProtocol): ) ) - # take resume value from parent - if scratchpad := cast( - Optional[PregelScratchpad], configurable.get(CONFIG_KEY_SCRATCHPAD) - ): - if isinstance(scratchpad, PregelScratchpad): - null_resume = scratchpad.get_null_resume(False) - if null_resume is not None: - self.put_writes(NULL_TASK_ID, [(RESUME, null_resume)]) # map command to writes if isinstance(self.input, Command): if self.input.resume is not None and not self.checkpointer: diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 4c3e84c57..5efb89385 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -130,7 +130,7 @@ class Interrupt: value: Any resumable: bool = False ns: Optional[Sequence[str]] = None - when: Literal["during"] = "during" + when: Literal["during"] = dataclasses.field(default="during", repr=False) class PregelTask(NamedTuple): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index c0bc65b69..ccd899ce9 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7317,3 +7317,297 @@ def test_empty_invoke() -> None: "111": 111, "222": 222, } + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_parallel_interrupts( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + from pydantic import BaseModel, Field + + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + # --- CHILD GRAPH --- + + class ChildState(BaseModel): + prompt: str = Field(..., description="What is going to be asked to the user?") + human_input: Optional[str] = Field(None, description="What the human said") + human_inputs: Annotated[List[str], operator.add] = Field( + default_factory=list, description="All of my messages" + ) + + def get_human_input(state: ChildState): + human_input = interrupt(state.prompt) + + return dict( + human_input=human_input, # update child state + human_inputs=[human_input], # update parent state + ) + + child_graph_builder = StateGraph(ChildState) + child_graph_builder.add_node("get_human_input", get_human_input) + child_graph_builder.add_edge(START, "get_human_input") + child_graph_builder.add_edge("get_human_input", END) + child_graph = child_graph_builder.compile() + + # --- PARENT GRAPH --- + + class ParentState(BaseModel): + prompts: List[str] = Field( + ..., description="What is going to be asked to the user?" + ) + human_inputs: Annotated[List[str], operator.add] = Field( + default_factory=list, description="All of my messages" + ) + + def assign_workers(state: ParentState): + return [ + Send( + "child_graph", + dict( + prompt=prompt, + ), + ) + for prompt in state.prompts + ] + + def cleanup(state: ParentState): + assert len(state.human_inputs) == len(state.prompts) + + parent_graph_builder = StateGraph(ParentState) + parent_graph_builder.add_node("child_graph", child_graph) + parent_graph_builder.add_node("cleanup", cleanup) + + parent_graph_builder.add_conditional_edges(START, assign_workers, ["child_graph"]) + parent_graph_builder.add_edge("child_graph", "cleanup") + parent_graph_builder.add_edge("cleanup", END) + + parent_graph = parent_graph_builder.compile(checkpointer=checkpointer) + + # --- CLIENT INVOCATION --- + + thread_config = dict( + configurable=dict( + thread_id=str(uuid.uuid4()), + ) + ) + current_input = dict( + prompts=["a", "b"], + ) + + invokes = 0 + events: dict[int, list[dict]] = {} + while invokes < 10: + # reset interrupt + invokes += 1 + events[invokes] = [] + current_interrupts: list[Interrupt] = [] + + # start / resume the graph + for event in parent_graph.stream( + input=current_input, + config=thread_config, + stream_mode="updates", + ): + events[invokes].append(event) + # handle the interrupt + if "__interrupt__" in event: + current_interrupts.extend(event["__interrupt__"]) + # assume that it breaks here, because it is an interrupt + + # get human input and resume + if any(i.resumable for i in current_interrupts): + current_input = Command(resume=f"Resume #{invokes}") + + # not more human input required, must be completed + else: + break + else: + assert False, "Detected infinite loop" + + assert invokes == 3 + assert len(events) == 3 + + assert events[1] == UnsortedSequence( + { + "__interrupt__": ( + Interrupt( + value="a", + resumable=True, + ns=[ + AnyStr("child_graph:"), + AnyStr("get_human_input:"), + ], + ), + ) + }, + { + "__interrupt__": ( + Interrupt( + value="b", + resumable=True, + ns=[ + AnyStr("child_graph:"), + AnyStr("get_human_input:"), + ], + ), + ) + }, + ) + assert events[2] in ( + UnsortedSequence( + { + "__interrupt__": ( + Interrupt( + value="a", + resumable=True, + ns=[ + AnyStr("child_graph:"), + AnyStr("get_human_input:"), + ], + ), + ) + }, + {"child_graph": {"human_inputs": ["Resume #1"]}}, + ), + UnsortedSequence( + { + "__interrupt__": ( + Interrupt( + value="b", + resumable=True, + ns=[ + AnyStr("child_graph:"), + AnyStr("get_human_input:"), + ], + ), + ) + }, + {"child_graph": {"human_inputs": ["Resume #1"]}}, + ), + ) + assert events[3] == UnsortedSequence( + { + "child_graph": {"human_inputs": ["Resume #1"]}, + "__metadata__": {"cached": True}, + }, + {"child_graph": {"human_inputs": ["Resume #2"]}}, + {"cleanup": None}, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_parallel_interrupts_double( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + from pydantic import BaseModel, Field + + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + # --- CHILD GRAPH --- + + class ChildState(BaseModel): + prompt: str = Field(..., description="What is going to be asked to the user?") + human_input: Optional[str] = Field(None, description="What the human said") + human_inputs: Annotated[List[str], operator.add] = Field( + default_factory=list, description="All of my messages" + ) + + def get_human_input(state: ChildState): + human_input = interrupt(state.prompt) + + return dict( + human_inputs=[human_input], # update parent state + ) + + def get_dolphin_input(state: ChildState): + human_input = interrupt(state.prompt) + + return dict( + human_inputs=[human_input], # update parent state + ) + + child_graph_builder = StateGraph(ChildState) + child_graph_builder.add_node("get_human_input", get_human_input) + child_graph_builder.add_node("get_dolphin_input", get_dolphin_input) + child_graph_builder.add_edge(START, "get_human_input") + child_graph_builder.add_edge(START, "get_dolphin_input") + child_graph = child_graph_builder.compile() + + # --- PARENT GRAPH --- + + class ParentState(BaseModel): + prompts: List[str] = Field( + ..., description="What is going to be asked to the user?" + ) + human_inputs: Annotated[List[str], operator.add] = Field( + default_factory=list, description="All of my messages" + ) + + def assign_workers(state: ParentState): + return [ + Send( + "child_graph", + dict( + prompt=prompt, + ), + ) + for prompt in state.prompts + ] + + def cleanup(state: ParentState): + assert len(state.human_inputs) == len(state.prompts) * 2 + + parent_graph_builder = StateGraph(ParentState) + parent_graph_builder.add_node("child_graph", child_graph) + parent_graph_builder.add_node("cleanup", cleanup) + + parent_graph_builder.add_conditional_edges(START, assign_workers, ["child_graph"]) + parent_graph_builder.add_edge("child_graph", "cleanup") + parent_graph_builder.add_edge("cleanup", END) + + parent_graph = parent_graph_builder.compile(checkpointer=checkpointer) + + # --- CLIENT INVOCATION --- + + thread_config = dict( + configurable=dict( + thread_id=str(uuid.uuid4()), + ) + ) + current_input = dict( + prompts=["a", "b"], + ) + + invokes = 0 + events: dict[int, list[dict]] = {} + while invokes < 10: + # reset interrupt + invokes += 1 + events[invokes] = [] + current_interrupts: list[Interrupt] = [] + + # start / resume the graph + for event in parent_graph.stream( + input=current_input, + config=thread_config, + stream_mode="updates", + ): + events[invokes].append(event) + # handle the interrupt + if "__interrupt__" in event: + current_interrupts.extend(event["__interrupt__"]) + # assume that it breaks here, because it is an interrupt + + # get human input and resume + if any(i.resumable for i in current_interrupts): + current_input = Command(resume=f"Resume #{invokes}") + + # not more human input required, must be completed + else: + break + else: + assert False, "Detected infinite loop" + + assert invokes == 5 + assert len(events) == 5