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

---------

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

libs/community/tests/integration_tests/storage/test_astradb.py

@@ -1,9 +1,16 @@
 """Implement integration tests for AstraDB storage."""
+from __future__ import annotations
+
 import os
+from typing import TYPE_CHECKING
 
 import pytest
 
 from langchain_community.storage.astradb import AstraDBByteStore, AstraDBStore
+from langchain_community.utilities.astradb import SetupMode
+
+if TYPE_CHECKING:
+    from astrapy.db import AstraDB, AsyncAstraDB
 
 
 def _has_env_vars() -> bool:
@@ -16,7 +23,7 @@ def _has_env_vars() -> bool:
 
 
 @pytest.fixture
-def astra_db():  # type: ignore[no-untyped-def]
+def astra_db() -> AstraDB:
     from astrapy.db import AstraDB
 
     return AstraDB(
@@ -26,24 +33,45 @@ def astra_db():  # type: ignore[no-untyped-def]
     )
 
 
-def init_store(astra_db, collection_name: str):  # type: ignore[no-untyped-def, no-untyped-def]
-    astra_db.create_collection(collection_name)
+@pytest.fixture
+def async_astra_db() -> AsyncAstraDB:
+    from astrapy.db import AsyncAstraDB
+
+    return AsyncAstraDB(
+        token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
+        api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"],
+        namespace=os.environ.get("ASTRA_DB_KEYSPACE"),
+    )
+
+
+def init_store(astra_db: AstraDB, collection_name: str) -> AstraDBStore:
     store = AstraDBStore(collection_name=collection_name, astra_db_client=astra_db)
     store.mset([("key1", [0.1, 0.2]), ("key2", "value2")])
     return store
 
 
-def init_bytestore(astra_db, collection_name: str):  # type: ignore[no-untyped-def, no-untyped-def]
-    astra_db.create_collection(collection_name)
+def init_bytestore(astra_db: AstraDB, collection_name: str) -> AstraDBByteStore:
     store = AstraDBByteStore(collection_name=collection_name, astra_db_client=astra_db)
     store.mset([("key1", b"value1"), ("key2", b"value2")])
     return store
 
 
+async def init_async_store(
+    async_astra_db: AsyncAstraDB, collection_name: str
+) -> AstraDBStore:
+    store = AstraDBStore(
+        collection_name=collection_name,
+        async_astra_db_client=async_astra_db,
+        setup_mode=SetupMode.ASYNC,
+    )
+    await store.amset([("key1", [0.1, 0.2]), ("key2", "value2")])
+    return store
+
+
 @pytest.mark.requires("astrapy")
 @pytest.mark.skipif(not _has_env_vars(), reason="Missing Astra DB env. vars")
 class TestAstraDBStore:
-    def test_mget(self, astra_db) -> None:  # type: ignore[no-untyped-def]
+    def test_mget(self, astra_db: AstraDB) -> None:
         """Test AstraDBStore mget method."""
         collection_name = "lc_test_store_mget"
         try:
@@ -52,7 +80,16 @@ class TestAstraDBStore:
         finally:
             astra_db.delete_collection(collection_name)
 
-    def test_mset(self, astra_db) -> None:  # type: ignore[no-untyped-def]
+    async def test_amget(self, async_astra_db: AsyncAstraDB) -> None:
+        """Test AstraDBStore amget method."""
+        collection_name = "lc_test_store_mget"
+        try:
+            store = await init_async_store(async_astra_db, collection_name)
+            assert await store.amget(["key1", "key2"]) == [[0.1, 0.2], "value2"]
+        finally:
+            await async_astra_db.delete_collection(collection_name)
+
+    def test_mset(self, astra_db: AstraDB) -> None:
         """Test that multiple keys can be set with AstraDBStore."""
         collection_name = "lc_test_store_mset"
         try:
@@ -64,7 +101,19 @@ class TestAstraDBStore:
         finally:
             astra_db.delete_collection(collection_name)
 
-    def test_mdelete(self, astra_db) -> None:  # type: ignore[no-untyped-def]
+    async def test_amset(self, async_astra_db: AsyncAstraDB) -> None:
+        """Test that multiple keys can be set with AstraDBStore."""
+        collection_name = "lc_test_store_mset"
+        try:
+            store = await init_async_store(async_astra_db, collection_name)
+            result = await store.async_collection.find_one({"_id": "key1"})
+            assert result["data"]["document"]["value"] == [0.1, 0.2]
+            result = await store.async_collection.find_one({"_id": "key2"})
+            assert result["data"]["document"]["value"] == "value2"
+        finally:
+            await async_astra_db.delete_collection(collection_name)
+
+    def test_mdelete(self, astra_db: AstraDB) -> None:
         """Test that deletion works as expected."""
         collection_name = "lc_test_store_mdelete"
         try:
@@ -75,7 +124,18 @@ class TestAstraDBStore:
         finally:
             astra_db.delete_collection(collection_name)
 
-    def test_yield_keys(self, astra_db) -> None:  # type: ignore[no-untyped-def]
+    async def test_amdelete(self, async_astra_db: AsyncAstraDB) -> None:
+        """Test that deletion works as expected."""
+        collection_name = "lc_test_store_mdelete"
+        try:
+            store = await init_async_store(async_astra_db, collection_name)
+            await store.amdelete(["key1", "key2"])
+            result = await store.amget(["key1", "key2"])
+            assert result == [None, None]
+        finally:
+            await async_astra_db.delete_collection(collection_name)
+
+    def test_yield_keys(self, astra_db: AstraDB) -> None:
         collection_name = "lc_test_store_yield_keys"
         try:
             store = init_store(astra_db, collection_name)
@@ -85,7 +145,20 @@ class TestAstraDBStore:
         finally:
             astra_db.delete_collection(collection_name)
 
-    def test_bytestore_mget(self, astra_db) -> None:  # type: ignore[no-untyped-def]
+    async def test_ayield_keys(self, async_astra_db: AsyncAstraDB) -> None:
+        collection_name = "lc_test_store_yield_keys"
+        try:
+            store = await init_async_store(async_astra_db, collection_name)
+            assert {key async for key in store.ayield_keys()} == {"key1", "key2"}
+            assert {key async for key in store.ayield_keys(prefix="key")} == {
+                "key1",
+                "key2",
+            }
+            assert {key async for key in store.ayield_keys(prefix="lang")} == set()
+        finally:
+            await async_astra_db.delete_collection(collection_name)
+
+    def test_bytestore_mget(self, astra_db: AstraDB) -> None:
         """Test AstraDBByteStore mget method."""
         collection_name = "lc_test_bytestore_mget"
         try:
@@ -94,7 +167,7 @@ class TestAstraDBStore:
         finally:
             astra_db.delete_collection(collection_name)
 
-    def test_bytestore_mset(self, astra_db) -> None:  # type: ignore[no-untyped-def]
+    def test_bytestore_mset(self, astra_db: AstraDB) -> None:
         """Test that multiple keys can be set with AstraDBByteStore."""
         collection_name = "lc_test_bytestore_mset"
         try: