|
|
@@ -0,0 +1,303 @@
|
|
|
+"""极简 Redis 异步 checkpoint saver（不依赖 RediSearch / RedisJSON）。
|
|
|
+
|
|
|
+只使用 Redis String、Hash 与 SCAN，适配原生 Redis 5+。
|
|
|
+"""
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+from collections.abc import AsyncIterator, Sequence
|
|
|
+from types import TracebackType
|
|
|
+from typing import Any
|
|
|
+from urllib.parse import quote, unquote
|
|
|
+
|
|
|
+from langchain_core.runnables import RunnableConfig
|
|
|
+from langgraph.checkpoint.base import (
|
|
|
+ WRITES_IDX_MAP,
|
|
|
+ BaseCheckpointSaver,
|
|
|
+ ChannelVersions,
|
|
|
+ Checkpoint,
|
|
|
+ CheckpointMetadata,
|
|
|
+ CheckpointTuple,
|
|
|
+ get_checkpoint_id,
|
|
|
+ get_checkpoint_metadata,
|
|
|
+)
|
|
|
+from redis.asyncio import Redis
|
|
|
+
|
|
|
+
|
|
|
+class PlainRedisSaver(BaseCheckpointSaver[str]):
|
|
|
+ """原生 Redis 异步 checkpoint saver。
|
|
|
+
|
|
|
+ 每个 checkpoint 存为一个 Redis Hash，字段为 type/data（JsonPlusSerializer 的
|
|
|
+ 序列化结果），latest 指针存为普通 String。这样不要求 Redis 安装 RedisJSON 或
|
|
|
+ RediSearch 模块。
|
|
|
+ """
|
|
|
+
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ redis_url: str,
|
|
|
+ *,
|
|
|
+ key_prefix: str = "ka",
|
|
|
+ ttl_seconds: int | None = None,
|
|
|
+ serde=None,
|
|
|
+ ) -> None:
|
|
|
+ super().__init__(serde=serde)
|
|
|
+ self.redis_url = redis_url
|
|
|
+ self.key_prefix = key_prefix
|
|
|
+ self.ttl_seconds = ttl_seconds
|
|
|
+ self._redis: Redis | None = None
|
|
|
+
|
|
|
+ @property
|
|
|
+ def redis(self) -> Redis:
|
|
|
+ if self._redis is None:
|
|
|
+ raise RuntimeError("Redis client 尚未初始化")
|
|
|
+ return self._redis
|
|
|
+
|
|
|
+ async def asetup(self) -> "PlainRedisSaver":
|
|
|
+ self._redis = Redis.from_url(self.redis_url, decode_responses=False)
|
|
|
+ await self._redis.ping()
|
|
|
+ return self
|
|
|
+
|
|
|
+ async def __aenter__(self) -> "PlainRedisSaver":
|
|
|
+ await self.asetup()
|
|
|
+ return self
|
|
|
+
|
|
|
+ async def __aexit__(
|
|
|
+ self,
|
|
|
+ exc_type: type[BaseException] | None,
|
|
|
+ exc: BaseException | None,
|
|
|
+ tb: TracebackType | None,
|
|
|
+ ) -> None:
|
|
|
+ if self._redis is not None:
|
|
|
+ await self._redis.aclose()
|
|
|
+ self._redis = None
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def _part(value: str | int) -> str:
|
|
|
+ return quote(str(value), safe="")
|
|
|
+
|
|
|
+ def _threads_key(self) -> str:
|
|
|
+ return f"{self.key_prefix}:threads"
|
|
|
+
|
|
|
+ def _cp_key(self, thread_id: str, checkpoint_ns: str, checkpoint_id: str) -> str:
|
|
|
+ return (
|
|
|
+ f"{self.key_prefix}:cp:"
|
|
|
+ f"{self._part(thread_id)}:{self._part(checkpoint_ns)}:"
|
|
|
+ f"{self._part(checkpoint_id)}"
|
|
|
+ )
|
|
|
+
|
|
|
+ def _writes_key(self, thread_id: str, checkpoint_ns: str, checkpoint_id: str) -> str:
|
|
|
+ return (
|
|
|
+ f"{self.key_prefix}:writes:"
|
|
|
+ f"{self._part(thread_id)}:{self._part(checkpoint_ns)}:"
|
|
|
+ f"{self._part(checkpoint_id)}"
|
|
|
+ )
|
|
|
+
|
|
|
+ def _latest_key(self, thread_id: str, checkpoint_ns: str) -> str:
|
|
|
+ return f"{self.key_prefix}:latest:{self._part(thread_id)}:{self._part(checkpoint_ns)}"
|
|
|
+
|
|
|
+ async def _save_obj(self, key: str, obj: Any) -> None:
|
|
|
+ type_, data = self.serde.dumps_typed(obj)
|
|
|
+ await self.redis.hset(key, mapping={"type": type_, "data": data})
|
|
|
+ if self.ttl_seconds is not None:
|
|
|
+ await self.redis.expire(key, self.ttl_seconds)
|
|
|
+
|
|
|
+ async def _load_obj(self, key: str) -> Any:
|
|
|
+ pipe = self.redis.pipeline()
|
|
|
+ pipe.hget(key, "type")
|
|
|
+ pipe.hget(key, "data")
|
|
|
+ type_, data = await pipe.execute()
|
|
|
+ if type_ is None or data is None:
|
|
|
+ return None
|
|
|
+ return self.serde.loads_typed((type_.decode("utf-8"), data))
|
|
|
+
|
|
|
+ async def _get_writes(
|
|
|
+ self, thread_id: str, checkpoint_ns: str, checkpoint_id: str
|
|
|
+ ) -> list[tuple[str, str, Any]]:
|
|
|
+ stored = await self._load_obj(self._writes_key(thread_id, checkpoint_ns, checkpoint_id))
|
|
|
+ if not isinstance(stored, list):
|
|
|
+ return []
|
|
|
+ pending: list[tuple[str, str, Any]] = []
|
|
|
+ for row in stored:
|
|
|
+ if not isinstance(row, (list, tuple)) or len(row) != 4:
|
|
|
+ continue
|
|
|
+ task_id, channel, _idx, value = row
|
|
|
+ pending.append((task_id, channel, value))
|
|
|
+ return pending
|
|
|
+
|
|
|
+ async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
|
|
+ thread_id: str = config["configurable"]["thread_id"]
|
|
|
+ checkpoint_ns: str = config["configurable"].get("checkpoint_ns", "")
|
|
|
+ requested_id = get_checkpoint_id(config)
|
|
|
+ checkpoint_id = requested_id
|
|
|
+
|
|
|
+ if not checkpoint_id:
|
|
|
+ raw_latest = await self.redis.get(self._latest_key(thread_id, checkpoint_ns))
|
|
|
+ checkpoint_id = raw_latest.decode("utf-8") if raw_latest else None
|
|
|
+
|
|
|
+ if not checkpoint_id:
|
|
|
+ return None
|
|
|
+
|
|
|
+ payload = await self._load_obj(self._cp_key(thread_id, checkpoint_ns, checkpoint_id))
|
|
|
+ if not isinstance(payload, dict):
|
|
|
+ return None
|
|
|
+
|
|
|
+ checkpoint = payload.get("checkpoint")
|
|
|
+ metadata = payload.get("metadata", {})
|
|
|
+ parent_checkpoint_id = payload.get("parent_checkpoint_id")
|
|
|
+ pending_writes = await self._get_writes(thread_id, checkpoint_ns, checkpoint_id)
|
|
|
+
|
|
|
+ if requested_id:
|
|
|
+ return_config: RunnableConfig = config
|
|
|
+ else:
|
|
|
+ return_config = {
|
|
|
+ "configurable": {
|
|
|
+ "thread_id": thread_id,
|
|
|
+ "checkpoint_ns": checkpoint_ns,
|
|
|
+ "checkpoint_id": checkpoint_id,
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ parent_config: RunnableConfig | None = None
|
|
|
+ if parent_checkpoint_id:
|
|
|
+ parent_config = {
|
|
|
+ "configurable": {
|
|
|
+ "thread_id": thread_id,
|
|
|
+ "checkpoint_ns": checkpoint_ns,
|
|
|
+ "checkpoint_id": parent_checkpoint_id,
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ return CheckpointTuple(
|
|
|
+ config=return_config,
|
|
|
+ checkpoint=checkpoint,
|
|
|
+ metadata=metadata,
|
|
|
+ parent_config=parent_config,
|
|
|
+ pending_writes=pending_writes,
|
|
|
+ )
|
|
|
+
|
|
|
+ async def aput(
|
|
|
+ self,
|
|
|
+ config: RunnableConfig,
|
|
|
+ checkpoint: Checkpoint,
|
|
|
+ metadata: CheckpointMetadata,
|
|
|
+ new_versions: ChannelVersions,
|
|
|
+ ) -> RunnableConfig:
|
|
|
+ thread_id: str = config["configurable"]["thread_id"]
|
|
|
+ checkpoint_ns: str = config["configurable"].get("checkpoint_ns", "")
|
|
|
+ config_checkpoint_id = config["configurable"].get("checkpoint_id")
|
|
|
+ checkpoint_id = checkpoint.get("id") or config_checkpoint_id
|
|
|
+
|
|
|
+ if not checkpoint_id:
|
|
|
+ raise RuntimeError("checkpoint 缺少 checkpoint_id")
|
|
|
+
|
|
|
+ parent_checkpoint_id = None
|
|
|
+ if config_checkpoint_id and config_checkpoint_id != checkpoint_id:
|
|
|
+ parent_checkpoint_id = config_checkpoint_id
|
|
|
+
|
|
|
+ payload = {
|
|
|
+ "checkpoint": checkpoint,
|
|
|
+ "metadata": get_checkpoint_metadata(config, metadata),
|
|
|
+ "parent_checkpoint_id": parent_checkpoint_id,
|
|
|
+ }
|
|
|
+ await self._save_obj(self._cp_key(thread_id, checkpoint_ns, checkpoint_id), payload)
|
|
|
+ if self.ttl_seconds is not None:
|
|
|
+ await self.redis.set(
|
|
|
+ self._latest_key(thread_id, checkpoint_ns), checkpoint_id, ex=self.ttl_seconds
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ await self.redis.set(self._latest_key(thread_id, checkpoint_ns), checkpoint_id)
|
|
|
+ await self.redis.sadd(self._threads_key(), thread_id)
|
|
|
+
|
|
|
+ return {
|
|
|
+ "configurable": {
|
|
|
+ "thread_id": thread_id,
|
|
|
+ "checkpoint_ns": checkpoint_ns,
|
|
|
+ "checkpoint_id": checkpoint_id,
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ async def aput_writes(
|
|
|
+ self,
|
|
|
+ config: RunnableConfig,
|
|
|
+ writes: Sequence[tuple[str, Any]],
|
|
|
+ task_id: str,
|
|
|
+ task_path: str = "",
|
|
|
+ ) -> None:
|
|
|
+ thread_id: str = config["configurable"]["thread_id"]
|
|
|
+ checkpoint_ns: str = config["configurable"].get("checkpoint_ns", "")
|
|
|
+ checkpoint_id = config["configurable"].get("checkpoint_id")
|
|
|
+ if not checkpoint_id:
|
|
|
+ raw_latest = await self.redis.get(self._latest_key(thread_id, checkpoint_ns))
|
|
|
+ checkpoint_id = raw_latest.decode("utf-8") if raw_latest else None
|
|
|
+ if not checkpoint_id:
|
|
|
+ return
|
|
|
+ key = self._writes_key(thread_id, checkpoint_ns, checkpoint_id)
|
|
|
+
|
|
|
+ stored = await self._load_obj(key)
|
|
|
+ records: list[list[Any]] = [
|
|
|
+ list(row) for row in stored if isinstance(row, (list, tuple)) and len(row) == 4
|
|
|
+ ] if isinstance(stored, list) else []
|
|
|
+
|
|
|
+ for idx, (channel, value) in enumerate(writes):
|
|
|
+ write_idx = WRITES_IDX_MAP.get(channel, idx)
|
|
|
+ if write_idx >= 0 and any(
|
|
|
+ row[0] == task_id and row[2] == write_idx for row in records
|
|
|
+ ):
|
|
|
+ continue
|
|
|
+ records.append([task_id, channel, write_idx, value])
|
|
|
+
|
|
|
+ await self._save_obj(key, records)
|
|
|
+
|
|
|
+ async def alist(
|
|
|
+ self,
|
|
|
+ config: RunnableConfig | None,
|
|
|
+ *,
|
|
|
+ filter: dict[str, Any] | None = None,
|
|
|
+ before: RunnableConfig | None = None,
|
|
|
+ limit: int | None = None,
|
|
|
+ ) -> AsyncIterator[CheckpointTuple]:
|
|
|
+ if config is None:
|
|
|
+ return
|
|
|
+ thread_id: str = config["configurable"]["thread_id"]
|
|
|
+ checkpoint_ns: str = config["configurable"].get("checkpoint_ns", "")
|
|
|
+
|
|
|
+ match = f"{self.key_prefix}:cp:{self._part(thread_id)}:{self._part(checkpoint_ns)}:*"
|
|
|
+ keys = [key async for key in self.redis.scan_iter(match=match)]
|
|
|
+ checkpoint_ids: list[str] = []
|
|
|
+ for key in keys:
|
|
|
+ try:
|
|
|
+ checkpoint_ids.append(unquote(key.decode("utf-8").rsplit(":", 1)[-1]))
|
|
|
+ except Exception:
|
|
|
+ continue
|
|
|
+
|
|
|
+ count = 0
|
|
|
+ for checkpoint_id in sorted(checkpoint_ids, reverse=True):
|
|
|
+ tup = await self.aget_tuple(
|
|
|
+ {
|
|
|
+ "configurable": {
|
|
|
+ "thread_id": thread_id,
|
|
|
+ "checkpoint_ns": checkpoint_ns,
|
|
|
+ "checkpoint_id": checkpoint_id,
|
|
|
+ }
|
|
|
+ }
|
|
|
+ )
|
|
|
+ if tup is None:
|
|
|
+ continue
|
|
|
+ yield tup
|
|
|
+ count += 1
|
|
|
+ if limit is not None and count >= limit:
|
|
|
+ return
|
|
|
+
|
|
|
+ async def adelete_thread(self, thread_id: str) -> None:
|
|
|
+ tid = self._part(thread_id)
|
|
|
+ patterns = (
|
|
|
+ f"{self.key_prefix}:cp:{tid}:*",
|
|
|
+ f"{self.key_prefix}:writes:{tid}:*",
|
|
|
+ f"{self.key_prefix}:latest:{tid}:*",
|
|
|
+ )
|
|
|
+ keys: list[bytes] = []
|
|
|
+ for pattern in patterns:
|
|
|
+ async for key in self.redis.scan_iter(match=pattern):
|
|
|
+ keys.append(key)
|
|
|
+ if keys:
|
|
|
+ await self.redis.delete(*keys)
|
|
|
+ await self.redis.srem(self._threads_key(), thread_id)
|
wangxi