mirror of
https://github.com/hwchase17/langchain.git
synced 2025-05-25 17:08:46 +00:00
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:
parent
411aa9a41e
commit
de86b84a70
@ -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)
|
||||
|
@ -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)
|
||||
|
Loading…
Reference in New Issue
Block a user