Add a BaseStore backed by AstraDB (#15812)

- **Description:** this change adds a `BaseStore` backed by AstraDB
  - **Twitter handle:** cbornet_

libs/community/langchain_community/storage/__init__.py

@@ -6,6 +6,10 @@ to a simple key-value interface.
 The primary goal of these storages is to support implementation of caching.
 """
 
+from langchain_community.storage.astradb import (
+    AstraDBByteStore,
+    AstraDBStore,
+)
 from langchain_community.storage.redis import RedisStore
 from langchain_community.storage.upstash_redis import (
     UpstashRedisByteStore,
@@ -13,6 +17,8 @@ from langchain_community.storage.upstash_redis import (
 )
 
 __all__ = [
+    "AstraDBStore",
+    "AstraDBByteStore",
     "RedisStore",
     "UpstashRedisByteStore",
     "UpstashRedisStore",

libs/community/langchain_community/storage/astradb.py

@@ -0,0 +1,120 @@
+import base64
+from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+)
+
+from langchain_core.stores import BaseStore, ByteStore
+
+V = TypeVar("V")
+
+
+class AstraDBBaseStore(Generic[V], BaseStore[str, V], ABC):
+    def __init__(
+        self,
+        collection_name: str,
+        token: Optional[str] = None,
+        api_endpoint: Optional[str] = None,
+        astra_db_client: Optional[Any] = None,  # 'astrapy.db.AstraDB' if passed
+        namespace: Optional[str] = None,
+    ) -> None:
+        try:
+            from astrapy.db import AstraDB, AstraDBCollection
+        except (ImportError, ModuleNotFoundError):
+            raise ImportError(
+                "Could not import a recent astrapy python package. "
+                "Please install it with `pip install --upgrade astrapy`."
+            )
+
+        # Conflicting-arg checks:
+        if astra_db_client is not None:
+            if token is not None or api_endpoint is not None:
+                raise ValueError(
+                    "You cannot pass 'astra_db_client' to AstraDB if passing "
+                    "'token' and 'api_endpoint'."
+                )
+
+        astra_db = astra_db_client or AstraDB(
+            token=token,
+            api_endpoint=api_endpoint,
+            namespace=namespace,
+        )
+        self.collection = AstraDBCollection(collection_name, astra_db=astra_db)
+
+    @abstractmethod
+    def decode_value(self, value: Any) -> Optional[V]:
+        """Decodes value from Astra DB"""
+
+    @abstractmethod
+    def encode_value(self, value: Optional[V]) -> Any:
+        """Encodes value for Astra DB"""
+
+    def mget(self, keys: Sequence[str]) -> List[Optional[V]]:
+        """Get the values associated with the given keys."""
+        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]
+
+    def mset(self, key_value_pairs: Sequence[Tuple[str, V]]) -> None:
+        """Set the given key-value pairs."""
+        for k, v in key_value_pairs:
+            self.collection.upsert({"_id": k, "value": self.encode_value(v)})
+
+    def mdelete(self, keys: Sequence[str]) -> None:
+        """Delete the given keys."""
+        self.collection.delete_many(filter={"_id": {"$in": list(keys)}})
+
+    def yield_keys(self, *, prefix: Optional[str] = None) -> Iterator[str]:
+        """Yield keys in the store."""
+        docs = self.collection.paginated_find()
+        for doc in docs:
+            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.
+    The value type can be any type serializable by json.dumps.
+    Can be used to store embeddings with the CacheBackedEmbeddings.
+    Documents in the AstraDB collection will have the format
+    {
+      "_id": "<key>",
+      "value": <value>
+    }
+    """
+
+    def decode_value(self, value: Any) -> Any:
+        return value
+
+    def encode_value(self, value: Any) -> Any:
+        return value
+
+
+class AstraDBByteStore(AstraDBBaseStore[bytes], ByteStore):
+    """ByteStore implementation using DataStax AstraDB as the underlying store.
+    The bytes values are converted to base64 encoded strings
+    Documents in the AstraDB collection will have the format
+    {
+      "_id": "<key>",
+      "value": "<byte64 string value>"
+    }
+    """
+
+    def decode_value(self, value: Any) -> Optional[bytes]:
+        if value is None:
+            return None
+        return base64.b64decode(value)
+
+    def encode_value(self, value: Optional[bytes]) -> Any:
+        if value is None:
+            return None
+        return base64.b64encode(value).decode("ascii")

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

@@ -0,0 +1,107 @@
+"""Implement integration tests for AstraDB storage."""
+import os
+
+import pytest
+
+from langchain_community.storage.astradb import AstraDBByteStore, AstraDBStore
+
+
+def _has_env_vars() -> bool:
+    return all(
+        [
+            "ASTRA_DB_APPLICATION_TOKEN" in os.environ,
+            "ASTRA_DB_API_ENDPOINT" in os.environ,
+        ]
+    )
+
+
+@pytest.fixture
+def astra_db():
+    from astrapy.db import AstraDB
+
+    return AstraDB(
+        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, collection_name: str):
+    astra_db.create_collection(collection_name)
+    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):
+    astra_db.create_collection(collection_name)
+    store = AstraDBByteStore(collection_name=collection_name, astra_db_client=astra_db)
+    store.mset([("key1", b"value1"), ("key2", b"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:
+        """Test AstraDBStore mget method."""
+        collection_name = "lc_test_store_mget"
+        try:
+            store = init_store(astra_db, collection_name)
+            assert store.mget(["key1", "key2"]) == [[0.1, 0.2], "value2"]
+        finally:
+            astra_db.delete_collection(collection_name)
+
+    def test_mset(self, astra_db) -> None:
+        """Test that multiple keys can be set with AstraDBStore."""
+        collection_name = "lc_test_store_mset"
+        try:
+            store = init_store(astra_db, collection_name)
+            result = store.collection.find_one({"_id": "key1"})
+            assert result["data"]["document"]["value"] == [0.1, 0.2]
+            result = store.collection.find_one({"_id": "key2"})
+            assert result["data"]["document"]["value"] == "value2"
+        finally:
+            astra_db.delete_collection(collection_name)
+
+    def test_mdelete(self, astra_db) -> None:
+        """Test that deletion works as expected."""
+        collection_name = "lc_test_store_mdelete"
+        try:
+            store = init_store(astra_db, collection_name)
+            store.mdelete(["key1", "key2"])
+            result = store.mget(["key1", "key2"])
+            assert result == [None, None]
+        finally:
+            astra_db.delete_collection(collection_name)
+
+    def test_yield_keys(self, astra_db) -> None:
+        collection_name = "lc_test_store_yield_keys"
+        try:
+            store = init_store(astra_db, collection_name)
+            assert set(store.yield_keys()) == {"key1", "key2"}
+            assert set(store.yield_keys(prefix="key")) == {"key1", "key2"}
+            assert set(store.yield_keys(prefix="lang")) == set()
+        finally:
+            astra_db.delete_collection(collection_name)
+
+    def test_bytestore_mget(self, astra_db) -> None:
+        """Test AstraDBByteStore mget method."""
+        collection_name = "lc_test_bytestore_mget"
+        try:
+            store = init_bytestore(astra_db, collection_name)
+            assert store.mget(["key1", "key2"]) == [b"value1", b"value2"]
+        finally:
+            astra_db.delete_collection(collection_name)
+
+    def test_bytestore_mset(self, astra_db) -> None:
+        """Test that multiple keys can be set with AstraDBByteStore."""
+        collection_name = "lc_test_bytestore_mset"
+        try:
+            store = init_bytestore(astra_db, collection_name)
+            result = store.collection.find_one({"_id": "key1"})
+            assert result["data"]["document"]["value"] == "dmFsdWUx"
+            result = store.collection.find_one({"_id": "key2"})
+            assert result["data"]["document"]["value"] == "dmFsdWUy"
+        finally:
+            astra_db.delete_collection(collection_name)

libs/community/tests/unit_tests/storage/test_imports.py

@@ -1,6 +1,8 @@
 from langchain_community.storage import __all__
 
 EXPECTED_ALL = [
+    "AstraDBStore",
+    "AstraDBByteStore",
     "RedisStore",
     "UpstashRedisByteStore",
     "UpstashRedisStore",