Add async methods to BaseStore (#16669)

- **Description:**

The BaseStore methods are currently blocking. Some implementations
(AstraDBStore, RedisStore) would benefit from having async methods.
Also once we have async methods for BaseStore, we can implement the
async `aembed_documents` in CacheBackedEmbeddings to cache the
embeddings asynchronously.

* adds async methods amget, amset, amedelete and ayield_keys to
BaseStore
  * implements the async methods for InMemoryStore
  * adds tests for InMemoryStore async methods

- **Twitter handle:** cbornet_

libs/core/langchain_core/stores.py

@@ -1,5 +1,17 @@
 from abc import ABC, abstractmethod
-from typing import Generic, Iterator, List, Optional, Sequence, Tuple, TypeVar, Union
+from typing import (
+    AsyncIterator,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+from langchain_core.runnables import run_in_executor
 
 K = TypeVar("K")
 V = TypeVar("V")
@@ -20,6 +32,18 @@ class BaseStore(Generic[K, V], ABC):
             If a key is not found, the corresponding value will be None.
         """
 
+    async def amget(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.
+        """
+        return await run_in_executor(None, self.mget, keys)
+
     @abstractmethod
     def mset(self, key_value_pairs: Sequence[Tuple[K, V]]) -> None:
         """Set the values for the given keys.
@@ -28,6 +52,14 @@ class BaseStore(Generic[K, V], ABC):
             key_value_pairs (Sequence[Tuple[K, V]]): A sequence of key-value pairs.
         """
 
+    async def amset(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.
+        """
+        return await run_in_executor(None, self.mset, key_value_pairs)
+
     @abstractmethod
     def mdelete(self, keys: Sequence[K]) -> None:
         """Delete the given keys and their associated values.
@@ -36,6 +68,14 @@ class BaseStore(Generic[K, V], ABC):
             keys (Sequence[K]): A sequence of keys to delete.
         """
 
+    async def amdelete(self, keys: Sequence[K]) -> None:
+        """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)
+
     @abstractmethod
     def yield_keys(
         self, *, prefix: Optional[str] = None
@@ -52,5 +92,27 @@ class BaseStore(Generic[K, V], ABC):
             depending on what makes more sense for the given store.
         """
 
+    async def ayield_keys(
+        self, *, prefix: Optional[str] = None
+    ) -> Union[AsyncIterator[K], AsyncIterator[str]]:
+        """Get an iterator over keys that match the given prefix.
+
+        Args:
+            prefix (str): The prefix to match.
+
+        Returns:
+            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
+
 
 ByteStore = BaseStore[str, bytes]