Prefer byte store interface for Upstash BaseStore to match other Redis (

langchain-ai#14201)

If we are not going to make the existing Docstore class also implement
`BaseStore[str, Document]`, IMO all base store implementations should
always be `[str, bytes]` so that they are more interchangeable.

CC @rlancemartin @eyurtsev

libs/langchain/langchain/storage/upstash_redis.py

@@ -1,9 +1,10 @@
 from typing import Any, Iterator, List, Optional, Sequence, Tuple, cast
 
+from langchain_core._api.deprecation import deprecated
 from langchain_core.stores import BaseStore
 
 
-class UpstashRedisStore(BaseStore[str, str]):
+class _UpstashRedisStore(BaseStore[str, str]):
     """BaseStore implementation using Upstash Redis as the underlying store."""
 
     def __init__(
@@ -117,3 +118,57 @@ def yield_keys(self, *, prefix: Optional[str] = None) -> Iterator[str]:
                     yield relative_key
                 else:
                     yield key
+
+
+@deprecated("0.0.335", alternative="UpstashRedisByteStore")
+class UpstashRedisStore(_UpstashRedisStore):
+    """
+    BaseStore implementation using Upstash Redis
+    as the underlying store to store strings.
+
+    Deprecated in favor of the more generic UpstashRedisByteStore.
+    """
+
+
+class UpstashRedisByteStore(BaseStore[str, bytes]):
+    """
+    BaseStore implementation using Upstash Redis
+    as the underlying store to store raw bytes.
+    """
+
+    def __init__(
+        self,
+        *,
+        client: Any = None,
+        url: Optional[str] = None,
+        token: Optional[str] = None,
+        ttl: Optional[int] = None,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self.underlying_store = _UpstashRedisStore(
+            client=client, url=url, token=token, ttl=ttl, namespace=namespace
+        )
+
+    def mget(self, keys: Sequence[str]) -> List[Optional[bytes]]:
+        """Get the values associated with the given keys."""
+        return [
+            value.encode("utf-8") if value is not None else None
+            for value in self.underlying_store.mget(keys)
+        ]
+
+    def mset(self, key_value_pairs: Sequence[Tuple[str, bytes]]) -> None:
+        """Set the given key-value pairs."""
+        self.underlying_store.mset(
+            [
+                (k, v.decode("utf-8")) if v is not None else None
+                for k, v in key_value_pairs
+            ]
+        )
+
+    def mdelete(self, keys: Sequence[str]) -> None:
+        """Delete the given keys."""
+        self.underlying_store.mdelete(keys)
+
+    def yield_keys(self, *, prefix: Optional[str] = None) -> Iterator[str]:
+        """Yield keys in the store."""
+        yield from self.underlying_store.yield_keys(prefix=prefix)

libs/langchain/tests/integration_tests/storage/test_upstash_redis.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from langchain.storage.upstash_redis import UpstashRedisStore
+from langchain.storage.upstash_redis import UpstashRedisByteStore
 
 if TYPE_CHECKING:
     from upstash_redis import Redis
@@ -34,24 +34,24 @@ def redis_client() -> Redis:
 
 
 def test_mget(redis_client: Redis) -> None:
-    store = UpstashRedisStore(client=redis_client, ttl=None)
+    store = UpstashRedisByteStore(client=redis_client, ttl=None)
     keys = ["key1", "key2"]
     redis_client.mset({"key1": "value1", "key2": "value2"})
     result = store.mget(keys)
-    assert result == ["value1", "value2"]
+    assert result == [b"value1", b"value2"]
 
 
 def test_mset(redis_client: Redis) -> None:
-    store = UpstashRedisStore(client=redis_client, ttl=None)
-    key_value_pairs = [("key1", "value1"), ("key2", "value2")]
+    store = UpstashRedisByteStore(client=redis_client, ttl=None)
+    key_value_pairs = [("key1", b"value1"), ("key2", b"value2")]
     store.mset(key_value_pairs)
     result = redis_client.mget("key1", "key2")
     assert result == ["value1", "value2"]
 
 
 def test_mdelete(redis_client: Redis) -> None:
     """Test that deletion works as expected."""
-    store = UpstashRedisStore(client=redis_client, ttl=None)
+    store = UpstashRedisByteStore(client=redis_client, ttl=None)
     keys = ["key1", "key2"]
     redis_client.mset({"key1": "value1", "key2": "value2"})
     store.mdelete(keys)
@@ -60,16 +60,16 @@ def test_mdelete(redis_client: Redis) -> None:
 
 
 def test_yield_keys(redis_client: Redis) -> None:
-    store = UpstashRedisStore(client=redis_client, ttl=None)
+    store = UpstashRedisByteStore(client=redis_client, ttl=None)
     redis_client.mset({"key1": "value2", "key2": "value2"})
     assert sorted(store.yield_keys()) == ["key1", "key2"]
     assert sorted(store.yield_keys(prefix="key*")) == ["key1", "key2"]
     assert sorted(store.yield_keys(prefix="lang*")) == []
 
 
 def test_namespace(redis_client: Redis) -> None:
-    store = UpstashRedisStore(client=redis_client, ttl=None, namespace="meow")
-    key_value_pairs = [("key1", "value1"), ("key2", "value2")]
+    store = UpstashRedisByteStore(client=redis_client, ttl=None, namespace="meow")
+    key_value_pairs = [("key1", b"value1"), ("key2", b"value2")]
     store.mset(key_value_pairs)
 
     cursor, all_keys = redis_client.scan(0)