codekingpro/portable-devtools

1from __future__ import annotations2 3import datetime4import threading5from collections.abc import Mapping, Sequence6 7from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT8from langgraph.checkpoint.serde.base import SerializerProtocol9 10 11class InMemoryCache(BaseCache[ValueT]):12    def __init__(self, *, serde: SerializerProtocol | None = None):13        super().__init__(serde=serde)14        self._cache: dict[Namespace, dict[str, tuple[str, bytes, float | None]]] = {}15        self._lock = threading.RLock()16 17    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:18        """Get the cached values for the given keys."""19        with self._lock:20            if not keys:21                return {}22            now = datetime.datetime.now(datetime.timezone.utc).timestamp()23            values: dict[FullKey, ValueT] = {}24            for ns_tuple, key in keys:25                ns = Namespace(ns_tuple)26                if ns in self._cache and key in self._cache[ns]:27                    enc, val, expiry = self._cache[ns][key]28                    if expiry is None or now < expiry:29                        values[(ns, key)] = self.serde.loads_typed((enc, val))30                    else:31                        del self._cache[ns][key]32            return values33 34    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:35        """Asynchronously get the cached values for the given keys."""36        return self.get(keys)37 38    def set(self, keys: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:39        """Set the cached values for the given keys."""40        with self._lock:41            now = datetime.datetime.now(datetime.timezone.utc)42            for (ns, key), (value, ttl) in keys.items():43                if ttl is not None:44                    delta = datetime.timedelta(seconds=ttl)45                    expiry: float | None = (now + delta).timestamp()46                else:47                    expiry = None48                if ns not in self._cache:49                    self._cache[ns] = {}50                self._cache[ns][key] = (51                    *self.serde.dumps_typed(value),52                    expiry,53                )54 55    async def aset(self, keys: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:56        """Asynchronously set the cached values for the given keys."""57        self.set(keys)58 59    def clear(self, namespaces: Sequence[Namespace] | None = None) -> None:60        """Delete the cached values for the given namespaces.61        If no namespaces are provided, clear all cached values."""62        with self._lock:63            if namespaces is None:64                self._cache.clear()65            else:66                for ns in namespaces:67                    if ns in self._cache:68                        del self._cache[ns]69 70    async def aclear(self, namespaces: Sequence[Namespace] | None = None) -> None:71        """Asynchronously delete the cached values for the given namespaces.72        If no namespaces are provided, clear all cached values."""73        self.clear(namespaces)74