mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 14:42:28 +02:00
When using global resume value, ensure subgraphs consume it (#3889)
- Previously the global resume value was passed to subgraphs without being consumed - This would result in two parallel subgraph calls being able to use the same resume value - Note this behavior can't be implemented over the wire, that will be fixed in future PR Closes #3398
This commit is contained in:
Nuno Campos
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user