From def55a5ac561774b2d669c95b761cf851c2b2883 Mon Sep 17 00:00:00 2001 From: Will Fu-Hinthorn Date: Tue, 21 Apr 2026 14:31:43 -0700 Subject: [PATCH] fun --- .../langgraph/checkpoint/base/__init__.py | 57 +++++++++---------- libs/langgraph/langgraph/channels/delta.py | 4 +- libs/langgraph/langgraph/pregel/_loop.py | 6 +- libs/langgraph/langgraph/pregel/main.py | 11 ++-- 4 files changed, 40 insertions(+), 38 deletions(-) diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index ea54475f8..86b4ef6ab 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import copy import dataclasses import logging @@ -567,22 +568,34 @@ class BaseCheckpointSaver(Generic[V]): thread_id = str(config["configurable"]["thread_id"]) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") current_checkpoint_id = checkpoint.get("id") - assembled: dict[str, Any] = {} - for spec in plan.channels: - value = checkpoint["channel_values"].get(spec.name) - if spec.kind != "delta" or not isinstance(value, DeltaValue): - continue + targets = [ + (spec, checkpoint["channel_values"][spec.name]) + for spec in plan.channels + if spec.kind == "delta" + and isinstance(checkpoint["channel_values"].get(spec.name), DeltaValue) + ] + if not targets: + return checkpoint - assembled_value = await self._amaterialize_delta_channel( - thread_id=thread_id, - checkpoint_ns=checkpoint_ns, - current_checkpoint_id=current_checkpoint_id, - channel=spec.name, - value=value, + # Walks for independent channels can run concurrently — each has its own chain. + results = await asyncio.gather( + *( + self._amaterialize_delta_channel( + thread_id=thread_id, + checkpoint_ns=checkpoint_ns, + current_checkpoint_id=current_checkpoint_id, + channel=spec.name, + value=value, + ) + for spec, value in targets ) - if assembled_value is not None: - assembled[spec.name] = assembled_value + ) + assembled = { + spec.name: result + for (spec, _), result in zip(targets, results, strict=True) + if result is not None + } if not assembled: return checkpoint @@ -683,18 +696,6 @@ class BaseCheckpointSaver(Generic[V]): clone.serde = maybe_add_typed_methods(serde) return clone - def _get_checkpoint_tuple_for_materialization( - self, config: RunnableConfig - ) -> CheckpointTuple | None: - """Internal raw checkpoint lookup used by fallback materialization.""" - return self.get_tuple(config) - - async def _aget_checkpoint_tuple_for_materialization( - self, config: RunnableConfig - ) -> CheckpointTuple | None: - """Async internal raw checkpoint lookup used by fallback materialization.""" - return await self.aget_tuple(config) - def _materialize_delta_channel( self, *, @@ -738,7 +739,7 @@ class BaseCheckpointSaver(Generic[V]): "checkpoint_id": prev_id, } } - parent_tuple = self._get_checkpoint_tuple_for_materialization(parent_config) + parent_tuple = self.get_tuple(parent_config) if parent_tuple is None: logger.warning( "DeltaChannel chain broken: checkpoint %r not found for channel %r", @@ -805,9 +806,7 @@ class BaseCheckpointSaver(Generic[V]): "checkpoint_id": prev_id, } } - parent_tuple = await self._aget_checkpoint_tuple_for_materialization( - parent_config - ) + parent_tuple = await self.aget_tuple(parent_config) if parent_tuple is None: logger.warning( "DeltaChannel chain broken: checkpoint %r not found for channel %r", diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index 7408e7b80..99aad844f 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -2,7 +2,7 @@ from __future__ import annotations import collections.abc from collections.abc import Callable, Sequence -from typing import Any, Generic +from typing import Any, Generic, Literal from langgraph.checkpoint.base import DeltaChainValue, DeltaValue from typing_extensions import Self @@ -99,7 +99,7 @@ class DeltaChannel(Generic[Value], BaseChannel[list[Value], Value, DeltaValue]): return self.typ | list[self.typ] # type: ignore[name-defined] @property - def checkpoint_hydration_kind(self) -> str: + def checkpoint_hydration_kind(self) -> Literal["delta"]: return "delta" def copy(self) -> Self: diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 1986c1d20..6b15096a8 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -12,6 +12,7 @@ from contextlib import ( ExitStack, ) from datetime import datetime, timezone +from functools import cached_property from inspect import signature from types import TracebackType from typing import ( @@ -316,6 +317,7 @@ class PregelLoop: ) self.prev_checkpoint_config = None + @cached_property def _checkpoint_hydration_plan(self) -> CheckpointHydrationPlan | None: """Build the saver hydration plan from this loop's channel specs.""" return checkpoint_hydration_plan(self.specs) @@ -326,7 +328,7 @@ class PregelLoop: if saved is None or self.checkpointer is None: return saved return self.checkpointer.materialize_checkpoint_tuple( - saved, self._checkpoint_hydration_plan() + saved, self._checkpoint_hydration_plan ) async def _amaterialize_saved_checkpoint( @@ -335,7 +337,7 @@ class PregelLoop: if saved is None or self.checkpointer is None: return saved return await self.checkpointer.amaterialize_checkpoint_tuple( - saved, self._checkpoint_hydration_plan() + saved, self._checkpoint_hydration_plan ) def _push_graph_lifecycle_event( diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 9f965e709..7440a7348 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -17,7 +17,7 @@ from collections.abc import ( Sequence, ) from dataclasses import is_dataclass, replace -from functools import partial +from functools import cached_property, partial from inspect import isclass from typing import ( Any, @@ -730,6 +730,7 @@ class Pregel( return checkpointer return _serde.apply_checkpointer_allowlist(checkpointer, self._serde_allowlist) + @cached_property def _checkpoint_hydration_plan(self) -> CheckpointHydrationPlan | None: return checkpoint_hydration_plan(self.channels) @@ -741,7 +742,7 @@ class Pregel( if saved is None or checkpointer is None: return saved return checkpointer.materialize_checkpoint_tuple( - saved, self._checkpoint_hydration_plan() + saved, self._checkpoint_hydration_plan ) async def _amaterialize_saved_checkpoint( @@ -752,7 +753,7 @@ class Pregel( if saved is None or checkpointer is None: return saved return await checkpointer.amaterialize_checkpoint_tuple( - saved, self._checkpoint_hydration_plan() + saved, self._checkpoint_hydration_plan ) def _materialize_saved_checkpoints( @@ -763,7 +764,7 @@ class Pregel( if checkpointer is None or not saved: return list(saved) return checkpointer.materialize_checkpoint_tuples( - saved, self._checkpoint_hydration_plan() + saved, self._checkpoint_hydration_plan ) async def _amaterialize_saved_checkpoints( @@ -774,7 +775,7 @@ class Pregel( if checkpointer is None or not saved: return list(saved) return await checkpointer.amaterialize_checkpoint_tuples( - saved, self._checkpoint_hydration_plan() + saved, self._checkpoint_hydration_plan ) def get_graph(