codekingpro/portable-devtools

1from __future__ import annotations2 3from collections.abc import Mapping, Sequence4from typing import Any5 6from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT7from langgraph.checkpoint.serde.base import SerializerProtocol8 9 10class RedisCache(BaseCache[ValueT]):11    """Redis-based cache implementation with TTL support."""12 13    def __init__(14        self,15        redis: Any,16        *,17        serde: SerializerProtocol | None = None,18        prefix: str = "langgraph:cache:",19    ) -> None:20        """Initialize the cache with a Redis client.21 22        Args:23            redis: Redis client instance (sync or async)24            serde: Serializer to use for values25            prefix: Key prefix for all cached values26        """27        super().__init__(serde=serde)28        self.redis = redis29        self.prefix = prefix30 31    def _make_key(self, ns: Namespace, key: str) -> str:32        """Create a Redis key from namespace and key."""33        ns_str = ":".join(ns) if ns else ""34        return f"{self.prefix}{ns_str}:{key}" if ns_str else f"{self.prefix}{key}"35 36    def _parse_key(self, redis_key: str) -> tuple[Namespace, str]:37        """Parse a Redis key back to namespace and key."""38        if not redis_key.startswith(self.prefix):39            raise ValueError(40                f"Key {redis_key} does not start with prefix {self.prefix}"41            )42 43        remaining = redis_key[len(self.prefix) :]44        if ":" in remaining:45            parts = remaining.split(":")46            key = parts[-1]47            ns_parts = parts[:-1]48            return (tuple(ns_parts), key)49        else:50            return (tuple(), remaining)51 52    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:53        """Get the cached values for the given keys."""54        if not keys:55            return {}56 57        # Build Redis keys58        redis_keys = [self._make_key(ns, key) for ns, key in keys]59 60        # Get values from Redis using MGET61        try:62            raw_values = self.redis.mget(redis_keys)63        except Exception:64            # If Redis is unavailable, return empty dict65            return {}66 67        values: dict[FullKey, ValueT] = {}68        for i, raw_value in enumerate(raw_values):69            if raw_value is not None:70                try:71                    # Deserialize the value72                    encoding, data = raw_value.split(b":", 1)73                    values[keys[i]] = self.serde.loads_typed((encoding.decode(), data))74                except Exception:75                    # Skip corrupted entries76                    continue77 78        return values79 80    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:81        """Asynchronously get the cached values for the given keys."""82        return self.get(keys)83 84    def set(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:85        """Set the cached values for the given keys and TTLs."""86        if not mapping:87            return88 89        # Use pipeline for efficient batch operations90        pipe = self.redis.pipeline()91 92        for (ns, key), (value, ttl) in mapping.items():93            redis_key = self._make_key(ns, key)94            encoding, data = self.serde.dumps_typed(value)95 96            # Store as "encoding:data" format97            serialized_value = f"{encoding}:".encode() + data98 99            if ttl is not None:100                pipe.setex(redis_key, ttl, serialized_value)101            else:102                pipe.set(redis_key, serialized_value)103 104        try:105            pipe.execute()106        except Exception:107            # Silently fail if Redis is unavailable108            pass109 110    async def aset(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:111        """Asynchronously set the cached values for the given keys and TTLs."""112        self.set(mapping)113 114    def clear(self, namespaces: Sequence[Namespace] | None = None) -> None:115        """Delete the cached values for the given namespaces.116        If no namespaces are provided, clear all cached values."""117        try:118            if namespaces is None:119                # Clear all keys with our prefix120                pattern = f"{self.prefix}*"121                keys = self.redis.keys(pattern)122                if keys:123                    self.redis.delete(*keys)124            else:125                # Clear specific namespaces126                keys_to_delete = []127                for ns in namespaces:128                    ns_str = ":".join(ns) if ns else ""129                    pattern = (130                        f"{self.prefix}{ns_str}:*" if ns_str else f"{self.prefix}*"131                    )132                    keys = self.redis.keys(pattern)133                    keys_to_delete.extend(keys)134 135                if keys_to_delete:136                    self.redis.delete(*keys_to_delete)137        except Exception:138            # Silently fail if Redis is unavailable139            pass140 141    async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None:142        """Asynchronously delete the cached values for the given namespaces.143        If no namespaces are provided, clear all cached values."""144        self.clear(namespaces)145