Compare commits

...
Author SHA1 Message Date
Harrison Chase 67f6ab27b8 add overwrite command 2025-10-13 16:50:08 -07:00
3 changed files with 63 additions and 10 deletions
+2 -1
View File
@@ -1,3 +1,4 @@
from langgraph.pregel._write import Overwrite, overwrite
from langgraph.pregel.main import NodeBuilder, Pregel from langgraph.pregel.main import NodeBuilder, Pregel
__all__ = ("Pregel", "NodeBuilder") __all__ = ("Pregel", "NodeBuilder", "overwrite", "Overwrite")
+35 -9
View File
@@ -71,6 +71,7 @@ from langgraph.pregel._call import get_runnable_for_task, identifier
from langgraph.pregel._io import read_channels from langgraph.pregel._io import read_channels
from langgraph.pregel._log import logger from langgraph.pregel._log import logger
from langgraph.pregel._read import INPUT_CACHE_KEY_TYPE, PregelNode from langgraph.pregel._read import INPUT_CACHE_KEY_TYPE, PregelNode
from langgraph.pregel._write import _Overwrite
from langgraph.runtime import DEFAULT_RUNTIME, Runtime from langgraph.runtime import DEFAULT_RUNTIME, Runtime
from langgraph.types import ( from langgraph.types import (
All, All,
@@ -198,8 +199,16 @@ def local_read(
# apply writes # apply writes
local_channels: dict[str, BaseChannel] = {} local_channels: dict[str, BaseChannel] = {}
for k in channels: for k in channels:
cc = channels[k].copy() if updated[k]:
cc.update(updated[k]) # If any overwrite is present for this channel, reflect it directly
ow = next((v for v in updated[k] if isinstance(v, _Overwrite)), None)
if ow is not None:
cc = channels[k].from_checkpoint(ow.value)
else:
cc = channels[k].copy()
cc.update(updated[k])
else:
cc = channels[k].copy()
local_channels[k] = cc local_channels[k] = cc
# read fresh values # read fresh values
values = read_channels(local_channels, select) values = read_channels(local_channels, select)
@@ -277,12 +286,16 @@ def apply_writes(
# Group writes by channel # Group writes by channel
pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
overwrite_by_channel: dict[str, Any] = {}
for task in tasks: for task in tasks:
for chan, val in task.writes: for chan, val in task.writes:
if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR): if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR):
pass continue
elif chan in channels: if chan in channels:
pending_writes_by_channel[chan].append(val) if isinstance(val, _Overwrite):
overwrite_by_channel[chan] = val.value
else:
pending_writes_by_channel[chan].append(val)
else: else:
logger.warning( logger.warning(
f"Task {task.name} with path {task.path} wrote to unknown channel {chan}, ignoring it." f"Task {task.name} with path {task.path} wrote to unknown channel {chan}, ignoring it."
@@ -290,13 +303,26 @@ def apply_writes(
# Apply writes to channels # Apply writes to channels
updated_channels: set[str] = set() updated_channels: set[str] = set()
for chan, vals in pending_writes_by_channel.items(): for chan in set(pending_writes_by_channel.keys()) | set(
overwrite_by_channel.keys()
):
if chan in channels: if chan in channels:
if channels[chan].update(vals) and next_version is not None: if chan in overwrite_by_channel:
checkpoint["channel_versions"][chan] = next_version # Overwrite the entire channel value, bypassing reducers.
# unavailable channels can't trigger tasks, so don't add them channels[chan] = channels[chan].from_checkpoint(
overwrite_by_channel[chan]
)
if next_version is not None:
checkpoint["channel_versions"][chan] = next_version
if channels[chan].is_available(): if channels[chan].is_available():
updated_channels.add(chan) updated_channels.add(chan)
else:
vals = pending_writes_by_channel.get(chan, [])
if channels[chan].update(vals) and next_version is not None:
checkpoint["channel_versions"][chan] = next_version
# unavailable channels can't trigger tasks, so don't add them
if channels[chan].is_available():
updated_channels.add(chan)
# Channels that weren't updated in this step are notified of a new step # Channels that weren't updated in this step are notified of a new step
if bump_step: if bump_step:
+26
View File
@@ -26,6 +26,32 @@ SKIP_WRITE = object()
PASSTHROUGH = object() PASSTHROUGH = object()
class _Overwrite:
"""Marker wrapper indicating a direct channel overwrite.
Use via `overwrite(channel, value)` or `Overwrite(value)`.
"""
__slots__ = ("value",)
def __init__(self, value: Any):
self.value = value
def Overwrite(value: Any) -> _Overwrite:
"""Wrap a value to force overwrite a channel, bypassing reducers."""
return _Overwrite(value)
def overwrite(channel: str, value: Any) -> ChannelWriteEntry:
"""Convenience factory for a write that overwrites the target channel.
Example:
NodeBuilder().write_to(overwrite("foo", 123))
"""
return ChannelWriteEntry(channel, Overwrite(value))
class ChannelWriteEntry(NamedTuple): class ChannelWriteEntry(NamedTuple):
channel: str channel: str
"""Channel name to write to.""" """Channel name to write to."""