community[minor]: Add async methods to the AstraDB BaseStore (#16872)

---------

Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com>

libs/community/langchain_community/storage/astradb.py

@@ -5,6 +5,7 @@ from abc import ABC, abstractmethod
 from typing import (
     TYPE_CHECKING,
     Any,
+    AsyncIterator,
     Generic,
     Iterator,
     List,
@@ -16,10 +17,13 @@ from typing import (
 
 from langchain_core.stores import BaseStore, ByteStore
 
-from langchain_community.utilities.astradb import _AstraDBEnvironment
+from langchain_community.utilities.astradb import (
+    SetupMode,
+    _AstraDBCollectionEnvironment,
+)
 
 if TYPE_CHECKING:
-    from astrapy.db import AstraDB
+    from astrapy.db import AstraDB, AsyncAstraDB
 
 V = TypeVar("V")
 
@@ -34,17 +38,23 @@ class AstraDBBaseStore(Generic[V], BaseStore[str, V], ABC):
         api_endpoint: Optional[str] = None,
         astra_db_client: Optional[AstraDB] = None,
         namespace: Optional[str] = None,
+        *,
+        async_astra_db_client: Optional[AsyncAstraDB] = None,
+        pre_delete_collection: bool = False,
+        setup_mode: SetupMode = SetupMode.SYNC,
     ) -> None:
-        astra_env = _AstraDBEnvironment(
+        self.astra_env = _AstraDBCollectionEnvironment(
+            collection_name=collection_name,
             token=token,
             api_endpoint=api_endpoint,
             astra_db_client=astra_db_client,
+            async_astra_db_client=async_astra_db_client,
             namespace=namespace,
+            setup_mode=setup_mode,
+            pre_delete_collection=pre_delete_collection,
         )
-        self.astra_db = astra_env.astra_db
-        self.collection = self.astra_db.create_collection(
-            collection_name=collection_name,
-        )
+        self.collection = self.astra_env.collection
+        self.async_collection = self.astra_env.async_collection
 
     @abstractmethod
     def decode_value(self, value: Any) -> Optional[V]:
@@ -56,28 +66,63 @@ class AstraDBBaseStore(Generic[V], BaseStore[str, V], ABC):
 
     def mget(self, keys: Sequence[str]) -> List[Optional[V]]:
         """Get the values associated with the given keys."""
+        self.astra_env.ensure_db_setup()
         docs_dict = {}
         for doc in self.collection.paginated_find(filter={"_id": {"$in": list(keys)}}):
             docs_dict[doc["_id"]] = doc.get("value")
         return [self.decode_value(docs_dict.get(key)) for key in keys]
 
+    async def amget(self, keys: Sequence[str]) -> List[Optional[V]]:
+        """Get the values associated with the given keys."""
+        await self.astra_env.aensure_db_setup()
+        docs_dict = {}
+        async for doc in self.async_collection.paginated_find(
+            filter={"_id": {"$in": list(keys)}}
+        ):
+            docs_dict[doc["_id"]] = doc.get("value")
+        return [self.decode_value(docs_dict.get(key)) for key in keys]
+
     def mset(self, key_value_pairs: Sequence[Tuple[str, V]]) -> None:
         """Set the given key-value pairs."""
+        self.astra_env.ensure_db_setup()
         for k, v in key_value_pairs:
             self.collection.upsert({"_id": k, "value": self.encode_value(v)})
 
+    async def amset(self, key_value_pairs: Sequence[Tuple[str, V]]) -> None:
+        """Set the given key-value pairs."""
+        await self.astra_env.aensure_db_setup()
+        for k, v in key_value_pairs:
+            await self.async_collection.upsert(
+                {"_id": k, "value": self.encode_value(v)}
+            )
+
     def mdelete(self, keys: Sequence[str]) -> None:
         """Delete the given keys."""
+        self.astra_env.ensure_db_setup()
         self.collection.delete_many(filter={"_id": {"$in": list(keys)}})
 
+    async def amdelete(self, keys: Sequence[str]) -> None:
+        """Delete the given keys."""
+        await self.astra_env.aensure_db_setup()
+        await self.async_collection.delete_many(filter={"_id": {"$in": list(keys)}})
+
     def yield_keys(self, *, prefix: Optional[str] = None) -> Iterator[str]:
         """Yield keys in the store."""
+        self.astra_env.ensure_db_setup()
         docs = self.collection.paginated_find()
         for doc in docs:
             key = doc["_id"]
             if not prefix or key.startswith(prefix):
                 yield key
 
+    async def ayield_keys(self, *, prefix: Optional[str] = None) -> AsyncIterator[str]:
+        """Yield keys in the store."""
+        await self.astra_env.aensure_db_setup()
+        async for doc in self.async_collection.paginated_find():
+            key = doc["_id"]
+            if not prefix or key.startswith(prefix):
+                yield key
+
 
 class AstraDBStore(AstraDBBaseStore[Any]):
     """BaseStore implementation using DataStax AstraDB as the underlying store.