WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-24 08:32:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

Implement map-reduce api for StateGraph/Pregel

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