diff --git a/libs/langchain/langchain/storage/upstash_redis.py b/libs/langchain/langchain/storage/upstash_redis.py index 193b8c7c4b5c3d..7dc436ce33e2cb 100644 --- a/libs/langchain/langchain/storage/upstash_redis.py +++ b/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) diff --git a/libs/langchain/tests/integration_tests/storage/test_upstash_redis.py b/libs/langchain/tests/integration_tests/storage/test_upstash_redis.py index 183e09515b6a47..01ab5831fe3ed2 100644 --- a/libs/langchain/tests/integration_tests/storage/test_upstash_redis.py +++ b/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,16 +34,16 @@ 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"] @@ -51,7 +51,7 @@ def test_mset(redis_client: Redis) -> None: 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,7 +60,7 @@ 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"] @@ -68,8 +68,8 @@ def test_yield_keys(redis_client: Redis) -> None: 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)