Prefer byte store interface for Upstash BaseStore to match other Redis (#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
This commit is contained in:
Jacob Lee 2023-12-04 14:17:33 -08:00 committed by GitHub
parent 411aa9a41e
commit de86b84a70
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 65 additions and 10 deletions

View File

@ -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 @@ class UpstashRedisStore(BaseStore[str, 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)

View File

@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
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)