|
|
from typing import ( |
|
|
Any, |
|
|
AsyncIterator, |
|
|
Callable, |
|
|
Iterator, |
|
|
List, |
|
|
Optional, |
|
|
Sequence, |
|
|
Tuple, |
|
|
TypeVar, |
|
|
Union, |
|
|
) |
|
|
|
|
|
from langchain_core.stores import BaseStore |
|
|
|
|
|
K = TypeVar("K") |
|
|
V = TypeVar("V") |
|
|
|
|
|
|
|
|
class EncoderBackedStore(BaseStore[K, V]): |
|
|
"""Wraps a store with key and value encoders/decoders. |
|
|
|
|
|
Examples that uses JSON for encoding/decoding: |
|
|
|
|
|
.. code-block:: python |
|
|
|
|
|
import json |
|
|
|
|
|
def key_encoder(key: int) -> str: |
|
|
return json.dumps(key) |
|
|
|
|
|
def value_serializer(value: float) -> str: |
|
|
return json.dumps(value) |
|
|
|
|
|
def value_deserializer(serialized_value: str) -> float: |
|
|
return json.loads(serialized_value) |
|
|
|
|
|
# Create an instance of the abstract store |
|
|
abstract_store = MyCustomStore() |
|
|
|
|
|
# Create an instance of the encoder-backed store |
|
|
store = EncoderBackedStore( |
|
|
store=abstract_store, |
|
|
key_encoder=key_encoder, |
|
|
value_serializer=value_serializer, |
|
|
value_deserializer=value_deserializer |
|
|
) |
|
|
|
|
|
# Use the encoder-backed store methods |
|
|
store.mset([(1, 3.14), (2, 2.718)]) |
|
|
values = store.mget([1, 2]) # Retrieves [3.14, 2.718] |
|
|
store.mdelete([1, 2]) # Deletes the keys 1 and 2 |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
store: BaseStore[str, Any], |
|
|
key_encoder: Callable[[K], str], |
|
|
value_serializer: Callable[[V], bytes], |
|
|
value_deserializer: Callable[[Any], V], |
|
|
) -> None: |
|
|
"""Initialize an EncodedStore.""" |
|
|
self.store = store |
|
|
self.key_encoder = key_encoder |
|
|
self.value_serializer = value_serializer |
|
|
self.value_deserializer = value_deserializer |
|
|
|
|
|
def mget(self, keys: Sequence[K]) -> List[Optional[V]]: |
|
|
"""Get the values associated with the given keys.""" |
|
|
encoded_keys: List[str] = [self.key_encoder(key) for key in keys] |
|
|
values = self.store.mget(encoded_keys) |
|
|
return [ |
|
|
self.value_deserializer(value) if value is not None else value |
|
|
for value in values |
|
|
] |
|
|
|
|
|
async def amget(self, keys: Sequence[K]) -> List[Optional[V]]: |
|
|
"""Get the values associated with the given keys.""" |
|
|
encoded_keys: List[str] = [self.key_encoder(key) for key in keys] |
|
|
values = await self.store.amget(encoded_keys) |
|
|
return [ |
|
|
self.value_deserializer(value) if value is not None else value |
|
|
for value in values |
|
|
] |
|
|
|
|
|
def mset(self, key_value_pairs: Sequence[Tuple[K, V]]) -> None: |
|
|
"""Set the values for the given keys.""" |
|
|
encoded_pairs = [ |
|
|
(self.key_encoder(key), self.value_serializer(value)) |
|
|
for key, value in key_value_pairs |
|
|
] |
|
|
self.store.mset(encoded_pairs) |
|
|
|
|
|
async def amset(self, key_value_pairs: Sequence[Tuple[K, V]]) -> None: |
|
|
"""Set the values for the given keys.""" |
|
|
encoded_pairs = [ |
|
|
(self.key_encoder(key), self.value_serializer(value)) |
|
|
for key, value in key_value_pairs |
|
|
] |
|
|
await self.store.amset(encoded_pairs) |
|
|
|
|
|
def mdelete(self, keys: Sequence[K]) -> None: |
|
|
"""Delete the given keys and their associated values.""" |
|
|
encoded_keys = [self.key_encoder(key) for key in keys] |
|
|
self.store.mdelete(encoded_keys) |
|
|
|
|
|
async def amdelete(self, keys: Sequence[K]) -> None: |
|
|
"""Delete the given keys and their associated values.""" |
|
|
encoded_keys = [self.key_encoder(key) for key in keys] |
|
|
await self.store.amdelete(encoded_keys) |
|
|
|
|
|
def yield_keys( |
|
|
self, *, prefix: Optional[str] = None |
|
|
) -> Union[Iterator[K], Iterator[str]]: |
|
|
"""Get an iterator over keys that match the given prefix.""" |
|
|
|
|
|
|
|
|
yield from self.store.yield_keys(prefix=prefix) |
|
|
|
|
|
async def ayield_keys( |
|
|
self, *, prefix: Optional[str] = None |
|
|
) -> Union[AsyncIterator[K], AsyncIterator[str]]: |
|
|
"""Get an iterator over keys that match the given prefix.""" |
|
|
|
|
|
|
|
|
async for key in self.store.ayield_keys(prefix=prefix): |
|
|
yield key |
|
|
|
