diff --git a/libs/checkpoint/langgraph/cache/base/__init__.py b/libs/checkpoint/langgraph/cache/base/__init__.py index 0dd5aca8e..e98525f81 100644 --- a/libs/checkpoint/langgraph/cache/base/__init__.py +++ b/libs/checkpoint/langgraph/cache/base/__init__.py @@ -5,10 +5,10 @@ from typing import Generic, Sequence, TypeVar from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -T = TypeVar("T") +ValueT = TypeVar("ValueT") -class BaseCache(ABC, Generic[T]): +class BaseCache(ABC, Generic[ValueT]): """Base class for a cache.""" serde: SerializerProtocol = JsonPlusSerializer(pickle_fallback=True) @@ -18,19 +18,19 @@ class BaseCache(ABC, Generic[T]): self.serde = serde or self.serde @abstractmethod - def get(self, keys: Sequence[str]) -> dict[str, T]: + def get(self, keys: Sequence[str]) -> dict[str, ValueT]: """Get the cached values for the given keys.""" @abstractmethod - async def aget(self, keys: Sequence[str]) -> dict[str, T]: + async def aget(self, keys: Sequence[str]) -> dict[str, ValueT]: """Asynchronously get the cached values for the given keys.""" @abstractmethod - def set(self, mapping: Mapping[str, tuple[T, int | None]]) -> None: + def set(self, mapping: Mapping[str, tuple[ValueT, int | None]]) -> None: """Set the cached values for the given keys and TTLs.""" @abstractmethod - async def aset(self, mapping: Mapping[str, tuple[T, int | None]]) -> None: + async def aset(self, mapping: Mapping[str, tuple[ValueT, int | None]]) -> None: """Asynchronously set the cached values for the given keys and TTLs.""" @abstractmethod 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/file/py.typed b/libs/checkpoint/langgraph/cache/file/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 00cb3220b..c0b15fd18 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -585,6 +585,28 @@ 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() + if call.cache_policy: + args_key = call.cache_policy.key(*call.input[0], **call.input[1]) + cache_key: Optional[CacheKey] = CacheKey( + 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(), + args_key.encode() + if isinstance(args_key, str) + else args_key, + ) + ) + ), + call.cache_policy.ttl, + call.cache_policy.refresh, + ) + else: + cache_key = None return PregelExecutableTask( name, call.input, @@ -630,24 +652,7 @@ def prepare_single_task( ), triggers, call.retry, - CacheKey( - 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_policy.key(*call.input[0], **call.input[1]), - ) - ) - ), - call.cache_policy.ttl, - call.cache_policy.refresh, - ) - if call.cache_policy - else None, + cache_key, task_id, task_path, ) @@ -709,6 +714,28 @@ def prepare_single_task( if proc.metadata: metadata.update(proc.metadata) writes = deque() + if proc.cache_policy: + args_key = proc.cache_policy.key(packet.arg) + cache_key = CacheKey( + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + packet.node.encode(), + args_key.encode() + if isinstance(args_key, str) + else args_key, + ) + ) + ), + proc.cache_policy.ttl, + proc.cache_policy.refresh, + ) + else: + cache_key = None return PregelExecutableTask( packet.node, packet.arg, @@ -758,24 +785,7 @@ def prepare_single_task( ), triggers, proc.retry_policy, - CacheKey( - 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, - proc.cache_policy.refresh, - ) - if proc.cache_policy - else None, + cache_key, task_id, task_path, writers=proc.flat_writers, @@ -842,6 +852,28 @@ def prepare_single_task( if proc.metadata: metadata.update(proc.metadata) writes = deque() + if proc.cache_policy: + args_key = proc.cache_policy.key(val) + cache_key = CacheKey( + xxh3_128_hexdigest( + b"".join( + ( + b"__pregel_cache", + recast_checkpoint_ns(parent_ns).encode() + if parent_ns + else b"", + name.encode(), + args_key.encode() + if isinstance(args_key, str) + else args_key, + ) + ) + ), + proc.cache_policy.ttl, + proc.cache_policy.refresh, + ) + else: + cache_key = None return PregelExecutableTask( name, val, @@ -903,24 +935,7 @@ def prepare_single_task( ), triggers, proc.retry_policy, - CacheKey( - 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, - proc.cache_policy.refresh, - ) - if proc.cache_policy - else None, + cache_key, task_id, task_path[:3], writers=proc.flat_writers, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7ab11dbb2..e46a7fcc6 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -134,6 +134,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: @@ -148,7 +149,7 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol: class PregelLoop(LoopProtocol): input: Optional[Any] input_model: Optional[type[BaseModel]] - cache: Optional[BaseCache[Sequence[tuple[str, Any]]]] + cache: Optional[BaseCache[WritesT]] checkpointer: Optional[BaseCheckpointSaver] nodes: Mapping[str, PregelNode] specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]] @@ -163,16 +164,14 @@ class PregelLoop(LoopProtocol): debug: bool 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, ], @@ -303,7 +302,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 @@ -923,7 +922,7 @@ class PregelLoop(LoopProtocol): self.stream((self.checkpoint_ns, mode, v)) def _output_writes( - self, task_id: str, writes: Sequence[tuple[str, Any]], *, cached: bool = False + 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( @@ -1057,10 +1056,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: - for key, values in self.cache.get(cached.keys()).items(): + for key, values in self.cache.get(tuple(cached)).items(): cached[key].writes.extend(values) - 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.""" super().put_writes(task_id, writes) if not writes or self.cache is None or not hasattr(self, "tasks"): @@ -1233,10 +1232,10 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): for t in self.tasks.values() if t.cache_key and not t.cache_key.refresh and not t.writes }: - for key, values in (await self.cache.aget(cached.keys())).items(): + for key, values in (await self.cache.aget(tuple(cached))).items(): cached[key].writes.extend(values) - 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.""" super().put_writes(task_id, writes) if not writes or self.cache is None or not hasattr(self, "tasks"): diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 8b42c0e10..9fbb92217 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -32,7 +32,7 @@ if TYPE_CHECKING: try: - from langchain_core.messages.tool import ToolOutputMixin # type: ignore + from langchain_core.messages.tool import ToolOutputMixin except ImportError: class ToolOutputMixin: # type: ignore[no-redef]