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