diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 3fdb4bec5..b8b87b516 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -65,7 +65,7 @@ from langgraph.constants import ( ) from langgraph.errors import InvalidUpdateError from langgraph.managed.base import ManagedValueMapping -from langgraph.pregel.call import get_runnable_for_task +from langgraph.pregel.call import get_runnable_for_task, identifier from langgraph.pregel.io import read_channels from langgraph.pregel.log import logger from langgraph.pregel.read import INPUT_CACHE_KEY_TYPE, PregelNode @@ -79,7 +79,7 @@ from langgraph.types import ( PregelTask, RetryPolicy, ) -from langgraph.utils.config import merge_configs, patch_config +from langgraph.utils.config import merge_configs, patch_config, recast_checkpoint_ns GetNextVersion = Callable[[Optional[V], BaseChannel], V] SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) @@ -631,7 +631,18 @@ def prepare_single_task( triggers, call.retry, CacheKey( - xxh3_128_hexdigest(call.cache.key(*call.input[0], **call.input[1])), + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + (identifier(call.func) or "__dynamic__").encode(), + call.cache.key(*call.input[0], **call.input[1]), + ) + ) + ), call.cache.ttl, ) if call.cache @@ -747,7 +758,18 @@ def prepare_single_task( triggers, proc.retry_policy, CacheKey( - xxh3_128_hexdigest(proc.cache_policy.key(packet.arg)), + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + packet.node.encode(), + proc.cache_policy.key(packet.arg), + ) + ) + ), proc.cache_policy.ttl, ) if proc.cache_policy @@ -880,7 +902,18 @@ def prepare_single_task( triggers, proc.retry_policy, CacheKey( - xxh3_128_hexdigest(proc.cache_policy.key(val)), + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + name.encode(), + proc.cache_policy.key(val), + ) + ) + ), proc.cache_policy.ttl, ) if proc.cache_policy diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index e3719d2b6..5458d3972 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -74,6 +74,25 @@ def _whichmodule(obj: Any, name: str) -> Optional[str]: return None +def identifier(obj: Any, name: Optional[str] = None) -> Optional[str]: + if name is None: + name = getattr(obj, "__qualname__", None) + if name is None: # pragma: no cover + # This used to be needed for Python 2.7 support but is probably not + # needed anymore. However we keep the __name__ introspection in case + # users of cloudpickle rely on this old behavior for unknown reasons. + name = getattr(obj, "__name__", None) + if name is None: + return None + + module_name = getattr(obj, "__module__", None) + if module_name is None: + # In this case, obj.__module__ is None. obj is thus treated as dynamic. + return None + + return f"{module_name}.{name}" + + def _lookup_module_and_qualname( obj: Any, name: Optional[str] = None ) -> Optional[tuple[types.ModuleType, str]]: diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 7902ebfbd..805523195 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -586,6 +586,7 @@ def _call( retry=retry, callbacks=callbacks, schedule_task=schedule_task, + match_cached_writes=match_cached_writes, submit=submit, reraise=reraise, ),