91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
"""Replay state for subgraph checkpoint loading during time-travel."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import TYPE_CHECKING
|
||
|
||
from langgraph._internal._constants import NS_END
|
||
|
||
if TYPE_CHECKING:
|
||
from langchain_core.runnables import RunnableConfig
|
||
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointTuple
|
||
|
||
|
||
class ReplayState:
|
||
"""Tracks which subgraphs have already loaded their pre-replay checkpoint.
|
||
|
||
During a parent replay, each subgraph's first invocation should restore the
|
||
checkpoint from before the replay point. Subsequent invocations of the same
|
||
subgraph (e.g. in a loop) should use normal checkpoint loading so they pick
|
||
up freshly created checkpoints.
|
||
|
||
The single `ReplayState` instance is shared by reference across all derived
|
||
configs within one parent execution.
|
||
"""
|
||
|
||
__slots__ = ("checkpoint_id", "_visited_ns")
|
||
|
||
def __init__(self, checkpoint_id: str) -> None:
|
||
self.checkpoint_id = checkpoint_id
|
||
# DO NOT CHANGE THIS VARIABLE – it may need to be rehydrated
|
||
# in other runtimes
|
||
self._visited_ns: set[str] = set()
|
||
|
||
def _is_first_visit(self, checkpoint_ns: str) -> bool:
|
||
"""Return True the first time a subgraph namespace is seen.
|
||
|
||
The task-id suffix is stripped so that the same logical subgraph
|
||
(e.g. ``"sub_node"``) is recognized across loop iterations even
|
||
though each iteration has a different task id.
|
||
"""
|
||
# "sub_node:task_id" -> "sub_node"
|
||
stable_ns = (
|
||
checkpoint_ns.rsplit(NS_END, 1)[0]
|
||
if NS_END in checkpoint_ns
|
||
else checkpoint_ns
|
||
)
|
||
if stable_ns in self._visited_ns:
|
||
return False
|
||
self._visited_ns.add(stable_ns)
|
||
return True
|
||
|
||
def get_checkpoint(
|
||
self,
|
||
checkpoint_ns: str,
|
||
checkpointer: BaseCheckpointSaver,
|
||
checkpoint_config: RunnableConfig,
|
||
) -> CheckpointTuple | None:
|
||
"""Load the right checkpoint for a subgraph during replay.
|
||
|
||
On the first call for a given subgraph namespace, returns the latest
|
||
checkpoint created *before* the replay point. On subsequent calls
|
||
(e.g. the same subgraph in a later loop iteration), falls back to
|
||
normal latest-checkpoint loading.
|
||
"""
|
||
if self._is_first_visit(checkpoint_ns):
|
||
for saved in checkpointer.list(
|
||
checkpoint_config,
|
||
before={"configurable": {"checkpoint_id": self.checkpoint_id}},
|
||
limit=1,
|
||
):
|
||
return saved
|
||
return None
|
||
return checkpointer.get_tuple(checkpoint_config)
|
||
|
||
async def aget_checkpoint(
|
||
self,
|
||
checkpoint_ns: str,
|
||
checkpointer: BaseCheckpointSaver,
|
||
checkpoint_config: RunnableConfig,
|
||
) -> CheckpointTuple | None:
|
||
"""Async version of `get_checkpoint`."""
|
||
if self._is_first_visit(checkpoint_ns):
|
||
async for saved in checkpointer.alist(
|
||
checkpoint_config,
|
||
before={"configurable": {"checkpoint_id": self.checkpoint_id}},
|
||
limit=1,
|
||
):
|
||
return saved
|
||
return None
|
||
return await checkpointer.aget_tuple(checkpoint_config)
|