diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d2a66a0b9..9fa0770a0 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -92,7 +92,7 @@ from langgraph.pregel.algo import ( prepare_next_tasks, ) from langgraph.pregel.debug import tasks_w_writes -from langgraph.pregel.io import read_channels +from langgraph.pregel.io import map_input, read_channels from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.messages import StreamMessagesHandler @@ -1332,6 +1332,58 @@ class Pregel(PregelProtocol): return patch_checkpoint_map( next_config, saved.metadata if saved else None ) + + # act as an input + if as_node == INPUT: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when updating as input" + ) + + if input_writes := deque(map_input(self.input_channels, values)): + apply_writes( + checkpoint, + channels, + [PregelTaskWrites((), INPUT, input_writes, [])], + checkpointer.get_next_version, + ) + + # apply input write to channels + next_step = ( + step + 1 + if saved and saved.metadata.get("step") is not None + else -1 + ) + next_config = checkpointer.put( + checkpoint_config, + create_checkpoint(checkpoint, channels, next_step), + { + **checkpoint_metadata, + "source": "input", + "step": next_step, + "writes": dict(input_writes), + }, + get_new_channel_versions( + checkpoint_previous_versions, + checkpoint["channel_versions"], + ), + ) + + # store the writes + checkpointer.put_writes( + next_config, + input_writes, + str(uuid5(UUID(checkpoint["id"]), INPUT)), + ) + + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + else: + raise InvalidUpdateError( + f"Received no input writes for {self.input_channels}" + ) + # no values, copy checkpoint if values is None and as_node == "__copy__": if len(updates) > 1: @@ -1417,119 +1469,107 @@ class Pregel(PregelProtocol): isinstance(self.input_channels, str) and self.input_channels in self.nodes ): - if ( - isinstance(self.input_channels, str) - and self.input_channels in self.nodes - ): - as_node = self.input_channels - elif as_node is None: - last_seen_by_node = sorted( - (v, n) - for n, seen in checkpoint["versions_seen"].items() - if n in self.nodes - for v in seen.values() - ) - # if two nodes updated the state at the same time, it's ambiguous - if last_seen_by_node: - if len(last_seen_by_node) == 1: - as_node = last_seen_by_node[0][1] - elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: - as_node = last_seen_by_node[-1][1] + as_node = self.input_channels + elif as_node is None: + last_seen_by_node = sorted( + (v, n) + for n, seen in checkpoint["versions_seen"].items() + if n in self.nodes + for v in seen.values() + ) + # if two nodes updated the state at the same time, it's ambiguous + if last_seen_by_node: + if len(last_seen_by_node) == 1: + as_node = last_seen_by_node[0][1] + elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: + as_node = last_seen_by_node[-1][1] + if as_node is None: + raise InvalidUpdateError("Ambiguous update, specify as_node") + if as_node not in self.nodes: + raise InvalidUpdateError(f"Node {as_node} does not exist") + valid_updates.append((as_node, values)) + else: + for values, as_node in updates: if as_node is None: - raise InvalidUpdateError("Ambiguous update, specify as_node") + raise InvalidUpdateError( + "as_node is required when applying multiple updates" + ) if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") valid_updates.append((as_node, values)) - else: - for values, as_node in updates: - if as_node is None: - raise InvalidUpdateError( - "as_node is required when applying multiple updates" - ) - if as_node not in self.nodes: - raise InvalidUpdateError(f"Node {as_node} does not exist") + run_tasks: list[PregelTaskWrites] = [] + run_task_ids: list[str] = [] - valid_updates.append((as_node, values)) - - run_tasks: list[PregelTaskWrites] = [] - run_task_ids: list[str] = [] - - for as_node, values in valid_updates: - # create task to run all writers of the chosen node - writers = self.nodes[as_node].flat_writers - if not writers: - raise InvalidUpdateError(f"Node {as_node} has no writers") - writes: deque[tuple[str, Any]] = deque() - task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) - task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) - run_tasks.append(task) - run_task_ids.append(task_id) - run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] - # execute task - run.invoke( - values, - patch_config( - config, - run_name=self.name + "UpdateState", - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - local_write, - writes.extend, - self.nodes.keys(), - ), - CONFIG_KEY_READ: partial( - local_read, - step + 1, - checkpoint, - channels, - managed, - task, - config, - ), - }, - ), - ) - # save task writes - for task_id, task in zip(run_task_ids, run_tasks): - channel_writes = [w for w in task.writes if w[0] != PUSH] - # channel writes are saved to current checkpoint - if saved and channel_writes: - checkpointer.put_writes( - checkpoint_config, channel_writes, task_id - ) - # apply to checkpoint and save - mv_writes = apply_writes( - checkpoint, channels, run_tasks, checkpointer.get_next_version - ) - assert not mv_writes, "Can't write to SharedValues from update_state" - checkpoint = create_checkpoint(checkpoint, channels, step + 1) - next_config = checkpointer.put( - checkpoint_config, - checkpoint, - { - **checkpoint_metadata, - "source": "update", - "step": step + 1, - "writes": { - as_node: values for as_node, values in valid_updates + for as_node, values in valid_updates: + # create task to run all writers of the chosen node + writers = self.nodes[as_node].flat_writers + if not writers: + raise InvalidUpdateError(f"Node {as_node} has no writers") + writes: deque[tuple[str, Any]] = deque() + task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) + task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) + run_tasks.append(task) + run_task_ids.append(task_id) + run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] + # execute task + run.invoke( + values, + patch_config( + config, + run_name=self.name + "UpdateState", + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, + writes.extend, + self.nodes.keys(), + ), + CONFIG_KEY_READ: partial( + local_read, + step + 1, + checkpoint, + channels, + managed, + task, + config, + ), }, - "parents": saved.metadata.get("parents", {}) if saved else {}, - }, - get_new_channel_versions( - checkpoint_previous_versions, checkpoint["channel_versions"] ), ) - for task_id, task in zip(run_task_ids, run_tasks): - # save push writes - if push_writes := [w for w in task.writes if w[0] == PUSH]: - checkpointer.put_writes(next_config, push_writes, task_id) + # save task writes + for task_id, task in zip(run_task_ids, run_tasks): + # channel writes are saved to current checkpoint + channel_writes = [w for w in task.writes if w[0] != PUSH] + if saved and channel_writes: + checkpointer.put_writes(checkpoint_config, channel_writes, task_id) + # apply to checkpoint and save + mv_writes = apply_writes( + checkpoint, channels, run_tasks, checkpointer.get_next_version + ) + assert not mv_writes, "Can't write to SharedValues from update_state" + checkpoint = create_checkpoint(checkpoint, channels, step + 1) + next_config = checkpointer.put( + checkpoint_config, + checkpoint, + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {as_node: values for as_node, values in valid_updates}, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + get_new_channel_versions( + checkpoint_previous_versions, checkpoint["channel_versions"] + ), + ) + for task_id, task in zip(run_task_ids, run_tasks): + # save push writes + if push_writes := [w for w in task.writes if w[0] == PUSH]: + checkpointer.put_writes(next_config, push_writes, task_id) - return patch_checkpoint_map( - next_config, saved.metadata if saved else None - ) + return patch_checkpoint_map(next_config, saved.metadata if saved else None) current_config = config for superstep in supersteps: @@ -1706,6 +1746,58 @@ class Pregel(PregelProtocol): return patch_checkpoint_map( next_config, saved.metadata if saved else None ) + + # act as an input + if as_node == INPUT: + if len(updates) > 1: + raise InvalidUpdateError( + "Cannot apply multiple updates when updating as input" + ) + + if input_writes := deque(map_input(self.input_channels, values)): + apply_writes( + checkpoint, + channels, + [PregelTaskWrites((), INPUT, input_writes, [])], + checkpointer.get_next_version, + ) + + # apply input write to channels + next_step = ( + step + 1 + if saved and saved.metadata.get("step") is not None + else -1 + ) + next_config = await checkpointer.aput( + checkpoint_config, + create_checkpoint(checkpoint, channels, next_step), + { + **checkpoint_metadata, + "source": "input", + "step": next_step, + "writes": dict(input_writes), + }, + get_new_channel_versions( + checkpoint_previous_versions, + checkpoint["channel_versions"], + ), + ) + + # store the writes + await checkpointer.aput_writes( + next_config, + input_writes, + str(uuid5(UUID(checkpoint["id"]), INPUT)), + ) + + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + else: + raise InvalidUpdateError( + f"Received no input writes for {self.input_channels}" + ) + # no values, copy checkpoint if values is None and as_node == "__copy__": if len(updates) > 1: @@ -1811,9 +1903,6 @@ class Pregel(PregelProtocol): raise InvalidUpdateError( "as_node is required when applying multiple updates" ) - if as_node is None: - raise InvalidUpdateError("Ambiguous update, specify as_node") - if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 201e230f7..0eeb492e2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7734,3 +7734,166 @@ def test_bulk_state_updates( ], ], ) + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) +def test_update_as_input( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + foo: str + + def agent(state: State) -> State: + return {"foo": "agent"} + + def tool(state: State) -> State: + return {"foo": "tool"} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("tool", tool) + .add_edge(START, "agent") + .add_edge("agent", "tool") + .compile(checkpointer=checkpointer) + ) + + assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == { + "foo": "tool" + } + + assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == { + "foo": "tool" + } + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + } + + history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "1"}}) + ] + + graph.bulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + # First turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + # Second turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + ], + ) + + state = graph.get_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "tool"} + + new_history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "2"}}) + ] + + assert new_history == history + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) +def test_batch_update_as_input( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + foo: str + tasks: Annotated[list[int], operator.add] + + def agent(state: State) -> State: + return {"foo": "agent"} + + def map(state: State) -> Command["task"]: + return Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ) + + def task(state: dict) -> State: + return {"tasks": [state["index"]]} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("map", map) + .add_node("task", task) + .add_edge(START, "agent") + .add_edge("agent", "map") + .compile(checkpointer=checkpointer) + ) + + assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == { + "foo": "map", + "tasks": [0, 1, 2], + } + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + "tasks": [t.name for t in i.tasks], + } + + history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "1"}}) + ] + + graph.bulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent", "tasks": []}, "agent")], + [ + StateUpdate( + Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ), + "map", + ) + ], + [ + StateUpdate({"tasks": [0]}, "task"), + StateUpdate({"tasks": [1]}, "task"), + StateUpdate({"tasks": [2]}, "task"), + ], + ], + ) + + state = graph.get_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "map", "tasks": [0, 1, 2]} + + new_history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "2"}}) + ] + + assert new_history == history diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 8b83a5cf2..438d669db 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7961,3 +7961,169 @@ async def test_bulk_state_updates(checkpointer_name: str) -> None: ], ], ) + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_update_as_input(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + class State(TypedDict): + foo: str + + def agent(state: State) -> State: + return {"foo": "agent"} + + def tool(state: State) -> State: + return {"foo": "tool"} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("tool", tool) + .add_edge(START, "agent") + .add_edge("agent", "tool") + .compile(checkpointer=checkpointer) + ) + + assert await graph.ainvoke( + {"foo": "input"}, {"configurable": {"thread_id": "1"}} + ) == {"foo": "tool"} + + assert await graph.ainvoke( + {"foo": "input"}, {"configurable": {"thread_id": "1"}} + ) == {"foo": "tool"} + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + } + + history = [ + map_snapshot(s) + async for s in graph.aget_state_history( + {"configurable": {"thread_id": "1"}} + ) + ] + + await graph.abulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + # First turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + # Second turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + ], + ) + + state = await graph.aget_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "tool"} + + new_history = [ + map_snapshot(s) + async for s in graph.aget_state_history( + {"configurable": {"thread_id": "2"}} + ) + ] + + assert new_history == history + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_batch_update_as_input(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + class State(TypedDict): + foo: str + tasks: Annotated[list[int], operator.add] + + def agent(state: State) -> State: + return {"foo": "agent"} + + def map(state: State) -> Command["task"]: + return Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ) + + def task(state: dict) -> State: + return {"tasks": [state["index"]]} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("map", map) + .add_node("task", task) + .add_edge(START, "agent") + .add_edge("agent", "map") + .compile(checkpointer=checkpointer) + ) + + assert await graph.ainvoke( + {"foo": "input"}, {"configurable": {"thread_id": "1"}} + ) == {"foo": "map", "tasks": [0, 1, 2]} + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + "tasks": [t.name for t in i.tasks], + } + + history = [ + map_snapshot(s) + async for s in graph.aget_state_history( + {"configurable": {"thread_id": "1"}} + ) + ] + + await graph.abulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent", "tasks": []}, "agent")], + [ + StateUpdate( + Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ), + "map", + ) + ], + [ + StateUpdate({"tasks": [0]}, "task"), + StateUpdate({"tasks": [1]}, "task"), + StateUpdate({"tasks": [2]}, "task"), + ], + ], + ) + + state = await graph.aget_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "map", "tasks": [0, 1, 2]} + + new_history = [ + map_snapshot(s) + async for s in graph.aget_state_history( + {"configurable": {"thread_id": "2"}} + ) + ] + + assert new_history == history