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..cd327bd48 --- /dev/null +++ b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py @@ -0,0 +1,120 @@ +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, ValueT +from langgraph.checkpoint.serde.base import SerializerProtocol + + +class SqliteCache(BaseCache[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 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: + 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 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.clear, namespaces) + + def __del__(self) -> None: + try: + self._conn.close() + except Exception: + pass diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py new file mode 100644 index 000000000..c23e859fd --- /dev/null +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +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 + +ValueT = TypeVar("ValueT") +Namespace = tuple[str, ...] +FullKey = tuple[Namespace, str] + + +class BaseCache(ABC, Generic[ValueT]): + """Base class for a cache.""" + + serde: SerializerProtocol = JsonPlusSerializer(pickle_fallback=True) + + 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[FullKey]) -> dict[FullKey, ValueT]: + """Get the cached values for the given keys.""" + + @abstractmethod + async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]: + """Asynchronously get the cached values for the given keys.""" + + @abstractmethod + 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, pairs: Mapping[FullKey, tuple[ValueT, int | None]]) -> None: + """Asynchronously set the cached values for the given keys and TTLs.""" + + @abstractmethod + 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 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/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/memory/__init__.py b/libs/checkpoint/langgraph/cache/memory/__init__.py new file mode 100644 index 000000000..7b10db051 --- /dev/null +++ b/libs/checkpoint/langgraph/cache/memory/__init__.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import datetime +import threading +from collections.abc import Mapping, Sequence + +from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT +from langgraph.checkpoint.serde.base import SerializerProtocol + + +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]]] = {} + 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: 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) + 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: Mapping[FullKey, tuple[ValueT, int | None]]) -> None: + """Asynchronously set the cached values for the given keys.""" + self.set(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.""" + with self._lock: + if namespaces is None: + self._cache.clear() + else: + for ns in namespaces: + if ns in self._cache: + del self._cache[ns] + + 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/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..1a7752608 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/base.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/base.py @@ -1,7 +1,17 @@ +from __future__ import annotations + 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 +22,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 +44,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/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 3a9ab50e2..aef405a95 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 self.pickle_fallback and type_ == "pickle": + return pickle.loads(data_) else: raise NotImplementedError(f"Unknown serialization type: {type_}") diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 167c340f1..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 @@ -61,6 +65,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") diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index a951ef749..143b1a2de 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -16,10 +16,11 @@ 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 -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, @@ -27,11 +28,55 @@ 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 from langgraph.store.base import BaseStore -from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode +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: + 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 [attr-defined] + 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) + + 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.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.aclear( + ((CACHE_NS_WRITES, identifier(self.func) or "__dynamic__"),) + ) @overload @@ -39,16 +84,17 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, + 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( @@ -56,12 +102,13 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, + cache_policy: Optional[CachePolicy[Callable[P, Union[str, bytes]]]] = None, ) -> 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. @@ -129,21 +176,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) - object.__setattr__(call_func, "_is_pregel_task", True) - return functools.update_wrapper(call_func, func) + return TaskFunction( + func, retry=retry_policies, cache_policy=cache_policy, name=name + ) if __func_or_none__ is not None: return decorator(__func_or_none__) @@ -316,11 +351,17 @@ class entrypoint: self, checkpointer: Optional[BaseCheckpointSaver] = None, 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) @@ -450,5 +491,8 @@ class entrypoint: stream_eager=True, 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..f6a29198b 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 ( @@ -79,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 @@ -113,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 @@ -259,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. @@ -276,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.""" @@ -290,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. @@ -303,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. @@ -431,6 +437,7 @@ class StateGraph(Graph): metadata, input=input or self.schema, retry_policy=retry, + cache_policy=cache_policy, ends=ends, defer=defer, ) @@ -571,6 +578,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 +663,7 @@ class StateGraph(Graph): auto_validate=False, debug=debug, store=store, + cache=cache, name=name or "LangGraph", ) @@ -811,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 3abe2092f..fe4f2c833 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, ) @@ -46,7 +47,9 @@ from langgraph.checkpoint.base import ( copy_checkpoint, ) from langgraph.constants import ( + CACHE_NS_WRITES, CONF, + CONFIG_KEY_CACHE, CONFIG_KEY_CHECKPOINT_DURING, CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_NS, @@ -85,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 @@ -102,6 +106,7 @@ from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore from langgraph.types import ( All, + CachePolicy, Checkpointer, Interrupt, LoopProtocol, @@ -495,8 +500,15 @@ class Pregel(PregelProtocol): store: BaseStore | None = None """Memory store to use for SharedValues. Defaults to None.""" - retry_policy: Sequence[RetryPolicy] | None = None - """Retry policies to use when running tasks. Set to None to disable.""" + cache: BaseCache | None = None + """Cache to use for storing node results. Defaults to None.""" + + 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 @@ -525,7 +537,9 @@ class Pregel(PregelProtocol): debug: bool | None = None, checkpointer: BaseCheckpointSaver | None = None, store: BaseStore | None = None, - retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, + cache: BaseCache | 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, @@ -545,10 +559,11 @@ class Pregel(PregelProtocol): self.debug = debug if debug is not None else get_debug() self.checkpointer = checkpointer self.store = store - if isinstance(retry_policy, RetryPolicy): - self.retry_policy: Sequence[RetryPolicy] = (retry_policy,) - else: - self.retry_policy = retry_policy + self.cache = cache + 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 @@ -2193,6 +2208,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 +2241,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 +2253,7 @@ class Pregel(PregelProtocol): interrupt_after, checkpointer, store, + cache, ) def stream( @@ -2405,6 +2426,7 @@ class Pregel(PregelProtocol): interrupt_after_, checkpointer, store, + cache, ) = self._defaults( config, stream_mode=stream_mode, @@ -2436,6 +2458,7 @@ class Pregel(PregelProtocol): stream=StreamProtocol(stream.put, stream_modes), config=config, store=store, + cache=cache, checkpointer=checkpointer, nodes=self.nodes, specs=self.channels, @@ -2450,6 +2473,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( @@ -2494,11 +2519,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): + 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, + match_cached_writes=loop.match_cached_writes, ): # emit output yield from output() @@ -2710,6 +2737,7 @@ class Pregel(PregelProtocol): interrupt_after_, checkpointer, store, + cache, ) = self._defaults( config, stream_mode=stream_mode, @@ -2743,6 +2771,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, @@ -2757,6 +2786,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( @@ -2792,11 +2823,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): + 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, + match_cached_writes=loop.amatch_cached_writes, ): # emit output for o in output(): @@ -2958,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.clear(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.aclear(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/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c51afd7d3..33c3bf2bb 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, @@ -65,13 +66,15 @@ 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 from langgraph.store.base import BaseStore from langgraph.types import ( All, + CacheKey, + CachePolicy, PregelExecutableTask, PregelScratchpad, PregelTask, @@ -111,24 +114,27 @@ class PregelTaskWrites(NamedTuple): class Call: - __slots__ = ("func", "input", "retry", "callbacks") + __slots__ = ("func", "input", "retry", "cache_policy", "callbacks") func: Callable - input: Any + input: tuple[tuple[Any, ...], dict[str, Any]] retry: Optional[Sequence[RetryPolicy]] + cache_policy: Optional[CachePolicy] callbacks: Callbacks def __init__( self, func: Callable, - input: Any, + input: tuple[tuple[Any, ...], dict[str, Any]], *, retry: Optional[Sequence[RetryPolicy]], + cache_policy: Optional[CachePolicy], callbacks: Callbacks, ) -> None: self.func = func self.input = input self.retry = retry + self.cache_policy = cache_policy self.callbacks = callbacks @@ -382,6 +388,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]: ... @@ -401,6 +409,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]: ... @@ -419,6 +429,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. @@ -469,6 +481,8 @@ def prepare_next_tasks( checkpointer=checkpointer, manager=manager, input_cache=input_cache, + cache_policy=cache_policy, + retry_policy=retry_policy, ): tasks.append(task) @@ -512,6 +526,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} @@ -538,6 +554,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.""" @@ -580,6 +598,21 @@ 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() + cache_policy = call.cache_policy or cache_policy + 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( + args_key.encode() if isinstance(args_key, str) else args_key, + ), + cache_policy.ttl, + ) + else: + cache_key = None return PregelExecutableTask( name, call.input, @@ -624,8 +657,8 @@ def prepare_single_task( }, ), triggers, - call.retry, - None, + call.retry or retry_policy, + cache_key, task_id, task_path, ) @@ -649,6 +682,11 @@ 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 = ( @@ -679,73 +717,80 @@ 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() + cache_policy = proc.cache_policy or cache_policy + if cache_policy: + args_key = cache_policy.key_func(packet.arg) + cache_key = CacheKey( + ( + CACHE_NS_WRITES, + (identifier(proc) or "__dynamic__"), + packet.node, ), - triggers, - proc.retry_policy, - None, - task_id, - task_path, - writers=proc.flat_writers, - subgraphs=proc.subgraphs, + xxh3_128_hexdigest( + args_key.encode() if isinstance(args_key, str) else args_key, + ), + cache_policy.ttl, ) + else: + cache_key = None + 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 + ), + 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 or retry_policy, + cache_key, + 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: @@ -807,6 +852,24 @@ def prepare_single_task( if proc.metadata: metadata.update(proc.metadata) writes = deque() + cache_policy = proc.cache_policy or cache_policy + if cache_policy: + args_key = cache_policy.key_func(val) + cache_key = CacheKey( + ( + CACHE_NS_WRITES, + (identifier(proc) or "__dynamic__"), + name, + ), + xxh3_128_hexdigest( + args_key.encode() + if isinstance(args_key, str) + else args_key, + ), + cache_policy.ttl, + ) + else: + cache_key = None return PregelExecutableTask( name, val, @@ -867,8 +930,8 @@ def prepare_single_task( }, ), triggers, - proc.retry_policy, - None, + proc.retry_policy or retry_policy, + cache_key, task_id, task_path[:3], writers=proc.flat_writers, @@ -1000,7 +1063,8 @@ def _proc_input( val = channels[chan].get() break else: - val[k] = managed[k]() + val = managed[chan]() + break else: return MISSING else: @@ -1019,18 +1083,20 @@ def _proc_input( return val -def _uuid5_str(namespace: bytes, *parts: str) -> 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) - 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: 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() 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..d18d09b4b 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}") @@ -73,6 +74,35 @@ def _whichmodule(obj: Any, name: str) -> Optional[str]: return None +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 + # 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]]: @@ -135,7 +165,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 +190,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 +252,16 @@ def call( func: Callable[P, T], *args: Any, retry: Optional[Sequence[RetryPolicy]] = 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, callbacks=config["callbacks"]) + fut = impl( + func, + (args, kwargs), + retry=retry, + cache_policy=cache_policy, + callbacks=config["callbacks"], + ) return fut diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 805779d47..204dbdd3e 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, @@ -117,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, ) @@ -133,6 +136,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: @@ -147,6 +151,7 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol: class PregelLoop(LoopProtocol): input: Optional[Any] input_model: Optional[type[BaseModel]] + cache: Optional[BaseCache[WritesT]] checkpointer: Optional[BaseCheckpointSaver] nodes: Mapping[str, PregelNode] specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]] @@ -159,18 +164,18 @@ 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, 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, ], @@ -206,6 +211,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]], @@ -218,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__( @@ -230,6 +238,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 @@ -244,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]: @@ -299,7 +310,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 @@ -346,7 +357,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: @@ -418,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 @@ -547,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 = [] @@ -609,10 +624,16 @@ 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) -> Sequence[PregelExecutableTask]: + raise NotImplementedError + + async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]: + raise NotImplementedError + # private def _match_writes(self, tasks: Mapping[str, PregelExecutableTask]) -> None: @@ -912,8 +933,8 @@ class PregelLoop(LoopProtocol): for v in values(*args, **kwargs): self.stream((self.checkpoint_ns, mode, v)) - def _output_writes( - self, task_id: str, writes: Sequence[tuple[str, Any]], *, cached: bool = False + def output_writes( + 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( @@ -967,6 +988,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]], @@ -979,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__( @@ -987,6 +1011,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): stream=stream, config=config, checkpointer=checkpointer, + cache=cache, store=store, nodes=nodes, specs=specs, @@ -998,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() @@ -1037,6 +1064,39 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): return self.submit(cast(WritableManagedValue, managed_value).update, values) + def match_cached_writes(self) -> Sequence[PregelExecutableTask]: + if self.cache is None: + return () + matched: list[PregelExecutableTask] = [] + if cached := { + (t.cache_key.ns, 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(tuple(cached)).items(): + 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.""" + 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.ns, task.cache_key.key): ( + task.writes, + task.cache_key.ttl, + ) + }, + ) + # context manager def __enter__(self) -> Self: @@ -1117,6 +1177,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]], @@ -1129,6 +1190,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__( @@ -1137,6 +1200,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): stream=stream, config=config, checkpointer=checkpointer, + cache=cache, store=store, nodes=nodes, specs=specs, @@ -1148,6 +1212,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() @@ -1187,6 +1253,42 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): return self.submit(cast(WritableManagedValue, managed_value).aupdate, values) + async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]: + if self.cache is None: + return [] + matched: list[PregelExecutableTask] = [] + if cached := { + (t.cache_key.ns, 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(tuple(cached))).items(): + 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.""" + 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.ns, task.cache_key.key): ( + task.writes, + task.cache_key.ttl, + ) + }, + ) + # context manager async def __aenter__(self) -> Self: 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/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 78c2f8b0c..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, @@ -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 @@ -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 @@ -149,7 +157,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/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 6a5992d04..b0b2f31a9 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,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[[], Sequence[PregelExecutableTask]] + ] = None, ) -> Iterator[None]: tasks = tuple(tasks) futures = FuturesDict( @@ -147,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( @@ -160,6 +170,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, ), @@ -191,25 +202,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, - 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 @@ -266,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) @@ -277,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( @@ -292,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, @@ -324,33 +341,33 @@ 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, + match_cached_writes=match_cached_writes, + 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 @@ -515,6 +532,7 @@ def _call( input: Any, *, retry: Optional[Sequence[RetryPolicy]] = None, + cache_policy: Optional[CachePolicy] = None, callbacks: Callbacks = None, futures: weakref.ref[FuturesDict], schedule_task: weakref.ref[ @@ -522,6 +540,7 @@ def _call( [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] ] ], + match_cached_writes: Optional[Callable[[], Sequence[PregelExecutableTask]]], submit: weakref.ref[Submit], reraise: bool, ) -> concurrent.futures.Future[Any]: @@ -535,8 +554,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_policy=cache_policy, callbacks=callbacks), ): + if match_cached_writes: + match_cached_writes() if fut := next( ( f @@ -574,6 +595,7 @@ def _call( retry=retry, callbacks=callbacks, schedule_task=schedule_task, + match_cached_writes=match_cached_writes, submit=submit, reraise=reraise, ), @@ -596,6 +618,7 @@ def _acall( input: Any, *, retry: Optional[Sequence[RetryPolicy]] = None, + cache_policy: Optional[CachePolicy] = None, callbacks: Callbacks = None, # injected dependencies futures: weakref.ref[FuturesDict], @@ -604,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, @@ -616,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, callbacks=callbacks), + Call(func, input, retry=retry, cache_policy=cache_policy, callbacks=callbacks), ): if fut := next( ( @@ -652,6 +678,7 @@ def _acall( next_task, retry, stream=stream, + match_cached_writes=match_cached_writes, configurable={ CONFIG_KEY_CALL: partial( _acall, @@ -659,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/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3d34b35fd..06bc20cfd 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -22,6 +22,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: @@ -122,13 +123,19 @@ 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): - """Configuration for caching nodes. +KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., Union[str, bytes]]) - !!! version-added "Added in version 0.2.24." - """ - pass +@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. + 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) @@ -174,6 +181,17 @@ else: _T_DC_KWARGS = {"frozen": True} +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] + """Time to live for the cache entry in seconds.""" + + @dataclasses.dataclass(**_T_DC_KWARGS) class PregelExecutableTask: name: str @@ -182,8 +200,8 @@ class PregelExecutableTask: writes: deque[tuple[str, Any]] config: RunnableConfig triggers: Sequence[str] - retry_policy: Optional[Sequence[RetryPolicy]] - cache_policy: Optional[CachePolicy] + retry_policy: Sequence[RetryPolicy] + cache_key: Optional[CacheKey] id: str path: tuple[Union[str, int, tuple], ...] scheduled: bool = False @@ -313,7 +331,7 @@ class Command(Generic[N], ToolOutputMixin): graph: Optional[str] = None update: Optional[Any] = None - resume: Optional[Union[Any, dict[str, Any]]] = None + resume: Optional[Union[dict[str, Any], Any]] = None goto: Union[Send, Sequence[Union[Send, N]], N] = () def __repr__(self) -> str: diff --git a/libs/langgraph/langgraph/utils/cache.py b/libs/langgraph/langgraph/utils/cache.py new file mode 100644 index 000000000..c9c451474 --- /dev/null +++ b/libs/langgraph/langgraph/utils/cache.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from collections.abc import Hashable, Mapping, Sequence +from typing import Any + + +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, 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 ( + 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) -> str | 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) diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 5bfec63bd..7a4829e03 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,5 +1,5 @@ import sys -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Iterator from contextlib import asynccontextmanager from typing import Optional from uuid import UUID, uuid4 @@ -11,6 +11,9 @@ 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.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 ( @@ -361,6 +364,16 @@ async def _store_postgres_aio_pool(): await conn.execute(f"DROP DATABASE {database}") +@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") 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..d26a6e408 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 @@ -38,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 @@ -64,6 +58,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, @@ -3571,7 +3566,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, cache: BaseCache +) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3586,7 +3584,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: @@ -3613,7 +3615,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) @@ -3628,7 +3634,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=cache) assert app.invoke({"query": "what is weather in sf"}) == { "query": "analyzed: query: analyzed: query: what is weather in sf", @@ -3637,12 +3643,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" @@ -3653,6 +3671,18 @@ 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 + + # 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: @@ -6559,7 +6589,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 +6626,81 @@ 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, cache: BaseCache +): + """Test multiple interrupts with functional API.""" + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + 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=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 + + # 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(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( request: pytest.FixtureRequest, checkpointer_name: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 38c3ed52d..0c815f49e 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, 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=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,18 @@ 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 + + # 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: @@ -7510,6 +7548,80 @@ 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, 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=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 + + 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(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) async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: 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,