WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-10 03:37:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

optimizations i sure hope

Browse Source
This commit is contained in:
Sydney Runkle
2026-04-22 20:47:51 -04:00
parent 5b7fdf5655
commit acc7eda8c5
7 changed files with 192 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+3 -2
View File
@@ -15,6 +15,7 @@ from langgraph.checkpoint.base import (
CheckpointMetadata,
CheckpointTuple,
DeltaChannelWrites,
_overwrite_types,
get_checkpoint_id,
get_serializable_checkpoint_metadata,
)
@@ -402,7 +403,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
cur: Any,
) -> list[Any]:
"""Async version of _get_channel_writes_cur — see sync version for rationale."""
from langgraph.types import Overwrite # type: ignore[import-untyped]
overwrite_types = _overwrite_types()
await cur.execute(
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
@@ -435,7 +436,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
for type_tag, blob in writes_by_cp.get(cid, []):
val = self.serde.loads_typed((type_tag, blob))
collected.append(val)
if isinstance(val, Overwrite):
if isinstance(val, overwrite_types):
collected.reverse()
return collected
collected.reverse()
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+3 -2
View File
@@ -14,6 +14,7 @@ from langgraph.checkpoint.base import (
BaseCheckpointSaver,
ChannelVersions,
DeltaChannelWrites,
_overwrite_types,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.types import TASKS
@@ -235,7 +236,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread.
2. Walk ancestry in Python, then fetch writes with a plain ANY() filter.
"""
from langgraph.types import Overwrite # type: ignore[import-untyped]
overwrite_types = _overwrite_types()
cur.execute(
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
@@ -268,7 +269,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
for type_tag, blob in writes_by_cp.get(cid, []):
val = self.serde.loads_typed((type_tag, blob))
collected.append(val)
if isinstance(val, Overwrite):
if isinstance(val, overwrite_types):
collected.reverse()
return collected
collected.reverse()
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+26 -10
View File
@@ -55,6 +55,26 @@ def _overwrite_types() -> tuple[type, ...]:
return (Overwrite,)
def _split_list_config(
config: RunnableConfig,
) -> tuple[RunnableConfig, RunnableConfig | None]:
"""Split a `get_channel_writes` config into `(list_config, before_config)`.
Most savers collapse `list(config)` to a single row when `config` carries a
`checkpoint_id`. For ancestor traversal we need the opposite: every tuple
strictly earlier than the target. Drop `checkpoint_id` from `list_config`
and pass the original as `before=` (which filters `checkpoint_id < target`).
"""
configurable = config.get("configurable", {}) or {}
target_id = configurable.get("checkpoint_id")
if target_id is None:
return config, None
list_config: RunnableConfig = {
"configurable": {k: v for k, v in configurable.items() if k != "checkpoint_id"}
}
return list_config, config
logger = logging.getLogger(__name__)
@@ -503,14 +523,12 @@ class BaseCheckpointSaver(Generic[V]):
if getattr(_DELTA_RECONSTRUCTION, "active", False):
return []
overwrite_types = _overwrite_types()
list_config, before_config = _split_list_config(config)
_DELTA_RECONSTRUCTION.active = True
try:
collected: list[Any] = [] # newest first
target_id = config["configurable"].get("checkpoint_id")
for tup in self.list(config): # newest → oldest
if tup.config["configurable"].get("checkpoint_id") == target_id:
continue
for tup in self.list(list_config, before=before_config): # newest → oldest
if not tup.pending_writes:
continue
# Within a superstep, pending_writes are oldest→newest; reverse
@@ -531,17 +549,15 @@ class BaseCheckpointSaver(Generic[V]):
self, config: RunnableConfig, channel: str
) -> List[Any]: # noqa: UP006
"""Async version of get_channel_writes."""
overwrite_types = _overwrite_types()
if getattr(_DELTA_RECONSTRUCTION, "active", False):
return []
overwrite_types = _overwrite_types()
list_config, before_config = _split_list_config(config)
_DELTA_RECONSTRUCTION.active = True
try:
collected: list[Any] = []
target_id = config["configurable"].get("checkpoint_id")
async for tup in self.alist(config):
if tup.config["configurable"].get("checkpoint_id") == target_id:
continue
async for tup in self.alist(list_config, before=before_config):
if not tup.pending_writes:
continue
for _, ch, value in reversed(tup.pending_writes):
libs/checkpoint/tests/test_memory.py
+122
View File
@@ -390,3 +390,125 @@ class TestInMemorySaverDeltaChannel:
}
result = saver.get_channel_writes(config, channel)
assert result == [{"content": "hi"}, {"content": "bye"}]
class TestBaseFallbackGetChannelWrites:
"""Exercises the `BaseCheckpointSaver.get_channel_writes` default
implementation — the path third-party savers inherit when they don't
override `get_channel_writes` themselves.
Regression guard for a bug where the fallback passed the caller's config
(with `checkpoint_id`) straight to `self.list()`, which most savers
collapse to a single row — causing the fallback to return `[]`.
"""
def _build_saver_with_chain(self) -> tuple[InMemorySaver, str, str]:
"""Build an InMemorySaver with a 3-checkpoint chain and per-step writes
for a `messages` channel.
Returns `(saver, thread_id, namespace)`. The saver subclass deletes the
InMemorySaver override so the base class fallback is exercised.
"""
class _ThirdPartyStyleSaver(InMemorySaver):
get_channel_writes = (
InMemorySaver.__mro__[1].get_channel_writes # type: ignore[attr-defined]
)
aget_channel_writes = (
InMemorySaver.__mro__[1].aget_channel_writes # type: ignore[attr-defined]
)
saver = _ThirdPartyStyleSaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
cp0 = empty_checkpoint()
cp0["id"] = "00000000000000000000000000000001.0000000000000000"
cp1 = empty_checkpoint()
cp1["id"] = "00000000000000000000000000000002.0000000000000000"
cp2 = empty_checkpoint()
cp2["id"] = "00000000000000000000000000000003.0000000000000000"
saver.storage[thread_id][ns] = {
cp0["id"]: (serde.dumps_typed(cp0), serde.dumps_typed({}), None),
cp1["id"]: (serde.dumps_typed(cp1), serde.dumps_typed({}), cp0["id"]),
cp2["id"]: (serde.dumps_typed(cp2), serde.dumps_typed({}), cp1["id"]),
}
# Writes under cp0 produced cp1's state; writes under cp1 produced cp2's.
saver.writes[(thread_id, ns, cp0["id"])][("task1", 0)] = (
"task1",
channel,
serde.dumps_typed({"content": "first"}),
"",
)
saver.writes[(thread_id, ns, cp1["id"])][("task2", 0)] = (
"task2",
channel,
serde.dumps_typed({"content": "second"}),
"",
)
return saver, thread_id, ns
def test_fallback_returns_ancestor_writes_oldest_first(self) -> None:
saver, thread_id, ns = self._build_saver_with_chain()
target_id = "00000000000000000000000000000003.0000000000000000"
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target_id,
}
}
result = saver.get_channel_writes(config, "messages")
assert result == [{"content": "first"}, {"content": "second"}]
async def test_async_fallback_returns_ancestor_writes_oldest_first(self) -> None:
saver, thread_id, ns = self._build_saver_with_chain()
target_id = "00000000000000000000000000000003.0000000000000000"
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target_id,
}
}
result = await saver.aget_channel_writes(config, "messages")
assert result == [{"content": "first"}, {"content": "second"}]
def test_fallback_stops_at_first_overwrite(self) -> None:
"""An `Overwrite` dominates older history: scan newest→oldest stops at
the first one (so `snapshot_every` / user Overwrites bound replay cost).
"""
langgraph_types = pytest.importorskip(
"langgraph.types", reason="langgraph core not installed"
)
Overwrite = langgraph_types.Overwrite
saver, thread_id, ns = self._build_saver_with_chain()
serde = JsonPlusSerializer()
cp1_id = "00000000000000000000000000000002.0000000000000000"
# Replace cp1's write with an Overwrite — cp0's write must be dropped.
saver.writes[(thread_id, ns, cp1_id)][("task2", 0)] = (
"task2",
"messages",
serde.dumps_typed(Overwrite([{"content": "reset"}])),
"",
)
target_id = "00000000000000000000000000000003.0000000000000000"
config: RunnableConfig = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": ns,
"checkpoint_id": target_id,
}
}
result = saver.get_channel_writes(config, "messages")
assert len(result) == 1
assert isinstance(result[0], Overwrite)
assert result[0].value == [{"content": "reset"}]
libs/langgraph/langgraph/channels/delta.py
+28 -27
View File
@@ -95,6 +95,24 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
new._writes_since_snapshot = self._writes_since_snapshot
return new
def _apply_write(self, value: Any, write: Any, counter: int) -> tuple[Any, int]:
"""Apply one write to `value`; return (new_value, new_counter).
An `Overwrite` resets the counter to 0; any other write increments it.
Centralizes the Overwrite/reducer branching used by both `update` (live
super-step) and `from_checkpoint` (ancestor replay).
"""
is_overwrite, overwrite_value = _get_overwrite(write)
if is_overwrite:
new_value = (
_copy.copy(overwrite_value)
if overwrite_value is not None
else _empty(self.typ)
)
return new_value, 0
base = _empty(self.typ) if value is MISSING else value
return self.operator(base, write), counter + 1
def from_checkpoint(self, checkpoint: Any) -> Self:
new: DeltaChannel[Value] = DeltaChannel(
self.operator, snapshot_every=self.snapshot_every
@@ -106,22 +124,12 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
new._writes_since_snapshot = 0
elif isinstance(checkpoint, DeltaChannelWrites):
# Saver reconstructed per-step writes; replay through the operator.
# Count writes since last Overwrite so the snapshot cadence stays
# accurate across reloads.
# Counter tracks writes since the last Overwrite so snapshot cadence
# stays accurate across reloads.
value: Any = _empty(new.typ)
counter = 0
for write in checkpoint.writes:
is_overwrite, overwrite_value = _get_overwrite(write)
if is_overwrite:
value = (
_copy.copy(overwrite_value)
if overwrite_value is not None
else _empty(new.typ)
)
counter = 0
else:
value = new.operator(value, write)
counter += 1
value, counter = new._apply_write(value, write, counter)
new.value = value
new._writes_since_snapshot = counter
else:
@@ -135,9 +143,8 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
if not values:
return False
seen_overwrite = False
applied = 0
for value in values:
is_overwrite, overwrite_value = _get_overwrite(value)
is_overwrite, _ = _get_overwrite(value)
if is_overwrite:
if seen_overwrite:
from langgraph.errors import (
@@ -151,19 +158,13 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]):
error_code=ErrorCode.INVALID_CONCURRENT_GRAPH_UPDATE,
)
raise InvalidUpdateError(msg)
self.value = (
_copy.copy(overwrite_value)
if overwrite_value is not None
else _empty(self.typ)
)
seen_overwrite = True
self._writes_since_snapshot = 0
elif not seen_overwrite:
base = _empty(self.typ) if self.value is MISSING else self.value
self.value = self.operator(base, value)
applied += 1
if not seen_overwrite:
self._writes_since_snapshot += applied
elif seen_overwrite:
# Post-Overwrite writes within the same super-step are dropped.
continue
self.value, self._writes_since_snapshot = self._apply_write(
self.value, value, self._writes_since_snapshot
)
return True
def get(self) -> Any:
libs/langgraph/langgraph/pregel/main.py
+10 -10
View File
@@ -1049,14 +1049,13 @@ class Pregel(
step = saved.metadata.get("step", -1) + 1
stop = step + 2
checkpoint = saved.checkpoint
channels, managed = channels_from_checkpoint(
self.channels,
checkpoint,
saved.checkpoint,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
@@ -1169,14 +1168,13 @@ class Pregel(
step = saved.metadata.get("step", -1) + 1
stop = step + 2
checkpoint = saved.checkpoint
channels, managed = channels_from_checkpoint(
self.channels,
checkpoint,
saved.checkpoint,
)
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
checkpoint,
saved.checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
@@ -1522,8 +1520,9 @@ class Pregel(
saved = checkpointer.get_tuple(config)
if saved is not None:
self._migrate_checkpoint(saved.checkpoint)
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
checkpoint = (
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
)
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"].copy() if saved else {}
)
@@ -1967,8 +1966,9 @@ class Pregel(
saved = await checkpointer.aget_tuple(config)
if saved is not None:
self._migrate_checkpoint(saved.checkpoint)
base_checkpoint = saved.checkpoint if saved else empty_checkpoint()
checkpoint = copy_checkpoint(base_checkpoint) if saved else base_checkpoint
checkpoint = (
copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
)
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"].copy() if saved else {}
)
pr_delta_channel.md
-133
View File
@@ -1,133 +0,0 @@
# feat(channels): DeltaChannel — O(N) incremental checkpoint storage
## The problem
LangGraph checkpoints store the **full accumulated value** of every channel on every step. For a `messages` channel backed by `add_messages`, that means each checkpoint blob contains the entire conversation history up to that point.
Storage cost grows **O(N²)** in the number of turns:
| Step | Checkpoint blob |
|------|----------------|
| 1 | [msg_1] |
| 2 | [msg_1, msg_2] |
| N | [msg_1, ..., msg_N] |
At 100K tokens of conversation data, a single thread accumulates ~250 MB; with large messages or file attachments costs scale even faster.
## The fix: `DeltaChannel`
`DeltaChannel` is an opt-in wrapper around any binary reducer that stores only a **sentinel marker** in `checkpoint_blobs` rather than the full accumulated value. The actual per-step writes stay in `checkpoint_writes` (which every checkpointer already writes unconditionally). At read time the saver walks the ancestor chain, collects all writes for the channel, and replays them through the reducer.
Storage scales **O(N)** — the sentinel blob is effectively zero bytes, and the writes table already exists.
```python
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import add_messages
class State(TypedDict):
# Before: O(N²) storage
messages: Annotated[list[AnyMessage], add_messages]
# After: O(N) storage
messages: Annotated[list[AnyMessage], DeltaChannel(add_messages)]
```
## Benchmarks
Simulated with realistic paragraph-length messages (~100 tokens each, ~400 chars). Each turn = one human + one AI message (~200 tokens total).
### Storage (InMemorySaver)
| turns | ctx | add_msgs | delta | savings |
|------:|----:|---------:|------:|--------:|
| 10 | ~2K tok | 108.6 KB | 4.0 KB | 27x |
| 25 | ~5K tok | 649.0 KB | 10.1 KB | 64x |
| 50 | ~10K tok | 2.6 MB | 20.2 KB | 126x |
| 100 | ~20K tok | 10.2 MB | 40.5 KB | 251x |
| 500 | ~100K tok | 252.6 MB | 202.8 KB | 1245x |
Savings grow with N because `add_messages` is O(N²) while `DeltaChannel` is O(N). The sentinel blob itself is essentially zero bytes.
### Read latency (avg of 5 `get_state` calls = cost per `invoke`)
| turns | ctx | add_msgs | delta |
|------:|----:|---------:|------:|
| 10 | ~2K tok | 0.1ms | 0.2ms |
| 25 | ~5K tok | 0.2ms | 0.5ms |
| 50 | ~10K tok | 0.4ms | 1.5ms |
| 100 | ~20K tok | 0.7ms | 4.8ms |
| 500 | ~100K tok | 5.8ms | 114.9ms |
**This cost is paid once per `invoke`/`stream` call, not per node.** Within a single invocation, all channels are loaded into memory once at the start and shared across every node — there is no per-node reconstruction. The 114.9ms at 500 turns is what you pay each time a user sends a new message, not on each step of the graph.
## How it works
**Write:** `DeltaChannel.checkpoint()` always emits `DeltaChannelSentinel()` — a tiny marker (zero payload bytes) stored in `checkpoint_blobs`. Per-step writes flow into `checkpoint_writes` as they normally do for every channel.
**Read:** The saver detects `DeltaChannelSentinel` values in `channel_values` and replaces them by calling `get_channel_writes` / `aget_channel_writes`, which walks the ancestor checkpoint chain and collects all writes for that channel (oldest→newest). `DeltaChannel.from_checkpoint()` replays those writes through the operator to reconstruct the full value.
**Saver implementations:**
- `InMemorySaver` — direct dict traversal of `self.storage` and `self.writes`, no I/O
- `PostgresSaver` (sync + async) — two queries: one cheap ID walk across the thread, one `ANY()` fetch of writes; no recursive CTE
- All other savers — `BaseCheckpointSaver.get_channel_writes` fallback via `list()`, with a re-entrancy guard to prevent infinite recursion
## Changes
**`libs/checkpoint`**
- `base/__init__.py` — add `DeltaChannelSentinel` marker dataclass; add `get_channel_writes` / `aget_channel_writes` to `BaseCheckpointSaver` with a `list()`-based fallback and re-entrancy guard
**`libs/checkpoint/memory`**
- `memory/__init__.py` — `get_channel_writes` via direct dict traversal; `_resolve_delta_channels` helper called in `get_tuple` / `aget_tuple` to replace sentinels with reconstructed write lists
**`libs/langgraph`**
- `channels/delta.py` — `DeltaChannel` implementation: `checkpoint()` always emits sentinel, `from_checkpoint()` replays writes list
- `channels/__init__.py` — export `DeltaChannel`
- `graph/state.py` — recognize `DeltaChannel` as a valid channel annotation
- `pregel/_checkpoint.py` / `pregel/_loop.py` — wire `after_checkpoint` hook; call it after each checkpointing step so `DeltaChannel` can advance internal state
**`libs/checkpoint-postgres`**
- `postgres/base.py` — `_get_channel_writes_cur` two-query ancestor walk (sync); `_resolve_delta_channels` called after `_load_blobs`
- `postgres/aio.py` — `_aget_channel_writes_cur` (async counterpart)
## Open questions
**Should we add a compile-time capability check?**
Currently misconfiguring `DeltaChannel` with an unsupported saver only errors at runtime on first reload. A protocol-based check at `compile()` time would give an early warning without requiring a manual boolean flag.
**`snapshot_every` for bounded reconstruction cost?**
Both per-invoke read latency and total write wall time grow O(N) per invoke / O(N²) total as the conversation lengthens. A `snapshot_every` parameter — periodically store a full snapshot in `checkpoint_blobs` to cap chain depth — would bound reconstruction cost and is a natural follow-up once the core design is stable.
## Backwards compatibility
| Scenario | Behaviour |
|----------|-----------|
| Existing graph using `add_messages` | Unaffected — no code or schema changes |
| `DeltaChannel` loading an old full-list checkpoint blob | Handled via backwards-compat path in `from_checkpoint` |
| `DeltaChannel` with `InMemorySaver` or `PostgresSaver` | Fully supported |
| Time-travel to a past checkpoint | Ancestor walk uses the version at that checkpoint — correct by construction |
| `Overwrite` value | Resets the effective chain; reconstruction starts from that step |
## Test plan
- [x] `DeltaChannel` unit tests: `update` → `checkpoint` lifecycle, `from_checkpoint` chain replay, backwards-compat with plain list, `Overwrite` resets chain
- [x] `InMemorySaver` `get_channel_writes`: assembles write list from dict storage
- [x] Serde round-trip for `DeltaChannelSentinel`
- [x] End-to-end graph tests: multi-turn conversations accumulate correctly, time-travel reconstructs correct partial history
- [x] `PostgresSaver` two-query chain reconstruction (sync + async)
- [x] `BaseCheckpointSaver` fallback path via `list()` with re-entrancy guard
- [x] Storage benchmark: `DeltaChannel` uses strictly less storage than `add_messages` at all measured turn counts
---
## Changes from previous base branch
The previous version stored `DeltaValue` objects (containing the per-step writes) directly in `checkpoint_blobs` and used a `DeltaChainValue` to represent the assembled chain. Reconstruction required a dedicated `get_delta_chain` / `aget_delta_chain` protocol and a recursive CTE in Postgres.
This version pivots to a simpler design:
- **Sentinel in blobs, writes in `checkpoint_writes`** — `checkpoint_blobs` stores only a zero-byte `DeltaChannelSentinel` marker. The actual per-step data already lives in `checkpoint_writes` (written unconditionally by every checkpointer), so blob storage is essentially free. This is why storage savings jump to 1245x at 500 turns.
- **No custom serde type for the delta payload** — `DeltaValue` / `DeltaChainValue` and the `"delta"` serde type tag are gone. Writes are deserialized with the same serde path they were originally written with.
- **Postgres: two queries instead of a recursive CTE** — fetch all `(checkpoint_id, parent_checkpoint_id)` pairs for the thread, walk the ancestor chain in Python, then fetch writes with a plain `ANY()` filter.
- **Universal fallback on `BaseCheckpointSaver`** — the base class now provides `get_channel_writes` via `list()`, so any third-party saver works without modification.
- **`snapshot_every` removed** — deferred as a follow-up; the simpler design is easier to reason about and delivers larger storage savings.