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 (