This commit is contained in:
Nuno Campos
2025-05-08 16:49:01 -07:00
parent 057da43cd0
commit 64491a2b29
6 changed files with 86 additions and 72 deletions
+6 -6
View File
@@ -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
View File
View File
+69 -54
View File
@@ -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,
+10 -11
View File
@@ -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"):
+1 -1
View File
@@ -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]