Source code for langchain_core.stores

"""**Store** implements the key-value stores and storage helpers.

Module provides implementations of various key-value stores that conform
to a simple key-value interface.

The primary goal of these storages is to support implementation of caching.
"""

from abc import ABC, abstractmethod
from collections.abc import AsyncIterator, Iterator, Sequence
from typing import (
    Any,
    Generic,
    Optional,
    TypeVar,
    Union,
)

from langchain_core.exceptions import LangChainException
from langchain_core.runnables import run_in_executor

K = TypeVar("K")
V = TypeVar("V")



[docs]
class BaseStore(Generic[K, V], ABC):
    """Abstract interface for a key-value store.

    This is an interface that's meant to abstract away the details of
    different key-value stores. It provides a simple interface for
    getting, setting, and deleting key-value pairs.

    The basic methods are `mget`, `mset`, and `mdelete` for getting,
    setting, and deleting multiple key-value pairs at once. The `yield_keys`
    method is used to iterate over keys that match a given prefix.

    The async versions of these methods are also provided, which are
    meant to be used in async contexts. The async methods are named with
    an `a` prefix, e.g., `amget`, `amset`, `amdelete`, and `ayield_keys`.

    By default, the `amget`, `amset`, `amdelete`, and `ayield_keys` methods
    are implemented using the synchronous methods. If the store can natively
    support async  operations, it should override these methods.

    By design the methods only accept batches of keys and values, and not
    single keys or values. This is done to force user code to work with batches
    which will usually be more efficient by saving on round trips to the store.

    Examples:

        .. code-block:: python

            from langchain.storage import BaseStore

            class MyInMemoryStore(BaseStore[str, int]):

                def __init__(self):
                    self.store = {}

                def mget(self, keys):
                    return [self.store.get(key) for key in keys]

                def mset(self, key_value_pairs):
                    for key, value in key_value_pairs:
                        self.store[key] = value

                def mdelete(self, keys):
                    for key in keys:
                        if key in self.store:
                            del self.store[key]

                def yield_keys(self, prefix=None):
                    if prefix is None:
                        yield from self.store.keys()
                    else:
                        for key in self.store.keys():
                            if key.startswith(prefix):
                                yield key
    """


[docs]
    @abstractmethod
    def mget(self, keys: Sequence[K]) -> list[Optional[V]]:
        """Get the values associated with the given keys.

        Args:
            keys (Sequence[K]): A sequence of keys.

        Returns:
            A sequence of optional values associated with the keys.
            If a key is not found, the corresponding value will be None.
        """



[docs]
    async def amget(self, keys: Sequence[K]) -> list[Optional[V]]:
        """Async get the values associated with the given keys.

        Args:
            keys (Sequence[K]): A sequence of keys.

        Returns:
            A sequence of optional values associated with the keys.
            If a key is not found, the corresponding value will be None.
        """
        return await run_in_executor(None, self.mget, keys)



[docs]
    @abstractmethod
    def mset(self, key_value_pairs: Sequence[tuple[K, V]]) -> None:
        """Set the values for the given keys.

        Args:
            key_value_pairs (Sequence[Tuple[K, V]]): A sequence of key-value pairs.
        """



[docs]
    async def amset(self, key_value_pairs: Sequence[tuple[K, V]]) -> None:
        """Async set the values for the given keys.

        Args:
            key_value_pairs (Sequence[Tuple[K, V]]): A sequence of key-value pairs.
        """
        return await run_in_executor(None, self.mset, key_value_pairs)



[docs]
    @abstractmethod
    def mdelete(self, keys: Sequence[K]) -> None:
        """Delete the given keys and their associated values.

        Args:
            keys (Sequence[K]): A sequence of keys to delete.
        """



[docs]
    async def amdelete(self, keys: Sequence[K]) -> None:
        """Async delete the given keys and their associated values.

        Args:
            keys (Sequence[K]): A sequence of keys to delete.
        """
        return await run_in_executor(None, self.mdelete, keys)



[docs]
    @abstractmethod
    def yield_keys(
        self, *, prefix: Optional[str] = None
    ) -> Union[Iterator[K], Iterator[str]]:
        """Get an iterator over keys that match the given prefix.

        Args:
            prefix (str): The prefix to match.

        Yields:
            Iterator[K | str]: An iterator over keys that match the given prefix.
            This method is allowed to return an iterator over either K or str
            depending on what makes more sense for the given store.
        """



[docs]
    async def ayield_keys(
        self, *, prefix: Optional[str] = None
    ) -> Union[AsyncIterator[K], AsyncIterator[str]]:
        """Async get an iterator over keys that match the given prefix.

        Args:
            prefix (str): The prefix to match.

        Yields:
            Iterator[K | str]: An iterator over keys that match the given prefix.
            This method is allowed to return an iterator over either K or str
            depending on what makes more sense for the given store.
        """
        iterator = await run_in_executor(None, self.yield_keys, prefix=prefix)
        done = object()
        while True:
            item = await run_in_executor(None, lambda it: next(it, done), iterator)
            if item is done:
                break
            yield item  # type: ignore[misc]




