From a446f34ed91ac6ca3b8f9078d9d16ebc4160cc16 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 29 Apr 2025 11:44:31 -0700 Subject: [PATCH 01/41] WIP --- .../langgraph/checkpoint/base/__init__.py | 34 +++ libs/langgraph/langgraph/constants.py | 1 + libs/langgraph/langgraph/func/__init__.py | 6 +- libs/langgraph/langgraph/pregel/algo.py | 204 ++++++++++-------- libs/langgraph/langgraph/pregel/call.py | 12 +- libs/langgraph/langgraph/pregel/loop.py | 5 +- libs/langgraph/langgraph/pregel/read.py | 12 +- libs/langgraph/langgraph/types.py | 11 +- libs/langgraph/langgraph/utils/cache.py | 26 +++ 9 files changed, 206 insertions(+), 105 deletions(-) create mode 100644 libs/langgraph/langgraph/utils/cache.py diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 9aa051f9a..b5f5a6a7e 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -254,6 +254,23 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError + def get_writes( + self, + task_ids: Sequence[str], + ) -> dict[str, list[tuple[str, Any]]]: + """Fetch writes associated with the given task IDs. + + Args: + task_ids (Sequence[str]): List of task IDs to fetch writes for. + + Returns: + dict[str, list[tuple[str, Any]]]: Mapping from task ID to list of writes. + + Raises: + NotImplementedError: Implement this method in your custom checkpoint saver. + """ + raise NotImplementedError + def list( self, config: Optional[RunnableConfig], @@ -358,6 +375,23 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError + async def aget_writes( + self, + task_ids: Sequence[str], + ) -> dict[str, list[tuple[str, Any]]]: + """Asynchronously fetch writes associated with the given task IDs. + + Args: + task_ids (Sequence[str]): List of task IDs to fetch writes for. + + Returns: + dict[str, list[tuple[str, Any]]]: Mapping from task ID to list of writes. + + Raises: + NotImplementedError: Implement this method in your custom checkpoint saver. + """ + raise NotImplementedError + async def alist( self, config: Optional[RunnableConfig], diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 167c340f1..bcb614ca4 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -106,6 +106,7 @@ NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000") # the task_id to use for writes that are not associated with a task CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map") # holds a mapping of task ns -> resume value for resuming tasks +TASK_CACHE_NAMESPACE = b"__pregel_task_cache" RESERVED = { TAG_HIDDEN, diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index a951ef749..a3afe33ae 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -31,7 +31,7 @@ from langgraph.pregel.call import ( from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode +from langgraph.types import _DC_KWARGS, CachePolicy, RetryPolicy, StreamMode @overload @@ -39,6 +39,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, + cache: Optional[CachePolicy[P]] = None, ) -> Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], Callable[P, SyncAsyncFuture[T]], @@ -56,6 +57,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, + cache: Optional[CachePolicy[P]] = None, ) -> Union[ Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], @@ -141,7 +143,7 @@ def task( # handle regular functions / partials / callable classes, etc. func.__name__ = name - call_func = functools.partial(call, func, retry=retry_policies) + call_func = functools.partial(call, func, retry=retry_policies, cache=cache) object.__setattr__(call_func, "_is_pregel_task", True) return functools.update_wrapper(call_func, func) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c51afd7d3..a546bc8b4 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -60,6 +60,7 @@ from langgraph.constants import ( RESUME, RETURN, TAG_HIDDEN, + TASK_CACHE_NAMESPACE, TASKS, Send, ) @@ -72,6 +73,7 @@ from langgraph.pregel.read import INPUT_CACHE_KEY_TYPE, PregelNode from langgraph.store.base import BaseStore from langgraph.types import ( All, + CachePolicy, PregelExecutableTask, PregelScratchpad, PregelTask, @@ -111,24 +113,27 @@ class PregelTaskWrites(NamedTuple): class Call: - __slots__ = ("func", "input", "retry", "callbacks") + __slots__ = ("func", "input", "retry", "cache", "callbacks") func: Callable - input: Any + input: tuple[tuple[Any, ...], dict[str, Any]] retry: Optional[Sequence[RetryPolicy]] + cache: Optional[CachePolicy] callbacks: Callbacks def __init__( self, func: Callable, - input: Any, + input: tuple[tuple[Any, ...], dict[str, Any]], *, retry: Optional[Sequence[RetryPolicy]], + cache: Optional[CachePolicy], callbacks: Callbacks, ) -> None: self.func = func self.input = input self.retry = retry + self.cache = cache self.callbacks = callbacks @@ -556,15 +561,23 @@ def prepare_single_task( # create task id triggers: Sequence[str] = PUSH_TRIGGER checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name - task_id = task_id_func( - checkpoint_id_bytes, - checkpoint_ns, - str(step), - name, - PUSH, - task_path_str(task_path[1]), - str(task_path[2]), - ) + if call.cache: + task_id = task_id_func( + TASK_CACHE_NAMESPACE, + name, + call.cache.key(*call.input[0], **call.input[1]), + # TODO add truncated iso timestamp here for ttl + ) + else: + task_id = task_id_func( + checkpoint_id_bytes, + checkpoint_ns, + str(step), + name, + PUSH, + task_path_str(task_path[1]), + str(task_path[2]), + ) task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" # we append True to the task path to indicate that a call is being # made, so we should not return interrupts from this task (responsibility lies with the parent) @@ -625,7 +638,7 @@ def prepare_single_task( ), triggers, call.retry, - None, + call.cache, task_id, task_path, ) @@ -649,19 +662,31 @@ def prepare_single_task( f"Ignoring unknown node name {packet.node} in pending sends" ) return + # find process + proc = processes[packet.node] + proc_node = proc.node + if proc_node is None: + return # create task id triggers = PUSH_TRIGGER checkpoint_ns = ( f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node ) - task_id = task_id_func( - checkpoint_id_bytes, - checkpoint_ns, - str(step), - packet.node, - PUSH, - str(idx), - ) + if proc.cache_policy: + task_id = task_id_func( + TASK_CACHE_NAMESPACE, + packet.node, + proc.cache_policy.key(packet.arg), + ) + else: + task_id = task_id_func( + checkpoint_id_bytes, + checkpoint_ns, + str(step), + packet.node, + PUSH, + str(idx), + ) else: logger.warning(f"Ignoring invalid PUSH task path {task_path}") return @@ -679,73 +704,64 @@ def prepare_single_task( if task_id_checksum is not None: assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}" if for_execution: - proc = processes[packet.node] - if node := proc.node: - if proc.metadata: - metadata.update(proc.metadata) - writes = deque() - return PregelExecutableTask( - packet.node, - packet.arg, - node, - writes, - patch_config( - merge_configs( - config, {"metadata": metadata, "tags": proc.tags} - ), - run_name=packet.node, - callbacks=( - manager.get_child(f"graph:step:{step}") if manager else None - ), - configurable={ - CONFIG_KEY_TASK_ID: task_id, - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - local_write, - writes.extend, - processes.keys(), - ), - CONFIG_KEY_READ: partial( - local_read, - channels, - managed, - PregelTaskWrites( - task_path, packet.node, writes, triggers - ), - ), - CONFIG_KEY_STORE: ( - store or configurable.get(CONFIG_KEY_STORE) - ), - CONFIG_KEY_CHECKPOINTER: ( - checkpointer - or configurable.get(CONFIG_KEY_CHECKPOINTER) - ), - CONFIG_KEY_CHECKPOINT_MAP: { - **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}), - parent_ns: checkpoint["id"], - }, - CONFIG_KEY_CHECKPOINT_ID: None, - CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, - CONFIG_KEY_SCRATCHPAD: _scratchpad( - config[CONF].get(CONFIG_KEY_SCRATCHPAD), - pending_writes, - task_id, - xxh3_128_hexdigest(task_checkpoint_ns.encode()), - config[CONF].get(CONFIG_KEY_RESUME_MAP), - ), - CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get( - PREVIOUS, None - ), - }, + if proc.metadata: + metadata.update(proc.metadata) + writes = deque() + return PregelExecutableTask( + packet.node, + packet.arg, + proc_node, + writes, + patch_config( + merge_configs(config, {"metadata": metadata, "tags": proc.tags}), + run_name=packet.node, + callbacks=( + manager.get_child(f"graph:step:{step}") if manager else None ), - triggers, - proc.retry_policy, - None, - task_id, - task_path, - writers=proc.flat_writers, - subgraphs=proc.subgraphs, - ) + configurable={ + CONFIG_KEY_TASK_ID: task_id, + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, + writes.extend, + processes.keys(), + ), + CONFIG_KEY_READ: partial( + local_read, + channels, + managed, + PregelTaskWrites(task_path, packet.node, writes, triggers), + ), + CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)), + CONFIG_KEY_CHECKPOINTER: ( + checkpointer or configurable.get(CONFIG_KEY_CHECKPOINTER) + ), + CONFIG_KEY_CHECKPOINT_MAP: { + **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}), + parent_ns: checkpoint["id"], + }, + CONFIG_KEY_CHECKPOINT_ID: None, + CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, + CONFIG_KEY_SCRATCHPAD: _scratchpad( + config[CONF].get(CONFIG_KEY_SCRATCHPAD), + pending_writes, + task_id, + xxh3_128_hexdigest(task_checkpoint_ns.encode()), + config[CONF].get(CONFIG_KEY_RESUME_MAP), + ), + CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get( + PREVIOUS, None + ), + }, + ), + triggers, + proc.retry_policy, + proc.cache_policy, + task_id, + task_path, + writers=proc.flat_writers, + subgraphs=proc.subgraphs, + ) else: return PregelTask(task_id, packet.node, task_path) elif task_path[0] == PULL: @@ -784,6 +800,7 @@ def prepare_single_task( # create task id checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name + # TODO implement cache task_id = task_id_func( checkpoint_id_bytes, checkpoint_ns, @@ -1000,7 +1017,8 @@ def _proc_input( val = channels[chan].get() break else: - val[k] = managed[k]() + val = managed[chan]() + break else: return MISSING else: @@ -1019,18 +1037,20 @@ def _proc_input( return val -def _uuid5_str(namespace: bytes, *parts: str) -> str: +def _uuid5_str(namespace: bytes, *parts: str | bytes) -> str: """Generate a UUID from the SHA-1 hash of a namespace and str parts.""" sha = sha1(namespace, usedforsecurity=False) - sha.update(b"".join(p.encode() for p in parts)) + sha.update(b"".join(p.encode() if isinstance(p, str) else p for p in parts)) hex = sha.hexdigest() return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}" -def _xxhash_str(namespace: bytes, *parts: str) -> str: +def _xxhash_str(namespace: bytes, *parts: str | bytes) -> str: """Generate a UUID from the XXH3 hash of a namespace and str parts.""" - hex = xxh3_128_hexdigest(namespace + b"".join(p.encode() for p in parts)) + hex = xxh3_128_hexdigest( + namespace + b"".join(p.encode() if isinstance(p, str) else p for p in parts) + ) return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}" diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index e8dde073a..e3719d2b6 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -13,7 +13,7 @@ from typing_extensions import ParamSpec from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry -from langgraph.types import RetryPolicy +from langgraph.types import CachePolicy, RetryPolicy from langgraph.utils.config import get_config from langgraph.utils.runnable import ( RunnableCallable, @@ -28,6 +28,7 @@ from langgraph.utils.runnable import ( def _getattribute(obj: Any, name: str) -> Any: + parent = None for subpath in name.split("."): if subpath == "": raise AttributeError(f"Can't get local attribute {name!r} on {obj!r}") @@ -135,7 +136,7 @@ def _explode_args_trace_inputs( return arguments -def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq: +def get_runnable_for_entrypoint(func: Callable[..., Any]) -> Runnable: key = (func, False) if key in CACHE: return CACHE[key] @@ -160,7 +161,7 @@ def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq: return CACHE.setdefault(key, run) -def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq: +def get_runnable_for_task(func: Callable[..., Any]) -> Runnable: key = (func, True) if key in CACHE: return CACHE[key] @@ -222,9 +223,12 @@ def call( func: Callable[P, T], *args: Any, retry: Optional[Sequence[RetryPolicy]] = None, + cache: Optional[CachePolicy] = None, **kwargs: Any, ) -> SyncAsyncFuture[T]: config = get_config() impl = config[CONF][CONFIG_KEY_CALL] - fut = impl(func, (args, kwargs), retry=retry, callbacks=config["callbacks"]) + fut = impl( + func, (args, kwargs), retry=retry, cache=cache, callbacks=config["callbacks"] + ) return fut diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 805779d47..34d1d9bb9 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -434,7 +434,10 @@ class PregelLoop(LoopProtocol): # save the new task self.tasks[pushed.id] = pushed # match any pending writes to the new task - if self.skip_done_tasks: + if call.cache: + pass + # how to call async method here... + elif self.skip_done_tasks: self._match_writes({pushed.id: pushed}) # return the new task, to be started if not run before return pushed diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 7eb2a530c..028b08437 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -14,14 +14,15 @@ from langchain_core.runnables import ( RunnablePassthrough, RunnableSerializable, ) -from langchain_core.runnables.base import Input, Other, coerce_to_runnable -from langchain_core.runnables.utils import ConfigurableFieldSpec +from langchain_core.runnables.base import Other, coerce_to_runnable +from langchain_core.runnables.utils import ConfigurableFieldSpec, Input from langgraph.constants import CONF, CONFIG_KEY_READ from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.utils import find_subgraph_pregel from langgraph.pregel.write import ChannelWrite +from langgraph.types import CachePolicy from langgraph.utils.config import merge_configs from langgraph.utils.runnable import RunnableCallable, RunnableSeq @@ -143,6 +144,9 @@ class PregelNode(Runnable): retry_policy: Sequence[RetryPolicy] | None """The retry policies to use when invoking the node.""" + cache_policy: CachePolicy | None + """The cache policy to use when invoking the node.""" + tags: Sequence[str] | None """Tags to attach to the node for tracing.""" @@ -163,6 +167,7 @@ class PregelNode(Runnable): metadata: Mapping[str, Any] | None = None, bound: Runnable[Any, Any] | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, + cache_policy: CachePolicy | None = None, subgraphs: Sequence[PregelProtocol] | None = None, ) -> None: self.channels = channels @@ -170,8 +175,9 @@ class PregelNode(Runnable): self.mapper = mapper self.writers = writers or [] self.bound = bound if bound is not None else DEFAULT_BOUND + self.cache_policy = cache_policy if isinstance(retry_policy, RetryPolicy): - self.retry_policy: Sequence[RetryPolicy] = (retry_policy,) + self.retry_policy = (retry_policy,) else: self.retry_policy = retry_policy self.tags = tags diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3d34b35fd..51a3068ac 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -11,6 +11,7 @@ from typing import ( Literal, NamedTuple, Optional, + ParamSpec, TypeVar, Union, cast, @@ -22,6 +23,7 @@ from typing_extensions import Self from xxhash import xxh3_128_hexdigest from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata +from langgraph.utils.cache import default_cache_key from langgraph.utils.fields import get_update_as_tuples if TYPE_CHECKING: @@ -30,7 +32,7 @@ if TYPE_CHECKING: try: - from langchain_core.messages.tool import ToolOutputMixin + from langchain_core.messages.tool import ToolOutputMixin # type: ignore except ImportError: class ToolOutputMixin: # type: ignore[no-redef] @@ -122,13 +124,16 @@ class RetryPolicy(NamedTuple): """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" -class CachePolicy(NamedTuple): +P = ParamSpec("P") + + +class CachePolicy(NamedTuple, Generic[P]): """Configuration for caching nodes. !!! version-added "Added in version 0.2.24." """ - pass + key: Callable[P, str | bytes] = default_cache_key @dataclasses.dataclass(**_DC_KWARGS) diff --git a/libs/langgraph/langgraph/utils/cache.py b/libs/langgraph/langgraph/utils/cache.py new file mode 100644 index 000000000..ddae896c8 --- /dev/null +++ b/libs/langgraph/langgraph/utils/cache.py @@ -0,0 +1,26 @@ +from collections.abc import Hashable +from typing import Any + + +def _freeze(obj: Any) -> Hashable: + if isinstance(obj, dict): + # sort keys so {"a":1,"b":2} == {"b":2,"a":1} + return tuple(sorted((k, _freeze(v)) for k, v in obj.items())) + elif isinstance(obj, (list, tuple, set, frozenset)): + return tuple(_freeze(x) for x in obj) + # numpy / pandas etc. can provide their own .tobytes() + elif hasattr(obj, "tobytes"): + return ( + type(obj).__name__, + obj.tobytes(), + obj.shape if hasattr(obj, "shape") else None, + ) + return obj # strings, ints, dataclasses with frozen=True, etc. + + +def default_cache_key(*args: Any, **kwargs: Any) -> bytes: + """Default cache key function that uses the arguments and keyword arguments to generate a hashable key.""" + import pickle + + # protocol 5 strikes a good balance between speed and size + return pickle.dumps((_freeze(args), _freeze(kwargs)), protocol=5, fix_imports=False) From 1aecde3cd8e4177335cf4981ce46a33f3f6ad72d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Apr 2025 15:06:23 -0700 Subject: [PATCH 02/41] Re-do with separate cache interface --- .../langgraph/cache/base/__init__.py | 42 +++++++++++ .../langgraph/cache/file/__init__.py | 53 ++++++++++++++ .../langgraph/checkpoint/base/__init__.py | 34 --------- .../langgraph/checkpoint/memory/__init__.py | 3 +- .../langgraph/checkpoint/serde/base.py | 20 +++-- libs/langgraph/langgraph/constants.py | 3 +- libs/langgraph/langgraph/func/__init__.py | 3 + libs/langgraph/langgraph/pregel/__init__.py | 21 ++++++ libs/langgraph/langgraph/pregel/algo.py | 73 +++++++++---------- libs/langgraph/langgraph/pregel/loop.py | 70 +++++++++++++++++- libs/langgraph/langgraph/pregel/runner.py | 19 ++++- libs/langgraph/langgraph/types.py | 16 +++- libs/langgraph/tests/conftest.py | 20 ++++- libs/langgraph/tests/test_pregel.py | 46 +++++++++++- 14 files changed, 333 insertions(+), 90 deletions(-) create mode 100644 libs/checkpoint/langgraph/cache/base/__init__.py create mode 100644 libs/checkpoint/langgraph/cache/file/__init__.py diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py new file mode 100644 index 000000000..a81b48cf6 --- /dev/null +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -0,0 +1,42 @@ +from abc import ABC, abstractmethod +from collections.abc import Mapping +from typing import Generic, Sequence, TypeVar + +from langgraph.checkpoint.serde.base import SerializerProtocol +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer + +T = TypeVar("T") + + +class BaseCache(ABC, Generic[T]): + """Base class for a cache.""" + + serde: SerializerProtocol = JsonPlusSerializer() + + def __init__(self, *, serde: SerializerProtocol | None = None) -> None: + """Initialize the cache with a serializer.""" + self.serde = serde or self.serde + + @abstractmethod + def get(self, keys: Sequence[str]) -> dict[str, T]: + """Get the cached values for the given keys.""" + + @abstractmethod + def set(self, mapping: Mapping[str, tuple[T, int | None]]) -> None: + """Set the cached values for the given keys and TTLs.""" + + @abstractmethod + def delete(self, keys: Sequence[str]) -> None: + """Delete the cached values for the given keys.""" + + @abstractmethod + async def aget(self, keys: Sequence[str]) -> dict[str, T]: + """Asynchronously get the cached values for the given keys.""" + + @abstractmethod + async def aset(self, mapping: Mapping[str, tuple[T, int | None]]) -> None: + """Asynchronously set the cached values for the given keys and TTLs.""" + + @abstractmethod + async def adelete(self, keys: Sequence[str]) -> None: + """Asynchronously delete the cached values for the given keys.""" diff --git a/libs/checkpoint/langgraph/cache/file/__init__.py b/libs/checkpoint/langgraph/cache/file/__init__.py new file mode 100644 index 000000000..cc6499543 --- /dev/null +++ b/libs/checkpoint/langgraph/cache/file/__init__.py @@ -0,0 +1,53 @@ +import asyncio +import dbm + +import ormsgpack + +from langgraph.cache.base import BaseCache +from langgraph.checkpoint.serde.base import SerializerProtocol + + +class FileCache(BaseCache): + """File-based cache using dbm.""" + + def __init__( + self, + *, + path: str, + serde: SerializerProtocol | None = None, + ) -> None: + """Initialize the cache with a file path.""" + super().__init__(serde=serde) + self._db = dbm.open(path, "c") + + def get(self, keys: list[str]) -> dict[str, bytes]: + """Get the cached values for the given keys.""" + return { + key: self.serde.loads_typed(ormsgpack.unpackb(self._db[key])) + for key in keys + if key in self._db + } + + def set(self, mapping: dict[str, tuple[bytes, int | None]]) -> None: + """Set the cached values for the given keys and TTLs.""" + for key, (value, _) in mapping.items(): + # File-based caches do not support TTLs, so we ignore them. + self._db[key] = ormsgpack.packb(self.serde.dumps_typed(value)) + + def delete(self, keys: list[str]) -> None: + """Delete the cached values for the given keys.""" + for key in keys: + if key in self._db: + del self._db[key] + + async def aget(self, keys: list[str]) -> dict[str, bytes]: + """Asynchronously get the cached values for the given keys.""" + return await asyncio.to_thread(self.get, keys) + + async def aset(self, mapping: dict[str, tuple[bytes, int | None]]) -> None: + """Asynchronously set the cached values for the given keys and TTLs.""" + await asyncio.to_thread(self.set, mapping) + + async def adelete(self, keys: list[str]) -> None: + """Asynchronously delete the cached values for the given keys.""" + await asyncio.to_thread(self.delete, keys) diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index b5f5a6a7e..9aa051f9a 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -254,23 +254,6 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError - def get_writes( - self, - task_ids: Sequence[str], - ) -> dict[str, list[tuple[str, Any]]]: - """Fetch writes associated with the given task IDs. - - Args: - task_ids (Sequence[str]): List of task IDs to fetch writes for. - - Returns: - dict[str, list[tuple[str, Any]]]: Mapping from task ID to list of writes. - - Raises: - NotImplementedError: Implement this method in your custom checkpoint saver. - """ - raise NotImplementedError - def list( self, config: Optional[RunnableConfig], @@ -375,23 +358,6 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError - async def aget_writes( - self, - task_ids: Sequence[str], - ) -> dict[str, list[tuple[str, Any]]]: - """Asynchronously fetch writes associated with the given task IDs. - - Args: - task_ids (Sequence[str]): List of task IDs to fetch writes for. - - Returns: - dict[str, list[tuple[str, Any]]]: Mapping from task ID to list of writes. - - Raises: - NotImplementedError: Implement this method in your custom checkpoint saver. - """ - raise NotImplementedError - async def alist( self, config: Optional[RunnableConfig], diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index f475606dc..3e916196f 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -68,8 +68,9 @@ class InMemorySaver( str, dict[str, tuple[tuple[str, bytes], tuple[str, bytes], Optional[str]]] ], ] + # (thread ID, checkpoint NS, checkpoint ID) -> (task ID, write idx) writes: defaultdict[ - tuple[str, str, str], # thread ID, checkpoint NS, checkpoint ID + tuple[str, str, str], dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]], ] blobs: dict[ diff --git a/libs/checkpoint/langgraph/checkpoint/serde/base.py b/libs/checkpoint/langgraph/checkpoint/serde/base.py index f593dd999..b007341eb 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/base.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/base.py @@ -1,7 +1,15 @@ from typing import Any, Protocol -class SerializerProtocol(Protocol): +class UntypedSerializerProtocol(Protocol): + """Protocol for serialization and deserialization of objects.""" + + def dumps(self, obj: Any) -> bytes: ... + + def loads(self, data: bytes) -> Any: ... + + +class SerializerProtocol(UntypedSerializerProtocol, Protocol): """Protocol for serialization and deserialization of objects. - `dumps`: Serialize an object to bytes. @@ -12,17 +20,13 @@ class SerializerProtocol(Protocol): Valid implementations include the `pickle`, `json` and `orjson` modules. """ - def dumps(self, obj: Any) -> bytes: ... - def dumps_typed(self, obj: Any) -> tuple[str, bytes]: ... - def loads(self, data: bytes) -> Any: ... - def loads_typed(self, data: tuple[str, bytes]) -> Any: ... class SerializerCompat(SerializerProtocol): - def __init__(self, serde: SerializerProtocol) -> None: + def __init__(self, serde: UntypedSerializerProtocol) -> None: self.serde = serde def dumps(self, obj: Any) -> bytes: @@ -38,7 +42,9 @@ class SerializerCompat(SerializerProtocol): return self.serde.loads(data[1]) -def maybe_add_typed_methods(serde: SerializerProtocol) -> SerializerProtocol: +def maybe_add_typed_methods( + serde: SerializerProtocol | UntypedSerializerProtocol, +) -> SerializerProtocol: """Wrap serde old serde implementations in a class with loads_typed and dumps_typed for backwards compatibility.""" if not hasattr(serde, "loads_typed") or not hasattr(serde, "dumps_typed"): diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index bcb614ca4..c3ca2821b 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -61,6 +61,8 @@ CONFIG_KEY_STREAM_WRITER = sys.intern("__pregel_stream_writer") # holds a `StreamWriter` for stream_mode=custom CONFIG_KEY_STORE = sys.intern("__pregel_store") # holds a `BaseStore` made available to managed values +CONFIG_KEY_CACHE = sys.intern("__pregel_cache") +# holds a `BaseCache` made available to subgraphs CONFIG_KEY_RESUMING = sys.intern("__pregel_resuming") # holds a boolean indicating if subgraphs should resume from a previous checkpoint CONFIG_KEY_TASK_ID = sys.intern("__pregel_task_id") @@ -106,7 +108,6 @@ NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000") # the task_id to use for writes that are not associated with a task CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map") # holds a mapping of task ns -> resume value for resuming tasks -TASK_CACHE_NAMESPACE = b"__pregel_task_cache" RESERVED = { TAG_HIDDEN, diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index a3afe33ae..5992eb5de 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -318,11 +318,13 @@ class entrypoint: self, checkpointer: Optional[BaseCheckpointSaver] = None, store: Optional[BaseStore] = None, + cache: Optional[CachePolicy] = None, config_schema: Optional[type[Any]] = None, ) -> None: """Initialize the entrypoint decorator.""" self.checkpointer = checkpointer self.store = store + self.cache = cache self.config_schema = config_schema @dataclass(**_DC_KWARGS) @@ -452,5 +454,6 @@ class entrypoint: stream_eager=True, checkpointer=self.checkpointer, store=self.store, + cache=self.cache, config_type=self.config_schema, ) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 3abe2092f..fde60fedd 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -36,6 +36,7 @@ from langchain_core.runnables.utils import ( from pydantic import BaseModel from typing_extensions import Self +from langgraph.cache.base import BaseCache from langgraph.channels.base import ( BaseChannel, ) @@ -47,6 +48,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.constants import ( CONF, + CONFIG_KEY_CACHE, CONFIG_KEY_CHECKPOINT_DURING, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, @@ -495,6 +497,9 @@ class Pregel(PregelProtocol): store: BaseStore | None = None """Memory store to use for SharedValues. Defaults to None.""" + cache: BaseCache | None = None + """Cache to use for storing node results. Defaults to None.""" + retry_policy: Sequence[RetryPolicy] | None = None """Retry policies to use when running tasks. Set to None to disable.""" @@ -525,6 +530,7 @@ class Pregel(PregelProtocol): debug: bool | None = None, checkpointer: BaseCheckpointSaver | None = None, store: BaseStore | None = None, + cache: BaseCache | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, config_type: type[Any] | None = None, input_model: type[BaseModel] | None = None, @@ -545,6 +551,7 @@ class Pregel(PregelProtocol): self.debug = debug if debug is not None else get_debug() self.checkpointer = checkpointer self.store = store + self.cache = cache if isinstance(retry_policy, RetryPolicy): self.retry_policy: Sequence[RetryPolicy] = (retry_policy,) else: @@ -2193,6 +2200,7 @@ class Pregel(PregelProtocol): All | Sequence[str], BaseCheckpointSaver | None, BaseStore | None, + BaseCache | None, ]: if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") @@ -2225,6 +2233,10 @@ class Pregel(PregelProtocol): store: BaseStore | None = config[CONF][CONFIG_KEY_STORE] else: store = self.store + if CONFIG_KEY_CACHE in config.get(CONF, {}): + cache: BaseCache | None = config[CONF][CONFIG_KEY_CACHE] + else: + cache = self.cache return ( debug, set(stream_mode), @@ -2233,6 +2245,7 @@ class Pregel(PregelProtocol): interrupt_after, checkpointer, store, + cache, ) def stream( @@ -2405,6 +2418,7 @@ class Pregel(PregelProtocol): interrupt_after_, checkpointer, store, + cache, ) = self._defaults( config, stream_mode=stream_mode, @@ -2436,6 +2450,7 @@ class Pregel(PregelProtocol): stream=StreamProtocol(stream.put, stream_modes), config=config, store=store, + cache=cache, checkpointer=checkpointer, nodes=self.nodes, specs=self.channels, @@ -2494,11 +2509,13 @@ class Pregel(PregelProtocol): # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps. while loop.tick(input_keys=self.input_channels): + loop.match_cached_writes() for _ in runner.tick( loop.tasks.values(), timeout=self.step_timeout, retry_policy=self.retry_policy, get_waiter=get_waiter, + match_cached_writes=loop.match_cached_writes, ): # emit output yield from output() @@ -2710,6 +2727,7 @@ class Pregel(PregelProtocol): interrupt_after_, checkpointer, store, + cache, ) = self._defaults( config, stream_mode=stream_mode, @@ -2743,6 +2761,7 @@ class Pregel(PregelProtocol): stream=StreamProtocol(stream.put_nowait, stream_modes), config=config, store=store, + cache=cache, checkpointer=checkpointer, nodes=self.nodes, specs=self.channels, @@ -2792,11 +2811,13 @@ class Pregel(PregelProtocol): # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps while loop.tick(input_keys=self.input_channels): + await loop.amatch_cached_writes() async for _ in runner.atick( loop.tasks.values(), timeout=self.step_timeout, retry_policy=self.retry_policy, get_waiter=get_waiter, + # TODO pass match_cached_writes ): # emit output for o in output(): diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index a546bc8b4..3fdb4bec5 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -60,7 +60,6 @@ from langgraph.constants import ( RESUME, RETURN, TAG_HIDDEN, - TASK_CACHE_NAMESPACE, TASKS, Send, ) @@ -73,6 +72,7 @@ from langgraph.pregel.read import INPUT_CACHE_KEY_TYPE, PregelNode from langgraph.store.base import BaseStore from langgraph.types import ( All, + CacheKey, CachePolicy, PregelExecutableTask, PregelScratchpad, @@ -561,23 +561,15 @@ def prepare_single_task( # create task id triggers: Sequence[str] = PUSH_TRIGGER checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name - if call.cache: - task_id = task_id_func( - TASK_CACHE_NAMESPACE, - name, - call.cache.key(*call.input[0], **call.input[1]), - # TODO add truncated iso timestamp here for ttl - ) - else: - task_id = task_id_func( - checkpoint_id_bytes, - checkpoint_ns, - str(step), - name, - PUSH, - task_path_str(task_path[1]), - str(task_path[2]), - ) + task_id = task_id_func( + checkpoint_id_bytes, + checkpoint_ns, + str(step), + name, + PUSH, + task_path_str(task_path[1]), + str(task_path[2]), + ) task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" # we append True to the task path to indicate that a call is being # made, so we should not return interrupts from this task (responsibility lies with the parent) @@ -638,7 +630,12 @@ def prepare_single_task( ), triggers, call.retry, - call.cache, + CacheKey( + xxh3_128_hexdigest(call.cache.key(*call.input[0], **call.input[1])), + call.cache.ttl, + ) + if call.cache + else None, task_id, task_path, ) @@ -672,21 +669,14 @@ def prepare_single_task( checkpoint_ns = ( f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node ) - if proc.cache_policy: - task_id = task_id_func( - TASK_CACHE_NAMESPACE, - packet.node, - proc.cache_policy.key(packet.arg), - ) - else: - task_id = task_id_func( - checkpoint_id_bytes, - checkpoint_ns, - str(step), - packet.node, - PUSH, - str(idx), - ) + task_id = task_id_func( + checkpoint_id_bytes, + checkpoint_ns, + str(step), + packet.node, + PUSH, + str(idx), + ) else: logger.warning(f"Ignoring invalid PUSH task path {task_path}") return @@ -756,7 +746,12 @@ def prepare_single_task( ), triggers, proc.retry_policy, - proc.cache_policy, + CacheKey( + xxh3_128_hexdigest(proc.cache_policy.key(packet.arg)), + proc.cache_policy.ttl, + ) + if proc.cache_policy + else None, task_id, task_path, writers=proc.flat_writers, @@ -800,7 +795,6 @@ def prepare_single_task( # create task id checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name - # TODO implement cache task_id = task_id_func( checkpoint_id_bytes, checkpoint_ns, @@ -885,7 +879,12 @@ def prepare_single_task( ), triggers, proc.retry_policy, - None, + CacheKey( + xxh3_128_hexdigest(proc.cache_policy.key(val)), + proc.cache_policy.ttl, + ) + if proc.cache_policy + else None, task_id, task_path[:3], writers=proc.flat_writers, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 34d1d9bb9..4a5f5c74d 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -27,6 +27,7 @@ from langchain_core.runnables import RunnableConfig from pydantic import BaseModel from typing_extensions import ParamSpec, Self +from langgraph.cache.base import BaseCache from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import ( EXCLUDED_METADATA_KEYS, @@ -147,6 +148,7 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol: class PregelLoop(LoopProtocol): input: Optional[Any] input_model: Optional[type[BaseModel]] + cache: Optional[BaseCache[Sequence[tuple[str, Any]]]] checkpointer: Optional[BaseCheckpointSaver] nodes: Mapping[str, PregelNode] specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]] @@ -206,6 +208,7 @@ class PregelLoop(LoopProtocol): stream: Optional[StreamProtocol], config: RunnableConfig, store: Optional[BaseStore], + cache: Optional[BaseCache], checkpointer: Optional[BaseCheckpointSaver], nodes: Mapping[str, PregelNode], specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], @@ -230,6 +233,7 @@ class PregelLoop(LoopProtocol): self.input = input self.input_model = input_model self.checkpointer = checkpointer + self.cache = cache self.nodes = nodes self.specs = specs self.output_keys = output_keys @@ -434,10 +438,7 @@ class PregelLoop(LoopProtocol): # save the new task self.tasks[pushed.id] = pushed # match any pending writes to the new task - if call.cache: - pass - # how to call async method here... - elif self.skip_done_tasks: + if self.skip_done_tasks: self._match_writes({pushed.id: pushed}) # return the new task, to be started if not run before return pushed @@ -616,6 +617,12 @@ class PregelLoop(LoopProtocol): return True + def match_cached_writes(self) -> None: + raise NotImplementedError + + async def amatch_cached_writes(self) -> None: + raise NotImplementedError + # private def _match_writes(self, tasks: Mapping[str, PregelExecutableTask]) -> None: @@ -970,6 +977,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): stream: Optional[StreamProtocol], config: RunnableConfig, store: Optional[BaseStore], + cache: Optional[BaseCache], checkpointer: Optional[BaseCheckpointSaver], nodes: Mapping[str, PregelNode], specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], @@ -990,6 +998,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): stream=stream, config=config, checkpointer=checkpointer, + cache=cache, store=store, nodes=nodes, specs=specs, @@ -1040,6 +1049,30 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): return self.submit(cast(WritableManagedValue, managed_value).update, values) + def match_cached_writes(self) -> None: + if self.cache is None: + return + if cached := { + t.cache_key.key: t + for t in self.tasks.values() + if t.cache_key and not t.writes + }: + for key, values in self.cache.get(cached.keys()).items(): + cached[key].writes.extend(values) + + def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: + """Put writes for a task, to be read by the next tick.""" + super().put_writes(task_id, writes) + if not writes or self.cache is None or not hasattr(self, "tasks"): + return + task = self.tasks.get(task_id) + if task is None or task.cache_key is None: + return + self.submit( + self.cache.set, + {task.cache_key.key: (task.writes, task.cache_key.ttl)}, + ) + # context manager def __enter__(self) -> Self: @@ -1120,6 +1153,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): stream: Optional[StreamProtocol], config: RunnableConfig, store: Optional[BaseStore], + cache: Optional[BaseCache], checkpointer: Optional[BaseCheckpointSaver], nodes: Mapping[str, PregelNode], specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], @@ -1140,6 +1174,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): stream=stream, config=config, checkpointer=checkpointer, + cache=cache, store=store, nodes=nodes, specs=specs, @@ -1190,6 +1225,33 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): return self.submit(cast(WritableManagedValue, managed_value).aupdate, values) + async def amatch_cached_writes(self) -> None: + if self.cache is None: + return + if cached := { + t.cache_key.key: t + for t in self.tasks.values() + if t.cache_key and not t.writes + }: + for key, values in (await self.cache.aget(cached.keys())).items(): + cached[key].writes.extend(values) + + def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: + """Put writes for a task, to be read by the next tick.""" + super().put_writes(task_id, writes) + if not writes or self.cache is None or not hasattr(self, "tasks"): + return + task = self.tasks.get(task_id) + if task is None or task.cache_key is None: + return + if writes[0][0] in (INTERRUPT, ERROR): + # only cache successful tasks + return + self.submit( + self.cache.aset, + {task.cache_key.key: (task.writes, task.cache_key.ttl)}, + ) + # context manager async def __aenter__(self) -> Self: diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 6a5992d04..7902ebfbd 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -33,7 +33,12 @@ from langgraph.errors import GraphBubbleUp, GraphInterrupt from langgraph.pregel.algo import Call from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry -from langgraph.types import PregelExecutableTask, PregelScratchpad, RetryPolicy +from langgraph.types import ( + CachePolicy, + PregelExecutableTask, + PregelScratchpad, + RetryPolicy, +) from langgraph.utils.future import chain_future F = TypeVar("F", concurrent.futures.Future, asyncio.Future) @@ -137,6 +142,7 @@ class PregelRunner: timeout: Optional[float] = None, retry_policy: Optional[Sequence[RetryPolicy]] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, + match_cached_writes: Optional[Callable[[], None]] = None, ) -> Iterator[None]: tasks = tuple(tasks) futures = FuturesDict( @@ -160,6 +166,7 @@ class PregelRunner: retry=retry_policy, futures=weakref.ref(futures), schedule_task=self.schedule_task, + match_cached_writes=match_cached_writes, submit=self.submit, reraise=reraise, ), @@ -203,6 +210,7 @@ class PregelRunner: retry=retry_policy, futures=weakref.ref(futures), schedule_task=self.schedule_task, + match_cached_writes=match_cached_writes, submit=self.submit, reraise=reraise, ), @@ -515,6 +523,7 @@ def _call( input: Any, *, retry: Optional[Sequence[RetryPolicy]] = None, + cache: Optional[CachePolicy] = None, callbacks: Callbacks = None, futures: weakref.ref[FuturesDict], schedule_task: weakref.ref[ @@ -522,6 +531,7 @@ def _call( [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] ] ], + match_cached_writes: Optional[Callable[[], None]], submit: weakref.ref[Submit], reraise: bool, ) -> concurrent.futures.Future[Any]: @@ -535,8 +545,10 @@ def _call( if next_task := schedule_task()( # type: ignore[misc] task(), # type: ignore[arg-type] scratchpad.call_counter(), - Call(func, input, retry=retry, callbacks=callbacks), + Call(func, input, retry=retry, cache=cache, callbacks=callbacks), ): + if match_cached_writes: + match_cached_writes() if fut := next( ( f @@ -596,6 +608,7 @@ def _acall( input: Any, *, retry: Optional[Sequence[RetryPolicy]] = None, + cache: Optional[CachePolicy] = None, callbacks: Callbacks = None, # injected dependencies futures: weakref.ref[FuturesDict], @@ -616,7 +629,7 @@ def _acall( if next_task := schedule_task()( # type: ignore[misc] task(), # type: ignore[arg-type] scratchpad.call_counter(), - Call(func, input, retry=retry, callbacks=callbacks), + Call(func, input, retry=retry, cache=cache, callbacks=callbacks), ): if fut := next( ( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 51a3068ac..bd012b729 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -134,6 +134,11 @@ class CachePolicy(NamedTuple, Generic[P]): """ key: Callable[P, str | bytes] = default_cache_key + """Function to generate a cache key from the node's input. + Defaults to hashing the input with pickle.""" + + ttl: Optional[int] = None + """Time to live for the cache entry in seconds. If None, the entry never expires.""" @dataclasses.dataclass(**_DC_KWARGS) @@ -179,6 +184,15 @@ else: _T_DC_KWARGS = {"frozen": True} +class CacheKey(NamedTuple): + """Cache key for a task.""" + + key: str + """Key for the cache entry.""" + ttl: Optional[int] + """Time to live for the cache entry in seconds.""" + + @dataclasses.dataclass(**_T_DC_KWARGS) class PregelExecutableTask: name: str @@ -188,7 +202,7 @@ class PregelExecutableTask: config: RunnableConfig triggers: Sequence[str] retry_policy: Optional[Sequence[RetryPolicy]] - cache_policy: Optional[CachePolicy] + cache_key: Optional[CacheKey] id: str path: tuple[Union[str, int, tuple], ...] scheduled: bool = False diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 5bfec63bd..6de17548a 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,5 +1,7 @@ +import os import sys -from collections.abc import AsyncIterator +import tempfile +from collections.abc import AsyncIterator, Iterator from contextlib import asynccontextmanager from typing import Optional from uuid import UUID, uuid4 @@ -11,6 +13,8 @@ from psycopg import AsyncConnection, Connection from psycopg_pool import AsyncConnectionPool, ConnectionPool from pytest_mock import MockerFixture +from langgraph.cache.base import BaseCache +from langgraph.cache.file import FileCache from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver from langgraph.checkpoint.postgres.aio import ( @@ -361,6 +365,20 @@ async def _store_postgres_aio_pool(): await conn.execute(f"DROP DATABASE {database}") +@pytest.fixture(scope="function") +def file_cache() -> Iterator[BaseCache]: + _, path = tempfile.mkstemp() + os.remove(path) + try: + yield FileCache(path=path) + finally: + # Cleanup the file + try: + os.remove(path) + except OSError: + pass + + @pytest.fixture(scope="function") def store_postgres(): database = f"test_{uuid4().hex[:16]}" diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a5baeeece..93ad840db 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -38,6 +38,7 @@ from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import TypedDict +from langgraph.cache.base import BaseCache from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context @@ -64,6 +65,7 @@ from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner from langgraph.store.base import BaseStore from langgraph.types import ( + CachePolicy, Command, Interrupt, PregelTask, @@ -6559,7 +6561,7 @@ def test_falsy_return_from_task( @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_interrupts_functional( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + request: pytest.FixtureRequest, checkpointer_name: str ): """Test multiple interrupts with functional API.""" checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @@ -6596,6 +6598,48 @@ def test_multiple_interrupts_functional( assert counter == 3 +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_multiple_interrupts_functional_cache( + request: pytest.FixtureRequest, checkpointer_name: str, file_cache: BaseCache +): + """Test multiple interrupts with functional API.""" + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + counter = 0 + + @task(cache=CachePolicy()) + def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x + + @entrypoint(checkpointer=checkpointer, cache=file_cache) + def graph(state: dict) -> dict: + """React tool.""" + + values = [] + + for idx in [1, 1, 2, 2, 3, 3]: + values.extend([double(idx).result(), interrupt({"a": "boo"})]) + + return {"values": values} + + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + graph.invoke({}, configurable) + graph.invoke(Command(resume="a"), configurable) + graph.invoke(Command(resume="b"), configurable) + graph.invoke(Command(resume="c"), configurable) + graph.invoke(Command(resume="d"), configurable) + graph.invoke(Command(resume="e"), configurable) + result = graph.invoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_double_interrupt_subgraph( request: pytest.FixtureRequest, checkpointer_name: str From 2a1c63ff9cdc948d101d070d561ad41f2c79f9fa Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Apr 2025 18:01:20 -0700 Subject: [PATCH 03/41] Fixes --- libs/langgraph/langgraph/pregel/algo.py | 43 ++++++++++++++++++++--- libs/langgraph/langgraph/pregel/call.py | 19 ++++++++++ libs/langgraph/langgraph/pregel/runner.py | 1 + 3 files changed, 58 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 3fdb4bec5..b8b87b516 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -65,7 +65,7 @@ from langgraph.constants import ( ) from langgraph.errors import InvalidUpdateError from langgraph.managed.base import ManagedValueMapping -from langgraph.pregel.call import get_runnable_for_task +from langgraph.pregel.call import get_runnable_for_task, identifier from langgraph.pregel.io import read_channels from langgraph.pregel.log import logger from langgraph.pregel.read import INPUT_CACHE_KEY_TYPE, PregelNode @@ -79,7 +79,7 @@ from langgraph.types import ( PregelTask, RetryPolicy, ) -from langgraph.utils.config import merge_configs, patch_config +from langgraph.utils.config import merge_configs, patch_config, recast_checkpoint_ns GetNextVersion = Callable[[Optional[V], BaseChannel], V] SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) @@ -631,7 +631,18 @@ def prepare_single_task( triggers, call.retry, CacheKey( - xxh3_128_hexdigest(call.cache.key(*call.input[0], **call.input[1])), + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + (identifier(call.func) or "__dynamic__").encode(), + call.cache.key(*call.input[0], **call.input[1]), + ) + ) + ), call.cache.ttl, ) if call.cache @@ -747,7 +758,18 @@ def prepare_single_task( triggers, proc.retry_policy, CacheKey( - xxh3_128_hexdigest(proc.cache_policy.key(packet.arg)), + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + packet.node.encode(), + proc.cache_policy.key(packet.arg), + ) + ) + ), proc.cache_policy.ttl, ) if proc.cache_policy @@ -880,7 +902,18 @@ def prepare_single_task( triggers, proc.retry_policy, CacheKey( - xxh3_128_hexdigest(proc.cache_policy.key(val)), + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + name.encode(), + proc.cache_policy.key(val), + ) + ) + ), proc.cache_policy.ttl, ) if proc.cache_policy diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index e3719d2b6..5458d3972 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -74,6 +74,25 @@ def _whichmodule(obj: Any, name: str) -> Optional[str]: return None +def identifier(obj: Any, name: Optional[str] = None) -> Optional[str]: + if name is None: + name = getattr(obj, "__qualname__", None) + if name is None: # pragma: no cover + # This used to be needed for Python 2.7 support but is probably not + # needed anymore. However we keep the __name__ introspection in case + # users of cloudpickle rely on this old behavior for unknown reasons. + name = getattr(obj, "__name__", None) + if name is None: + return None + + module_name = getattr(obj, "__module__", None) + if module_name is None: + # In this case, obj.__module__ is None. obj is thus treated as dynamic. + return None + + return f"{module_name}.{name}" + + def _lookup_module_and_qualname( obj: Any, name: Optional[str] = None ) -> Optional[tuple[types.ModuleType, str]]: diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 7902ebfbd..805523195 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -586,6 +586,7 @@ def _call( retry=retry, callbacks=callbacks, schedule_task=schedule_task, + match_cached_writes=match_cached_writes, submit=submit, reraise=reraise, ), From b3ea406e813f1e17eae90b6fbc8a28fcdfa221f1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Apr 2025 18:01:31 -0700 Subject: [PATCH 04/41] Add pickle_fallback for json plus serializer --- libs/checkpoint/langgraph/cache/base/__init__.py | 2 +- .../checkpoint/langgraph/checkpoint/serde/jsonplus.py | 11 ++++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py index a81b48cf6..9c8653d23 100644 --- a/libs/checkpoint/langgraph/cache/base/__init__.py +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -11,7 +11,7 @@ T = TypeVar("T") class BaseCache(ABC, Generic[T]): """Base class for a cache.""" - serde: SerializerProtocol = JsonPlusSerializer() + serde: SerializerProtocol = JsonPlusSerializer(pickle_fallback=True) def __init__(self, *, serde: SerializerProtocol | None = None) -> None: """Initialize the cache with a serializer.""" diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 3a9ab50e2..53b96065c 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -3,6 +3,7 @@ import decimal import importlib import json import pathlib +import pickle import re from collections import deque from collections.abc import Sequence @@ -37,8 +38,12 @@ class JsonPlusSerializer(SerializerProtocol): """Serializer that uses ormsgpack, with a fallback to extended JSON serializer.""" def __init__( - self, *, __unpack_ext_hook__: Optional[Callable[[int, bytes], Any]] = None + self, + *, + pickle_fallback: bool = False, + __unpack_ext_hook__: Optional[Callable[[int, bytes], Any]] = None, ) -> None: + self.pickle_fallback = pickle_fallback self._unpack_ext_hook = ( __unpack_ext_hook__ if __unpack_ext_hook__ is not None @@ -209,6 +214,8 @@ class JsonPlusSerializer(SerializerProtocol): except ormsgpack.MsgpackEncodeError as exc: if "valid UTF-8" in str(exc): return "json", self.dumps(obj) + elif self.pickle_fallback: + return "pickle", pickle.dumps(obj) raise exc def loads(self, data: bytes) -> Any: @@ -228,6 +235,8 @@ class JsonPlusSerializer(SerializerProtocol): return ormsgpack.unpackb( data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS ) + elif type_ == "pickle": + return pickle.loads(data_) else: raise NotImplementedError(f"Unknown serialization type: {type_}") From ed3f05260a5896a4a4064c01dfa32c4e46c096c4 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Apr 2025 18:05:37 -0700 Subject: [PATCH 05/41] Secure! --- libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 53b96065c..aef405a95 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -235,7 +235,7 @@ class JsonPlusSerializer(SerializerProtocol): return ormsgpack.unpackb( data_, ext_hook=self._unpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS ) - elif type_ == "pickle": + elif self.pickle_fallback and type_ == "pickle": return pickle.loads(data_) else: raise NotImplementedError(f"Unknown serialization type: {type_}") From 5cff35d1c36dc025d5be695d3ea235382d93d131 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 1 May 2025 08:05:38 -0700 Subject: [PATCH 06/41] Fix type annotation --- libs/langgraph/langgraph/func/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 5992eb5de..120d44205 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -16,6 +16,7 @@ from typing import ( overload, ) +from langgraph.cache.base import BaseCache from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import BaseCheckpointSaver @@ -318,7 +319,7 @@ class entrypoint: self, checkpointer: Optional[BaseCheckpointSaver] = None, store: Optional[BaseStore] = None, - cache: Optional[CachePolicy] = None, + cache: Optional[BaseCache] = None, config_schema: Optional[type[Any]] = None, ) -> None: """Initialize the entrypoint decorator.""" From 1d977f1c094e4b1cf8f613a341acf343ca6ed6cd Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 1 May 2025 09:57:48 -0700 Subject: [PATCH 07/41] Rename --- libs/langgraph/langgraph/func/__init__.py | 8 +++++--- libs/langgraph/langgraph/pregel/algo.py | 17 ++++++++++------- libs/langgraph/langgraph/pregel/call.py | 8 ++++++-- libs/langgraph/langgraph/pregel/loop.py | 4 ++-- libs/langgraph/langgraph/pregel/runner.py | 6 +++--- libs/langgraph/langgraph/types.py | 5 +++++ libs/langgraph/tests/test_pregel.py | 2 +- 7 files changed, 32 insertions(+), 18 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 120d44205..689ac05ab 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -40,7 +40,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, - cache: Optional[CachePolicy[P]] = None, + cache_policy: Optional[CachePolicy[P]] = None, ) -> Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], Callable[P, SyncAsyncFuture[T]], @@ -58,7 +58,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, - cache: Optional[CachePolicy[P]] = None, + cache_policy: Optional[CachePolicy[P]] = None, ) -> Union[ Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], @@ -144,7 +144,9 @@ def task( # handle regular functions / partials / callable classes, etc. func.__name__ = name - call_func = functools.partial(call, func, retry=retry_policies, cache=cache) + call_func = functools.partial( + call, func, retry=retry_policies, cache=cache_policy + ) object.__setattr__(call_func, "_is_pregel_task", True) return functools.update_wrapper(call_func, func) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index b8b87b516..00cb3220b 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -113,12 +113,12 @@ class PregelTaskWrites(NamedTuple): class Call: - __slots__ = ("func", "input", "retry", "cache", "callbacks") + __slots__ = ("func", "input", "retry", "cache_policy", "callbacks") func: Callable input: tuple[tuple[Any, ...], dict[str, Any]] retry: Optional[Sequence[RetryPolicy]] - cache: Optional[CachePolicy] + cache_policy: Optional[CachePolicy] callbacks: Callbacks def __init__( @@ -127,13 +127,13 @@ class Call: input: tuple[tuple[Any, ...], dict[str, Any]], *, retry: Optional[Sequence[RetryPolicy]], - cache: Optional[CachePolicy], + cache_policy: Optional[CachePolicy], callbacks: Callbacks, ) -> None: self.func = func self.input = input self.retry = retry - self.cache = cache + self.cache_policy = cache_policy self.callbacks = callbacks @@ -639,13 +639,14 @@ def prepare_single_task( if parent_ns else b"", (identifier(call.func) or "__dynamic__").encode(), - call.cache.key(*call.input[0], **call.input[1]), + call.cache_policy.key(*call.input[0], **call.input[1]), ) ) ), - call.cache.ttl, + call.cache_policy.ttl, + call.cache_policy.refresh, ) - if call.cache + if call.cache_policy else None, task_id, task_path, @@ -771,6 +772,7 @@ def prepare_single_task( ) ), proc.cache_policy.ttl, + proc.cache_policy.refresh, ) if proc.cache_policy else None, @@ -915,6 +917,7 @@ def prepare_single_task( ) ), proc.cache_policy.ttl, + proc.cache_policy.refresh, ) if proc.cache_policy else None, diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index 5458d3972..bc41eafb5 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -242,12 +242,16 @@ def call( func: Callable[P, T], *args: Any, retry: Optional[Sequence[RetryPolicy]] = None, - cache: Optional[CachePolicy] = None, + cache_policy: Optional[CachePolicy] = None, **kwargs: Any, ) -> SyncAsyncFuture[T]: config = get_config() impl = config[CONF][CONFIG_KEY_CALL] fut = impl( - func, (args, kwargs), retry=retry, cache=cache, callbacks=config["callbacks"] + func, + (args, kwargs), + retry=retry, + cache=cache_policy, + callbacks=config["callbacks"], ) return fut diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 4a5f5c74d..7ab11dbb2 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1055,7 +1055,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): if cached := { t.cache_key.key: t for t in self.tasks.values() - if t.cache_key and not t.writes + if t.cache_key and not t.cache_key.refresh and not t.writes }: for key, values in self.cache.get(cached.keys()).items(): cached[key].writes.extend(values) @@ -1231,7 +1231,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): if cached := { t.cache_key.key: t for t in self.tasks.values() - if t.cache_key and not t.writes + if t.cache_key and not t.cache_key.refresh and not t.writes }: for key, values in (await self.cache.aget(cached.keys())).items(): cached[key].writes.extend(values) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 805523195..42c83c7ca 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -523,7 +523,7 @@ def _call( input: Any, *, retry: Optional[Sequence[RetryPolicy]] = None, - cache: Optional[CachePolicy] = None, + cache_policy: Optional[CachePolicy] = None, callbacks: Callbacks = None, futures: weakref.ref[FuturesDict], schedule_task: weakref.ref[ @@ -545,7 +545,7 @@ def _call( if next_task := schedule_task()( # type: ignore[misc] task(), # type: ignore[arg-type] scratchpad.call_counter(), - Call(func, input, retry=retry, cache=cache, callbacks=callbacks), + Call(func, input, retry=retry, cache_policy=cache_policy, callbacks=callbacks), ): if match_cached_writes: match_cached_writes() @@ -630,7 +630,7 @@ def _acall( if next_task := schedule_task()( # type: ignore[misc] task(), # type: ignore[arg-type] scratchpad.call_counter(), - Call(func, input, retry=retry, cache=cache, callbacks=callbacks), + Call(func, input, retry=retry, cache_policy=cache, callbacks=callbacks), ): if fut := next( ( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index bd012b729..8b42c0e10 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -140,6 +140,9 @@ class CachePolicy(NamedTuple, Generic[P]): ttl: Optional[int] = None """Time to live for the cache entry in seconds. If None, the entry never expires.""" + refresh: bool = False + """Whether to force a refresh of the cache entry when it is accessed.""" + @dataclasses.dataclass(**_DC_KWARGS) class Interrupt: @@ -191,6 +194,8 @@ class CacheKey(NamedTuple): """Key for the cache entry.""" ttl: Optional[int] """Time to live for the cache entry in seconds.""" + refresh: bool + """Whether to force a refresh of the cache entry when it is accessed.""" @dataclasses.dataclass(**_T_DC_KWARGS) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 93ad840db..9ede99db1 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6607,7 +6607,7 @@ def test_multiple_interrupts_functional_cache( counter = 0 - @task(cache=CachePolicy()) + @task(cache_policy=CachePolicy()) def double(x: int) -> int: """Increment the counter.""" nonlocal counter From a11a62e68fdfbe41f80373e7552a7a58b46024de Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 1 May 2025 09:59:07 -0700 Subject: [PATCH 08/41] Thread-safe delete --- libs/checkpoint/langgraph/cache/file/__init__.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/libs/checkpoint/langgraph/cache/file/__init__.py b/libs/checkpoint/langgraph/cache/file/__init__.py index cc6499543..482d8d357 100644 --- a/libs/checkpoint/langgraph/cache/file/__init__.py +++ b/libs/checkpoint/langgraph/cache/file/__init__.py @@ -37,8 +37,7 @@ class FileCache(BaseCache): def delete(self, keys: list[str]) -> None: """Delete the cached values for the given keys.""" for key in keys: - if key in self._db: - del self._db[key] + self._db.pop(key, None) async def aget(self, keys: list[str]) -> dict[str, bytes]: """Asynchronously get the cached values for the given keys.""" From 057da43cd0a95c259d75d72db87a1adf704c999a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 1 May 2025 09:59:46 -0700 Subject: [PATCH 09/41] Lint --- libs/checkpoint/langgraph/cache/base/__init__.py | 16 ++++++++-------- libs/checkpoint/langgraph/cache/file/__init__.py | 16 ++++++++-------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py index 9c8653d23..0dd5aca8e 100644 --- a/libs/checkpoint/langgraph/cache/base/__init__.py +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -21,22 +21,22 @@ class BaseCache(ABC, Generic[T]): def get(self, keys: Sequence[str]) -> dict[str, T]: """Get the cached values for the given keys.""" - @abstractmethod - def set(self, mapping: Mapping[str, tuple[T, int | None]]) -> None: - """Set the cached values for the given keys and TTLs.""" - - @abstractmethod - def delete(self, keys: Sequence[str]) -> None: - """Delete the cached values for the given keys.""" - @abstractmethod async def aget(self, keys: Sequence[str]) -> dict[str, T]: """Asynchronously get the cached values for the given keys.""" + @abstractmethod + def set(self, mapping: Mapping[str, tuple[T, int | None]]) -> None: + """Set the cached values for the given keys and TTLs.""" + @abstractmethod async def aset(self, mapping: Mapping[str, tuple[T, int | None]]) -> None: """Asynchronously set the cached values for the given keys and TTLs.""" + @abstractmethod + def delete(self, keys: Sequence[str]) -> None: + """Delete the cached values for the given keys.""" + @abstractmethod async def adelete(self, keys: Sequence[str]) -> None: """Asynchronously delete the cached values for the given keys.""" diff --git a/libs/checkpoint/langgraph/cache/file/__init__.py b/libs/checkpoint/langgraph/cache/file/__init__.py index 482d8d357..2ed206abf 100644 --- a/libs/checkpoint/langgraph/cache/file/__init__.py +++ b/libs/checkpoint/langgraph/cache/file/__init__.py @@ -28,25 +28,25 @@ class FileCache(BaseCache): if key in self._db } + async def aget(self, keys: list[str]) -> dict[str, bytes]: + """Asynchronously get the cached values for the given keys.""" + return await asyncio.to_thread(self.get, keys) + def set(self, mapping: dict[str, tuple[bytes, int | None]]) -> None: """Set the cached values for the given keys and TTLs.""" for key, (value, _) in mapping.items(): # File-based caches do not support TTLs, so we ignore them. self._db[key] = ormsgpack.packb(self.serde.dumps_typed(value)) + async def aset(self, mapping: dict[str, tuple[bytes, int | None]]) -> None: + """Asynchronously set the cached values for the given keys and TTLs.""" + await asyncio.to_thread(self.set, mapping) + def delete(self, keys: list[str]) -> None: """Delete the cached values for the given keys.""" for key in keys: self._db.pop(key, None) - async def aget(self, keys: list[str]) -> dict[str, bytes]: - """Asynchronously get the cached values for the given keys.""" - return await asyncio.to_thread(self.get, keys) - - async def aset(self, mapping: dict[str, tuple[bytes, int | None]]) -> None: - """Asynchronously set the cached values for the given keys and TTLs.""" - await asyncio.to_thread(self.set, mapping) - async def adelete(self, keys: list[str]) -> None: """Asynchronously delete the cached values for the given keys.""" await asyncio.to_thread(self.delete, keys) From 64491a2b298c3d8611bdc347c16fdace60cbf07b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 1 May 2025 10:13:04 -0700 Subject: [PATCH 10/41] Lint --- .../langgraph/cache/base/__init__.py | 12 +- libs/checkpoint/langgraph/cache/base/py.typed | 0 libs/checkpoint/langgraph/cache/file/py.typed | 0 libs/langgraph/langgraph/pregel/algo.py | 123 ++++++++++-------- libs/langgraph/langgraph/pregel/loop.py | 21 ++- libs/langgraph/langgraph/types.py | 2 +- 6 files changed, 86 insertions(+), 72 deletions(-) create mode 100644 libs/checkpoint/langgraph/cache/base/py.typed create mode 100644 libs/checkpoint/langgraph/cache/file/py.typed diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py index 0dd5aca8e..e98525f81 100644 --- a/libs/checkpoint/langgraph/cache/base/__init__.py +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -5,10 +5,10 @@ from typing import Generic, Sequence, TypeVar from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -T = TypeVar("T") +ValueT = TypeVar("ValueT") -class BaseCache(ABC, Generic[T]): +class BaseCache(ABC, Generic[ValueT]): """Base class for a cache.""" serde: SerializerProtocol = JsonPlusSerializer(pickle_fallback=True) @@ -18,19 +18,19 @@ class BaseCache(ABC, Generic[T]): self.serde = serde or self.serde @abstractmethod - def get(self, keys: Sequence[str]) -> dict[str, T]: + def get(self, keys: Sequence[str]) -> dict[str, ValueT]: """Get the cached values for the given keys.""" @abstractmethod - async def aget(self, keys: Sequence[str]) -> dict[str, T]: + async def aget(self, keys: Sequence[str]) -> dict[str, ValueT]: """Asynchronously get the cached values for the given keys.""" @abstractmethod - def set(self, mapping: Mapping[str, tuple[T, int | None]]) -> None: + def set(self, mapping: Mapping[str, tuple[ValueT, int | None]]) -> None: """Set the cached values for the given keys and TTLs.""" @abstractmethod - async def aset(self, mapping: Mapping[str, tuple[T, int | None]]) -> None: + async def aset(self, mapping: Mapping[str, tuple[ValueT, int | None]]) -> None: """Asynchronously set the cached values for the given keys and TTLs.""" @abstractmethod diff --git a/libs/checkpoint/langgraph/cache/base/py.typed b/libs/checkpoint/langgraph/cache/base/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/libs/checkpoint/langgraph/cache/file/py.typed b/libs/checkpoint/langgraph/cache/file/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 00cb3220b..c0b15fd18 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -585,6 +585,28 @@ def prepare_single_task( assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}" if for_execution: writes: deque[tuple[str, Any]] = deque() + if call.cache_policy: + args_key = call.cache_policy.key(*call.input[0], **call.input[1]) + cache_key: Optional[CacheKey] = CacheKey( + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + (identifier(call.func) or "__dynamic__").encode(), + args_key.encode() + if isinstance(args_key, str) + else args_key, + ) + ) + ), + call.cache_policy.ttl, + call.cache_policy.refresh, + ) + else: + cache_key = None return PregelExecutableTask( name, call.input, @@ -630,24 +652,7 @@ def prepare_single_task( ), triggers, call.retry, - CacheKey( - xxh3_128_hexdigest( - b"".join( - ( - b"__pregel_cache", - recast_checkpoint_ns(parent_ns).encode() - if parent_ns - else b"", - (identifier(call.func) or "__dynamic__").encode(), - call.cache_policy.key(*call.input[0], **call.input[1]), - ) - ) - ), - call.cache_policy.ttl, - call.cache_policy.refresh, - ) - if call.cache_policy - else None, + cache_key, task_id, task_path, ) @@ -709,6 +714,28 @@ def prepare_single_task( if proc.metadata: metadata.update(proc.metadata) writes = deque() + if proc.cache_policy: + args_key = proc.cache_policy.key(packet.arg) + cache_key = CacheKey( + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + packet.node.encode(), + args_key.encode() + if isinstance(args_key, str) + else args_key, + ) + ) + ), + proc.cache_policy.ttl, + proc.cache_policy.refresh, + ) + else: + cache_key = None return PregelExecutableTask( packet.node, packet.arg, @@ -758,24 +785,7 @@ def prepare_single_task( ), triggers, proc.retry_policy, - CacheKey( - xxh3_128_hexdigest( - b"".join( - ( - b"__pregel_cache", - recast_checkpoint_ns(parent_ns).encode() - if parent_ns - else b"", - packet.node.encode(), - proc.cache_policy.key(packet.arg), - ) - ) - ), - proc.cache_policy.ttl, - proc.cache_policy.refresh, - ) - if proc.cache_policy - else None, + cache_key, task_id, task_path, writers=proc.flat_writers, @@ -842,6 +852,28 @@ def prepare_single_task( if proc.metadata: metadata.update(proc.metadata) writes = deque() + if proc.cache_policy: + args_key = proc.cache_policy.key(val) + cache_key = CacheKey( + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + name.encode(), + args_key.encode() + if isinstance(args_key, str) + else args_key, + ) + ) + ), + proc.cache_policy.ttl, + proc.cache_policy.refresh, + ) + else: + cache_key = None return PregelExecutableTask( name, val, @@ -903,24 +935,7 @@ def prepare_single_task( ), triggers, proc.retry_policy, - CacheKey( - xxh3_128_hexdigest( - b"".join( - ( - b"__pregel_cache", - recast_checkpoint_ns(parent_ns).encode() - if parent_ns - else b"", - name.encode(), - proc.cache_policy.key(val), - ) - ) - ), - proc.cache_policy.ttl, - proc.cache_policy.refresh, - ) - if proc.cache_policy - else None, + cache_key, task_id, task_path[:3], writers=proc.flat_writers, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7ab11dbb2..e46a7fcc6 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -134,6 +134,7 @@ INPUT_DONE = object() INPUT_RESUMING = object() INPUT_SHOULD_VALIDATE = object() SPECIAL_CHANNELS = (ERROR, INTERRUPT, SCHEDULED) +WritesT = Sequence[tuple[str, Any]] def DuplexStream(*streams: StreamProtocol) -> StreamProtocol: @@ -148,7 +149,7 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol: class PregelLoop(LoopProtocol): input: Optional[Any] input_model: Optional[type[BaseModel]] - cache: Optional[BaseCache[Sequence[tuple[str, Any]]]] + cache: Optional[BaseCache[WritesT]] checkpointer: Optional[BaseCheckpointSaver] nodes: Mapping[str, PregelNode] specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]] @@ -163,16 +164,14 @@ class PregelLoop(LoopProtocol): debug: bool checkpointer_get_next_version: GetNextVersion - checkpointer_put_writes: Optional[ - Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any] - ] + checkpointer_put_writes: Optional[Callable[[RunnableConfig, WritesT, str], Any]] checkpointer_put_writes_accepts_task_path: bool _checkpointer_put_after_previous: Optional[ Callable[ [ Optional[concurrent.futures.Future], RunnableConfig, - Sequence[tuple[str, Any]], + Checkpoint, str, ChannelVersions, ], @@ -303,7 +302,7 @@ class PregelLoop(LoopProtocol): ) self.prev_checkpoint_config = None - def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: + def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" if not writes: return @@ -923,7 +922,7 @@ class PregelLoop(LoopProtocol): self.stream((self.checkpoint_ns, mode, v)) def _output_writes( - self, task_id: str, writes: Sequence[tuple[str, Any]], *, cached: bool = False + self, task_id: str, writes: WritesT, *, cached: bool = False ) -> None: if task := self.tasks.get(task_id): if task.config is not None and TAG_HIDDEN in task.config.get( @@ -1057,10 +1056,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: - for key, values in self.cache.get(cached.keys()).items(): + for key, values in self.cache.get(tuple(cached)).items(): cached[key].writes.extend(values) - def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: + def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" super().put_writes(task_id, writes) if not writes or self.cache is None or not hasattr(self, "tasks"): @@ -1233,10 +1232,10 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: - for key, values in (await self.cache.aget(cached.keys())).items(): + for key, values in (await self.cache.aget(tuple(cached))).items(): cached[key].writes.extend(values) - def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: + def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" super().put_writes(task_id, writes) if not writes or self.cache is None or not hasattr(self, "tasks"): diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 8b42c0e10..9fbb92217 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -32,7 +32,7 @@ if TYPE_CHECKING: try: - from langchain_core.messages.tool import ToolOutputMixin # type: ignore + from langchain_core.messages.tool import ToolOutputMixin except ImportError: class ToolOutputMixin: # type: ignore[no-redef] From 09fdc14d0aa7fd37a2b68a2c3fb84e71341becf6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 1 May 2025 10:23:19 -0700 Subject: [PATCH 11/41] Output cached writes --- libs/langgraph/langgraph/pregel/__init__.py | 10 ++- libs/langgraph/langgraph/pregel/loop.py | 24 ++++-- libs/langgraph/langgraph/pregel/runner.py | 92 ++++++++++----------- 3 files changed, 67 insertions(+), 59 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index fde60fedd..6d8e7adb1 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2509,9 +2509,10 @@ class Pregel(PregelProtocol): # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps. while loop.tick(input_keys=self.input_channels): - loop.match_cached_writes() + for task in loop.match_cached_writes(): + loop.output_writes(task.id, task.writes, cached=True) for _ in runner.tick( - loop.tasks.values(), + [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, retry_policy=self.retry_policy, get_waiter=get_waiter, @@ -2811,9 +2812,10 @@ class Pregel(PregelProtocol): # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps while loop.tick(input_keys=self.input_channels): - await loop.amatch_cached_writes() + for task in await loop.amatch_cached_writes(): + loop.output_writes(task.id, task.writes, cached=True) async for _ in runner.atick( - loop.tasks.values(), + [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, retry_policy=self.retry_policy, get_waiter=get_waiter, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index e46a7fcc6..0bc710e04 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -349,7 +349,7 @@ class PregelLoop(LoopProtocol): ) # output writes if hasattr(self, "tasks"): - self._output_writes(task_id, writes) + self.output_writes(task_id, writes) def _put_pending_writes(self) -> None: if self.checkpointer_put_writes is None: @@ -612,14 +612,14 @@ class PregelLoop(LoopProtocol): # print output for any tasks we applied previous writes to for task in self.tasks.values(): if task.writes: - self._output_writes(task.id, task.writes, cached=True) + self.output_writes(task.id, task.writes, cached=True) return True - def match_cached_writes(self) -> None: + def match_cached_writes(self) -> list[PregelExecutableTask]: raise NotImplementedError - async def amatch_cached_writes(self) -> None: + async def amatch_cached_writes(self) -> list[PregelExecutableTask]: raise NotImplementedError # private @@ -921,7 +921,7 @@ class PregelLoop(LoopProtocol): for v in values(*args, **kwargs): self.stream((self.checkpoint_ns, mode, v)) - def _output_writes( + def output_writes( self, task_id: str, writes: WritesT, *, cached: bool = False ) -> None: if task := self.tasks.get(task_id): @@ -1048,16 +1048,20 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): return self.submit(cast(WritableManagedValue, managed_value).update, values) - def match_cached_writes(self) -> None: + def match_cached_writes(self) -> list[PregelExecutableTask]: if self.cache is None: return + matched: list[PregelExecutableTask] = [] if cached := { t.cache_key.key: t for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: for key, values in self.cache.get(tuple(cached)).items(): - cached[key].writes.extend(values) + task = cached[key] + task.writes.extend(values) + matched.append(task) + return matched def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" @@ -1227,13 +1231,17 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): async def amatch_cached_writes(self) -> None: if self.cache is None: return + matched: list[PregelExecutableTask] = [] if cached := { t.cache_key.key: t for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: for key, values in (await self.cache.aget(tuple(cached))).items(): - cached[key].writes.extend(values) + task = cached[key] + task.writes.extend(values) + matched.append(task) + return matched def put_writes(self, task_id: str, writes: WritesT) -> None: """Put writes for a task, to be read by the next tick.""" diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 42c83c7ca..dcf32f96e 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -198,26 +198,25 @@ class PregelRunner: futures[get_waiter()] = None # schedule tasks for t in tasks: - if not t.writes: - fut = self.submit()( # type: ignore[misc] - run_with_retry, - t, - retry_policy, - configurable={ - CONFIG_KEY_CALL: partial( - _call, - weakref.ref(t), - retry=retry_policy, - futures=weakref.ref(futures), - schedule_task=self.schedule_task, - match_cached_writes=match_cached_writes, - submit=self.submit, - reraise=reraise, - ), - }, - __reraise_on_exit__=reraise, - ) - futures[fut] = t + fut = self.submit()( # type: ignore[misc] + run_with_retry, + t, + retry_policy, + configurable={ + CONFIG_KEY_CALL: partial( + _call, + weakref.ref(t), + retry=retry_policy, + futures=weakref.ref(futures), + schedule_task=self.schedule_task, + match_cached_writes=match_cached_writes, + submit=self.submit, + reraise=reraise, + ), + }, + __reraise_on_exit__=reraise, + ) + futures[fut] = t # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes @@ -332,33 +331,32 @@ class PregelRunner: futures[get_waiter()] = None # schedule tasks for t in tasks: - if not t.writes: - fut = cast( - asyncio.Future, - self.submit()( # type: ignore[misc] - arun_with_retry, - t, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_CALL: partial( - _acall, - weakref.ref(t), - retry=retry_policy, - stream=self.use_astream, - futures=weakref.ref(futures), - schedule_task=self.schedule_task, - submit=self.submit, - reraise=reraise, - loop=loop, - ), - }, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, - ), - ) - futures[fut] = t + fut = cast( + asyncio.Future, + self.submit()( # type: ignore[misc] + arun_with_retry, + t, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_CALL: partial( + _acall, + weakref.ref(t), + retry=retry_policy, + stream=self.use_astream, + futures=weakref.ref(futures), + schedule_task=self.schedule_task, + submit=self.submit, + reraise=reraise, + loop=loop, + ), + }, + __name__=t.name, + __cancel_on_exit__=True, + __reraise_on_exit__=reraise, + ), + ) + futures[fut] = t # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes From d33c5a20e417aa8f60edc170b7a7d46b55362c99 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 1 May 2025 12:14:09 -0700 Subject: [PATCH 12/41] Fix --- libs/langgraph/langgraph/func/__init__.py | 2 +- libs/langgraph/langgraph/pregel/call.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 689ac05ab..dd6f5f4a5 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -145,7 +145,7 @@ def task( func.__name__ = name call_func = functools.partial( - call, func, retry=retry_policies, cache=cache_policy + call, func, retry=retry_policies, cache_policy=cache_policy ) object.__setattr__(call_func, "_is_pregel_task", True) return functools.update_wrapper(call_func, func) diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index bc41eafb5..d9abc6cf3 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -251,7 +251,7 @@ def call( func, (args, kwargs), retry=retry, - cache=cache_policy, + cache_policy=cache_policy, callbacks=config["callbacks"], ) return fut From 6b78bcd857d0da07632b4ff98333fc5bd56ac052 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 2 May 2025 10:39:35 -0700 Subject: [PATCH 13/41] Implement ttl in FileCache --- .../langgraph/cache/file/__init__.py | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/libs/checkpoint/langgraph/cache/file/__init__.py b/libs/checkpoint/langgraph/cache/file/__init__.py index 2ed206abf..e586c1355 100644 --- a/libs/checkpoint/langgraph/cache/file/__init__.py +++ b/libs/checkpoint/langgraph/cache/file/__init__.py @@ -1,4 +1,5 @@ import asyncio +import datetime import dbm import ormsgpack @@ -22,11 +23,16 @@ class FileCache(BaseCache): def get(self, keys: list[str]) -> dict[str, bytes]: """Get the cached values for the given keys.""" - return { - key: self.serde.loads_typed(ormsgpack.unpackb(self._db[key])) - for key in keys - if key in self._db - } + now = datetime.datetime.now(datetime.timezone.utc).timestamp() + values: dict[str, bytes] = {} + for key in keys: + if val := self._db.get(key): + expiry, *data = ormsgpack.unpackb(val) + if expiry is not None and now > expiry: + self._db.pop(key, None) + continue + values[key] = self.serde.loads_typed(data) + return values async def aget(self, keys: list[str]) -> dict[str, bytes]: """Asynchronously get the cached values for the given keys.""" @@ -34,9 +40,14 @@ class FileCache(BaseCache): def set(self, mapping: dict[str, tuple[bytes, int | None]]) -> None: """Set the cached values for the given keys and TTLs.""" - for key, (value, _) in mapping.items(): - # File-based caches do not support TTLs, so we ignore them. - self._db[key] = ormsgpack.packb(self.serde.dumps_typed(value)) + now = datetime.datetime.now(datetime.timezone.utc) + for key, (value, ttl) in mapping.items(): + if ttl is not None: + delta = datetime.timedelta(seconds=ttl) + expiry: float | None = (now + delta).timestamp() + else: + expiry = None + self._db[key] = ormsgpack.packb((expiry, *self.serde.dumps_typed(value))) async def aset(self, mapping: dict[str, tuple[bytes, int | None]]) -> None: """Asynchronously set the cached values for the given keys and TTLs.""" From a2a1a42c754c390f13e4b6c68abd0cbdc48d4eee Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 2 May 2025 10:41:53 -0700 Subject: [PATCH 14/41] Fix --- libs/langgraph/langgraph/pregel/loop.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 0bc710e04..95fd863b7 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -616,10 +616,10 @@ class PregelLoop(LoopProtocol): return True - def match_cached_writes(self) -> list[PregelExecutableTask]: + def match_cached_writes(self) -> Sequence[PregelExecutableTask]: raise NotImplementedError - async def amatch_cached_writes(self) -> list[PregelExecutableTask]: + async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]: raise NotImplementedError # private @@ -1048,9 +1048,9 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): return self.submit(cast(WritableManagedValue, managed_value).update, values) - def match_cached_writes(self) -> list[PregelExecutableTask]: + def match_cached_writes(self) -> Sequence[PregelExecutableTask]: if self.cache is None: - return + return () matched: list[PregelExecutableTask] = [] if cached := { t.cache_key.key: t @@ -1228,9 +1228,9 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): return self.submit(cast(WritableManagedValue, managed_value).aupdate, values) - async def amatch_cached_writes(self) -> None: + async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]: if self.cache is None: - return + return [] matched: list[PregelExecutableTask] = [] if cached := { t.cache_key.key: t From 0e81699fecc7edce8766c5c671ff197700e97a3d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 2 May 2025 10:59:18 -0700 Subject: [PATCH 15/41] Lint --- libs/checkpoint/langgraph/checkpoint/serde/base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/base.py b/libs/checkpoint/langgraph/checkpoint/serde/base.py index b007341eb..1a7752608 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/base.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/base.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Any, Protocol From 1edf5cee89d6a322722ef133031dc7e07d791e5e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 2 May 2025 10:59:34 -0700 Subject: [PATCH 16/41] Accept default cache_policy for graph/entrypoint/pregel --- libs/langgraph/langgraph/func/__init__.py | 6 +++ libs/langgraph/langgraph/graph/graph.py | 7 ++++ libs/langgraph/langgraph/graph/state.py | 3 ++ libs/langgraph/langgraph/pregel/__init__.py | 26 ++++++++----- libs/langgraph/langgraph/pregel/algo.py | 41 +++++++++++++-------- libs/langgraph/langgraph/pregel/loop.py | 20 ++++++++++ libs/langgraph/langgraph/pregel/retry.py | 4 +- libs/langgraph/langgraph/types.py | 2 +- 8 files changed, 82 insertions(+), 27 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index dd6f5f4a5..b2a4081fb 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -323,11 +323,15 @@ class entrypoint: store: Optional[BaseStore] = None, cache: Optional[BaseCache] = None, config_schema: Optional[type[Any]] = None, + cache_policy: Optional[CachePolicy] = None, + retry: Union[RetryPolicy, Sequence[RetryPolicy]] = (), ) -> None: """Initialize the entrypoint decorator.""" self.checkpointer = checkpointer self.store = store self.cache = cache + self.cache_policy = cache_policy + self.retry = retry self.config_schema = config_schema @dataclass(**_DC_KWARGS) @@ -458,5 +462,7 @@ class entrypoint: checkpointer=self.checkpointer, store=self.store, cache=self.cache, + cache_policy=self.cache_policy, + retry_policy=self.retry, config_type=self.config_schema, ) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index e8250ba57..6c4c069b7 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -14,6 +14,7 @@ from typing import ( from langchain_core.runnables import Runnable from typing_extensions import Self +from langgraph.cache.base import BaseCache from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.constants import ( EMPTY_SEQ, @@ -28,6 +29,7 @@ from langgraph.graph.branch import Branch from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry +from langgraph.store.base import BaseStore from langgraph.types import All, Checkpointer from langgraph.utils.runnable import RunnableLike, coerce_to_runnable @@ -316,6 +318,9 @@ class Graph: interrupt_after: Optional[Union[All, list[str]]] = None, debug: bool = False, name: Optional[str] = None, + *, + cache: Optional[BaseCache] = None, + store: Optional[BaseStore] = None, ) -> "CompiledGraph": """Compiles the graph into a `CompiledGraph` object. @@ -364,6 +369,8 @@ class Graph: auto_validate=False, debug=debug, name=name or "LangGraph", + cache=cache, + store=store, ) # attach nodes, edges, and branches diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index c86fc6b14..45138de0f 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -26,6 +26,7 @@ from pydantic import BaseModel from typing_extensions import Self from langgraph._api.deprecation import LangGraphDeprecationWarning +from langgraph.cache.base import BaseCache from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.dynamic_barrier_value import ( @@ -571,6 +572,7 @@ class StateGraph(Graph): self, checkpointer: Checkpointer = None, *, + cache: Optional[BaseCache] = None, store: Optional[BaseStore] = None, interrupt_before: Optional[Union[All, list[str]]] = None, interrupt_after: Optional[Union[All, list[str]]] = None, @@ -655,6 +657,7 @@ class StateGraph(Graph): auto_validate=False, debug=debug, store=store, + cache=cache, name=name or "LangGraph", ) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 6d8e7adb1..6aabc9a61 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -104,6 +104,7 @@ from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore from langgraph.types import ( All, + CachePolicy, Checkpointer, Interrupt, LoopProtocol, @@ -500,8 +501,12 @@ class Pregel(PregelProtocol): cache: BaseCache | None = None """Cache to use for storing node results. Defaults to None.""" - retry_policy: Sequence[RetryPolicy] | None = None - """Retry policies to use when running tasks. Set to None to disable.""" + retry_policy: Sequence[RetryPolicy] = () + """Retry policies to use when running tasks. Empty set disables retries.""" + + cache_policy: CachePolicy | None = None + """Cache policy to use for all nodes. Can be overridden by individual nodes. + Defaults to None.""" config_type: type[Any] | None = None @@ -531,7 +536,8 @@ class Pregel(PregelProtocol): checkpointer: BaseCheckpointSaver | None = None, store: BaseStore | None = None, cache: BaseCache | None = None, - retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, + retry_policy: RetryPolicy | Sequence[RetryPolicy] = (), + cache_policy: CachePolicy | None = None, config_type: type[Any] | None = None, input_model: type[BaseModel] | None = None, config: RunnableConfig | None = None, @@ -552,10 +558,10 @@ class Pregel(PregelProtocol): self.checkpointer = checkpointer self.store = store self.cache = cache - if isinstance(retry_policy, RetryPolicy): - self.retry_policy: Sequence[RetryPolicy] = (retry_policy,) - else: - self.retry_policy = retry_policy + self.retry_policy = ( + (retry_policy,) if isinstance(retry_policy, RetryPolicy) else retry_policy + ) + self.cache_policy = cache_policy self.config_type = config_type self.input_model = input_model self.config = config @@ -2465,6 +2471,8 @@ class Pregel(PregelProtocol): else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True), trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, + retry_policy=self.retry_policy, + cache_policy=self.cache_policy, ) as loop: # create runner runner = PregelRunner( @@ -2514,7 +2522,6 @@ class Pregel(PregelProtocol): for _ in runner.tick( [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, - retry_policy=self.retry_policy, get_waiter=get_waiter, match_cached_writes=loop.match_cached_writes, ): @@ -2777,6 +2784,8 @@ class Pregel(PregelProtocol): else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True), trigger_to_nodes=self.trigger_to_nodes, migrate_checkpoint=self._migrate_checkpoint, + retry_policy=self.retry_policy, + cache_policy=self.cache_policy, ) as loop: # create runner runner = PregelRunner( @@ -2817,7 +2826,6 @@ class Pregel(PregelProtocol): async for _ in runner.atick( [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, - retry_policy=self.retry_policy, get_waiter=get_waiter, # TODO pass match_cached_writes ): diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c0b15fd18..7449cd01d 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -424,6 +424,8 @@ def prepare_next_tasks( manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, updated_channels: Optional[set[str]] = None, + retry_policy: Sequence[RetryPolicy] = (), + cache_policy: Optional[CachePolicy] = None, ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]: """Prepare the set of tasks that will make up the next Pregel step. @@ -474,6 +476,8 @@ def prepare_next_tasks( checkpointer=checkpointer, manager=manager, input_cache=input_cache, + cache_policy=cache_policy, + retry_policy=retry_policy, ): tasks.append(task) @@ -517,6 +521,8 @@ def prepare_next_tasks( checkpointer=checkpointer, manager=manager, input_cache=input_cache, + cache_policy=cache_policy, + retry_policy=retry_policy, ): tasks.append(task) return {t.id: t for t in tasks} @@ -543,6 +549,8 @@ def prepare_single_task( checkpointer: Optional[BaseCheckpointSaver] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, input_cache: Optional[dict[INPUT_CACHE_KEY_TYPE, Any]] = None, + cache_policy: Optional[CachePolicy] = None, + retry_policy: Sequence[RetryPolicy] = (), ) -> Union[None, PregelTask, PregelExecutableTask]: """Prepares a single task for the next Pregel step, given a task path, which uniquely identifies a PUSH or PULL task within the graph.""" @@ -585,8 +593,9 @@ def prepare_single_task( assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}" if for_execution: writes: deque[tuple[str, Any]] = deque() - if call.cache_policy: - args_key = call.cache_policy.key(*call.input[0], **call.input[1]) + cache_policy = call.cache_policy or cache_policy + if cache_policy: + args_key = cache_policy.key(*call.input[0], **call.input[1]) cache_key: Optional[CacheKey] = CacheKey( xxh3_128_hexdigest( b"".join( @@ -602,8 +611,8 @@ def prepare_single_task( ) ) ), - call.cache_policy.ttl, - call.cache_policy.refresh, + cache_policy.ttl, + cache_policy.refresh, ) else: cache_key = None @@ -651,7 +660,7 @@ def prepare_single_task( }, ), triggers, - call.retry, + call.retry or retry_policy, cache_key, task_id, task_path, @@ -714,8 +723,9 @@ def prepare_single_task( if proc.metadata: metadata.update(proc.metadata) writes = deque() - if proc.cache_policy: - args_key = proc.cache_policy.key(packet.arg) + cache_policy = proc.cache_policy or cache_policy + if cache_policy: + args_key = cache_policy.key(packet.arg) cache_key = CacheKey( xxh3_128_hexdigest( b"".join( @@ -731,8 +741,8 @@ def prepare_single_task( ) ) ), - proc.cache_policy.ttl, - proc.cache_policy.refresh, + cache_policy.ttl, + cache_policy.refresh, ) else: cache_key = None @@ -784,7 +794,7 @@ def prepare_single_task( }, ), triggers, - proc.retry_policy, + proc.retry_policy or retry_policy, cache_key, task_id, task_path, @@ -852,8 +862,9 @@ def prepare_single_task( if proc.metadata: metadata.update(proc.metadata) writes = deque() - if proc.cache_policy: - args_key = proc.cache_policy.key(val) + cache_policy = proc.cache_policy or cache_policy + if cache_policy: + args_key = cache_policy.key(val) cache_key = CacheKey( xxh3_128_hexdigest( b"".join( @@ -869,8 +880,8 @@ def prepare_single_task( ) ) ), - proc.cache_policy.ttl, - proc.cache_policy.refresh, + cache_policy.ttl, + cache_policy.refresh, ) else: cache_key = None @@ -934,7 +945,7 @@ def prepare_single_task( }, ), triggers, - proc.retry_policy, + proc.retry_policy or retry_policy, cache_key, task_id, task_path[:3], diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 95fd863b7..00feb2650 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -118,10 +118,12 @@ from langgraph.pregel.utils import get_new_channel_versions, is_xxh3_128_hexdige from langgraph.store.base import BaseStore from langgraph.types import ( All, + CachePolicy, Command, LoopProtocol, PregelExecutableTask, PregelScratchpad, + RetryPolicy, StreamChunk, StreamProtocol, ) @@ -162,6 +164,8 @@ class PregelLoop(LoopProtocol): interrupt_before: Union[All, Sequence[str]] checkpoint_during: bool debug: bool + retry_policy: Sequence[RetryPolicy] + cache_policy: Optional[CachePolicy] checkpointer_get_next_version: GetNextVersion checkpointer_put_writes: Optional[Callable[[RunnableConfig, WritesT, str], Any]] @@ -220,6 +224,8 @@ class PregelLoop(LoopProtocol): input_model: Optional[type[BaseModel]] = None, debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, + retry_policy: Sequence[RetryPolicy] = (), + cache_policy: Optional[CachePolicy] = None, checkpoint_during: bool = True, ) -> None: super().__init__( @@ -247,6 +253,8 @@ class PregelLoop(LoopProtocol): ) self._migrate_checkpoint = migrate_checkpoint self.trigger_to_nodes = trigger_to_nodes + self.retry_policy = retry_policy + self.cache_policy = cache_policy self.checkpoint_during = checkpoint_during self.debug = debug if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: @@ -421,6 +429,8 @@ class PregelLoop(LoopProtocol): store=self.store, checkpointer=self.checkpointer, manager=self.manager, + retry_policy=self.retry_policy, + cache_policy=self.cache_policy, ), ): # don't start if we should interrupt *before* the new task @@ -550,6 +560,8 @@ class PregelLoop(LoopProtocol): checkpointer=self.checkpointer, trigger_to_nodes=self.trigger_to_nodes, updated_channels=updated_channels, + retry_policy=self.retry_policy, + cache_policy=self.cache_policy, ) self.to_interrupt = [] @@ -989,6 +1001,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): input_model: Optional[type[BaseModel]] = None, debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, + retry_policy: Sequence[RetryPolicy] = (), + cache_policy: Optional[CachePolicy] = None, checkpoint_during: bool = True, ) -> None: super().__init__( @@ -1009,6 +1023,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): debug=debug, migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, + retry_policy=retry_policy, + cache_policy=cache_policy, checkpoint_during=checkpoint_during, ) self.stack = ExitStack() @@ -1169,6 +1185,8 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): input_model: Optional[type[BaseModel]] = None, debug: bool = False, migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, + retry_policy: Sequence[RetryPolicy] = (), + cache_policy: Optional[CachePolicy] = None, checkpoint_during: bool = True, ) -> None: super().__init__( @@ -1189,6 +1207,8 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): debug=debug, migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, + retry_policy=retry_policy, + cache_policy=cache_policy, checkpoint_during=checkpoint_during, ) self.stack = AsyncExitStack() diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 78c2f8b0c..be5c0a91c 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -61,7 +61,7 @@ def run_with_retry( except Exception as exc: if SUPPORTS_EXC_NOTES: exc.add_note(f"During task with name '{task.name}' and id '{task.id}'") - if retry_policy is None: + if not retry_policy: raise # Check which retry policy applies to this exception @@ -149,7 +149,7 @@ async def arun_with_retry( except Exception as exc: if SUPPORTS_EXC_NOTES: exc.add_note(f"During task with name '{task.name}' and id '{task.id}'") - if retry_policies is None: + if not retry_policies: raise # Check which retry policy applies to this exception diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 9fbb92217..79da90b67 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -206,7 +206,7 @@ class PregelExecutableTask: writes: deque[tuple[str, Any]] config: RunnableConfig triggers: Sequence[str] - retry_policy: Optional[Sequence[RetryPolicy]] + retry_policy: Sequence[RetryPolicy] cache_key: Optional[CacheKey] id: str path: tuple[Union[str, int, tuple], ...] From 026436308365def6d6896e87a57478b38d71ceee Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 5 May 2025 10:21:24 -0700 Subject: [PATCH 17/41] Lint --- libs/checkpoint/langgraph/cache/file/__init__.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/libs/checkpoint/langgraph/cache/file/__init__.py b/libs/checkpoint/langgraph/cache/file/__init__.py index e586c1355..cb6765f0e 100644 --- a/libs/checkpoint/langgraph/cache/file/__init__.py +++ b/libs/checkpoint/langgraph/cache/file/__init__.py @@ -1,6 +1,7 @@ import asyncio import datetime import dbm +from collections.abc import Mapping, Sequence import ormsgpack @@ -21,7 +22,7 @@ class FileCache(BaseCache): super().__init__(serde=serde) self._db = dbm.open(path, "c") - def get(self, keys: list[str]) -> dict[str, bytes]: + def get(self, keys: Sequence[str]) -> dict[str, bytes]: """Get the cached values for the given keys.""" now = datetime.datetime.now(datetime.timezone.utc).timestamp() values: dict[str, bytes] = {} @@ -34,11 +35,11 @@ class FileCache(BaseCache): values[key] = self.serde.loads_typed(data) return values - async def aget(self, keys: list[str]) -> dict[str, bytes]: + async def aget(self, keys: Sequence[str]) -> dict[str, bytes]: """Asynchronously get the cached values for the given keys.""" return await asyncio.to_thread(self.get, keys) - def set(self, mapping: dict[str, tuple[bytes, int | None]]) -> None: + def set(self, mapping: Mapping[str, tuple[bytes, int | None]]) -> None: """Set the cached values for the given keys and TTLs.""" now = datetime.datetime.now(datetime.timezone.utc) for key, (value, ttl) in mapping.items(): @@ -49,15 +50,15 @@ class FileCache(BaseCache): expiry = None self._db[key] = ormsgpack.packb((expiry, *self.serde.dumps_typed(value))) - async def aset(self, mapping: dict[str, tuple[bytes, int | None]]) -> None: + async def aset(self, mapping: Mapping[str, tuple[bytes, int | None]]) -> None: """Asynchronously set the cached values for the given keys and TTLs.""" await asyncio.to_thread(self.set, mapping) - def delete(self, keys: list[str]) -> None: + def delete(self, keys: Sequence[str]) -> None: """Delete the cached values for the given keys.""" for key in keys: self._db.pop(key, None) - async def adelete(self, keys: list[str]) -> None: + async def adelete(self, keys: Sequence[str]) -> None: """Asynchronously delete the cached values for the given keys.""" await asyncio.to_thread(self.delete, keys) From 6c155f87c30453174ef7d5053359b853ce807a11 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 5 May 2025 10:22:05 -0700 Subject: [PATCH 18/41] Lint --- libs/checkpoint/langgraph/cache/base/__init__.py | 2 ++ libs/checkpoint/langgraph/cache/file/__init__.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py index e98525f81..abf14b1b6 100644 --- a/libs/checkpoint/langgraph/cache/base/__init__.py +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from abc import ABC, abstractmethod from collections.abc import Mapping from typing import Generic, Sequence, TypeVar diff --git a/libs/checkpoint/langgraph/cache/file/__init__.py b/libs/checkpoint/langgraph/cache/file/__init__.py index cb6765f0e..f37322ae2 100644 --- a/libs/checkpoint/langgraph/cache/file/__init__.py +++ b/libs/checkpoint/langgraph/cache/file/__init__.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import asyncio import datetime import dbm From d04570f1784f28f82f3b2fb2af62b4d08bae30b4 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 5 May 2025 12:16:25 -0700 Subject: [PATCH 19/41] Lint --- libs/langgraph/langgraph/pregel/algo.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 7449cd01d..ff94ec0f6 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -387,6 +387,8 @@ def prepare_next_tasks( manager: Literal[None] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, updated_channels: Optional[set[str]] = None, + retry_policy: Sequence[RetryPolicy] = (), + cache_policy: Literal[None] = None, ) -> dict[str, PregelTask]: ... @@ -406,6 +408,8 @@ def prepare_next_tasks( manager: Union[None, ParentRunManager, AsyncParentRunManager], trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, updated_channels: Optional[set[str]] = None, + retry_policy: Sequence[RetryPolicy] = (), + cache_policy: Optional[CachePolicy] = None, ) -> dict[str, PregelExecutableTask]: ... From 14b07d06fa20354ffeed4b711dc4d84d9d968ea5 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 11:12:16 -0700 Subject: [PATCH 20/41] Re-implement using sqlite --- .../langgraph/cache/file/__init__.py | 87 ++++++++++++++----- 1 file changed, 64 insertions(+), 23 deletions(-) diff --git a/libs/checkpoint/langgraph/cache/file/__init__.py b/libs/checkpoint/langgraph/cache/file/__init__.py index f37322ae2..080d24c66 100644 --- a/libs/checkpoint/langgraph/cache/file/__init__.py +++ b/libs/checkpoint/langgraph/cache/file/__init__.py @@ -2,17 +2,16 @@ from __future__ import annotations import asyncio import datetime -import dbm +import sqlite3 +import threading from collections.abc import Mapping, Sequence -import ormsgpack - from langgraph.cache.base import BaseCache from langgraph.checkpoint.serde.base import SerializerProtocol class FileCache(BaseCache): - """File-based cache using dbm.""" + """File-based cache using SQLite.""" def __init__( self, @@ -22,20 +21,46 @@ class FileCache(BaseCache): ) -> None: """Initialize the cache with a file path.""" super().__init__(serde=serde) - self._db = dbm.open(path, "c") + # SQLite backing store + self._conn = sqlite3.connect( + path, + check_same_thread=False, + ) + # Serialize access to the shared connection across threads + self._lock = threading.RLock() + # Better concurrency & atomicity + self._conn.execute("PRAGMA journal_mode=WAL;") + # Schema: key -> (expiry, encoding, value) + self._conn.execute( + """CREATE TABLE IF NOT EXISTS cache ( + key TEXT PRIMARY KEY, + expiry REAL, + encoding TEXT NOT NULL, + val BLOB NOT NULL + )""" + ) + self._conn.commit() def get(self, keys: Sequence[str]) -> dict[str, bytes]: """Get the cached values for the given keys.""" - now = datetime.datetime.now(datetime.timezone.utc).timestamp() - values: dict[str, bytes] = {} - for key in keys: - if val := self._db.get(key): - expiry, *data = ormsgpack.unpackb(val) + with self._lock, self._conn: + now = datetime.datetime.now(datetime.timezone.utc).timestamp() + if not keys: + return {} + placeholders = ",".join("?" for _ in keys) + cursor = self._conn.execute( + f"SELECT key, expiry, encoding, val FROM cache WHERE key IN ({placeholders})", + tuple(keys), + ) + values: dict[str, bytes] = {} + rows = cursor.fetchall() + for key, expiry, encoding, raw in rows: if expiry is not None and now > expiry: - self._db.pop(key, None) + # purge expired entry + self._conn.execute("DELETE FROM cache WHERE key = ?", (key,)) continue - values[key] = self.serde.loads_typed(data) - return values + values[key] = self.serde.loads_typed((encoding, raw)) + return values async def aget(self, keys: Sequence[str]) -> dict[str, bytes]: """Asynchronously get the cached values for the given keys.""" @@ -43,14 +68,19 @@ class FileCache(BaseCache): def set(self, mapping: Mapping[str, tuple[bytes, int | None]]) -> None: """Set the cached values for the given keys and TTLs.""" - now = datetime.datetime.now(datetime.timezone.utc) - for key, (value, ttl) in mapping.items(): - if ttl is not None: - delta = datetime.timedelta(seconds=ttl) - expiry: float | None = (now + delta).timestamp() - else: - expiry = None - self._db[key] = ormsgpack.packb((expiry, *self.serde.dumps_typed(value))) + with self._lock, self._conn: + now = datetime.datetime.now(datetime.timezone.utc) + for key, (value, ttl) in mapping.items(): + if ttl is not None: + delta = datetime.timedelta(seconds=ttl) + expiry: float | None = (now + delta).timestamp() + else: + expiry = None + encoding, raw = self.serde.dumps_typed(value) + self._conn.execute( + "INSERT OR REPLACE INTO cache (key, expiry, encoding, val) VALUES (?, ?, ?, ?)", + (key, expiry, encoding, raw), + ) async def aset(self, mapping: Mapping[str, tuple[bytes, int | None]]) -> None: """Asynchronously set the cached values for the given keys and TTLs.""" @@ -58,9 +88,20 @@ class FileCache(BaseCache): def delete(self, keys: Sequence[str]) -> None: """Delete the cached values for the given keys.""" - for key in keys: - self._db.pop(key, None) + if not keys: + return + with self._lock, self._conn: + placeholders = ",".join("?" for _ in keys) + self._conn.execute( + f"DELETE FROM cache WHERE key IN ({placeholders})", tuple(keys) + ) async def adelete(self, keys: Sequence[str]) -> None: """Asynchronously delete the cached values for the given keys.""" await asyncio.to_thread(self.delete, keys) + + def __del__(self) -> None: + try: + self._conn.close() + except Exception: + pass From 42d88a769a1a6ad5eac5199e97c8e78d9142d5db Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 14:54:07 -0700 Subject: [PATCH 21/41] Lint --- libs/langgraph/langgraph/func/__init__.py | 4 ++-- libs/langgraph/langgraph/pregel/algo.py | 6 +++--- libs/langgraph/langgraph/pregel/runner.py | 6 ++++-- libs/langgraph/langgraph/types.py | 7 +++---- libs/langgraph/langgraph/utils/cache.py | 2 +- 5 files changed, 13 insertions(+), 12 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index b2a4081fb..29338b02b 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -40,7 +40,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, - cache_policy: Optional[CachePolicy[P]] = None, + cache_policy: Optional[CachePolicy[Callable[P, str | bytes]]] = None, ) -> Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], Callable[P, SyncAsyncFuture[T]], @@ -58,7 +58,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, - cache_policy: Optional[CachePolicy[P]] = None, + cache_policy: Optional[CachePolicy[Callable[P, str | bytes]]] = None, ) -> Union[ Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index ff94ec0f6..225b0ccd3 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -599,7 +599,7 @@ def prepare_single_task( writes: deque[tuple[str, Any]] = deque() cache_policy = call.cache_policy or cache_policy if cache_policy: - args_key = cache_policy.key(*call.input[0], **call.input[1]) + args_key = cache_policy.key_func(*call.input[0], **call.input[1]) cache_key: Optional[CacheKey] = CacheKey( xxh3_128_hexdigest( b"".join( @@ -729,7 +729,7 @@ def prepare_single_task( writes = deque() cache_policy = proc.cache_policy or cache_policy if cache_policy: - args_key = cache_policy.key(packet.arg) + args_key = cache_policy.key_func(packet.arg) cache_key = CacheKey( xxh3_128_hexdigest( b"".join( @@ -868,7 +868,7 @@ def prepare_single_task( writes = deque() cache_policy = proc.cache_policy or cache_policy if cache_policy: - args_key = cache_policy.key(val) + args_key = cache_policy.key_func(val) cache_key = CacheKey( xxh3_128_hexdigest( b"".join( diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index dcf32f96e..916abef36 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -142,7 +142,9 @@ class PregelRunner: timeout: Optional[float] = None, retry_policy: Optional[Sequence[RetryPolicy]] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, - match_cached_writes: Optional[Callable[[], None]] = None, + match_cached_writes: Optional[ + Callable[[], Sequence[PregelExecutableTask]] + ] = None, ) -> Iterator[None]: tasks = tuple(tasks) futures = FuturesDict( @@ -529,7 +531,7 @@ def _call( [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] ] ], - match_cached_writes: Optional[Callable[[], None]], + match_cached_writes: Optional[Callable[[], Sequence[PregelExecutableTask]]], submit: weakref.ref[Submit], reraise: bool, ) -> concurrent.futures.Future[Any]: diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 79da90b67..4b7e36ee0 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -11,7 +11,6 @@ from typing import ( Literal, NamedTuple, Optional, - ParamSpec, TypeVar, Union, cast, @@ -124,16 +123,16 @@ class RetryPolicy(NamedTuple): """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" -P = ParamSpec("P") +KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., str | bytes]) -class CachePolicy(NamedTuple, Generic[P]): +class CachePolicy(NamedTuple, Generic[KeyFuncT]): """Configuration for caching nodes. !!! version-added "Added in version 0.2.24." """ - key: Callable[P, str | bytes] = default_cache_key + key_func: KeyFuncT = default_cache_key # type: ignore[assignment] """Function to generate a cache key from the node's input. Defaults to hashing the input with pickle.""" diff --git a/libs/langgraph/langgraph/utils/cache.py b/libs/langgraph/langgraph/utils/cache.py index ddae896c8..01ee2b2ea 100644 --- a/libs/langgraph/langgraph/utils/cache.py +++ b/libs/langgraph/langgraph/utils/cache.py @@ -18,7 +18,7 @@ def _freeze(obj: Any) -> Hashable: return obj # strings, ints, dataclasses with frozen=True, etc. -def default_cache_key(*args: Any, **kwargs: Any) -> bytes: +def default_cache_key(*args: Any, **kwargs: Any) -> str | bytes: """Default cache key function that uses the arguments and keyword arguments to generate a hashable key.""" import pickle From c4deb2c6216419ea10308a2c0b27af9965f9bb9f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 15:58:50 -0700 Subject: [PATCH 22/41] Finish implementation, add tests --- libs/langgraph/langgraph/graph/state.py | 9 ++- libs/langgraph/langgraph/pregel/__init__.py | 2 +- libs/langgraph/langgraph/pregel/retry.py | 12 +++- libs/langgraph/langgraph/pregel/runner.py | 22 ++++-- libs/langgraph/tests/test_pregel.py | 34 +++++++-- libs/langgraph/tests/test_pregel_async.py | 78 +++++++++++++++++++-- 6 files changed, 139 insertions(+), 18 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 45138de0f..f6a29198b 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -80,7 +80,7 @@ from langgraph.pregel.write import ( ChannelWriteTupleEntry, ) from langgraph.store.base import BaseStore -from langgraph.types import All, Checkpointer, Command, RetryPolicy +from langgraph.types import All, CachePolicy, Checkpointer, Command, RetryPolicy from langgraph.utils.fields import get_field_default, get_update_as_tuples from langgraph.utils.pydantic import create_model from langgraph.utils.runnable import RunnableLike, coerce_to_runnable @@ -114,6 +114,7 @@ class StateNodeSpec(NamedTuple): metadata: Optional[dict[str, Any]] input: type[Any] retry_policy: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] + cache_policy: Optional[CachePolicy] ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ defer: bool = False @@ -260,6 +261,7 @@ class StateGraph(Graph): metadata: Optional[dict[str, Any]] = None, input: Optional[type[Any]] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, + cache_policy: Optional[CachePolicy] = None, destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None, ) -> Self: """Add a new node to the state graph. @@ -277,6 +279,7 @@ class StateGraph(Graph): metadata: Optional[dict[str, Any]] = None, input: Optional[type[Any]] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, + cache_policy: Optional[CachePolicy] = None, destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None, ) -> Self: """Add a new node to the state graph.""" @@ -291,6 +294,7 @@ class StateGraph(Graph): metadata: Optional[dict[str, Any]] = None, input: Optional[type[Any]] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, + cache_policy: Optional[CachePolicy] = None, destinations: Optional[Union[dict[str, str], tuple[str, ...]]] = None, ) -> Self: """Add a new node to the state graph. @@ -304,6 +308,7 @@ class StateGraph(Graph): input: The input schema for the node. (default: the graph's input schema) retry: The policy for retrying the node. (default: None) If a sequence is provided, the first matching policy will be applied. + cache_policy: The cache policy for the node. (default: None) destinations: Destinations that indicate where a node can route to. This is useful for edgeless graphs with nodes that return `Command` objects. If a dict is provided, the keys will be used as the target node names and the values will be used as the labels for the edges. @@ -432,6 +437,7 @@ class StateGraph(Graph): metadata, input=input or self.schema, retry_policy=retry, + cache_policy=cache_policy, ends=ends, defer=defer, ) @@ -814,6 +820,7 @@ class CompiledStateGraph(CompiledGraph): writers=[ChannelWrite(write_entries)], metadata=node.metadata, retry_policy=node.retry_policy, + cache_policy=node.cache_policy, bound=node.runnable, ) else: diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 6aabc9a61..f59fdd951 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2827,7 +2827,7 @@ class Pregel(PregelProtocol): [t for t in loop.tasks.values() if not t.writes], timeout=self.step_timeout, get_waiter=get_waiter, - # TODO pass match_cached_writes + match_cached_writes=loop.amatch_cached_writes, ): # emit output for o in output(): diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index be5c0a91c..42b73de36 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -3,9 +3,9 @@ import logging import random import sys import time -from collections.abc import Sequence +from collections.abc import Awaitable, Sequence from dataclasses import replace -from typing import Any, Optional +from typing import Any, Callable, Optional from langgraph.constants import ( CONF, @@ -106,6 +106,9 @@ async def arun_with_retry( task: PregelExecutableTask, retry_policies: Optional[Sequence[RetryPolicy]], stream: bool = False, + match_cached_writes: Optional[ + Callable[[], Awaitable[Sequence[PregelExecutableTask]]] + ] = None, configurable: Optional[dict[str, Any]] = None, ) -> None: """Run a task asynchronously with retries.""" @@ -114,6 +117,11 @@ async def arun_with_retry( config = task.config if configurable is not None: config = patch_configurable(config, configurable) + if match_cached_writes is not None and task.cache_key is not None: + for t in await match_cached_writes(): + if t is task: + # if the task is already cached, return + return while True: try: # clear any writes from previous attempts diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 916abef36..b0b2f31a9 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -155,7 +155,9 @@ class PregelRunner: # give control back to the caller yield # fast path if single task with no timeout and no waiter - if len(tasks) == 1 and timeout is None and get_waiter is None: + if len(tasks) == 0: + return + elif len(tasks) == 1 and timeout is None and get_waiter is None: t = tasks[0] try: run_with_retry( @@ -275,6 +277,9 @@ class PregelRunner: timeout: Optional[float] = None, retry_policy: Optional[Sequence[RetryPolicy]] = None, get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, + match_cached_writes: Optional[ + Callable[[], Awaitable[Sequence[PregelExecutableTask]]] + ] = None, ) -> AsyncIterator[None]: loop = asyncio.get_event_loop() tasks = tuple(tasks) @@ -286,7 +291,9 @@ class PregelRunner: # give control back to the caller yield # fast path if single task with no waiter and no timeout - if len(tasks) == 1 and get_waiter is None and timeout is None: + if len(tasks) == 0: + return + elif len(tasks) == 1 and get_waiter is None and timeout is None: t = tasks[0] try: await arun_with_retry( @@ -301,6 +308,7 @@ class PregelRunner: retry=retry_policy, futures=weakref.ref(futures), schedule_task=self.schedule_task, + match_cached_writes=match_cached_writes, submit=self.submit, reraise=reraise, loop=loop, @@ -348,6 +356,7 @@ class PregelRunner: stream=self.use_astream, futures=weakref.ref(futures), schedule_task=self.schedule_task, + match_cached_writes=match_cached_writes, submit=self.submit, reraise=reraise, loop=loop, @@ -609,7 +618,7 @@ def _acall( input: Any, *, retry: Optional[Sequence[RetryPolicy]] = None, - cache: Optional[CachePolicy] = None, + cache_policy: Optional[CachePolicy] = None, callbacks: Callbacks = None, # injected dependencies futures: weakref.ref[FuturesDict], @@ -618,6 +627,9 @@ def _acall( [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] ] ], + match_cached_writes: Optional[ + Callable[[], Awaitable[Sequence[PregelExecutableTask]]] + ] = None, submit: weakref.ref[Submit], loop: asyncio.AbstractEventLoop, reraise: bool = False, @@ -630,7 +642,7 @@ def _acall( if next_task := schedule_task()( # type: ignore[misc] task(), # type: ignore[arg-type] scratchpad.call_counter(), - Call(func, input, retry=retry, cache_policy=cache, callbacks=callbacks), + Call(func, input, retry=retry, cache_policy=cache_policy, callbacks=callbacks), ): if fut := next( ( @@ -666,6 +678,7 @@ def _acall( next_task, retry, stream=stream, + match_cached_writes=match_cached_writes, configurable={ CONFIG_KEY_CALL: partial( _acall, @@ -673,6 +686,7 @@ def _acall( stream=stream, futures=futures, schedule_task=schedule_task, + match_cached_writes=match_cached_writes, submit=submit, loop=loop, reraise=reraise, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 9ede99db1..53233a0ce 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -3573,7 +3573,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( ] -def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: +@pytest.mark.parametrize("with_cache", [True, False]) +def test_in_one_fan_out_state_graph_waiting_edge_multiple( + with_cache: bool, file_cache: BaseCache +) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3588,7 +3591,11 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: answer: str docs: Annotated[list[str], sorted_add] + rewrite_query_count = 0 + def rewrite_query(data: State) -> State: + nonlocal rewrite_query_count + rewrite_query_count += 1 return {"query": f"query: {data['query']}"} def analyzer_one(data: State) -> State: @@ -3615,7 +3622,11 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: workflow = StateGraph(State) - workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node( + "rewrite_query", + rewrite_query, + cache_policy=CachePolicy() if with_cache else None, + ) workflow.add_node("analyzer_one", analyzer_one) workflow.add_node("retriever_one", retriever_one) workflow.add_node("retriever_two", retriever_two) @@ -3630,7 +3641,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: workflow.add_conditional_edges("decider", decider_cond) workflow.set_finish_point("qa") - app = workflow.compile() + app = workflow.compile(cache=file_cache) assert app.invoke({"query": "what is weather in sf"}) == { "query": "analyzed: query: analyzed: query: what is weather in sf", @@ -3639,12 +3650,24 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: } assert [*app.stream({"query": "what is weather in sf"})] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "rewrite_query": {"query": "query: what is weather in sf"}, + "__metadata__": {"cached": True}, + } + if with_cache + else {"rewrite_query": {"query": "query: what is weather in sf"}}, {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"decider": None}, - {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, + { + "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}, + "__metadata__": {"cached": True}, + } + if with_cache + else { + "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"} + }, { "analyzer_one": { "query": "analyzed: query: analyzed: query: what is weather in sf" @@ -3655,6 +3678,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: {"decider": None}, {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, ] + assert rewrite_query_count == 2 if with_cache else 4 def test_callable_in_conditional_edges_with_no_path_map() -> None: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 38c3ed52d..03d3f329a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -31,6 +31,7 @@ from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import TypedDict +from langgraph.cache.base import BaseCache from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context @@ -55,6 +56,7 @@ from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner from langgraph.store.base import BaseStore from langgraph.types import ( + CachePolicy, Command, Interrupt, PregelTask, @@ -5395,7 +5397,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( ] -async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: +@pytest.mark.parametrize("with_cache", [True, False]) +async def test_in_one_fan_out_state_graph_waiting_edge_multiple( + with_cache: bool, file_cache: BaseCache +) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -5410,7 +5415,11 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: answer: str docs: Annotated[list[str], sorted_add] + rewrite_query_count = 0 + async def rewrite_query(data: State) -> State: + nonlocal rewrite_query_count + rewrite_query_count += 1 return {"query": f"query: {data['query']}"} async def analyzer_one(data: State) -> State: @@ -5437,7 +5446,11 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: workflow = StateGraph(State) - workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node( + "rewrite_query", + rewrite_query, + cache_policy=CachePolicy() if with_cache else None, + ) workflow.add_node("analyzer_one", analyzer_one) workflow.add_node("retriever_one", retriever_one) workflow.add_node("retriever_two", retriever_two) @@ -5452,21 +5465,34 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: workflow.add_conditional_edges("decider", decider_cond) workflow.set_finish_point("qa") - app = workflow.compile() + app = workflow.compile(cache=file_cache) assert await app.ainvoke({"query": "what is weather in sf"}) == { "query": "analyzed: query: analyzed: query: what is weather in sf", "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], } + assert rewrite_query_count == 2 assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "rewrite_query": {"query": "query: what is weather in sf"}, + "__metadata__": {"cached": True}, + } + if with_cache + else {"rewrite_query": {"query": "query: what is weather in sf"}}, {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, {"retriever_two": {"docs": ["doc3", "doc4"]}}, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"decider": None}, - {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, + { + "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}, + "__metadata__": {"cached": True}, + } + if with_cache + else { + "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"} + }, { "analyzer_one": { "query": "analyzed: query: analyzed: query: what is weather in sf" @@ -5477,6 +5503,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: {"decider": None}, {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, ] + assert rewrite_query_count == 2 if with_cache else 4 async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None: @@ -7510,6 +7537,47 @@ async def test_multiple_interrupts_functional(checkpointer_name: str) -> None: assert counter == 3 +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_multiple_interrupts_functional_cache( + checkpointer_name: str, file_cache: BaseCache +): + """Test multiple interrupts with functional API.""" + async with awith_checkpointer(checkpointer_name) as checkpointer: + counter = 0 + + @task(cache_policy=CachePolicy()) + def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x + + @entrypoint(checkpointer=checkpointer, cache=file_cache) + def graph(state: dict) -> dict: + """React tool.""" + + values = [] + + for idx in [1, 1, 2, 2, 3, 3]: + values.extend([double(idx).result(), interrupt({"a": "boo"})]) + + return {"values": values} + + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({}, configurable) + await graph.ainvoke(Command(resume="a"), configurable) + await graph.ainvoke(Command(resume="b"), configurable) + await graph.ainvoke(Command(resume="c"), configurable) + await graph.ainvoke(Command(resume="d"), configurable) + await graph.ainvoke(Command(resume="e"), configurable) + result = await graph.ainvoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 + + @NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: From d38303494cfef39a7fd3abc11d9cd93b51eef6e1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 16:03:26 -0700 Subject: [PATCH 23/41] Lint --- libs/langgraph/langgraph/types.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 4b7e36ee0..d1b68b4eb 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -126,11 +126,9 @@ class RetryPolicy(NamedTuple): KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., str | bytes]) -class CachePolicy(NamedTuple, Generic[KeyFuncT]): - """Configuration for caching nodes. - - !!! version-added "Added in version 0.2.24." - """ +@dataclasses.dataclass(**_DC_KWARGS) +class CachePolicy(Generic[KeyFuncT]): + """Configuration for caching nodes.""" key_func: KeyFuncT = default_cache_key # type: ignore[assignment] """Function to generate a cache key from the node's input. From 7c9f9aa89d6793c8da8f998c665906126e899869 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 16:06:26 -0700 Subject: [PATCH 24/41] Lint --- libs/langgraph/langgraph/utils/cache.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/langgraph/langgraph/utils/cache.py b/libs/langgraph/langgraph/utils/cache.py index 01ee2b2ea..5555cf9bc 100644 --- a/libs/langgraph/langgraph/utils/cache.py +++ b/libs/langgraph/langgraph/utils/cache.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from collections.abc import Hashable from typing import Any From f5bf77b3eb1ab36231e2a15c44a67b50cec1b553 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 16:16:03 -0700 Subject: [PATCH 25/41] Lint --- libs/langgraph/langgraph/types.py | 60 +++++++++++++++---------------- 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index d1b68b4eb..031f63c1a 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import dataclasses import sys from collections import deque @@ -10,9 +12,7 @@ from typing import ( Generic, Literal, NamedTuple, - Optional, TypeVar, - Union, cast, get_type_hints, ) @@ -41,7 +41,7 @@ except ImportError: All = Literal["*"] """Special value to indicate that graph should interrupt on all nodes.""" -Checkpointer = Union[None, bool, BaseCheckpointSaver] +Checkpointer = None | bool | BaseCheckpointSaver """Type of the checkpointer to use for a subgraph. - True enables persistent checkpointing for this subgraph. - False disables checkpointing, even if the parent graph has a checkpointer. @@ -117,9 +117,9 @@ class RetryPolicy(NamedTuple): """Maximum number of attempts to make before giving up, including the first.""" jitter: bool = True """Whether to add random jitter to the interval between retries.""" - retry_on: Union[ - type[Exception], Sequence[type[Exception]], Callable[[Exception], bool] - ] = default_retry_on + retry_on: ( + type[Exception] | Sequence[type[Exception]] | Callable[[Exception], bool] + ) = default_retry_on """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" @@ -134,7 +134,7 @@ class CachePolicy(Generic[KeyFuncT]): """Function to generate a cache key from the node's input. Defaults to hashing the input with pickle.""" - ttl: Optional[int] = None + ttl: int | None = None """Time to live for the cache entry in seconds. If None, the entry never expires.""" refresh: bool = False @@ -150,7 +150,7 @@ class Interrupt: value: Any resumable: bool = False - ns: Optional[Sequence[str]] = None + ns: Sequence[str] | None = None when: Literal["during"] = dataclasses.field(default="during", repr=False) @property @@ -162,8 +162,8 @@ class Interrupt: class StateUpdate(NamedTuple): - values: Optional[dict[str, Any]] - as_node: Optional[str] = None + values: dict[str, Any] | None + as_node: str | None = None class PregelTask(NamedTuple): @@ -171,11 +171,11 @@ class PregelTask(NamedTuple): id: str name: str - path: tuple[Union[str, int, tuple], ...] - error: Optional[Exception] = None + path: tuple[str | int | tuple, ...] + error: Exception | None = None interrupts: tuple[Interrupt, ...] = () - state: Union[None, RunnableConfig, "StateSnapshot"] = None - result: Optional[Any] = None + state: RunnableConfig | StateSnapshot | None = None + result: Any | None = None if sys.version_info > (3, 11): @@ -189,7 +189,7 @@ class CacheKey(NamedTuple): key: str """Key for the cache entry.""" - ttl: Optional[int] + ttl: int | None """Time to live for the cache entry in seconds.""" refresh: bool """Whether to force a refresh of the cache entry when it is accessed.""" @@ -204,28 +204,28 @@ class PregelExecutableTask: config: RunnableConfig triggers: Sequence[str] retry_policy: Sequence[RetryPolicy] - cache_key: Optional[CacheKey] + cache_key: CacheKey | None id: str - path: tuple[Union[str, int, tuple], ...] + path: tuple[str | int | tuple, ...] scheduled: bool = False writers: Sequence[Runnable] = () - subgraphs: Sequence["PregelProtocol"] = () + subgraphs: Sequence[PregelProtocol] = () class StateSnapshot(NamedTuple): """Snapshot of the state of the graph at the beginning of a step.""" - values: Union[dict[str, Any], Any] + values: dict[str, Any] | Any """Current values of channels.""" next: tuple[str, ...] """The name of the node to execute in each task for this step.""" config: RunnableConfig """Config used to fetch this snapshot.""" - metadata: Optional[CheckpointMetadata] + metadata: CheckpointMetadata | None """Metadata associated with this snapshot.""" - created_at: Optional[str] + created_at: str | None """Timestamp of snapshot creation.""" - parent_config: Optional[RunnableConfig] + parent_config: RunnableConfig | None """Config used to fetch the parent snapshot, if any.""" tasks: tuple[PregelTask, ...] """Tasks to execute in this step. If already attempted, may contain an error.""" @@ -332,10 +332,10 @@ class Command(Generic[N], ToolOutputMixin): - sequence of `Send` objects """ - graph: Optional[str] = None - update: Optional[Any] = None - resume: Optional[Union[Any, dict[str, Any]]] = None - goto: Union[Send, Sequence[Union[Send, N]], N] = () + graph: str | None = None + update: Any | None = None + resume: dict[str, Any] | Any | None = None + goto: Send | N | Sequence[Send | N] = () def __repr__(self) -> str: # get all non-None values @@ -385,8 +385,8 @@ class StreamProtocol: class LoopProtocol: config: RunnableConfig - store: Optional["BaseStore"] - stream: Optional[StreamProtocol] + store: BaseStore | None + stream: StreamProtocol | None step: int stop: int @@ -396,8 +396,8 @@ class LoopProtocol: step: int, stop: int, config: RunnableConfig, - store: Optional["BaseStore"] = None, - stream: Optional[StreamProtocol] = None, + store: BaseStore | None = None, + stream: StreamProtocol | None = None, ) -> None: self.stream = stream self.config = config From 0dd9fba0af87745c1842ad3b40da7e8904ef20cb Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 16:18:23 -0700 Subject: [PATCH 26/41] Lint --- libs/langgraph/langgraph/types.py | 62 +++++++++++++++---------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 031f63c1a..eba426deb 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import dataclasses import sys from collections import deque @@ -12,7 +10,9 @@ from typing import ( Generic, Literal, NamedTuple, + Optional, TypeVar, + Union, cast, get_type_hints, ) @@ -41,7 +41,7 @@ except ImportError: All = Literal["*"] """Special value to indicate that graph should interrupt on all nodes.""" -Checkpointer = None | bool | BaseCheckpointSaver +Checkpointer = Union[None, bool, BaseCheckpointSaver] """Type of the checkpointer to use for a subgraph. - True enables persistent checkpointing for this subgraph. - False disables checkpointing, even if the parent graph has a checkpointer. @@ -117,13 +117,13 @@ class RetryPolicy(NamedTuple): """Maximum number of attempts to make before giving up, including the first.""" jitter: bool = True """Whether to add random jitter to the interval between retries.""" - retry_on: ( - type[Exception] | Sequence[type[Exception]] | Callable[[Exception], bool] - ) = default_retry_on + retry_on: Union[ + type[Exception], Sequence[type[Exception]], Callable[[Exception], bool] + ] = default_retry_on """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" -KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., str | bytes]) +KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., Union[str, bytes]]) @dataclasses.dataclass(**_DC_KWARGS) @@ -134,7 +134,7 @@ class CachePolicy(Generic[KeyFuncT]): """Function to generate a cache key from the node's input. Defaults to hashing the input with pickle.""" - ttl: int | None = None + ttl: Optional[int] = None """Time to live for the cache entry in seconds. If None, the entry never expires.""" refresh: bool = False @@ -150,7 +150,7 @@ class Interrupt: value: Any resumable: bool = False - ns: Sequence[str] | None = None + ns: Optional[Sequence[str]] = None when: Literal["during"] = dataclasses.field(default="during", repr=False) @property @@ -162,8 +162,8 @@ class Interrupt: class StateUpdate(NamedTuple): - values: dict[str, Any] | None - as_node: str | None = None + values: Optional[dict[str, Any]] + as_node: Optional[str] = None class PregelTask(NamedTuple): @@ -171,11 +171,11 @@ class PregelTask(NamedTuple): id: str name: str - path: tuple[str | int | tuple, ...] - error: Exception | None = None + path: tuple[Union[str, int, tuple], ...] + error: Optional[Exception] = None interrupts: tuple[Interrupt, ...] = () - state: RunnableConfig | StateSnapshot | None = None - result: Any | None = None + state: Union[None, RunnableConfig, "StateSnapshot"] = None + result: Optional[Any] = None if sys.version_info > (3, 11): @@ -189,7 +189,7 @@ class CacheKey(NamedTuple): key: str """Key for the cache entry.""" - ttl: int | None + ttl: Optional[int] """Time to live for the cache entry in seconds.""" refresh: bool """Whether to force a refresh of the cache entry when it is accessed.""" @@ -204,28 +204,28 @@ class PregelExecutableTask: config: RunnableConfig triggers: Sequence[str] retry_policy: Sequence[RetryPolicy] - cache_key: CacheKey | None + cache_key: Optional[CacheKey] id: str - path: tuple[str | int | tuple, ...] + path: tuple[Union[str, int, tuple], ...] scheduled: bool = False writers: Sequence[Runnable] = () - subgraphs: Sequence[PregelProtocol] = () + subgraphs: Sequence["PregelProtocol"] = () class StateSnapshot(NamedTuple): """Snapshot of the state of the graph at the beginning of a step.""" - values: dict[str, Any] | Any + values: Union[dict[str, Any], Any] """Current values of channels.""" next: tuple[str, ...] """The name of the node to execute in each task for this step.""" config: RunnableConfig """Config used to fetch this snapshot.""" - metadata: CheckpointMetadata | None + metadata: Optional[CheckpointMetadata] """Metadata associated with this snapshot.""" - created_at: str | None + created_at: Optional[str] """Timestamp of snapshot creation.""" - parent_config: RunnableConfig | None + parent_config: Optional[RunnableConfig] """Config used to fetch the parent snapshot, if any.""" tasks: tuple[PregelTask, ...] """Tasks to execute in this step. If already attempted, may contain an error.""" @@ -332,10 +332,10 @@ class Command(Generic[N], ToolOutputMixin): - sequence of `Send` objects """ - graph: str | None = None - update: Any | None = None - resume: dict[str, Any] | Any | None = None - goto: Send | N | Sequence[Send | N] = () + graph: Optional[str] = None + update: Optional[Any] = None + resume: Optional[Union[dict[str, Any], Any]] = None + goto: Union[Send, Sequence[Union[Send, N]], N] = () def __repr__(self) -> str: # get all non-None values @@ -385,8 +385,8 @@ class StreamProtocol: class LoopProtocol: config: RunnableConfig - store: BaseStore | None - stream: StreamProtocol | None + store: Optional["BaseStore"] + stream: Optional[StreamProtocol] step: int stop: int @@ -396,8 +396,8 @@ class LoopProtocol: step: int, stop: int, config: RunnableConfig, - store: BaseStore | None = None, - stream: StreamProtocol | None = None, + store: Optional["BaseStore"] = None, + stream: Optional[StreamProtocol] = None, ) -> None: self.stream = stream self.config = config From 0211886bf54af330e322adaadf9de8f005972fc0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 16:20:37 -0700 Subject: [PATCH 27/41] Lint --- libs/langgraph/langgraph/pregel/algo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 225b0ccd3..4d7d3075c 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -1102,7 +1102,7 @@ def _proc_input( return val -def _uuid5_str(namespace: bytes, *parts: str | bytes) -> str: +def _uuid5_str(namespace: bytes, *parts: Union[str, bytes]) -> str: """Generate a UUID from the SHA-1 hash of a namespace and str parts.""" sha = sha1(namespace, usedforsecurity=False) @@ -1111,7 +1111,7 @@ def _uuid5_str(namespace: bytes, *parts: str | bytes) -> str: return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}" -def _xxhash_str(namespace: bytes, *parts: str | bytes) -> str: +def _xxhash_str(namespace: bytes, *parts: Union[str, bytes]) -> str: """Generate a UUID from the XXH3 hash of a namespace and str parts.""" hex = xxh3_128_hexdigest( namespace + b"".join(p.encode() if isinstance(p, str) else p for p in parts) From f1c1eaf229697640ee349cf4e6e8202d2fb92a6f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 16:22:47 -0700 Subject: [PATCH 28/41] Lint --- libs/langgraph/langgraph/func/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 29338b02b..06aac25c2 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -40,7 +40,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, - cache_policy: Optional[CachePolicy[Callable[P, str | bytes]]] = None, + cache_policy: Optional[CachePolicy[Callable[P, Union[str, bytes]]]] = None, ) -> Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], Callable[P, SyncAsyncFuture[T]], @@ -58,7 +58,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, - cache_policy: Optional[CachePolicy[Callable[P, str | bytes]]] = None, + cache_policy: Optional[CachePolicy[Callable[P, Union[str, bytes]]]] = None, ) -> Union[ Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], From 4761eb76967391b34c696da1cd181f78d523d84f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 16:25:54 -0700 Subject: [PATCH 29/41] Lint --- libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 9f6d4110f..adb4ce04d 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -151,6 +151,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): msg["input"], config=ensure_config(msg["config"]), stream=None, + cache=self.graph.cache, store=self.graph.store, checkpointer=self.graph.checkpointer, nodes=graph.nodes, @@ -338,6 +339,7 @@ class KafkaOrchestrator(AbstractContextManager): msg["input"], config=ensure_config(msg["config"]), stream=None, + cache=self.graph.cache, store=self.graph.store, checkpointer=self.graph.checkpointer, nodes=graph.nodes, From bebb0e81644f01eb4c1f3be02fa8b023b27a4462 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 May 2025 16:31:27 -0700 Subject: [PATCH 30/41] Lint --- libs/langgraph/tests/test_pregel_async.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 03d3f329a..c43f4cdac 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7537,6 +7537,7 @@ async def test_multiple_interrupts_functional(checkpointer_name: str) -> None: assert counter == 3 +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_multiple_interrupts_functional_cache( checkpointer_name: str, file_cache: BaseCache From 31b135f75d27639ad198e89697ceb1bc35b37900 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 15:14:00 -0700 Subject: [PATCH 31/41] Add namespace to cache keys --- .../langgraph/cache/base/__init__.py | 14 +++-- .../langgraph/cache/file/__init__.py | 52 ++++++++++------- libs/langgraph/langgraph/constants.py | 4 ++ libs/langgraph/langgraph/pregel/algo.py | 58 +++++++------------ libs/langgraph/langgraph/pregel/call.py | 10 ++++ libs/langgraph/langgraph/pregel/loop.py | 18 ++++-- libs/langgraph/langgraph/types.py | 2 + 7 files changed, 90 insertions(+), 68 deletions(-) diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py index abf14b1b6..8f8660cca 100644 --- a/libs/checkpoint/langgraph/cache/base/__init__.py +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -8,6 +8,8 @@ from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer ValueT = TypeVar("ValueT") +Namespace = tuple[str, ...] +FullKey = tuple[Namespace, str] class BaseCache(ABC, Generic[ValueT]): @@ -20,25 +22,25 @@ class BaseCache(ABC, Generic[ValueT]): self.serde = serde or self.serde @abstractmethod - def get(self, keys: Sequence[str]) -> dict[str, ValueT]: + def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]: """Get the cached values for the given keys.""" @abstractmethod - async def aget(self, keys: Sequence[str]) -> dict[str, ValueT]: + async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]: """Asynchronously get the cached values for the given keys.""" @abstractmethod - def set(self, mapping: Mapping[str, tuple[ValueT, int | None]]) -> None: + def set(self, pairs: Mapping[FullKey, tuple[ValueT, int | None]]) -> None: """Set the cached values for the given keys and TTLs.""" @abstractmethod - async def aset(self, mapping: Mapping[str, tuple[ValueT, int | None]]) -> None: + async def aset(self, pairs: Mapping[FullKey, tuple[ValueT, int | None]]) -> None: """Asynchronously set the cached values for the given keys and TTLs.""" @abstractmethod - def delete(self, keys: Sequence[str]) -> None: + def delete(self, keys: Sequence[Namespace]) -> None: """Delete the cached values for the given keys.""" @abstractmethod - async def adelete(self, keys: Sequence[str]) -> None: + async def adelete(self, keys: Sequence[Namespace]) -> None: """Asynchronously delete the cached values for the given keys.""" diff --git a/libs/checkpoint/langgraph/cache/file/__init__.py b/libs/checkpoint/langgraph/cache/file/__init__.py index 080d24c66..29231b206 100644 --- a/libs/checkpoint/langgraph/cache/file/__init__.py +++ b/libs/checkpoint/langgraph/cache/file/__init__.py @@ -6,7 +6,7 @@ import sqlite3 import threading from collections.abc import Mapping, Sequence -from langgraph.cache.base import BaseCache +from langgraph.cache.base import BaseCache, FullKey, Namespace from langgraph.checkpoint.serde.base import SerializerProtocol @@ -33,40 +33,49 @@ class FileCache(BaseCache): # Schema: key -> (expiry, encoding, value) self._conn.execute( """CREATE TABLE IF NOT EXISTS cache ( - key TEXT PRIMARY KEY, + ns TEXT, + key TEXT, expiry REAL, encoding TEXT NOT NULL, - val BLOB NOT NULL + val BLOB NOT NULL, + PRIMARY KEY (ns, key) )""" ) self._conn.commit() - def get(self, keys: Sequence[str]) -> dict[str, bytes]: + def get(self, keys: Sequence[FullKey]) -> dict[FullKey, bytes]: """Get the cached values for the given keys.""" with self._lock, self._conn: now = datetime.datetime.now(datetime.timezone.utc).timestamp() if not keys: return {} - placeholders = ",".join("?" for _ in keys) + placeholders = ",".join("(?, ?)" for _ in keys) + params: list[str] = [] + for ns_tuple, key in keys: + params.extend((",".join(ns_tuple), key)) cursor = self._conn.execute( - f"SELECT key, expiry, encoding, val FROM cache WHERE key IN ({placeholders})", - tuple(keys), + f"SELECT ns, key, expiry, encoding, val FROM cache WHERE (ns, key) IN ({placeholders})", + tuple(params), ) - values: dict[str, bytes] = {} + values: dict[FullKey, bytes] = {} rows = cursor.fetchall() - for key, expiry, encoding, raw in rows: + for ns, key, expiry, encoding, raw in rows: if expiry is not None and now > expiry: # purge expired entry - self._conn.execute("DELETE FROM cache WHERE key = ?", (key,)) + self._conn.execute( + "DELETE FROM cache WHERE (ns, key) = (?, ?)", (ns, key) + ) continue - values[key] = self.serde.loads_typed((encoding, raw)) + values[(tuple(ns.split(",")), key)] = self.serde.loads_typed( + (encoding, raw) + ) return values - async def aget(self, keys: Sequence[str]) -> dict[str, bytes]: + async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, bytes]: """Asynchronously get the cached values for the given keys.""" return await asyncio.to_thread(self.get, keys) - def set(self, mapping: Mapping[str, tuple[bytes, int | None]]) -> None: + def set(self, mapping: Mapping[FullKey, tuple[bytes, int | None]]) -> None: """Set the cached values for the given keys and TTLs.""" with self._lock, self._conn: now = datetime.datetime.now(datetime.timezone.utc) @@ -78,26 +87,27 @@ class FileCache(BaseCache): expiry = None encoding, raw = self.serde.dumps_typed(value) self._conn.execute( - "INSERT OR REPLACE INTO cache (key, expiry, encoding, val) VALUES (?, ?, ?, ?)", - (key, expiry, encoding, raw), + "INSERT OR REPLACE INTO cache (ns, key, expiry, encoding, val) VALUES (?, ?, ?, ?, ?)", + (",".join(key[0]), key[1], expiry, encoding, raw), ) - async def aset(self, mapping: Mapping[str, tuple[bytes, int | None]]) -> None: + async def aset(self, mapping: Mapping[FullKey, tuple[bytes, int | None]]) -> None: """Asynchronously set the cached values for the given keys and TTLs.""" await asyncio.to_thread(self.set, mapping) - def delete(self, keys: Sequence[str]) -> None: - """Delete the cached values for the given keys.""" + def delete(self, keys: Sequence[Namespace]) -> None: + """Delete the cached values for the given namespaces.""" if not keys: return with self._lock, self._conn: placeholders = ",".join("?" for _ in keys) self._conn.execute( - f"DELETE FROM cache WHERE key IN ({placeholders})", tuple(keys) + f"DELETE FROM cache WHERE (ns) IN ({placeholders})", + tuple(",".join(key) for key in keys), ) - async def adelete(self, keys: Sequence[str]) -> None: - """Asynchronously delete the cached values for the given keys.""" + async def adelete(self, keys: Sequence[Namespace]) -> None: + """Asynchronously delete the cached values for the given namespaces.""" await asyncio.to_thread(self.delete, keys) def __del__(self) -> None: diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index c3ca2821b..e361c8f1a 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -46,6 +46,10 @@ TASKS = sys.intern("__pregel_tasks") RETURN = sys.intern("__return__") # for writes of a task where we simply record the return value +# --- Reserved cache namespaces --- +CACHE_NS_WRITES = sys.intern("__pregel_ns_writes") +# cache namespace for node writes + # --- Reserved config.configurable keys --- CONFIG_KEY_SEND = sys.intern("__pregel_send") # holds the `write` function that accepts writes to state/edges/reserved keys diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 4d7d3075c..a6b9ca52e 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -33,6 +33,7 @@ from langgraph.checkpoint.base import ( V, ) from langgraph.constants import ( + CACHE_NS_WRITES, CONF, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_MAP, @@ -79,7 +80,7 @@ from langgraph.types import ( PregelTask, RetryPolicy, ) -from langgraph.utils.config import merge_configs, patch_config, recast_checkpoint_ns +from langgraph.utils.config import merge_configs, patch_config GetNextVersion = Callable[[Optional[V], BaseChannel], V] SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) @@ -601,19 +602,12 @@ def prepare_single_task( if cache_policy: args_key = cache_policy.key_func(*call.input[0], **call.input[1]) cache_key: Optional[CacheKey] = CacheKey( + ( + CACHE_NS_WRITES, + (identifier(call.func) or "__dynamic__"), + ), xxh3_128_hexdigest( - b"".join( - ( - b"__pregel_cache", - recast_checkpoint_ns(parent_ns).encode() - if parent_ns - else b"", - (identifier(call.func) or "__dynamic__").encode(), - args_key.encode() - if isinstance(args_key, str) - else args_key, - ) - ) + args_key.encode() if isinstance(args_key, str) else args_key, ), cache_policy.ttl, cache_policy.refresh, @@ -731,19 +725,13 @@ def prepare_single_task( if cache_policy: args_key = cache_policy.key_func(packet.arg) cache_key = CacheKey( + ( + CACHE_NS_WRITES, + (identifier(proc) or "__dynamic__"), + packet.node, + ), xxh3_128_hexdigest( - b"".join( - ( - b"__pregel_cache", - recast_checkpoint_ns(parent_ns).encode() - if parent_ns - else b"", - packet.node.encode(), - args_key.encode() - if isinstance(args_key, str) - else args_key, - ) - ) + args_key.encode() if isinstance(args_key, str) else args_key, ), cache_policy.ttl, cache_policy.refresh, @@ -870,19 +858,15 @@ def prepare_single_task( if cache_policy: args_key = cache_policy.key_func(val) cache_key = CacheKey( + ( + CACHE_NS_WRITES, + (identifier(proc) or "__dynamic__"), + name, + ), xxh3_128_hexdigest( - b"".join( - ( - b"__pregel_cache", - recast_checkpoint_ns(parent_ns).encode() - if parent_ns - else b"", - name.encode(), - args_key.encode() - if isinstance(args_key, str) - else args_key, - ) - ) + args_key.encode() + if isinstance(args_key, str) + else args_key, ), cache_policy.ttl, cache_policy.refresh, diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index d9abc6cf3..d18d09b4b 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -75,6 +75,16 @@ def _whichmodule(obj: Any, name: str) -> Optional[str]: def identifier(obj: Any, name: Optional[str] = None) -> Optional[str]: + """Return the module and name of an object.""" + from langgraph.pregel.read import PregelNode + from langgraph.utils.runnable import RunnableCallable, RunnableSeq + + if isinstance(obj, PregelNode): + obj = obj.bound + if isinstance(obj, RunnableSeq): + obj = obj.steps[0] + if isinstance(obj, RunnableCallable): + obj = obj.func if name is None: name = getattr(obj, "__qualname__", None) if name is None: # pragma: no cover diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 00feb2650..2c3e26259 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1069,7 +1069,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): return () matched: list[PregelExecutableTask] = [] if cached := { - t.cache_key.key: t + (t.cache_key.ns, t.cache_key.key): t for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: @@ -1089,7 +1089,12 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): return self.submit( self.cache.set, - {task.cache_key.key: (task.writes, task.cache_key.ttl)}, + { + (task.cache_key.ns, task.cache_key.key): ( + task.writes, + task.cache_key.ttl, + ) + }, ) # context manager @@ -1253,7 +1258,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): return [] matched: list[PregelExecutableTask] = [] if cached := { - t.cache_key.key: t + (t.cache_key.ns, t.cache_key.key): t for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: @@ -1276,7 +1281,12 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): return self.submit( self.cache.aset, - {task.cache_key.key: (task.writes, task.cache_key.ttl)}, + { + (task.cache_key.ns, task.cache_key.key): ( + task.writes, + task.cache_key.ttl, + ) + }, ) # context manager diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index eba426deb..c77a2bc3b 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -187,6 +187,8 @@ else: class CacheKey(NamedTuple): """Cache key for a task.""" + ns: tuple[str, ...] + """Namespace for the cache entry.""" key: str """Key for the cache entry.""" ttl: Optional[int] From d5d6fc0fee4ac4f948d2272f3d7b3e3d302bcb20 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 15:55:15 -0700 Subject: [PATCH 32/41] Remove refresh --- libs/langgraph/langgraph/pregel/algo.py | 3 --- libs/langgraph/langgraph/pregel/loop.py | 4 ++-- libs/langgraph/langgraph/types.py | 5 ----- 3 files changed, 2 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index a6b9ca52e..33c3bf2bb 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -610,7 +610,6 @@ def prepare_single_task( args_key.encode() if isinstance(args_key, str) else args_key, ), cache_policy.ttl, - cache_policy.refresh, ) else: cache_key = None @@ -734,7 +733,6 @@ def prepare_single_task( args_key.encode() if isinstance(args_key, str) else args_key, ), cache_policy.ttl, - cache_policy.refresh, ) else: cache_key = None @@ -869,7 +867,6 @@ def prepare_single_task( else args_key, ), cache_policy.ttl, - cache_policy.refresh, ) else: cache_key = None diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 2c3e26259..204dbdd3e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1071,7 +1071,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): if cached := { (t.cache_key.ns, t.cache_key.key): t for t in self.tasks.values() - if t.cache_key and not t.cache_key.refresh and not t.writes + if t.cache_key and not t.writes }: for key, values in self.cache.get(tuple(cached)).items(): task = cached[key] @@ -1260,7 +1260,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): if cached := { (t.cache_key.ns, t.cache_key.key): t for t in self.tasks.values() - if t.cache_key and not t.cache_key.refresh and not t.writes + if t.cache_key and not t.writes }: for key, values in (await self.cache.aget(tuple(cached))).items(): task = cached[key] diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index c77a2bc3b..06bc20cfd 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -137,9 +137,6 @@ class CachePolicy(Generic[KeyFuncT]): ttl: Optional[int] = None """Time to live for the cache entry in seconds. If None, the entry never expires.""" - refresh: bool = False - """Whether to force a refresh of the cache entry when it is accessed.""" - @dataclasses.dataclass(**_DC_KWARGS) class Interrupt: @@ -193,8 +190,6 @@ class CacheKey(NamedTuple): """Key for the cache entry.""" ttl: Optional[int] """Time to live for the cache entry in seconds.""" - refresh: bool - """Whether to force a refresh of the cache entry when it is accessed.""" @dataclasses.dataclass(**_T_DC_KWARGS) From 6e0041529e44e547fa954a486cd409ff703acd1d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 16:25:59 -0700 Subject: [PATCH 33/41] Add clear cache methods --- libs/langgraph/langgraph/func/__init__.py | 63 ++++++++++++++------- libs/langgraph/langgraph/pregel/__init__.py | 40 +++++++++++++ libs/langgraph/tests/test_pregel.py | 53 ++++++++++++++--- libs/langgraph/tests/test_pregel_async.py | 43 ++++++++++++++ 4 files changed, 170 insertions(+), 29 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 06aac25c2..de0e25648 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -20,7 +20,7 @@ from langgraph.cache.base import BaseCache from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import END, PREVIOUS, START +from langgraph.constants import CACHE_NS_WRITES, END, PREVIOUS, START from langgraph.pregel import Pregel from langgraph.pregel.call import ( P, @@ -28,6 +28,7 @@ from langgraph.pregel.call import ( T, call, get_runnable_for_entrypoint, + identifier, ) from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry @@ -35,6 +36,40 @@ from langgraph.store.base import BaseStore from langgraph.types import _DC_KWARGS, CachePolicy, RetryPolicy, StreamMode +class TaskFunction(Generic[P, T]): + def __init__( + self, + func: Callable[P, T], + *, + retry: Optional[Sequence[RetryPolicy]] = (), + cache_policy: Optional[CachePolicy[Callable[P, Union[str, bytes]]]] = None, + name: Optional[str] = None, + ) -> None: + self.func = func + self.retry = retry + self.cache_policy = cache_policy + functools.update_wrapper(self, func) + if name is not None: + setattr(self, "__name__", name) + + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> SyncAsyncFuture[T]: + return call( + self.func, retry=self.retry, cache_policy=self.cache_policy, *args, **kwargs + ) + + def clear_cache(self, cache: BaseCache) -> None: + """Clear the cache for this task.""" + if self.cache_policy is not None: + cache.delete(((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),)) + + async def aclear_cache(self, cache: BaseCache) -> None: + """Clear the cache for this task.""" + if self.cache_policy is not None: + await cache.adelete( + ((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),) + ) + + @overload def task( *, @@ -43,14 +78,14 @@ def task( cache_policy: Optional[CachePolicy[Callable[P, Union[str, bytes]]]] = None, ) -> Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], - Callable[P, SyncAsyncFuture[T]], + TaskFunction[P, T], ]: ... @overload def task( __func_or_none__: Union[Callable[P, Awaitable[T]], Callable[P, T]], -) -> Callable[P, SyncAsyncFuture[T]]: ... +) -> TaskFunction[P, T]: ... def task( @@ -62,9 +97,9 @@ def task( ) -> Union[ Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], - Callable[P, SyncAsyncFuture[T]], + TaskFunction[P, T], ], - Callable[P, SyncAsyncFuture[T]], + TaskFunction[P, T], ]: """Define a LangGraph task using the `task` decorator. @@ -132,23 +167,9 @@ def task( ) -> Union[ Callable[P, concurrent.futures.Future[T]], Callable[P, asyncio.Future[T]] ]: - if name is not None: - if hasattr(func, "__func__"): - # handle class methods - # NOTE: we're modifying the instance method to avoid modifying - # the original class method in case it's shared across multiple tasks - instance_method = functools.partial(func.__func__, func.__self__) # type: ignore [union-attr] - instance_method.__name__ = name # type: ignore [attr-defined] - func = instance_method - else: - # handle regular functions / partials / callable classes, etc. - func.__name__ = name - - call_func = functools.partial( - call, func, retry=retry_policies, cache_policy=cache_policy + return TaskFunction( + func, retry=retry_policies, cache_policy=cache_policy, name=name ) - object.__setattr__(call_func, "_is_pregel_task", True) - return functools.update_wrapper(call_func, func) if __func_or_none__ is not None: return decorator(__func_or_none__) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f59fdd951..286a98ddc 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -47,6 +47,7 @@ from langgraph.checkpoint.base import ( copy_checkpoint, ) from langgraph.constants import ( + CACHE_NS_WRITES, CONF, CONFIG_KEY_CACHE, CONFIG_KEY_CHECKPOINT_DURING, @@ -87,6 +88,7 @@ from langgraph.pregel.algo import ( local_write, prepare_next_tasks, ) +from langgraph.pregel.call import identifier from langgraph.pregel.checkpoint import create_checkpoint, empty_checkpoint from langgraph.pregel.debug import tasks_w_writes from langgraph.pregel.draw import draw_graph @@ -2989,6 +2991,44 @@ class Pregel(PregelProtocol): else: return chunks + def clear_cache(self, nodes: Sequence[str] | None = None) -> None: + """Clear the cache for the given nodes.""" + if not self.cache: + raise ValueError("No cache is set for this graph. Cannot clear cache.") + nodes = nodes or self.nodes.keys() + # collect namespaces to clear + namespaces: list[tuple[str, ...]] = [] + for node in nodes: + if node in self.nodes: + namespaces.append( + ( + CACHE_NS_WRITES, + (identifier(self.nodes[node]) or "__dynamic__"), + node, + ), + ) + # clear cache + self.cache.delete(namespaces) + + async def aclear_cache(self, nodes: Sequence[str] | None = None) -> None: + """Asynchronously clear the cache for the given nodes.""" + if not self.cache: + raise ValueError("No cache is set for this graph. Cannot clear cache.") + nodes = nodes or self.nodes.keys() + # collect namespaces to clear + namespaces: list[tuple[str, ...]] = [] + for node in nodes: + if node in self.nodes: + namespaces.append( + ( + CACHE_NS_WRITES, + (identifier(self.nodes[node]) or "__dynamic__"), + node, + ), + ) + # clear cache + await self.cache.adelete(namespaces) + def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str]]: """Index from a trigger to nodes that depend on it.""" diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 53233a0ce..bca22ceb2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -14,14 +14,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, field from random import randrange -from typing import ( - Annotated, - Any, - Literal, - Optional, - Union, - get_type_hints, -) +from typing import Annotated, Any, Literal, Optional, Union, get_type_hints import httpx import pytest @@ -3680,6 +3673,17 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple( ] assert rewrite_query_count == 2 if with_cache else 4 + # clear the cache + if with_cache: + app.clear_cache() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + assert rewrite_query_count == 4 + def test_callable_in_conditional_edges_with_no_path_map() -> None: class State(TypedDict, total=False): @@ -6663,6 +6667,39 @@ def test_multiple_interrupts_functional_cache( } assert counter == 3 + # should all be cached now + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + graph.invoke({}, configurable) + graph.invoke(Command(resume="a"), configurable) + graph.invoke(Command(resume="b"), configurable) + graph.invoke(Command(resume="c"), configurable) + graph.invoke(Command(resume="d"), configurable) + graph.invoke(Command(resume="e"), configurable) + result = graph.invoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 + + # clear cache + double.clear_cache(file_cache) + + # should recompute now + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + graph.invoke({}, configurable) + graph.invoke(Command(resume="a"), configurable) + graph.invoke(Command(resume="b"), configurable) + graph.invoke(Command(resume="c"), configurable) + graph.invoke(Command(resume="d"), configurable) + graph.invoke(Command(resume="e"), configurable) + result = graph.invoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 6 + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_double_interrupt_subgraph( diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index c43f4cdac..63c0a4b3e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5505,6 +5505,17 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple( ] assert rewrite_query_count == 2 if with_cache else 4 + # clear the cache + if with_cache: + await app.aclear_cache() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + assert rewrite_query_count == 4 + async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None: def sorted_add( @@ -7578,6 +7589,38 @@ async def test_multiple_interrupts_functional_cache( } assert counter == 3 + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({}, configurable) + await graph.ainvoke(Command(resume="a"), configurable) + await graph.ainvoke(Command(resume="b"), configurable) + await graph.ainvoke(Command(resume="c"), configurable) + await graph.ainvoke(Command(resume="d"), configurable) + await graph.ainvoke(Command(resume="e"), configurable) + result = await graph.ainvoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 + + # clear the cache + await double.aclear_cache(file_cache) + + # now should recompute + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({}, configurable) + await graph.ainvoke(Command(resume="a"), configurable) + await graph.ainvoke(Command(resume="b"), configurable) + await graph.ainvoke(Command(resume="c"), configurable) + await graph.ainvoke(Command(resume="d"), configurable) + await graph.ainvoke(Command(resume="e"), configurable) + result = await graph.ainvoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 6 + @NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) From 1a6395fd07f45e06526cc5e47e97b6be09ca7fcd Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 16:48:04 -0700 Subject: [PATCH 34/41] Move FileCache to sqlite package, add InMemoryCache --- .../langgraph/cache/sqlite/__init__.py | 118 ++++++++++++++++++ .../langgraph/cache/memory/__init__.py | 69 ++++++++++ libs/langgraph/tests/conftest.py | 25 ++-- libs/langgraph/tests/test_pregel.py | 10 +- libs/langgraph/tests/test_pregel_async.py | 10 +- 5 files changed, 207 insertions(+), 25 deletions(-) create mode 100644 libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py create mode 100644 libs/checkpoint/langgraph/cache/memory/__init__.py diff --git a/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py new file mode 100644 index 000000000..780e88282 --- /dev/null +++ b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import asyncio +import datetime +import sqlite3 +import threading +from collections.abc import Mapping, Sequence +from typing import Generic + +from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT +from langgraph.checkpoint.serde.base import SerializerProtocol + + +class SqliteCache(BaseCache[ValueT], Generic[ValueT]): + """File-based cache using SQLite.""" + + def __init__( + self, + *, + path: str, + serde: SerializerProtocol | None = None, + ) -> None: + """Initialize the cache with a file path.""" + super().__init__(serde=serde) + # SQLite backing store + self._conn = sqlite3.connect( + path, + check_same_thread=False, + ) + # Serialize access to the shared connection across threads + self._lock = threading.RLock() + # Better concurrency & atomicity + self._conn.execute("PRAGMA journal_mode=WAL;") + # Schema: key -> (expiry, encoding, value) + self._conn.execute( + """CREATE TABLE IF NOT EXISTS cache ( + ns TEXT, + key TEXT, + expiry REAL, + encoding TEXT NOT NULL, + val BLOB NOT NULL, + PRIMARY KEY (ns, key) + )""" + ) + self._conn.commit() + + def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]: + """Get the cached values for the given keys.""" + with self._lock, self._conn: + now = datetime.datetime.now(datetime.timezone.utc).timestamp() + if not keys: + return {} + placeholders = ",".join("(?, ?)" for _ in keys) + params: list[str] = [] + for ns_tuple, key in keys: + params.extend((",".join(ns_tuple), key)) + cursor = self._conn.execute( + f"SELECT ns, key, expiry, encoding, val FROM cache WHERE (ns, key) IN ({placeholders})", + tuple(params), + ) + values: dict[FullKey, ValueT] = {} + rows = cursor.fetchall() + for ns, key, expiry, encoding, raw in rows: + if expiry is not None and now > expiry: + # purge expired entry + self._conn.execute( + "DELETE FROM cache WHERE (ns, key) = (?, ?)", (ns, key) + ) + continue + values[(tuple(ns.split(",")), key)] = self.serde.loads_typed( + (encoding, raw) + ) + return values + + async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]: + """Asynchronously get the cached values for the given keys.""" + return await asyncio.to_thread(self.get, keys) + + def set(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None: + """Set the cached values for the given keys and TTLs.""" + with self._lock, self._conn: + now = datetime.datetime.now(datetime.timezone.utc) + for key, (value, ttl) in mapping.items(): + if ttl is not None: + delta = datetime.timedelta(seconds=ttl) + expiry: float | None = (now + delta).timestamp() + else: + expiry = None + encoding, raw = self.serde.dumps_typed(value) + self._conn.execute( + "INSERT OR REPLACE INTO cache (ns, key, expiry, encoding, val) VALUES (?, ?, ?, ?, ?)", + (",".join(key[0]), key[1], expiry, encoding, raw), + ) + + async def aset(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None: + """Asynchronously set the cached values for the given keys and TTLs.""" + await asyncio.to_thread(self.set, mapping) + + def delete(self, keys: Sequence[Namespace]) -> None: + """Delete the cached values for the given namespaces.""" + if not keys: + return + with self._lock, self._conn: + placeholders = ",".join("?" for _ in keys) + self._conn.execute( + f"DELETE FROM cache WHERE (ns) IN ({placeholders})", + tuple(",".join(key) for key in keys), + ) + + async def adelete(self, keys: Sequence[Namespace]) -> None: + """Asynchronously delete the cached values for the given namespaces.""" + await asyncio.to_thread(self.delete, keys) + + def __del__(self) -> None: + try: + self._conn.close() + except Exception: + pass diff --git a/libs/checkpoint/langgraph/cache/memory/__init__.py b/libs/checkpoint/langgraph/cache/memory/__init__.py new file mode 100644 index 000000000..d93744e80 --- /dev/null +++ b/libs/checkpoint/langgraph/cache/memory/__init__.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import datetime +import threading +from collections.abc import Sequence +from typing import Generic + +from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT +from langgraph.checkpoint.serde.base import SerializerProtocol + + +class InMemoryCache(BaseCache[ValueT], Generic[ValueT]): + def __init__(self, *, serde: SerializerProtocol | None = None): + super().__init__(serde=serde) + self._cache: dict[Namespace, dict[str, tuple[str, bytes, int | None]]] = {} + self._lock = threading.RLock() + + def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]: + """Get the cached values for the given keys.""" + with self._lock: + if not keys: + return {} + now = datetime.datetime.now(datetime.timezone.utc).timestamp() + values: dict[FullKey, ValueT] = {} + for ns_tuple, key in keys: + ns = Namespace(ns_tuple) + if ns in self._cache and key in self._cache[ns]: + enc, val, expiry = self._cache[ns][key] + if expiry is None or now < expiry: + values[(ns, key)] = self.serde.loads_typed((enc, val)) + else: + del self._cache[ns][key] + return values + + async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]: + """Asynchronously get the cached values for the given keys.""" + return self.get(keys) + + def set(self, keys: dict[FullKey, tuple[ValueT, int | None]]) -> None: + """Set the cached values for the given keys.""" + with self._lock: + now = datetime.datetime.now(datetime.timezone.utc).timestamp() + for (ns, key), (value, ttl) in keys.items(): + if ttl is not None: + delta = datetime.timedelta(seconds=ttl) + expiry: float | None = (now + delta).timestamp() + else: + expiry = None + if ns not in self._cache: + self._cache[ns] = {} + self._cache[ns][key] = ( + *self.serde.dumps_typed(value), + expiry, + ) + + async def aset(self, keys: dict[FullKey, tuple[ValueT, int | None]]) -> None: + """Asynchronously set the cached values for the given keys.""" + self.set(keys) + + def delete(self, keys: Sequence[Namespace]) -> None: + """Delete the cached values for the given namespaces.""" + with self._lock: + for ns in keys: + if ns in self._cache: + del self._cache[ns] + + async def adelete(self, keys: Sequence[Namespace]) -> None: + """Asynchronously delete the cached values for the given namespaces.""" + self.delete(keys) diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 6de17548a..7a4829e03 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,6 +1,4 @@ -import os import sys -import tempfile from collections.abc import AsyncIterator, Iterator from contextlib import asynccontextmanager from typing import Optional @@ -14,7 +12,8 @@ from psycopg_pool import AsyncConnectionPool, ConnectionPool from pytest_mock import MockerFixture from langgraph.cache.base import BaseCache -from langgraph.cache.file import FileCache +from langgraph.cache.memory import InMemoryCache +from langgraph.cache.sqlite import SqliteCache from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver from langgraph.checkpoint.postgres.aio import ( @@ -365,18 +364,14 @@ async def _store_postgres_aio_pool(): await conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def file_cache() -> Iterator[BaseCache]: - _, path = tempfile.mkstemp() - os.remove(path) - try: - yield FileCache(path=path) - finally: - # Cleanup the file - try: - os.remove(path) - except OSError: - pass +@pytest.fixture(scope="function", params=["sqlite", "memory"]) +def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]: + if request.param == "sqlite": + yield SqliteCache(path=":memory:") + elif request.param == "memory": + yield InMemoryCache() + else: + raise ValueError(f"Unknown cache type: {request.param}") @pytest.fixture(scope="function") diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index bca22ceb2..d26a6e408 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -3568,7 +3568,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( @pytest.mark.parametrize("with_cache", [True, False]) def test_in_one_fan_out_state_graph_waiting_edge_multiple( - with_cache: bool, file_cache: BaseCache + with_cache: bool, cache: BaseCache ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -3634,7 +3634,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple( workflow.add_conditional_edges("decider", decider_cond) workflow.set_finish_point("qa") - app = workflow.compile(cache=file_cache) + app = workflow.compile(cache=cache) assert app.invoke({"query": "what is weather in sf"}) == { "query": "analyzed: query: analyzed: query: what is weather in sf", @@ -6628,7 +6628,7 @@ def test_multiple_interrupts_functional( @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_interrupts_functional_cache( - request: pytest.FixtureRequest, checkpointer_name: str, file_cache: BaseCache + request: pytest.FixtureRequest, checkpointer_name: str, cache: BaseCache ): """Test multiple interrupts with functional API.""" checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @@ -6642,7 +6642,7 @@ def test_multiple_interrupts_functional_cache( counter += 1 return 2 * x - @entrypoint(checkpointer=checkpointer, cache=file_cache) + @entrypoint(checkpointer=checkpointer, cache=cache) def graph(state: dict) -> dict: """React tool.""" @@ -6683,7 +6683,7 @@ def test_multiple_interrupts_functional_cache( assert counter == 3 # clear cache - double.clear_cache(file_cache) + double.clear_cache(cache) # should recompute now configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 63c0a4b3e..0c815f49e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5399,7 +5399,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( @pytest.mark.parametrize("with_cache", [True, False]) async def test_in_one_fan_out_state_graph_waiting_edge_multiple( - with_cache: bool, file_cache: BaseCache + with_cache: bool, cache: BaseCache ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -5465,7 +5465,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple( workflow.add_conditional_edges("decider", decider_cond) workflow.set_finish_point("qa") - app = workflow.compile(cache=file_cache) + app = workflow.compile(cache=cache) assert await app.ainvoke({"query": "what is weather in sf"}) == { "query": "analyzed: query: analyzed: query: what is weather in sf", @@ -7551,7 +7551,7 @@ async def test_multiple_interrupts_functional(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_multiple_interrupts_functional_cache( - checkpointer_name: str, file_cache: BaseCache + checkpointer_name: str, cache: BaseCache ): """Test multiple interrupts with functional API.""" async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -7564,7 +7564,7 @@ async def test_multiple_interrupts_functional_cache( counter += 1 return 2 * x - @entrypoint(checkpointer=checkpointer, cache=file_cache) + @entrypoint(checkpointer=checkpointer, cache=cache) def graph(state: dict) -> dict: """React tool.""" @@ -7604,7 +7604,7 @@ async def test_multiple_interrupts_functional_cache( assert counter == 3 # clear the cache - await double.aclear_cache(file_cache) + await double.aclear_cache(cache) # now should recompute configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} From c937d5f04863fb0f7f8c417d8a57ccb3a0738bb0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 16:54:06 -0700 Subject: [PATCH 35/41] Lint --- .../langgraph/cache/file/__init__.py | 117 ------------------ libs/checkpoint/langgraph/cache/file/py.typed | 0 .../langgraph/cache/memory/__init__.py | 10 +- 3 files changed, 5 insertions(+), 122 deletions(-) delete mode 100644 libs/checkpoint/langgraph/cache/file/__init__.py delete mode 100644 libs/checkpoint/langgraph/cache/file/py.typed diff --git a/libs/checkpoint/langgraph/cache/file/__init__.py b/libs/checkpoint/langgraph/cache/file/__init__.py deleted file mode 100644 index 29231b206..000000000 --- a/libs/checkpoint/langgraph/cache/file/__init__.py +++ /dev/null @@ -1,117 +0,0 @@ -from __future__ import annotations - -import asyncio -import datetime -import sqlite3 -import threading -from collections.abc import Mapping, Sequence - -from langgraph.cache.base import BaseCache, FullKey, Namespace -from langgraph.checkpoint.serde.base import SerializerProtocol - - -class FileCache(BaseCache): - """File-based cache using SQLite.""" - - def __init__( - self, - *, - path: str, - serde: SerializerProtocol | None = None, - ) -> None: - """Initialize the cache with a file path.""" - super().__init__(serde=serde) - # SQLite backing store - self._conn = sqlite3.connect( - path, - check_same_thread=False, - ) - # Serialize access to the shared connection across threads - self._lock = threading.RLock() - # Better concurrency & atomicity - self._conn.execute("PRAGMA journal_mode=WAL;") - # Schema: key -> (expiry, encoding, value) - self._conn.execute( - """CREATE TABLE IF NOT EXISTS cache ( - ns TEXT, - key TEXT, - expiry REAL, - encoding TEXT NOT NULL, - val BLOB NOT NULL, - PRIMARY KEY (ns, key) - )""" - ) - self._conn.commit() - - def get(self, keys: Sequence[FullKey]) -> dict[FullKey, bytes]: - """Get the cached values for the given keys.""" - with self._lock, self._conn: - now = datetime.datetime.now(datetime.timezone.utc).timestamp() - if not keys: - return {} - placeholders = ",".join("(?, ?)" for _ in keys) - params: list[str] = [] - for ns_tuple, key in keys: - params.extend((",".join(ns_tuple), key)) - cursor = self._conn.execute( - f"SELECT ns, key, expiry, encoding, val FROM cache WHERE (ns, key) IN ({placeholders})", - tuple(params), - ) - values: dict[FullKey, bytes] = {} - rows = cursor.fetchall() - for ns, key, expiry, encoding, raw in rows: - if expiry is not None and now > expiry: - # purge expired entry - self._conn.execute( - "DELETE FROM cache WHERE (ns, key) = (?, ?)", (ns, key) - ) - continue - values[(tuple(ns.split(",")), key)] = self.serde.loads_typed( - (encoding, raw) - ) - return values - - async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, bytes]: - """Asynchronously get the cached values for the given keys.""" - return await asyncio.to_thread(self.get, keys) - - def set(self, mapping: Mapping[FullKey, tuple[bytes, int | None]]) -> None: - """Set the cached values for the given keys and TTLs.""" - with self._lock, self._conn: - now = datetime.datetime.now(datetime.timezone.utc) - for key, (value, ttl) in mapping.items(): - if ttl is not None: - delta = datetime.timedelta(seconds=ttl) - expiry: float | None = (now + delta).timestamp() - else: - expiry = None - encoding, raw = self.serde.dumps_typed(value) - self._conn.execute( - "INSERT OR REPLACE INTO cache (ns, key, expiry, encoding, val) VALUES (?, ?, ?, ?, ?)", - (",".join(key[0]), key[1], expiry, encoding, raw), - ) - - async def aset(self, mapping: Mapping[FullKey, tuple[bytes, int | None]]) -> None: - """Asynchronously set the cached values for the given keys and TTLs.""" - await asyncio.to_thread(self.set, mapping) - - def delete(self, keys: Sequence[Namespace]) -> None: - """Delete the cached values for the given namespaces.""" - if not keys: - return - with self._lock, self._conn: - placeholders = ",".join("?" for _ in keys) - self._conn.execute( - f"DELETE FROM cache WHERE (ns) IN ({placeholders})", - tuple(",".join(key) for key in keys), - ) - - async def adelete(self, keys: Sequence[Namespace]) -> None: - """Asynchronously delete the cached values for the given namespaces.""" - await asyncio.to_thread(self.delete, keys) - - def __del__(self) -> None: - try: - self._conn.close() - except Exception: - pass diff --git a/libs/checkpoint/langgraph/cache/file/py.typed b/libs/checkpoint/langgraph/cache/file/py.typed deleted file mode 100644 index e69de29bb..000000000 diff --git a/libs/checkpoint/langgraph/cache/memory/__init__.py b/libs/checkpoint/langgraph/cache/memory/__init__.py index d93744e80..c39d4b40f 100644 --- a/libs/checkpoint/langgraph/cache/memory/__init__.py +++ b/libs/checkpoint/langgraph/cache/memory/__init__.py @@ -2,7 +2,7 @@ from __future__ import annotations import datetime import threading -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Generic from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT @@ -12,7 +12,7 @@ from langgraph.checkpoint.serde.base import SerializerProtocol class InMemoryCache(BaseCache[ValueT], Generic[ValueT]): def __init__(self, *, serde: SerializerProtocol | None = None): super().__init__(serde=serde) - self._cache: dict[Namespace, dict[str, tuple[str, bytes, int | None]]] = {} + self._cache: dict[Namespace, dict[str, tuple[str, bytes, float | None]]] = {} self._lock = threading.RLock() def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]: @@ -36,10 +36,10 @@ class InMemoryCache(BaseCache[ValueT], Generic[ValueT]): """Asynchronously get the cached values for the given keys.""" return self.get(keys) - def set(self, keys: dict[FullKey, tuple[ValueT, int | None]]) -> None: + def set(self, keys: Mapping[FullKey, tuple[ValueT, int | None]]) -> None: """Set the cached values for the given keys.""" with self._lock: - now = datetime.datetime.now(datetime.timezone.utc).timestamp() + now = datetime.datetime.now(datetime.timezone.utc) for (ns, key), (value, ttl) in keys.items(): if ttl is not None: delta = datetime.timedelta(seconds=ttl) @@ -53,7 +53,7 @@ class InMemoryCache(BaseCache[ValueT], Generic[ValueT]): expiry, ) - async def aset(self, keys: dict[FullKey, tuple[ValueT, int | None]]) -> None: + async def aset(self, keys: Mapping[FullKey, tuple[ValueT, int | None]]) -> None: """Asynchronously set the cached values for the given keys.""" self.set(keys) From 83d2f93566d4212d96c222bd40fa29d3fb6eefb7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 16:57:21 -0700 Subject: [PATCH 36/41] Lint --- libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py | 3 +-- libs/checkpoint/langgraph/cache/memory/__init__.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py index 780e88282..9942a1228 100644 --- a/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py @@ -5,13 +5,12 @@ import datetime import sqlite3 import threading from collections.abc import Mapping, Sequence -from typing import Generic from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT from langgraph.checkpoint.serde.base import SerializerProtocol -class SqliteCache(BaseCache[ValueT], Generic[ValueT]): +class SqliteCache(BaseCache[ValueT]): """File-based cache using SQLite.""" def __init__( diff --git a/libs/checkpoint/langgraph/cache/memory/__init__.py b/libs/checkpoint/langgraph/cache/memory/__init__.py index c39d4b40f..a6d858e80 100644 --- a/libs/checkpoint/langgraph/cache/memory/__init__.py +++ b/libs/checkpoint/langgraph/cache/memory/__init__.py @@ -3,13 +3,12 @@ from __future__ import annotations import datetime import threading from collections.abc import Mapping, Sequence -from typing import Generic from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT from langgraph.checkpoint.serde.base import SerializerProtocol -class InMemoryCache(BaseCache[ValueT], Generic[ValueT]): +class InMemoryCache(BaseCache[ValueT]): def __init__(self, *, serde: SerializerProtocol | None = None): super().__init__(serde=serde) self._cache: dict[Namespace, dict[str, tuple[str, bytes, float | None]]] = {} From 7850c8d799b31d62a2d3f419e54e53a862ba2321 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 17:08:47 -0700 Subject: [PATCH 37/41] Lint --- libs/langgraph/langgraph/func/__init__.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index de0e25648..72bc42bd2 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -45,12 +45,21 @@ class TaskFunction(Generic[P, T]): cache_policy: Optional[CachePolicy[Callable[P, Union[str, bytes]]]] = None, name: Optional[str] = None, ) -> None: + if name is not None: + if hasattr(func, "__func__"): + # handle class methods + # NOTE: we're modifying the instance method to avoid modifying + # the original class method in case it's shared across multiple tasks + instance_method = functools.partial(func.__func__, func.__self__) # type: ignore [union-attr] + instance_method.__name__ = name # type: ignore [attr-defined] + func = instance_method + else: + # handle regular functions / partials / callable classes, etc. + func.__name__ = name self.func = func self.retry = retry self.cache_policy = cache_policy functools.update_wrapper(self, func) - if name is not None: - setattr(self, "__name__", name) def __call__(self, *args: P.args, **kwargs: P.kwargs) -> SyncAsyncFuture[T]: return call( From 6fc1df90138fc6cb4e5ad8bb61ab6a4368c4df15 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 8 May 2025 17:12:22 -0700 Subject: [PATCH 38/41] Lint --- libs/langgraph/langgraph/func/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 72bc42bd2..76599c671 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -50,7 +50,7 @@ class TaskFunction(Generic[P, T]): # handle class methods # NOTE: we're modifying the instance method to avoid modifying # the original class method in case it's shared across multiple tasks - instance_method = functools.partial(func.__func__, func.__self__) # type: ignore [union-attr] + instance_method = functools.partial(func.__func__, func.__self__) # type: ignore [attr-defined] instance_method.__name__ = name # type: ignore [attr-defined] func = instance_method else: From 331d5b07cebc9e0f271802f9436fcdf0f17c9c84 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 9 May 2025 12:00:08 -0700 Subject: [PATCH 39/41] Limit depth --- libs/langgraph/langgraph/utils/cache.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/utils/cache.py b/libs/langgraph/langgraph/utils/cache.py index 5555cf9bc..c9c451474 100644 --- a/libs/langgraph/langgraph/utils/cache.py +++ b/libs/langgraph/langgraph/utils/cache.py @@ -1,15 +1,18 @@ from __future__ import annotations -from collections.abc import Hashable +from collections.abc import Hashable, Mapping, Sequence from typing import Any -def _freeze(obj: Any) -> Hashable: - if isinstance(obj, dict): +def _freeze(obj: Any, depth: int = 10) -> Hashable: + if isinstance(obj, Hashable) or depth <= 0: + # already hashable, no need to freeze + return obj + elif isinstance(obj, Mapping): # sort keys so {"a":1,"b":2} == {"b":2,"a":1} - return tuple(sorted((k, _freeze(v)) for k, v in obj.items())) - elif isinstance(obj, (list, tuple, set, frozenset)): - return tuple(_freeze(x) for x in obj) + return tuple(sorted((k, _freeze(v, depth - 1)) for k, v in obj.items())) + elif isinstance(obj, Sequence): + return tuple(_freeze(x, depth - 1) for x in obj) # numpy / pandas etc. can provide their own .tobytes() elif hasattr(obj, "tobytes"): return ( From 898f266f725038aa2fcfb22e0a6f595f957c1b7e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 9 May 2025 12:05:26 -0700 Subject: [PATCH 40/41] Overload clear method to delete all when called without args --- .../langgraph/cache/sqlite/__init__.py | 27 ++++++++++--------- .../langgraph/cache/base/__init__.py | 10 ++++--- .../langgraph/cache/memory/__init__.py | 21 +++++++++------ libs/langgraph/langgraph/func/__init__.py | 4 +-- libs/langgraph/langgraph/pregel/__init__.py | 4 +-- 5 files changed, 38 insertions(+), 28 deletions(-) diff --git a/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py index 9942a1228..258006c54 100644 --- a/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py @@ -95,20 +95,23 @@ class SqliteCache(BaseCache[ValueT]): """Asynchronously set the cached values for the given keys and TTLs.""" await asyncio.to_thread(self.set, mapping) - def delete(self, keys: Sequence[Namespace]) -> None: - """Delete the cached values for the given namespaces.""" - if not keys: - return + def clear(self, namespaces: Sequence[Namespace] | None = None) -> None: + """Delete the cached values for the given namespaces. + If no namespaces are provided, clear all cached values.""" with self._lock, self._conn: - placeholders = ",".join("?" for _ in keys) - self._conn.execute( - f"DELETE FROM cache WHERE (ns) IN ({placeholders})", - tuple(",".join(key) for key in keys), - ) + if namespaces is None: + self._conn.execute("DELETE FROM cache") + else: + placeholders = ",".join("?" for _ in namespaces) + self._conn.execute( + f"DELETE FROM cache WHERE (ns) IN ({placeholders})", + tuple(",".join(key) for key in namespaces), + ) - async def adelete(self, keys: Sequence[Namespace]) -> None: - """Asynchronously delete the cached values for the given namespaces.""" - await asyncio.to_thread(self.delete, keys) + async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None: + """Asynchronously delete the cached values for the given namespaces. + If no namespaces are provided, clear all cached values.""" + await asyncio.to_thread(self.delete, namespaces) def __del__(self) -> None: try: diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py index 8f8660cca..c23e859fd 100644 --- a/libs/checkpoint/langgraph/cache/base/__init__.py +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -38,9 +38,11 @@ class BaseCache(ABC, Generic[ValueT]): """Asynchronously set the cached values for the given keys and TTLs.""" @abstractmethod - def delete(self, keys: Sequence[Namespace]) -> None: - """Delete the cached values for the given keys.""" + def clear(self, namespaces: Sequence[Namespace] | None = None) -> None: + """Delete the cached values for the given namespaces. + If no namespaces are provided, clear all cached values.""" @abstractmethod - async def adelete(self, keys: Sequence[Namespace]) -> None: - """Asynchronously delete the cached values for the given keys.""" + async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None: + """Asynchronously delete the cached values for the given namespaces. + If no namespaces are provided, clear all cached values.""" diff --git a/libs/checkpoint/langgraph/cache/memory/__init__.py b/libs/checkpoint/langgraph/cache/memory/__init__.py index a6d858e80..7b10db051 100644 --- a/libs/checkpoint/langgraph/cache/memory/__init__.py +++ b/libs/checkpoint/langgraph/cache/memory/__init__.py @@ -56,13 +56,18 @@ class InMemoryCache(BaseCache[ValueT]): """Asynchronously set the cached values for the given keys.""" self.set(keys) - def delete(self, keys: Sequence[Namespace]) -> None: - """Delete the cached values for the given namespaces.""" + def clear(self, namespaces: Sequence[Namespace] | None = None) -> None: + """Delete the cached values for the given namespaces. + If no namespaces are provided, clear all cached values.""" with self._lock: - for ns in keys: - if ns in self._cache: - del self._cache[ns] + if namespaces is None: + self._cache.clear() + else: + for ns in namespaces: + if ns in self._cache: + del self._cache[ns] - async def adelete(self, keys: Sequence[Namespace]) -> None: - """Asynchronously delete the cached values for the given namespaces.""" - self.delete(keys) + async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None: + """Asynchronously delete the cached values for the given namespaces. + If no namespaces are provided, clear all cached values.""" + self.clear(namespaces) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 76599c671..143b1a2de 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -69,12 +69,12 @@ class TaskFunction(Generic[P, T]): def clear_cache(self, cache: BaseCache) -> None: """Clear the cache for this task.""" if self.cache_policy is not None: - cache.delete(((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),)) + cache.clear(((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),)) async def aclear_cache(self, cache: BaseCache) -> None: """Clear the cache for this task.""" if self.cache_policy is not None: - await cache.adelete( + await cache.aclear( ((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),) ) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 286a98ddc..fe4f2c833 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -3008,7 +3008,7 @@ class Pregel(PregelProtocol): ), ) # clear cache - self.cache.delete(namespaces) + self.cache.clear(namespaces) async def aclear_cache(self, nodes: Sequence[str] | None = None) -> None: """Asynchronously clear the cache for the given nodes.""" @@ -3027,7 +3027,7 @@ class Pregel(PregelProtocol): ), ) # clear cache - await self.cache.adelete(namespaces) + await self.cache.aclear(namespaces) def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str]]: From 562d64bbb7fba7947eb6ae3a04d82470032e7602 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 9 May 2025 12:07:59 -0700 Subject: [PATCH 41/41] Lint --- libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py index 258006c54..cd327bd48 100644 --- a/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py @@ -111,7 +111,7 @@ class SqliteCache(BaseCache[ValueT]): async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None: """Asynchronously delete the cached values for the given namespaces. If no namespaces are provided, clear all cached values.""" - await asyncio.to_thread(self.delete, namespaces) + await asyncio.to_thread(self.clear, namespaces) def __del__(self) -> None: try: