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..f4c372e71 100644 --- a/langgraph/constants.py +++ b/langgraph/constants.py @@ -1,6 +1,16 @@ +from typing import Any, NamedTuple + 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(NamedTuple): + node: str + arg: Any diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 5659b5031..8acc8f9b4 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 @@ -68,14 +69,16 @@ class Branch(NamedTuple): reader: Optional[Callable[[], Any]], 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] + if reader: + value = reader(config) + # passthrough additional keys from node to branch + # only doable when using dict states + if isinstance(value, dict) and isinstance(input, dict): + value = {**input, **value} else: - destinations = result - return writer(destinations) or input + value = input + result = self.path.invoke(value, config) + return self._finish(writer, input, result) async def _aroute( self, @@ -85,15 +88,32 @@ class Branch(NamedTuple): reader: Optional[Callable[[], Any]], writer: Callable[[list[str]], Optional[Runnable]], ) -> Runnable: - result = await self.path.ainvoke(reader(config) if reader else input, config) + if reader: + value = reader(config) + # passthrough additional keys from node to branch + # only doable when using dict states + if isinstance(value, dict) and isinstance(input, dict): + value = {**input, **value} + else: + value = input + result = await self.path.ainvoke(value, 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 to the END node") return writer(destinations) or input @@ -384,13 +404,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..49bd34e37 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 @@ -321,11 +321,15 @@ class CompiledStateGraph(CompiledGraph): triggers=[START], channels=[START], writers=[ - ChannelWrite(state_write_entries, tags=[TAG_HIDDEN]), + ChannelWrite( + state_write_entries, + tags=[TAG_HIDDEN], + require_at_least_one_of=state_keys, + ), ], ) else: - self.channels[key] = EphemeralValue(Any) + self.channels[key] = EphemeralValue(Any, guard=False) self.nodes[key] = PregelNode( triggers=[], # read state keys and managed values @@ -345,6 +349,7 @@ class CompiledStateGraph(CompiledGraph): ChannelWrite( [ChannelWriteEntry(key, key)] + state_write_entries, tags=[TAG_HIDDEN], + require_at_least_one_of=state_keys, ), ], ).pipe(node) @@ -377,17 +382,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..933392b46 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 ( @@ -1444,7 +1446,7 @@ 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 @@ -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"]: + checkpoint["pending_packets"].clear() + 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,50 @@ 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"]: + if for_execution: + if node := processes[packet.node].get_node(): + writes = deque() + tasks.append( + PregelExecutableTask( + packet.node, + packet.arg, + 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], + ) + ) + else: + tasks.append(PregelTaskDescription(packet.node, packet.arg)) + if for_execution: + 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 +1632,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 +1649,7 @@ def _prepare_next_tasks( if for_execution: if node := proc.get_node(): writes = deque() + triggers = sorted(triggers) tasks.append( PregelExecutableTask( name, @@ -1627,6 +1664,7 @@ def _prepare_next_tasks( "metadata": { "langgraph_step": step, "langgraph_node": name, + "langgraph_triggers": triggers, } }, ), @@ -1636,13 +1674,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 +1690,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/io.py b/langgraph/pregel/io.py index 6228167cc..0b6979606 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -1,9 +1,10 @@ +from collections import defaultdict 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 +103,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..a14c704e4 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -121,7 +121,12 @@ 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, + require_at_least_one_of=writers[-2].require_at_least_one_of, + ) writers.pop() return writers 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..82610c322 100644 --- a/langgraph/pregel/write.py +++ b/langgraph/pregel/write.py @@ -1,12 +1,23 @@ 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, + Union, +) from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.utils import ConfigurableFieldSpec -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,19 +36,28 @@ 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 - runnable to map input, or None to use the input, or any other value to use instead - whether to skip writing if the mapped value is None """ + require_at_least_one_of: Optional[Sequence[str]] + """ + If defined, at least one of these channels must be written to. + """ def __init__( - self, writes: Sequence[ChannelWriteEntry], *, tags: Optional[list[str]] = None + self, + writes: Sequence[Union[ChannelWriteEntry, Packet]], + *, + tags: Optional[list[str]] = None, + require_at_least_one_of: Optional[Sequence[str]] = None, ): super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags) self.writes = writes + self.require_at_least_one_of = require_at_least_one_of def __repr_args__(self) -> Any: return [("writes", self.writes)] @@ -46,7 +66,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 +82,87 @@ 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, + self.require_at_least_one_of if input is not None else None, + ) 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, + self.require_at_least_one_of if input is not None else None, + ) return input @staticmethod - def do_write(config: RunnableConfig, **values: Any) -> None: + def do_write( + config: RunnableConfig, + values: List[Tuple[str, Any]], + require_at_least_one_of: Optional[Sequence[str]] = None, + ) -> None: + filtered = [(chan, val) for chan, val in values if val is not SKIP_WRITE] + if require_at_least_one_of is not None: + if not {chan for chan, _ in filtered} & set(require_at_least_one_of): + raise InvalidUpdateError( + f"Must write to at least one of {require_at_least_one_of}" + ) write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] - write([(chan, val) for chan, val in values.items() if val is not SKIP_WRITE]) + write(filtered) @staticmethod def is_writer(runnable: Runnable) -> bool: diff --git a/langgraph/serde/jsonplus.py b/langgraph/serde/jsonplus.py index e2e836689..f4f49a7a4 100644 --- a/langgraph/serde/jsonplus.py +++ b/langgraph/serde/jsonplus.py @@ -3,12 +3,13 @@ import importlib import json from datetime import datetime, timedelta, timezone from enum import Enum -from typing import Any, Optional +from typing import Any, NamedTuple, Optional 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,8 @@ class JsonPlusSerializer(SerializerProtocol): ) elif isinstance(obj, Enum): return self._encode_constructor_args(obj.__class__, args=[obj.value]) + elif isinstance(obj, NamedTuple): + return self._encode_constructor_args(Packet, args=[*obj]) 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.py b/tests/test_pregel.py index c58ff7e7b..bfabb59b2 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -34,6 +34,7 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.sqlite import SqliteSaver +from langgraph.constants import Packet from langgraph.errors import InvalidUpdateError from langgraph.graph import END, Graph from langgraph.graph.graph import START @@ -140,6 +141,21 @@ def test_graph_validation() -> None: with pytest.raises(ValueError): # extra is dead-end workflow.compile() + class State(TypedDict): + hello: str + + def node_a(state: State) -> State: + # typo + return {"hell": "world"} + + builder = StateGraph(State) + builder.add_node("a", node_a) + builder.set_entry_point("a") + builder.set_finish_point("a") + graph = builder.compile() + with pytest.raises(InvalidUpdateError): + graph.invoke({"hello": "there"}) + def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) @@ -3427,6 +3443,523 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None: ] +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="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + id="ai2", + 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"), + ] + ) + + def agent(data: AgentState) -> AgentState: + return { + "messages": model.invoke(data["messages"]), + "something_extra": "hi there", + } + + # Define decision-making logic + def should_continue(data: AgentState) -> str: + assert ( + data["something_extra"] == "hi there" + ), "nodes can pass extra data to their cond edges, which isn't saved in state" + # Logic to decide whether to continue in the loop or exit + if tool_calls := data["messages"][-1].tool_calls: + return [Packet("tools", tool_call) for tool_call in tool_calls] + else: + return END + + def tools_node(tool_call: ToolCall, config: RunnableConfig) -> 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", agent) + 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 app.invoke({"messages": HumanMessage(content="what is weather in sf")}) == { + "messages": [ + HumanMessage(content="what is weather in sf", id=AnyStr()), + AIMessage( + id="ai1", + 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="ai2", + 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 + for c in app.stream( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + 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="ai2", + 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 + for c in app_w_interrupt.stream( + {"messages": HumanMessage(content="what is weather in sf")}, config + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + next=("tools",), + config=(app_w_interrupt.checkpointer.get_tuple(config)).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + }, + ) + + # modify ai message + last_message = (app_w_interrupt.get_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + app_w_interrupt.update_state( + config, {"messages": last_message, "something_extra": "hi there"} + ) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "source": "update", + "step": 2, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + "something_extra": "hi there", + } + }, + }, + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + { + "messages": ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } + ] + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + }, + }, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + ] + }, + next=("tools", "tools"), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 4, + "writes": { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + }, + }, + }, + ) + + app_w_interrupt.update_state( + config, + { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + }, + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + next=(), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "source": "update", + "step": 5, + "writes": { + "agent": { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + } + }, + }, + ) + + def test_message_graph( snapshot: SnapshotAssertion, deterministic_uuids: MockerFixture, diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 34e4a783c..0f9868b65 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -18,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 @@ -27,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 @@ -2484,7 +2490,13 @@ async def test_state_graph_few_shot() -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) - from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage + from langchain_core.messages import ( + AIMessage, + AnyMessage, + HumanMessage, + ToolCall, + ToolMessage, + ) from langchain_core.prompts import ChatPromptTemplate from langchain_core.tools import tool @@ -2510,6 +2522,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( [ @@ -2550,16 +2563,24 @@ Some examples of past conversations: # Define decision-making logic def should_continue(data: AgentState) -> str: # Logic to decide whether to continue in the loop or exit - if not data["messages"][-1].tool_calls: - return "exit" + if tool_calls := data["messages"][-1].tool_calls: + return [Packet("tools", tool_call) for tool_call in tool_calls] else: - return "continue" + return "exit" + + def tools_node(tool_call: ToolCall, config: RunnableConfig) -> 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) workflow.add_node("agent", agent) - workflow.add_node("tools", ToolNode(tools)) + workflow.add_node("tools", tools_node) workflow.set_entry_point("agent") workflow.add_conditional_edges( "agent", should_continue, {"continue": "tools", "exit": END} @@ -3046,6 +3067,513 @@ 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="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + id="ai2", + 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) for tool_call in tool_calls] + else: + return END + + def tools_node(tool_call: ToolCall, config: RunnableConfig) -> 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="ai1", + 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="ai2", + 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="ai1", + 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="ai2", + 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( + {"messages": HumanMessage(content="what is weather in sf")}, config + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + 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": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + }, + ) + + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + await app_w_interrupt.aupdate_state(config, {"messages": last_message}) + + # message was replaced instead of appended + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + 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": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ) + } + }, + }, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": [ + { + "messages": ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } + ] + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + }, + }, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + ] + }, + next=("tools", "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": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + }, + }, + }, + ) + + await app_w_interrupt.aupdate_state( + config, + {"messages": 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={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + 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": {"messages": 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 (