ByteStore = BaseStore[str, bytes]



[docs]
class InMemoryBaseStore(BaseStore[str, V], Generic[V]):
    """In-memory implementation of the BaseStore using a dictionary."""


[docs]
    def __init__(self) -> None:
        """Initialize an empty store."""
        self.store: dict[str, V] = {}



[docs]
    def mget(self, keys: Sequence[str]) -> list[Optional[V]]:
        """Get the values associated with the given keys.

        Args:
            keys (Sequence[str]): A sequence of keys.

        Returns:
            A sequence of optional values associated with the keys.
            If a key is not found, the corresponding value will be None.
        """
        return [self.store.get(key) for key in keys]



[docs]
    async def amget(self, keys: Sequence[str]) -> list[Optional[V]]:
        """Async get the values associated with the given keys.

        Args:
            keys (Sequence[str]): A sequence of keys.

        Returns:
            A sequence of optional values associated with the keys.
            If a key is not found, the corresponding value will be None.
        """
        return self.mget(keys)



[docs]
    def mset(self, key_value_pairs: Sequence[tuple[str, V]]) -> None:
        """Set the values for the given keys.

        Args:
            key_value_pairs (Sequence[Tuple[str, V]]): A sequence of key-value pairs.

        Returns:
            None
        """
        for key, value in key_value_pairs:
            self.store[key] = value



[docs]
    async def amset(self, key_value_pairs: Sequence[tuple[str, V]]) -> None:
        """Async set the values for the given keys.

        Args:
            key_value_pairs (Sequence[Tuple[str, V]]): A sequence of key-value pairs.

        Returns:
            None
        """
        return self.mset(key_value_pairs)



[docs]
    def mdelete(self, keys: Sequence[str]) -> None:
        """Delete the given keys and their associated values.

        Args:
            keys (Sequence[str]): A sequence of keys to delete.
        """
        for key in keys:
            if key in self.store:
                del self.store[key]



[docs]
    async def amdelete(self, keys: Sequence[str]) -> None:
        """Async delete the given keys and their associated values.

        Args:
            keys (Sequence[str]): A sequence of keys to delete.
        """
        self.mdelete(keys)



[docs]
    def yield_keys(self, prefix: Optional[str] = None) -> Iterator[str]:
        """Get an iterator over keys that match the given prefix.

        Args:
            prefix (str, optional): The prefix to match. Defaults to None.

        Yields:
            Iterator[str]: An iterator over keys that match the given prefix.
        """
        if prefix is None:
            yield from self.store.keys()
        else:
            for key in self.store:
                if key.startswith(prefix):
                    yield key



[docs]
    async def ayield_keys(self, prefix: Optional[str] = None) -> AsyncIterator[str]:
        """Async get an async iterator over keys that match the given prefix.

        Args:
            prefix (str, optional): The prefix to match. Defaults to None.

        Yields:
            AsyncIterator[str]: An async iterator over keys that match the given prefix.
        """
        if prefix is None:
            for key in self.store:
                yield key
        else:
            for key in self.store:
                if key.startswith(prefix):
                    yield key





[docs]
class InMemoryStore(InMemoryBaseStore[Any]):
    """In-memory store for any type of data.

    Attributes:
        store (Dict[str, Any]): The underlying dictionary that stores
            the key-value pairs.

    Examples:

        .. code-block:: python

            from langchain.storage import InMemoryStore

            store = InMemoryStore()
            store.mset([('key1', 'value1'), ('key2', 'value2')])
            store.mget(['key1', 'key2'])
            # ['value1', 'value2']
            store.mdelete(['key1'])
            list(store.yield_keys())
            # ['key2']
            list(store.yield_keys(prefix='k'))
            # ['key2']
    """




[docs]
class InMemoryByteStore(InMemoryBaseStore[bytes]):
    """In-memory store for bytes.

    Attributes:
        store (Dict[str, bytes]): The underlying dictionary that stores
            the key-value pairs.

    Examples:

        .. code-block:: python

            from langchain.storage import InMemoryByteStore

            store = InMemoryByteStore()
            store.mset([('key1', b'value1'), ('key2', b'value2')])
            store.mget(['key1', 'key2'])
            # [b'value1', b'value2']
            store.mdelete(['key1'])
            list(store.yield_keys())
            # ['key2']
            list(store.yield_keys(prefix='k'))
            # ['key2']
    """




[docs]
class InvalidKeyException(LangChainException):
    """Raised when a key is invalid; e.g., uses incorrect characters."""