WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-06 17:57:49 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #557 from langchain-ai/nc/30may/map-reduce

Browse Source 
Implement map-reduce api for StateGraph/Pregel
This commit is contained in:
Nuno Campos
2024-05-31 13:24:36 -07:00
committed by GitHub
parent 423ab4fe8b 9c513ee7e6
commit 1462b5e6b0
14 changed files with 1424 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/channels/base.py
+1
View File
@@ -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"],
)
langgraph/checkpoint/base.py
+7
View File
@@ -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(),
)
langgraph/constants.py
+10
View File
@@ -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
langgraph/graph/graph.py
+38 -17
View File
@@ -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:
langgraph/graph/state.py
+20 -8
View File
@@ -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])
langgraph/pregel/__init__.py
+135 -44
View File
@@ -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:
langgraph/pregel/io.py
+39 -18
View File
@@ -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")
langgraph/pregel/read.py
+6 -1
View File
@@ -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
langgraph/pregel/validate.py
+5 -5
View File
@@ -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:
langgraph/pregel/write.py
+77 -17
View File
@@ -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:
langgraph/serde/jsonplus.py
+4 -1
View File
@@ -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"
langgraph/utils.py
+15 -15
View File
@@ -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
tests/test_pregel.py
+533
View File
@@ -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,
tests/test_pregel_async.py
+534 -6
View File
@@ -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 (