codekingpro/portable-devtools

1"""**Store** implements the key-value stores and storage helpers.2 3Module provides implementations of various key-value stores that conform4to a simple key-value interface.5 6The primary goal of these storages is to support implementation of caching.7"""8 9from abc import ABC, abstractmethod10from collections.abc import AsyncIterator, Iterator, Sequence11from typing import (12    Any,13    Generic,14    TypeVar,15)16 17from typing_extensions import override18 19from langchain_core.exceptions import LangChainException20from langchain_core.runnables import run_in_executor21 22K = TypeVar("K")23V = TypeVar("V")24 25 26class BaseStore(ABC, Generic[K, V]):27    """Abstract interface for a key-value store.28 29    This is an interface that's meant to abstract away the details of different30    key-value stores. It provides a simple interface for getting, setting, and deleting31    key-value pairs.32 33    The basic methods are `mget`, `mset`, and `mdelete` for getting, setting, and34    deleting multiple key-value pairs at once. The `yield_keys` method is used to35    iterate over keys that match a given prefix.36 37    The async versions of these methods are also provided, which are meant to be used in38    async contexts. The async methods are named with an `a` prefix, e.g., `amget`,39    `amset`, `amdelete`, and `ayield_keys`.40 41    By default, the `amget`, `amset`, `amdelete`, and `ayield_keys` methods are42    implemented using the synchronous methods. If the store can natively support async43    operations, it should override these methods.44 45    By design the methods only accept batches of keys and values, and not single keys or46    values. This is done to force user code to work with batches which will usually be47    more efficient by saving on round trips to the store.48 49    Examples:50        ```python51        from langchain.storage import BaseStore52 53 54        class MyInMemoryStore(BaseStore[str, int]):55            def __init__(self) -> None:56                self.store: dict[str, int] = {}57 58            def mget(self, keys: Sequence[str]) -> list[int | None]:59                return [self.store.get(key) for key in keys]60 61            def mset(self, key_value_pairs: Sequence[tuple[str, int]]) -> None:62                for key, value in key_value_pairs:63                    self.store[key] = value64 65            def mdelete(self, keys: Sequence[str]) -> None:66                for key in keys:67                    if key in self.store:68                        del self.store[key]69 70            def yield_keys(self, prefix: str | None = None) -> Iterator[str]:71                if prefix is None:72                    yield from self.store.keys()73                else:74                    for key in self.store.keys():75                        if key.startswith(prefix):76                            yield key77        ```78    """79 80    @abstractmethod81    def mget(self, keys: Sequence[K]) -> list[V | None]:82        """Get the values associated with the given keys.83 84        Args:85            keys: A sequence of keys.86 87        Returns:88            A sequence of optional values associated with the keys.89                If a key is not found, the corresponding value will be `None`.90        """91 92    async def amget(self, keys: Sequence[K]) -> list[V | None]:93        """Async get the values associated with the given keys.94 95        Args:96            keys: A sequence of keys.97 98        Returns:99            A sequence of optional values associated with the keys.100                If a key is not found, the corresponding value will be `None`.101        """102        return await run_in_executor(None, self.mget, keys)103 104    @abstractmethod105    def mset(self, key_value_pairs: Sequence[tuple[K, V]]) -> None:106        """Set the values for the given keys.107 108        Args:109            key_value_pairs: A sequence of key-value pairs.110        """111 112    async def amset(self, key_value_pairs: Sequence[tuple[K, V]]) -> None:113        """Async set the values for the given keys.114 115        Args:116            key_value_pairs: A sequence of key-value pairs.117        """118        return await run_in_executor(None, self.mset, key_value_pairs)119 120    @abstractmethod121    def mdelete(self, keys: Sequence[K]) -> None:122        """Delete the given keys and their associated values.123 124        Args:125            keys: A sequence of keys to delete.126        """127 128    async def amdelete(self, keys: Sequence[K]) -> None:129        """Async delete the given keys and their associated values.130 131        Args:132            keys: A sequence of keys to delete.133        """134        return await run_in_executor(None, self.mdelete, keys)135 136    @abstractmethod137    def yield_keys(self, *, prefix: str | None = None) -> Iterator[K] | Iterator[str]:138        """Get an iterator over keys that match the given prefix.139 140        Args:141            prefix: The prefix to match.142 143        Yields:144            An iterator over keys that match the given prefix.145 146                This method is allowed to return an iterator over either K or str147                depending on what makes more sense for the given store.148        """149 150    async def ayield_keys(151        self, *, prefix: str | None = None152    ) -> AsyncIterator[K] | AsyncIterator[str]:153        """Async get an iterator over keys that match the given prefix.154 155        Args:156            prefix: The prefix to match.157 158        Yields:159            The keys that match the given prefix.160 161                This method is allowed to return an iterator over either K or str162                depending on what makes more sense for the given store.163        """164        iterator = await run_in_executor(None, self.yield_keys, prefix=prefix)165        done = object()166        while True:167            item = await run_in_executor(None, lambda it: next(it, done), iterator)168            if item is done:169                break170            yield item  # type: ignore[misc]171 172 173ByteStore = BaseStore[str, bytes]174 175 176class InMemoryBaseStore(BaseStore[str, V], Generic[V]):177    """In-memory implementation of the `BaseStore` using a dictionary."""178 179    def __init__(self) -> None:180        """Initialize an empty store."""181        self.store: dict[str, V] = {}182 183    @override184    def mget(self, keys: Sequence[str]) -> list[V | None]:185        return [self.store.get(key) for key in keys]186 187    @override188    async def amget(self, keys: Sequence[str]) -> list[V | None]:189        return self.mget(keys)190 191    @override192    def mset(self, key_value_pairs: Sequence[tuple[str, V]]) -> None:193        for key, value in key_value_pairs:194            self.store[key] = value195 196    @override197    async def amset(self, key_value_pairs: Sequence[tuple[str, V]]) -> None:198        return self.mset(key_value_pairs)199 200    @override201    def mdelete(self, keys: Sequence[str]) -> None:202        for key in keys:203            if key in self.store:204                del self.store[key]205 206    @override207    async def amdelete(self, keys: Sequence[str]) -> None:208        self.mdelete(keys)209 210    def yield_keys(self, *, prefix: str | None = None) -> Iterator[str]:211        """Get an iterator over keys that match the given prefix.212 213        Args:214            prefix: The prefix to match.215 216        Yields:217            The keys that match the given prefix.218        """219        if prefix is None:220            yield from self.store.keys()221        else:222            for key in self.store:223                if key.startswith(prefix):224                    yield key225 226    async def ayield_keys(self, *, prefix: str | None = None) -> AsyncIterator[str]:227        """Async get an async iterator over keys that match the given prefix.228 229        Args:230            prefix: The prefix to match.231 232        Yields:233            The keys that match the given prefix.234        """235        if prefix is None:236            for key in self.store:237                yield key238        else:239            for key in self.store:240                if key.startswith(prefix):241                    yield key242 243 244class InMemoryStore(InMemoryBaseStore[Any]):245    """In-memory store for any type of data.246 247    Attributes:248        store: The underlying dictionary that stores the key-value pairs.249 250    Examples:251        ```python252        from langchain.storage import InMemoryStore253 254        store = InMemoryStore()255        store.mset([("key1", "value1"), ("key2", "value2")])256        store.mget(["key1", "key2"])257        # ['value1', 'value2']258        store.mdelete(["key1"])259        list(store.yield_keys())260        # ['key2']261        list(store.yield_keys(prefix="k"))262        # ['key2']263        ```264    """265 266 267class InMemoryByteStore(InMemoryBaseStore[bytes]):268    """In-memory store for bytes.269 270    Attributes:271        store: The underlying dictionary that stores the key-value pairs.272 273    Examples:274        ```python275        from langchain.storage import InMemoryByteStore276 277        store = InMemoryByteStore()278        store.mset([("key1", b"value1"), ("key2", b"value2")])279        store.mget(["key1", "key2"])280        # [b'value1', b'value2']281        store.mdelete(["key1"])282        list(store.yield_keys())283        # ['key2']284        list(store.yield_keys(prefix="k"))285        # ['key2']286        ```287    """288 289 290class InvalidKeyException(LangChainException):291    """Raised when a key is invalid; e.g., uses incorrect characters."""292