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

Re-do with separate cache interface

Browse Source
This commit is contained in:
Nuno Campos
2025-05-08 16:49:01 -07:00
parent a446f34ed9
commit 1aecde3cd8
14 changed files with 333 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint/langgraph/cache/base/__init__.py
Vendored
+42
View File
@@ -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."""
libs/checkpoint/langgraph/cache/file/__init__.py
Vendored
+53
View File
@@ -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)
libs/checkpoint/langgraph/checkpoint/base/__init__.py
-34
View File
@@ -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],
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+2 -1
View File
@@ -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[
libs/checkpoint/langgraph/checkpoint/serde/base.py
+13 -7
View File
@@ -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"):
libs/langgraph/langgraph/constants.py
+2 -1
View File
@@ -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,
libs/langgraph/langgraph/func/__init__.py
+3
View File
@@ -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,
)
libs/langgraph/langgraph/pregel/__init__.py
+21
View File
@@ -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():
libs/langgraph/langgraph/pregel/algo.py
+36 -37
View File
@@ -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,
libs/langgraph/langgraph/pregel/loop.py
+66 -4
View File
@@ -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:
libs/langgraph/langgraph/pregel/runner.py
+16 -3
View File
@@ -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(
(
libs/langgraph/langgraph/types.py
+15 -1
View File
@@ -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
libs/langgraph/tests/conftest.py
+19 -1
View File
@@ -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]}"
libs/langgraph/tests/test_pregel.py
+45 -1
View File
@@ -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