From 7a1feb466b04d04b301e89e61af487df7e4d997b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 15:36:32 -0700 Subject: [PATCH] Implement map-reduce api for StateGraph/Pregel - current Pregel primitive is pull-based, ie. nodes write to channels, and it's up to other nodes to subscribe to those channels to "pull" updates - this adds a "push" primitive where a node can directly schedule a node (more than once if desired) for execution in the next step, with additional kwargs to be passed in. Nodes scheduled in this way get called with both the current state and any kwargs passed to Packet --- langgraph/channels/base.py | 1 + langgraph/checkpoint/base.py | 7 + langgraph/constants.py | 17 ++ langgraph/graph/graph.py | 39 +-- langgraph/graph/state.py | 21 +- langgraph/pregel/__init__.py | 199 +++++++++++---- langgraph/pregel/debug.py | 21 +- langgraph/pregel/io.py | 58 +++-- langgraph/pregel/read.py | 19 +- langgraph/pregel/types.py | 1 + langgraph/pregel/validate.py | 10 +- langgraph/pregel/write.py | 61 +++-- langgraph/serde/jsonplus.py | 5 + langgraph/utils.py | 30 +-- tests/test_pregel_async.py | 473 ++++++++++++++++++++++++++++++++++- 15 files changed, 813 insertions(+), 149 deletions(-) diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index a2a4c2d0d..9dc5fcbbc 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -128,6 +128,7 @@ def create_checkpoint( channel_values=values, channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], + pending_packets=checkpoint["pending_packets"], ) diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 2bdb52055..673a83278 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -5,6 +5,7 @@ from typing import ( Any, AsyncIterator, Iterator, + List, Literal, NamedTuple, Optional, @@ -14,6 +15,7 @@ from typing import ( from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig from langgraph.checkpoint.id import uuid6 +from langgraph.constants import Packet from langgraph.serde.base import SerializerProtocol from langgraph.serde.jsonplus import JsonPlusSerializer @@ -72,6 +74,9 @@ class Checkpoint(TypedDict): Used to determine which nodes to execute next. """ + pending_packets: List[Packet] + """List of packets sent to nodes but not yet processed. + Cleared by the next checkpoint.""" def _seen_dict(): @@ -86,6 +91,7 @@ def empty_checkpoint() -> Checkpoint: channel_values={}, channel_versions=defaultdict(int), versions_seen=defaultdict(_seen_dict), + pending_packets=[], ) @@ -100,6 +106,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: _seen_dict, {k: defaultdict(int, v) for k, v in checkpoint["versions_seen"].items()}, ), + pending_packets=checkpoint["pending_packets"].copy(), ) diff --git a/langgraph/constants.py b/langgraph/constants.py index e75e97452..d880aa733 100644 --- a/langgraph/constants.py +++ b/langgraph/constants.py @@ -1,6 +1,23 @@ +from typing import Any + +from langgraph.errors import InvalidUpdateError + CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" INTERRUPT = "__interrupt__" +TASKS = "__pregel_tasks" + +RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ} TAG_HIDDEN = "langsmith:hidden" + + +class Packet: + def __init__(self, /, __node__: str, **kwargs: Any) -> None: + if not kwargs: + raise InvalidUpdateError( + "Packet must have at least one keyword argument to pass to node" + ) + self.node = __node__ + self.kwargs = kwargs diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 5659b5031..4e00d148a 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -26,7 +26,8 @@ from langchain_core.runnables.graph import ( from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.checkpoint import BaseCheckpointSaver -from langgraph.constants import TAG_HIDDEN +from langgraph.constants import TAG_HIDDEN, Packet +from langgraph.errors import InvalidUpdateError from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import PregelNode from langgraph.pregel.types import All @@ -69,13 +70,7 @@ class Branch(NamedTuple): writer: Callable[[list[str]], Optional[Runnable]], ) -> Runnable: result = self.path.invoke(reader(config) if reader else input, config) - if not isinstance(result, list): - result = [result] - if self.ends: - destinations = [self.ends[r] for r in result] - else: - destinations = result - return writer(destinations) or input + return self._finish(writer, input, result) async def _aroute( self, @@ -86,14 +81,25 @@ class Branch(NamedTuple): writer: Callable[[list[str]], Optional[Runnable]], ) -> Runnable: result = await self.path.ainvoke(reader(config) if reader else input, config) + return self._finish(writer, input, result) + + def _finish( + self, writer: Callable[[list[str]], Optional[Runnable]], input: Any, result: Any + ): if not isinstance(result, list): result = [result] if self.ends: - destinations = [self.ends[r] for r in result] + destinations = [ + r if isinstance(r, Packet) else self.ends[r] for r in result + ] else: destinations = result - if any(dest is None for dest in destinations): + if any(dest is None or dest == START for dest in destinations): raise ValueError("Branch did not return a valid destination") + if any(p.node == END for p in destinations if isinstance(p, Packet)): + raise InvalidUpdateError( + "Cannot send a packet with keyword arguments to the END node" + ) return writer(destinations) or input @@ -384,13 +390,14 @@ class CompiledGraph(Pregel): self.nodes[end].channels.append(start) def attach_branch(self, start: str, name: str, branch: Branch) -> None: - def branch_writer(ends: list[str]) -> Optional[ChannelWrite]: - channels = [ - f"branch:{start}:{name}:{end}" if end != END else END for end in ends + def branch_writer(packets: list[Union[str, Packet]]) -> Optional[ChannelWrite]: + writes = [ + ChannelWriteEntry(f"branch:{start}:{name}:{p}" if p != END else END) + if not isinstance(p, Packet) + else p + for p in packets ] - return ChannelWrite( - [ChannelWriteEntry(ch) for ch in channels], tags=[TAG_HIDDEN] - ) + return ChannelWrite(writes, tags=[TAG_HIDDEN]) # add hidden start node if start == START and start not in self.nodes: diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 02df2cc3a..01b16aa69 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -27,7 +27,7 @@ from langgraph.channels.named_barrier_value import NamedBarrierValue from langgraph.checkpoint import BaseCheckpointSaver from langgraph.constants import TAG_HIDDEN from langgraph.errors import InvalidUpdateError -from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph +from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Packet from langgraph.managed.base import ManagedValue, is_managed_value from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All @@ -325,7 +325,7 @@ class CompiledStateGraph(CompiledGraph): ], ) else: - self.channels[key] = EphemeralValue(Any) + self.channels[key] = EphemeralValue(Any, guard=False) self.nodes[key] = PregelNode( triggers=[], # read state keys and managed values @@ -377,17 +377,24 @@ class CompiledStateGraph(CompiledGraph): ) def attach_branch(self, start: str, name: str, branch: Branch) -> None: - def branch_writer(ends: list[str]) -> Optional[ChannelWrite]: - if filtered_ends := [end for end in ends if end != END]: + def branch_writer(packets: list[Union[str, Packet]]) -> Optional[ChannelWrite]: + if filtered := [p for p in packets if p != END]: writes = [ - ChannelWriteEntry(f"branch:{start}:{name}:{end}", start) - for end in filtered_ends + ChannelWriteEntry(f"branch:{start}:{name}:{p}", start) + if not isinstance(p, Packet) + else p + for p in filtered ] if branch.then and branch.then != END: writes.append( ChannelWriteEntry( f"branch:{start}:{name}:then", - WaitForNames(set(filtered_ends)), + WaitForNames( + { + p.node if isinstance(p, Packet) else p + for p in filtered + } + ), ) ) return ChannelWrite(writes, tags=[TAG_HIDDEN]) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index bd4fea0c1..1b50065ea 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -66,6 +66,8 @@ from langgraph.constants import ( CONFIG_KEY_SEND, INTERRUPT, TAG_HIDDEN, + TASKS, + Packet, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ( @@ -877,7 +879,7 @@ class Pregel( # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() - for _, _, _, writes, _, _ in next_tasks: + for _, _, _, writes, _, _, _ in next_tasks: pending_writes.extend(writes) if debug: @@ -1180,7 +1182,7 @@ class Pregel( # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() - for _, _, _, writes, _, _ in next_tasks: + for _, _, _, writes, _, _, _ in next_tasks: pending_writes.extend(writes) if debug: @@ -1444,10 +1446,10 @@ def _should_interrupt( checkpoint["channel_versions"][chan] > seen[chan] for chan in snapshot_channels ) - # and any channel written to is in interrupt_nodes list + # and any triggered node is in interrupt_nodes list and any( node - for node, _, _, _, config, _ in tasks + for node, _, _, _, config, _, _ in tasks if ( (not config or TAG_HIDDEN not in config.get("tags")) if interrupt_nodes == "*" @@ -1473,15 +1475,40 @@ def _local_read( return read_channels(channels, select) +def _local_write( + commit: Callable[[Sequence[tuple[str, Any]]], None], + processes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + writes: Sequence[tuple[str, Any]], +) -> None: + for chan, value in writes: + if chan == TASKS: + if not isinstance(value, Packet): + raise InvalidUpdateError( + f"Invalid packet type, expected Packet, got {value}" + ) + if value.node not in processes: + raise InvalidUpdateError(f"Invalid node name {value.node} in packet") + elif chan not in channels: + logger.warning(f"Skipping write for channel '{chan}' which has no readers") + commit(writes) + + def _apply_writes( checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], pending_writes: Sequence[tuple[str, Any]], ) -> None: + if checkpoint["pending_packets"]: + raise RuntimeError("Cannot apply writes when there are pending packets") + pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) # Group writes by channel for chan, val in pending_writes: - pending_writes_by_channel[chan].append(val) + if chan == TASKS: + checkpoint["pending_packets"].append(val) + else: + pending_writes_by_channel[chan].append(val) # Find the highest version of all channels if checkpoint["channel_versions"]: @@ -1497,12 +1524,10 @@ def _apply_writes( channels[chan].update(vals) except InvalidUpdateError as e: raise InvalidUpdateError( - f"Invalid update for channel {chan}: {e}" + f"Invalid update for channel {chan} with values {vals}" ) from e checkpoint["channel_versions"][chan] = max_version + 1 updated_channels.add(chan) - else: - logger.warning(f"Skipping write for channel '{chan}' which has no readers") # Channels that weren't updated in this step are notified of a new step for chan in channels: if chan not in updated_channels: @@ -1550,6 +1575,64 @@ def _prepare_next_tasks( ) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]: checkpoint = copy_checkpoint(checkpoint) tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] + # Consume pending packets + for packet in checkpoint["pending_packets"]: + try: + val = next( + _proc_input( + step, + packet.node, + processes[packet.node], + managed, + channels, + catch=True, + ) + ) + except StopIteration: + logger.warn('No input for node "%s" in packet, skipping', packet.node) + continue + if for_execution: + if node := processes[packet.node].get_node(packet.kwargs): + writes = deque() + tasks.append( + PregelExecutableTask( + packet.node, + val, + node, + writes, + patch_config( + merge_configs( + config, + processes[packet.node].config, + { + "metadata": { + "langgraph_step": step, + "langgraph_node": packet.node, + "langgraph_triggers": [TASKS], + } + }, + ), + run_name=packet.node, + callbacks=manager.get_child(f"graph:step:{step}") + if manager + else None, + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + _local_write, writes.extend, processes, channels + ), + CONFIG_KEY_READ: partial( + _local_read, checkpoint, channels, tasks + ), + }, + ), + [TASKS], + packet.kwargs, + ) + ) + else: + tasks.append(PregelTaskDescription(packet.node, val)) + checkpoint["pending_packets"].clear() # Check if any processes should be run in next step # If so, prepare the values to be passed to them for name, proc in processes.items(): @@ -1563,43 +1646,10 @@ def _prepare_next_tasks( ) and checkpoint["channel_versions"][chan] > seen[chan] ]: - # If all trigger channels subscribed by this process are not empty - # then invoke the process with the values of all non-empty channels - if isinstance(proc.channels, dict): - try: - val: dict = { - k: read_channel(channels, chan, catch=chan not in proc.triggers) - for k, chan in proc.channels.items() - if isinstance(chan, str) - } - - managed_values = {} - for key, chan in proc.channels.items(): - if is_managed_value(chan): - managed_values[key] = managed[key]( - step, PregelTaskDescription(name, val) - ) - - val.update(managed_values) - except EmptyChannelError: - continue - elif isinstance(proc.channels, list): - for chan in proc.channels: - try: - val = read_channel(channels, chan, catch=False) - break - except EmptyChannelError: - pass - else: - continue - else: - raise RuntimeError( - "Invalid channels type, expected list or dict, got {proc.channels}" - ) - - # If the process has a mapper, apply it to the value - if proc.mapper is not None: - val = proc.mapper(val) + try: + val = next(_proc_input(step, name, proc, managed, channels)) + except StopIteration: + continue # update seen versions if for_execution: @@ -1613,6 +1663,7 @@ def _prepare_next_tasks( if for_execution: if node := proc.get_node(): writes = deque() + triggers = sorted(triggers) tasks.append( PregelExecutableTask( name, @@ -1627,6 +1678,7 @@ def _prepare_next_tasks( "metadata": { "langgraph_step": step, "langgraph_node": name, + "langgraph_triggers": triggers, } }, ), @@ -1636,13 +1688,15 @@ def _prepare_next_tasks( else None, configurable={ # deque.extend is thread-safe - CONFIG_KEY_SEND: writes.extend, + CONFIG_KEY_SEND: partial( + _local_write, writes.extend, processes, channels + ), CONFIG_KEY_READ: partial( _local_read, checkpoint, channels, writes ), }, ), - sorted(triggers), + triggers, ) ) else: @@ -1650,6 +1704,57 @@ def _prepare_next_tasks( return checkpoint, tasks +def _proc_input( + step: int, + name: str, + proc: PregelNode, + managed: ManagedValueMapping, + channels: Mapping[str, BaseChannel], + catch: bool = False, +) -> Iterator[Any]: + # If all trigger channels subscribed by this process are not empty + # then invoke the process with the values of all non-empty channels + if isinstance(proc.channels, dict): + try: + val: dict = { + k: read_channel( + channels, chan, catch=catch or chan not in proc.triggers + ) + for k, chan in proc.channels.items() + if isinstance(chan, str) + } + + managed_values = {} + for key, chan in proc.channels.items(): + if is_managed_value(chan): + managed_values[key] = managed[key]( + step, PregelTaskDescription(name, val) + ) + + val.update(managed_values) + except EmptyChannelError: + return + elif isinstance(proc.channels, list): + for chan in proc.channels: + try: + val = read_channel(channels, chan, catch=False) + break + except EmptyChannelError: + pass + else: + return + else: + raise RuntimeError( + "Invalid channels type, expected list or dict, got {proc.channels}" + ) + + # If the process has a mapper, apply it to the value + if proc.mapper is not None: + val = proc.mapper(val) + + yield val + + def _with_mode(mode: StreamMode, on: bool, iter: Iterator[Any]) -> Iterator[Any]: if on: for chunk in iter: diff --git a/langgraph/pregel/debug.py b/langgraph/pregel/debug.py index 0fd00bdc8..6dc86643f 100644 --- a/langgraph/pregel/debug.py +++ b/langgraph/pregel/debug.py @@ -64,20 +64,23 @@ def map_debug_tasks( step: int, tasks: list[PregelExecutableTask] ) -> Iterator[DebugOutputTask]: ts = datetime.now(timezone.utc).isoformat() - for idx, (name, input, _, _, config, triggers) in enumerate(tasks): + for idx, (name, input, _, _, config, triggers, kwargs) in enumerate(tasks): if config is not None and TAG_HIDDEN in config.get("tags", []): continue + payload = { + "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, idx)))), + "name": name, + "input": input, + "triggers": triggers, + } + if kwargs is not None: + payload["kwargs"] = kwargs yield { "type": "task", "timestamp": ts, "step": step, - "payload": { - "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, idx)))), - "name": name, - "input": input, - "triggers": triggers, - }, + "payload": payload, } @@ -87,7 +90,7 @@ def map_debug_task_results( stream_channels_list: Sequence[str], ) -> Iterator[DebugOutputTaskResult]: ts = datetime.now(timezone.utc).isoformat() - for idx, (name, _, _, writes, config, _) in enumerate(tasks): + for idx, (name, _, _, writes, config, _, _) in enumerate(tasks): if config is not None and TAG_HIDDEN in config.get("tags", []): continue @@ -130,7 +133,7 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None: ) + "\n".join( f"- {get_colored_text(name, 'green')} -> {pformat(val)}" - for name, val, _, _, _, _ in next_tasks + for name, val, _, _, _, _, _ in next_tasks ) ) diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 6228167cc..3e2bf173d 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -1,9 +1,11 @@ +from collections import defaultdict +from itertools import groupby from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union from langchain_core.runnables.utils import AddableDict from langgraph.channels.base import BaseChannel, EmptyChannelError -from langgraph.constants import TAG_HIDDEN +from langgraph.constants import TAG_HIDDEN, TASKS from langgraph.pregel.log import logger from langgraph.pregel.types import PregelExecutableTask @@ -102,24 +104,44 @@ def map_output_updates( t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags") ] if isinstance(output_channels, str): - if updated := AddableUpdatesDict( - { - node: value - for node, _, _, writes, _, _ in output_tasks - for chan, value in writes - if chan == output_channels - } - ): - yield updated + if updated := [ + (triggers == [TASKS], node, value) + for node, _, _, writes, _, triggers, _ in output_tasks + for chan, value in writes + if chan == output_channels + ]: + grouped = defaultdict(list) + for from_packet, node, value in updated: + if from_packet: + grouped[node].append(value) + for from_packet, node, value in updated: + if not from_packet: + if grouped[node]: + grouped[node].append(value) + else: + grouped[node] = value + yield AddableUpdatesDict(grouped) else: - if updated := AddableUpdatesDict( - { - node: {chan: value for chan, value in writes if chan in output_channels} - for node, _, _, writes, _, _ in output_tasks - if any(chan in output_channels for chan, _ in writes) - } - ): - yield updated + if updated := [ + ( + triggers == [TASKS], + node, + {chan: value for chan, value in writes if chan in output_channels}, + ) + for node, _, _, writes, _, triggers, _ in output_tasks + if any(chan in output_channels for chan, _ in writes) + ]: + grouped = defaultdict(list) + for from_packet, node, value in updated: + if from_packet: + grouped[node].append(value) + for from_packet, node, value in updated: + if not from_packet: + if grouped[node]: + grouped[node].append(value) + else: + grouped[node] = value + yield AddableUpdatesDict(grouped) T = TypeVar("T") diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index eefed2b47..2e90a505b 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Callable, Mapping, Optional, Sequence, Union +from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Union from langchain_core.pydantic_v1 import Field from langchain_core.runnables import ( @@ -121,11 +121,17 @@ class PregelNode(RunnableBindingBase): and isinstance(writers[-2], ChannelWrite) ): # we can combine writes if they are consecutive - writers[-2].writes += writers[-1].writes + # careful to not modify the original writers list or ChannelWrite + writers[-2] = ChannelWrite( + writes=writers[-2].writes + writers[-1].writes, + tags=writers[-2].config["tags"] if writers[-2].config else None, + ) writers.pop() return writers - def get_node(self) -> Optional[Runnable[Any, Any]]: + def get_node( + self, kwargs: Optional[Dict[str, Any]] = None + ) -> Optional[Runnable[Any, Any]]: writers = self.get_writers() if self.bound is DEFAULT_BOUND and not writers: return None @@ -134,9 +140,12 @@ class PregelNode(RunnableBindingBase): elif self.bound is DEFAULT_BOUND: return RunnableSequence(*writers) elif writers: - return RunnableSequence(self.bound, *writers) + return RunnableSequence( + self.bound.bind(**kwargs) if kwargs is not None else self.bound, + *writers, + ) else: - return self.bound + return self.bound.bind(**kwargs) if kwargs is not None else self.bound def __init__( self, diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index a3826314a..ea9cb69ef 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -18,6 +18,7 @@ class PregelExecutableTask(NamedTuple): writes: deque[tuple[str, Any]] config: Optional[RunnableConfig] triggers: list[str] + kwargs: Optional[dict[str, Any]] = None class StateSnapshot(NamedTuple): diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py index ad4959657..8627642e9 100644 --- a/langgraph/pregel/validate.py +++ b/langgraph/pregel/validate.py @@ -1,7 +1,7 @@ from typing import Mapping, Optional, Sequence, Union from langgraph.channels.base import BaseChannel -from langgraph.constants import INTERRUPT +from langgraph.constants import RESERVED from langgraph.pregel.read import PregelNode from langgraph.pregel.types import All @@ -16,13 +16,13 @@ def validate_graph( interrupt_before_nodes: Union[All, Sequence[str]], ) -> None: for chan in channels: - if chan == INTERRUPT: - raise ValueError(f"Channel name {INTERRUPT} is reserved") + if chan in RESERVED: + raise ValueError(f"Channel names {RESERVED} are reserved") subscribed_channels = set[str]() for name, node in nodes.items(): - if name == INTERRUPT: - raise ValueError(f"Node name {INTERRUPT} is reserved") + if name in RESERVED: + raise ValueError(f"Node names {RESERVED} are reserved") if isinstance(node, PregelNode): subscribed_channels.update(node.triggers) else: diff --git a/langgraph/pregel/write.py b/langgraph/pregel/write.py index 4b2ad9ea5..f9e9c1233 100644 --- a/langgraph/pregel/write.py +++ b/langgraph/pregel/write.py @@ -1,12 +1,14 @@ from __future__ import annotations import asyncio -from typing import Any, Callable, NamedTuple, Optional, Sequence, TypeVar +from typing import Any, Callable, List, NamedTuple, Optional, Sequence, Tuple, TypeVar from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.utils import ConfigurableFieldSpec +from sympy import Union -from langgraph.constants import CONFIG_KEY_SEND +from langgraph.constants import CONFIG_KEY_SEND, TASKS, Packet +from langgraph.errors import InvalidUpdateError from langgraph.utils import RunnableCallable TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None] @@ -25,7 +27,7 @@ class ChannelWriteEntry(NamedTuple): class ChannelWrite(RunnableCallable): - writes: Sequence[ChannelWriteEntry] + writes: Sequence[Union[ChannelWriteEntry, Packet]] """ Sequence of write entries, each of which is a tuple of: - channel name @@ -34,7 +36,10 @@ class ChannelWrite(RunnableCallable): """ def __init__( - self, writes: Sequence[ChannelWriteEntry], *, tags: Optional[list[str]] = None + self, + writes: Sequence[Union[ChannelWriteEntry, Packet]], + *, + tags: Optional[list[str]] = None, ): super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags) self.writes = writes @@ -46,7 +51,7 @@ class ChannelWrite(RunnableCallable): self, suffix: Optional[str] = None, *, name: Optional[str] = None ) -> str: if not name: - name = f"ChannelWrite<{','.join(chan for chan, _, _, _ in self.writes)}>" + name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else w.node for w in self.writes)}>" return super().get_name(suffix, name=name) @property @@ -62,47 +67,69 @@ class ChannelWrite(RunnableCallable): ] def _write(self, input: Any, config: RunnableConfig) -> None: + # split packets and entries + writes = [ + (TASKS, packet) for packet in self.writes if isinstance(packet, Packet) + ] + entries = [ + write for write in self.writes if isinstance(write, ChannelWriteEntry) + ] + for entry in entries: + if entry.channel == TASKS: + raise InvalidUpdateError("Cannot write to the reserved channel TASKS") + # process entries into values values = [ - input if write.value is PASSTHROUGH else write.value - for write in self.writes + input if write.value is PASSTHROUGH else write.value for write in entries ] values = [ val if write.mapper is None else write.mapper.invoke(val, config) - for val, write in zip(values, self.writes) + for val, write in zip(values, entries) ] values = [ (write.channel, val) - for val, write in zip(values, self.writes) + for val, write in zip(values, entries) if not write.skip_none or val is not None ] - self.do_write(config, **dict(values)) + # write packets and values + self.do_write(config, writes + values) return input async def _awrite(self, input: Any, config: RunnableConfig) -> None: + # split packets and entries + writes = [ + (TASKS, packet) for packet in self.writes if isinstance(packet, Packet) + ] + entries = [ + write for write in self.writes if isinstance(write, ChannelWriteEntry) + ] + for entry in entries: + if entry.channel == TASKS: + raise InvalidUpdateError("Cannot write to the reserved channel TASKS") + # process entries into values values = [ - input if write.value is PASSTHROUGH else write.value - for write in self.writes + input if write.value is PASSTHROUGH else write.value for write in entries ] values = await asyncio.gather( *( _mk_future(val) if write.mapper is None else write.mapper.ainvoke(val, config) - for val, write in zip(values, self.writes) + for val, write in zip(values, entries) ) ) values = [ (write.channel, val) - for val, write in zip(values, self.writes) + for val, write in zip(values, entries) if not write.skip_none or val is not None ] - self.do_write(config, **dict(values)) + # write packets and values + self.do_write(config, writes + values) return input @staticmethod - def do_write(config: RunnableConfig, **values: Any) -> None: + def do_write(config: RunnableConfig, values: List[Tuple[str, Any]]) -> None: write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] - write([(chan, val) for chan, val in values.items() if val is not SKIP_WRITE]) + write([(chan, val) for chan, val in values if val is not SKIP_WRITE]) @staticmethod def is_writer(runnable: Runnable) -> bool: diff --git a/langgraph/serde/jsonplus.py b/langgraph/serde/jsonplus.py index e2e836689..2b89f5af7 100644 --- a/langgraph/serde/jsonplus.py +++ b/langgraph/serde/jsonplus.py @@ -9,6 +9,7 @@ from uuid import UUID from langchain_core.load.load import Reviver from langchain_core.load.serializable import Serializable +from langgraph.constants import Packet from langgraph.serde.base import SerializerProtocol LC_REVIVER = Reviver() @@ -63,6 +64,10 @@ class JsonPlusSerializer(SerializerProtocol): ) elif isinstance(obj, Enum): return self._encode_constructor_args(obj.__class__, args=[obj.value]) + elif isinstance(obj, Packet): + return self._encode_constructor_args( + Packet, args=[obj.node], kwargs=obj.kwargs + ) else: raise TypeError( f"Object of type {obj.__class__.__name__} is not JSON serializable" diff --git a/langgraph/utils.py b/langgraph/utils.py index e766a7fb8..7d94b81e9 100644 --- a/langgraph/utils.py +++ b/langgraph/utils.py @@ -59,7 +59,7 @@ class RunnableCallable(Runnable): pass self.func = func self.afunc = afunc - self.config = {"tags": tags} if tags else None + self.config: Optional[RunnableConfig] = {"tags": tags} if tags else None self.kwargs = kwargs self.trace = trace self.recurse = recurse @@ -72,47 +72,47 @@ class RunnableCallable(Runnable): } return f"{self.get_name()}({', '.join(f'{k}={v!r}' for k, v in repr_args.items())})" - def invoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any: + def invoke( + self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any + ) -> Any: if self.func is None: raise TypeError( f'No synchronous function provided to "{self.name}".' "\nEither initialize with a synchronous function or invoke" " via the async API (ainvoke, astream, etc.)" ) + kwargs = {**self.kwargs, **kwargs} if self.trace: ret = self._call_with_config( - self.func, input, merge_configs(self.config, config), **self.kwargs + self.func, input, merge_configs(self.config, config), **kwargs ) else: config = merge_configs(self.config, config) context = copy_context() context.run(var_child_runnable_config.set, config) - kwargs = ( - {**self.kwargs, "config": config} - if accepts_config(self.func) - else self.kwargs - ) + if accepts_config(self.func): + kwargs["config"] = config ret = context.run(self.func, input, **kwargs) if isinstance(ret, Runnable) and self.recurse: return ret.invoke(input, config) return ret - async def ainvoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any: + async def ainvoke( + self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any + ) -> Any: if not self.afunc: return self.invoke(input, config) + kwargs = {**self.kwargs, **kwargs} if self.trace: ret = await self._acall_with_config( - self.afunc, input, merge_configs(self.config, config), **self.kwargs + self.afunc, input, merge_configs(self.config, config), **kwargs ) else: config = merge_configs(self.config, config) context = copy_context() context.run(var_child_runnable_config.set, config) - kwargs = ( - {**self.kwargs, "config": config} - if accepts_config(self.afunc) - else self.kwargs - ) + if accepts_config(self.afunc): + kwargs["config"] = config if sys.version_info >= (3, 11): ret = await asyncio.create_task( self.afunc(input, **kwargs), context=context diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 211e4ab6d..7515c0351 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -10,7 +10,6 @@ from typing import ( AsyncIterator, Dict, Generator, - NamedTuple, Optional, Sequence, TypedDict, @@ -19,7 +18,12 @@ from typing import ( from uuid import UUID import pytest -from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough +from langchain_core.runnables import ( + RunnableConfig, + RunnableLambda, + RunnablePassthrough, + RunnablePick, +) from pytest_mock import MockerFixture from syrupy import SnapshotAssertion @@ -28,6 +32,7 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver +from langgraph.constants import Packet from langgraph.errors import InvalidUpdateError from langgraph.graph import END, Graph, StateGraph from langgraph.graph.graph import START @@ -2503,7 +2508,7 @@ async def test_state_graph_few_shot() -> None: class BaseState(TypedDict): messages: Annotated[list[AnyMessage], add_messages] - tool_results: Annotated[list[str], operator.add] + # tool_results: Annotated[list[str], operator.add] class AgentState(BaseState): examples: Annotated[ @@ -2518,6 +2523,7 @@ async def test_state_graph_few_shot() -> None: return f"result for {query}" tools = [search_api] + tools_by_name = {t.name: t for t in tools} prompt = ChatPromptTemplate.from_messages( [ @@ -2555,20 +2561,23 @@ Some examples of past conversations: response = await model.ainvoke(formatted) return {"messages": response} - class GoTo: - def __init__(self, /, __node__: str, **kwargs: Any) -> None: - pass - # Define decision-making logic def should_continue(data: AgentState) -> str: # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: - return [GoTo("tools", tool_call=tool_call) for tool_call in tool_calls] + return [Packet("tools", tool_call=tool_call) for tool_call in tool_calls] else: return "exit" - def tools_node(data: AgentState, *, tool_call: ToolCall) -> AgentState: - return {"tool_results": ...} + def tools_node( + _: AgentState, config: RunnableConfig, *, tool_call: ToolCall + ) -> AgentState: + output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) + return { + "messages": ToolMessage( + content=output, name=tool_call["name"], tool_call_id=tool_call["id"] + ) + } # Define a new graph workflow = StateGraph(AgentState) @@ -3061,6 +3070,450 @@ async def test_prebuilt_chat() -> None: ] +async def test_state_graph_packets() -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + ToolCall, + ToolMessage, + ) + from langchain_core.tools import tool + + class AgentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + tools_by_name = {t.name: t for t in tools} + + model = FakeMessagesListChatModel( + responses=[ + AIMessage( + id="a1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + id="a2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + AIMessage(id="ai3", content="answer"), + ] + ) + + # Define decision-making logic + def should_continue(data: AgentState) -> str: + # Logic to decide whether to continue in the loop or exit + if tool_calls := data["messages"][-1].tool_calls: + return [Packet("tools", tool_call=tool_call) for tool_call in tool_calls] + else: + return END + + def tools_node( + _: AgentState, config: RunnableConfig, *, tool_call: ToolCall + ) -> AgentState: + output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) + return { + "messages": ToolMessage( + content=output, name=tool_call["name"], tool_call_id=tool_call["id"] + ) + } + + # Define a new graph + workflow = StateGraph(AgentState) + + # Define the two nodes we will cycle between + workflow.add_node("agent", {"messages": RunnablePick("messages") | model}) + workflow.add_node("tools", tools_node) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges("agent", should_continue) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("tools", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + assert await app.ainvoke( + {"messages": HumanMessage(content="what is weather in sf")} + ) == { + "messages": [ + HumanMessage(content="what is weather in sf", id=AnyStr()), + AIMessage( + id="a1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ToolMessage( + content="result for query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage( + id="a2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + ToolMessage( + content="result for another", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call234", + ), + ToolMessage( + content="result for a third one", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call567", + ), + AIMessage(content="answer", id="ai3"), + ] + } + + assert [ + c + async for c in app.astream( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="a1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + }, + }, + { + "tools": [ + { + "messages": ToolMessage( + content="result for query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } + ] + }, + { + "agent": { + "messages": AIMessage( + id="a2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + } + }, + { + "tools": [ + { + "messages": ToolMessage( + content="result for another", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call234", + ) + }, + { + "messages": ToolMessage( + content="result for a third one", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call567", + ), + }, + ] + }, + {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + HumanMessage(content="what is weather in sf"), config + ) + ] == [ + { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + id="ai1", + ) + }, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + id="ai1", + ), + ], + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ + "ts" + ], + metadata={ + "source": "loop", + "step": 1, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + id="ai1", + ) + }, + }, + ) + + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values[-1] + last_message.additional_kwargs["function_call"]["arguments"] = '"a different query"' + await app_w_interrupt.aupdate_state(config, last_message) + + # message was replaced instead of appended + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + ], + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ + "ts" + ], + metadata={ + "source": "update", + "step": 2, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ) + }, + }, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ) + }, + { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"another"'} + }, + id="ai2", + ) + }, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"another"'} + }, + id="ai2", + ), + ], + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ + "ts" + ], + metadata={ + "source": "loop", + "step": 4, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"another"', + } + }, + id="ai2", + ) + }, + }, + ) + + await app_w_interrupt.aupdate_state( + config, + AIMessage(content="answer", id="ai2"), + ) + + # replaces message even if object identity is different, as long as id is the same + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + ], + next=(), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ + "ts" + ], + metadata={ + "source": "update", + "step": 5, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + }, + ) + + async def test_message_graph() -> None: from langchain_core.agents import AgentAction from langchain_core.language_models.fake_chat_models import (