205 lines
7.7 KiB
Python
205 lines
7.7 KiB
Python
from __future__ import annotations
|
|
|
|
import collections.abc
|
|
import copy as _copy
|
|
from collections.abc import Callable, Sequence
|
|
from typing import Any, Generic
|
|
|
|
from langgraph.checkpoint.base import PendingWrite
|
|
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
|
from typing_extensions import Self
|
|
|
|
from langgraph._internal._typing import MISSING
|
|
from langgraph.channels.base import BaseChannel, Value
|
|
from langgraph.channels.binop import _get_overwrite, _operators_equal, _strip_extras
|
|
from langgraph.errors import (
|
|
EmptyChannelError,
|
|
ErrorCode,
|
|
InvalidUpdateError,
|
|
create_error_message,
|
|
)
|
|
|
|
__all__ = ("DeltaChannel",)
|
|
|
|
|
|
class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
|
|
"""Reducer channel that stores only a sentinel in checkpoint blobs and
|
|
reconstructs state by replaying ancestor writes through the reducer.
|
|
|
|
!!! warning "Beta"
|
|
|
|
`DeltaChannel` is in beta. The API and on-disk representation may
|
|
change in future releases. Threads written with `DeltaChannel` today
|
|
are expected to remain readable, but the surrounding contract
|
|
(`BaseCheckpointSaver.get_delta_channel_history`, the
|
|
`_DeltaSnapshot` blob shape, the `counters_since_delta_snapshot`
|
|
metadata field) is not yet stable.
|
|
|
|
The reducer receives the current accumulated value and a batch of writes
|
|
in one call: `reducer(state, [write1, write2, ...]) -> new_state`.
|
|
|
|
Reducers must be deterministic and batching-invariant (associative across
|
|
folds): applying two consecutive write batches separately must produce the
|
|
same state as applying their concatenation once:
|
|
|
|
reducer(reducer(state, xs), ys) == reducer(state, xs + ys)
|
|
|
|
This lets LangGraph replay checkpointed writes in larger batches than they
|
|
were originally produced without changing reconstructed state.
|
|
|
|
Snapshot cadence is driven by two counters: per-channel update count and
|
|
total supersteps since last snapshot. `create_checkpoint` writes a full
|
|
`_DeltaSnapshot` blob when EITHER the update count reaches
|
|
`snapshot_frequency` OR the supersteps count reaches the system-wide
|
|
`DELTA_MAX_SUPERSTEPS_SINCE_SNAPSHOT` bound (default 5000), bounding
|
|
replay depth even for channels that stop receiving writes.
|
|
|
|
Parameters:
|
|
reducer: `(state, list[writes]) -> new_state`. Must be deterministic
|
|
and batching-invariant as described above.
|
|
typ: The value type (e.g. `list`, `dict`). Inferred automatically
|
|
from the outer type when used inside `Annotated[T, DeltaChannel(...)]`.
|
|
snapshot_frequency: Every Nth update to this channel writes a snapshot
|
|
blob (default `1000`). Must be a positive int.
|
|
"""
|
|
|
|
__slots__ = ("value", "reducer", "snapshot_frequency")
|
|
value: Value | Any
|
|
|
|
def __init__(
|
|
self,
|
|
reducer: Callable[[Any, Sequence[Any]], Any],
|
|
typ: type[Value] | None = None,
|
|
*,
|
|
snapshot_frequency: int = 1000,
|
|
) -> None:
|
|
if snapshot_frequency <= 0:
|
|
raise ValueError(
|
|
f"snapshot_frequency must be a positive int, got {snapshot_frequency}"
|
|
)
|
|
if typ is None:
|
|
typ = list # type: ignore[assignment] # placeholder; overridden by _is_field_channel
|
|
super().__init__(typ)
|
|
self.reducer = reducer
|
|
self.snapshot_frequency = snapshot_frequency
|
|
typ = _strip_extras(typ)
|
|
if typ in (collections.abc.Sequence, collections.abc.MutableSequence):
|
|
typ = list
|
|
if typ in (collections.abc.Set, collections.abc.MutableSet):
|
|
typ = set
|
|
if typ in (collections.abc.Mapping, collections.abc.MutableMapping):
|
|
typ = dict
|
|
self.typ = typ
|
|
self.value: Any = MISSING
|
|
|
|
def __eq__(self, other: object) -> bool:
|
|
if not isinstance(other, DeltaChannel):
|
|
return False
|
|
if self.snapshot_frequency != other.snapshot_frequency:
|
|
return False
|
|
return _operators_equal(self.reducer, other.reducer)
|
|
|
|
@property
|
|
def ValueType(self) -> Any:
|
|
return self.typ
|
|
|
|
@property
|
|
def UpdateType(self) -> Any:
|
|
return self.typ
|
|
|
|
def copy(self) -> Self:
|
|
new = self.__class__(
|
|
self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency
|
|
)
|
|
new.key = self.key
|
|
new.value = self.value if self.value is MISSING else _copy.copy(self.value)
|
|
return new
|
|
|
|
def from_checkpoint(self, checkpoint: Any) -> Self:
|
|
"""Initialize from a stored blob.
|
|
|
|
Blob types:
|
|
* `MISSING`: start empty; caller replays writes.
|
|
* `_DeltaSnapshot(value)`: restore value directly from snapshot.
|
|
* plain value (migration from old `BinaryOperatorAggregate` blobs):
|
|
use directly.
|
|
"""
|
|
new = self.__class__(
|
|
self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency
|
|
)
|
|
new.key = self.key
|
|
if checkpoint is MISSING:
|
|
new.value = self.typ()
|
|
elif isinstance(checkpoint, _DeltaSnapshot):
|
|
new.value = checkpoint.value
|
|
else:
|
|
new.value = checkpoint
|
|
return new
|
|
|
|
def replay_writes(self, writes: Sequence[PendingWrite]) -> None:
|
|
"""Apply ancestor writes oldest-to-newest via a single reducer call.
|
|
|
|
If any write is an Overwrite, the last one in the sequence acts as
|
|
the reset point: its value becomes the new base and only writes
|
|
after it are passed to the reducer.
|
|
"""
|
|
values = [v for _, _, v in writes]
|
|
if not values:
|
|
return
|
|
base = self.value
|
|
start = 0
|
|
for i, v in enumerate(values):
|
|
is_ow, ow_value = _get_overwrite(v)
|
|
if is_ow:
|
|
base = _copy.copy(ow_value) if ow_value is not None else self.typ()
|
|
start = i + 1
|
|
remaining = values[start:]
|
|
self.value = self.reducer(base, remaining) if remaining else base
|
|
|
|
def update(self, values: Sequence[Any]) -> bool:
|
|
if not values:
|
|
return False
|
|
overwrite_idx: int | None = None
|
|
for i, v in enumerate(values):
|
|
is_ow, _ = _get_overwrite(v)
|
|
if is_ow:
|
|
if overwrite_idx is not None:
|
|
msg = create_error_message(
|
|
message="Can receive only one Overwrite value per super-step.",
|
|
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
|
|
)
|
|
raise InvalidUpdateError(msg)
|
|
overwrite_idx = i
|
|
if overwrite_idx is not None:
|
|
_, overwrite_value = _get_overwrite(values[overwrite_idx])
|
|
base = (
|
|
_copy.copy(overwrite_value)
|
|
if overwrite_value is not None
|
|
else self.typ()
|
|
)
|
|
remaining = [v for i, v in enumerate(values) if i != overwrite_idx]
|
|
self.value = self.reducer(base, remaining) if remaining else base
|
|
return True
|
|
base = self.typ() if self.value is MISSING else self.value
|
|
self.value = self.reducer(base, list(values))
|
|
return True
|
|
|
|
def get(self) -> Any:
|
|
if self.value is MISSING:
|
|
raise EmptyChannelError()
|
|
return self.value
|
|
|
|
def is_available(self) -> bool:
|
|
return self.value is not MISSING
|
|
|
|
def checkpoint(self) -> Any:
|
|
"""Return stored representation: always `MISSING`.
|
|
|
|
Snapshot decisions live in `create_checkpoint` (which has the channel
|
|
version) and write `_DeltaSnapshot(ch.get())` directly into
|
|
`channel_values`. For non-snapshot steps the channel does not appear
|
|
in `channel_values`; reconstruction walks ancestor writes via the
|
|
saver's `get_delta_channel_history`.
|
|
"""
|
|
return MISSING
|