diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 9aa051f9a..b5f5a6a7e 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -254,6 +254,23 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError + def get_writes( + self, + task_ids: Sequence[str], + ) -> dict[str, list[tuple[str, Any]]]: + """Fetch writes associated with the given task IDs. + + Args: + task_ids (Sequence[str]): List of task IDs to fetch writes for. + + Returns: + dict[str, list[tuple[str, Any]]]: Mapping from task ID to list of writes. + + Raises: + NotImplementedError: Implement this method in your custom checkpoint saver. + """ + raise NotImplementedError + def list( self, config: Optional[RunnableConfig], @@ -358,6 +375,23 @@ class BaseCheckpointSaver(Generic[V]): """ raise NotImplementedError + async def aget_writes( + self, + task_ids: Sequence[str], + ) -> dict[str, list[tuple[str, Any]]]: + """Asynchronously fetch writes associated with the given task IDs. + + Args: + task_ids (Sequence[str]): List of task IDs to fetch writes for. + + Returns: + dict[str, list[tuple[str, Any]]]: Mapping from task ID to list of writes. + + Raises: + NotImplementedError: Implement this method in your custom checkpoint saver. + """ + raise NotImplementedError + async def alist( self, config: Optional[RunnableConfig], diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 167c340f1..bcb614ca4 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -106,6 +106,7 @@ NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000") # the task_id to use for writes that are not associated with a task CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map") # holds a mapping of task ns -> resume value for resuming tasks +TASK_CACHE_NAMESPACE = b"__pregel_task_cache" RESERVED = { TAG_HIDDEN, diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index a951ef749..a3afe33ae 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -31,7 +31,7 @@ from langgraph.pregel.call import ( from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode +from langgraph.types import _DC_KWARGS, CachePolicy, RetryPolicy, StreamMode @overload @@ -39,6 +39,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, + cache: Optional[CachePolicy[P]] = None, ) -> Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], Callable[P, SyncAsyncFuture[T]], @@ -56,6 +57,7 @@ def task( *, name: Optional[str] = None, retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None, + cache: Optional[CachePolicy[P]] = None, ) -> Union[ Callable[ [Union[Callable[P, Awaitable[T]], Callable[P, T]]], @@ -141,7 +143,7 @@ def task( # handle regular functions / partials / callable classes, etc. func.__name__ = name - call_func = functools.partial(call, func, retry=retry_policies) + call_func = functools.partial(call, func, retry=retry_policies, cache=cache) object.__setattr__(call_func, "_is_pregel_task", True) return functools.update_wrapper(call_func, func) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c51afd7d3..a546bc8b4 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -60,6 +60,7 @@ from langgraph.constants import ( RESUME, RETURN, TAG_HIDDEN, + TASK_CACHE_NAMESPACE, TASKS, Send, ) @@ -72,6 +73,7 @@ from langgraph.pregel.read import INPUT_CACHE_KEY_TYPE, PregelNode from langgraph.store.base import BaseStore from langgraph.types import ( All, + CachePolicy, PregelExecutableTask, PregelScratchpad, PregelTask, @@ -111,24 +113,27 @@ class PregelTaskWrites(NamedTuple): class Call: - __slots__ = ("func", "input", "retry", "callbacks") + __slots__ = ("func", "input", "retry", "cache", "callbacks") func: Callable - input: Any + input: tuple[tuple[Any, ...], dict[str, Any]] retry: Optional[Sequence[RetryPolicy]] + cache: Optional[CachePolicy] callbacks: Callbacks def __init__( self, func: Callable, - input: Any, + input: tuple[tuple[Any, ...], dict[str, Any]], *, retry: Optional[Sequence[RetryPolicy]], + cache: Optional[CachePolicy], callbacks: Callbacks, ) -> None: self.func = func self.input = input self.retry = retry + self.cache = cache self.callbacks = callbacks @@ -556,15 +561,23 @@ def prepare_single_task( # create task id triggers: Sequence[str] = PUSH_TRIGGER checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name - task_id = task_id_func( - checkpoint_id_bytes, - checkpoint_ns, - str(step), - name, - PUSH, - task_path_str(task_path[1]), - str(task_path[2]), - ) + if call.cache: + task_id = task_id_func( + TASK_CACHE_NAMESPACE, + name, + call.cache.key(*call.input[0], **call.input[1]), + # TODO add truncated iso timestamp here for ttl + ) + else: + task_id = task_id_func( + checkpoint_id_bytes, + checkpoint_ns, + str(step), + name, + PUSH, + task_path_str(task_path[1]), + str(task_path[2]), + ) task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" # we append True to the task path to indicate that a call is being # made, so we should not return interrupts from this task (responsibility lies with the parent) @@ -625,7 +638,7 @@ def prepare_single_task( ), triggers, call.retry, - None, + call.cache, task_id, task_path, ) @@ -649,19 +662,31 @@ def prepare_single_task( f"Ignoring unknown node name {packet.node} in pending sends" ) return + # find process + proc = processes[packet.node] + proc_node = proc.node + if proc_node is None: + return # create task id triggers = PUSH_TRIGGER checkpoint_ns = ( f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node ) - task_id = task_id_func( - checkpoint_id_bytes, - checkpoint_ns, - str(step), - packet.node, - PUSH, - str(idx), - ) + if proc.cache_policy: + task_id = task_id_func( + TASK_CACHE_NAMESPACE, + packet.node, + proc.cache_policy.key(packet.arg), + ) + else: + task_id = task_id_func( + checkpoint_id_bytes, + checkpoint_ns, + str(step), + packet.node, + PUSH, + str(idx), + ) else: logger.warning(f"Ignoring invalid PUSH task path {task_path}") return @@ -679,73 +704,64 @@ def prepare_single_task( if task_id_checksum is not None: assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}" if for_execution: - proc = processes[packet.node] - if node := proc.node: - if proc.metadata: - metadata.update(proc.metadata) - writes = deque() - return PregelExecutableTask( - packet.node, - packet.arg, - node, - writes, - patch_config( - merge_configs( - config, {"metadata": metadata, "tags": proc.tags} - ), - run_name=packet.node, - callbacks=( - manager.get_child(f"graph:step:{step}") if manager else None - ), - configurable={ - CONFIG_KEY_TASK_ID: task_id, - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - local_write, - writes.extend, - processes.keys(), - ), - CONFIG_KEY_READ: partial( - local_read, - channels, - managed, - PregelTaskWrites( - task_path, packet.node, writes, triggers - ), - ), - CONFIG_KEY_STORE: ( - store or configurable.get(CONFIG_KEY_STORE) - ), - CONFIG_KEY_CHECKPOINTER: ( - checkpointer - or configurable.get(CONFIG_KEY_CHECKPOINTER) - ), - CONFIG_KEY_CHECKPOINT_MAP: { - **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}), - parent_ns: checkpoint["id"], - }, - CONFIG_KEY_CHECKPOINT_ID: None, - CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, - CONFIG_KEY_SCRATCHPAD: _scratchpad( - config[CONF].get(CONFIG_KEY_SCRATCHPAD), - pending_writes, - task_id, - xxh3_128_hexdigest(task_checkpoint_ns.encode()), - config[CONF].get(CONFIG_KEY_RESUME_MAP), - ), - CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get( - PREVIOUS, None - ), - }, + if proc.metadata: + metadata.update(proc.metadata) + writes = deque() + return PregelExecutableTask( + packet.node, + packet.arg, + proc_node, + writes, + patch_config( + merge_configs(config, {"metadata": metadata, "tags": proc.tags}), + run_name=packet.node, + callbacks=( + manager.get_child(f"graph:step:{step}") if manager else None ), - triggers, - proc.retry_policy, - None, - task_id, - task_path, - writers=proc.flat_writers, - subgraphs=proc.subgraphs, - ) + configurable={ + CONFIG_KEY_TASK_ID: task_id, + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, + writes.extend, + processes.keys(), + ), + CONFIG_KEY_READ: partial( + local_read, + channels, + managed, + PregelTaskWrites(task_path, packet.node, writes, triggers), + ), + CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)), + CONFIG_KEY_CHECKPOINTER: ( + checkpointer or configurable.get(CONFIG_KEY_CHECKPOINTER) + ), + CONFIG_KEY_CHECKPOINT_MAP: { + **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}), + parent_ns: checkpoint["id"], + }, + CONFIG_KEY_CHECKPOINT_ID: None, + CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, + CONFIG_KEY_SCRATCHPAD: _scratchpad( + config[CONF].get(CONFIG_KEY_SCRATCHPAD), + pending_writes, + task_id, + xxh3_128_hexdigest(task_checkpoint_ns.encode()), + config[CONF].get(CONFIG_KEY_RESUME_MAP), + ), + CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get( + PREVIOUS, None + ), + }, + ), + triggers, + proc.retry_policy, + proc.cache_policy, + task_id, + task_path, + writers=proc.flat_writers, + subgraphs=proc.subgraphs, + ) else: return PregelTask(task_id, packet.node, task_path) elif task_path[0] == PULL: @@ -784,6 +800,7 @@ def prepare_single_task( # create task id checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name + # TODO implement cache task_id = task_id_func( checkpoint_id_bytes, checkpoint_ns, @@ -1000,7 +1017,8 @@ def _proc_input( val = channels[chan].get() break else: - val[k] = managed[k]() + val = managed[chan]() + break else: return MISSING else: @@ -1019,18 +1037,20 @@ def _proc_input( return val -def _uuid5_str(namespace: bytes, *parts: str) -> str: +def _uuid5_str(namespace: bytes, *parts: str | bytes) -> str: """Generate a UUID from the SHA-1 hash of a namespace and str parts.""" sha = sha1(namespace, usedforsecurity=False) - sha.update(b"".join(p.encode() for p in parts)) + sha.update(b"".join(p.encode() if isinstance(p, str) else p for p in parts)) hex = sha.hexdigest() return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}" -def _xxhash_str(namespace: bytes, *parts: str) -> str: +def _xxhash_str(namespace: bytes, *parts: str | bytes) -> str: """Generate a UUID from the XXH3 hash of a namespace and str parts.""" - hex = xxh3_128_hexdigest(namespace + b"".join(p.encode() for p in parts)) + hex = xxh3_128_hexdigest( + namespace + b"".join(p.encode() if isinstance(p, str) else p for p in parts) + ) return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}" diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index e8dde073a..e3719d2b6 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -13,7 +13,7 @@ from typing_extensions import ParamSpec from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry -from langgraph.types import RetryPolicy +from langgraph.types import CachePolicy, RetryPolicy from langgraph.utils.config import get_config from langgraph.utils.runnable import ( RunnableCallable, @@ -28,6 +28,7 @@ from langgraph.utils.runnable import ( def _getattribute(obj: Any, name: str) -> Any: + parent = None for subpath in name.split("."): if subpath == "": raise AttributeError(f"Can't get local attribute {name!r} on {obj!r}") @@ -135,7 +136,7 @@ def _explode_args_trace_inputs( return arguments -def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq: +def get_runnable_for_entrypoint(func: Callable[..., Any]) -> Runnable: key = (func, False) if key in CACHE: return CACHE[key] @@ -160,7 +161,7 @@ def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq: return CACHE.setdefault(key, run) -def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq: +def get_runnable_for_task(func: Callable[..., Any]) -> Runnable: key = (func, True) if key in CACHE: return CACHE[key] @@ -222,9 +223,12 @@ def call( func: Callable[P, T], *args: Any, retry: Optional[Sequence[RetryPolicy]] = None, + cache: Optional[CachePolicy] = None, **kwargs: Any, ) -> SyncAsyncFuture[T]: config = get_config() impl = config[CONF][CONFIG_KEY_CALL] - fut = impl(func, (args, kwargs), retry=retry, callbacks=config["callbacks"]) + fut = impl( + func, (args, kwargs), retry=retry, cache=cache, callbacks=config["callbacks"] + ) return fut diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 805779d47..34d1d9bb9 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -434,7 +434,10 @@ class PregelLoop(LoopProtocol): # save the new task self.tasks[pushed.id] = pushed # match any pending writes to the new task - if self.skip_done_tasks: + if call.cache: + pass + # how to call async method here... + elif self.skip_done_tasks: self._match_writes({pushed.id: pushed}) # return the new task, to be started if not run before return pushed diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 7eb2a530c..028b08437 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -14,14 +14,15 @@ from langchain_core.runnables import ( RunnablePassthrough, RunnableSerializable, ) -from langchain_core.runnables.base import Input, Other, coerce_to_runnable -from langchain_core.runnables.utils import ConfigurableFieldSpec +from langchain_core.runnables.base import Other, coerce_to_runnable +from langchain_core.runnables.utils import ConfigurableFieldSpec, Input from langgraph.constants import CONF, CONFIG_KEY_READ from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.utils import find_subgraph_pregel from langgraph.pregel.write import ChannelWrite +from langgraph.types import CachePolicy from langgraph.utils.config import merge_configs from langgraph.utils.runnable import RunnableCallable, RunnableSeq @@ -143,6 +144,9 @@ class PregelNode(Runnable): retry_policy: Sequence[RetryPolicy] | None """The retry policies to use when invoking the node.""" + cache_policy: CachePolicy | None + """The cache policy to use when invoking the node.""" + tags: Sequence[str] | None """Tags to attach to the node for tracing.""" @@ -163,6 +167,7 @@ class PregelNode(Runnable): metadata: Mapping[str, Any] | None = None, bound: Runnable[Any, Any] | None = None, retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None, + cache_policy: CachePolicy | None = None, subgraphs: Sequence[PregelProtocol] | None = None, ) -> None: self.channels = channels @@ -170,8 +175,9 @@ class PregelNode(Runnable): self.mapper = mapper self.writers = writers or [] self.bound = bound if bound is not None else DEFAULT_BOUND + self.cache_policy = cache_policy if isinstance(retry_policy, RetryPolicy): - self.retry_policy: Sequence[RetryPolicy] = (retry_policy,) + self.retry_policy = (retry_policy,) else: self.retry_policy = retry_policy self.tags = tags diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3d34b35fd..51a3068ac 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -11,6 +11,7 @@ from typing import ( Literal, NamedTuple, Optional, + ParamSpec, TypeVar, Union, cast, @@ -22,6 +23,7 @@ from typing_extensions import Self from xxhash import xxh3_128_hexdigest from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata +from langgraph.utils.cache import default_cache_key from langgraph.utils.fields import get_update_as_tuples if TYPE_CHECKING: @@ -30,7 +32,7 @@ if TYPE_CHECKING: try: - from langchain_core.messages.tool import ToolOutputMixin + from langchain_core.messages.tool import ToolOutputMixin # type: ignore except ImportError: class ToolOutputMixin: # type: ignore[no-redef] @@ -122,13 +124,16 @@ class RetryPolicy(NamedTuple): """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" -class CachePolicy(NamedTuple): +P = ParamSpec("P") + + +class CachePolicy(NamedTuple, Generic[P]): """Configuration for caching nodes. !!! version-added "Added in version 0.2.24." """ - pass + key: Callable[P, str | bytes] = default_cache_key @dataclasses.dataclass(**_DC_KWARGS) diff --git a/libs/langgraph/langgraph/utils/cache.py b/libs/langgraph/langgraph/utils/cache.py new file mode 100644 index 000000000..ddae896c8 --- /dev/null +++ b/libs/langgraph/langgraph/utils/cache.py @@ -0,0 +1,26 @@ +from collections.abc import Hashable +from typing import Any + + +def _freeze(obj: Any) -> Hashable: + if isinstance(obj, dict): + # sort keys so {"a":1,"b":2} == {"b":2,"a":1} + return tuple(sorted((k, _freeze(v)) for k, v in obj.items())) + elif isinstance(obj, (list, tuple, set, frozenset)): + return tuple(_freeze(x) for x in obj) + # numpy / pandas etc. can provide their own .tobytes() + elif hasattr(obj, "tobytes"): + return ( + type(obj).__name__, + obj.tobytes(), + obj.shape if hasattr(obj, "shape") else None, + ) + return obj # strings, ints, dataclasses with frozen=True, etc. + + +def default_cache_key(*args: Any, **kwargs: Any) -> bytes: + """Default cache key function that uses the arguments and keyword arguments to generate a hashable key.""" + import pickle + + # protocol 5 strikes a good balance between speed and size + return pickle.dumps((_freeze(args), _freeze(kwargs)), protocol=5, fix_imports=False)