Source code for langchain_community.storage.upstash_redis
fromtypingimportAny,Iterator,List,Optional,Sequence,Tuple,castfromlangchain_core._api.deprecationimportdeprecatedfromlangchain_core.storesimportBaseStore,ByteStoreclass_UpstashRedisStore(BaseStore[str,str]):"""BaseStore implementation using Upstash Redis as the underlying store."""def__init__(self,*,client:Any=None,url:Optional[str]=None,token:Optional[str]=None,ttl:Optional[int]=None,namespace:Optional[str]=None,)->None:"""Initialize the UpstashRedisStore with HTTP API. Must provide either an Upstash Redis client or a url. Args: client: An Upstash Redis instance url: UPSTASH_REDIS_REST_URL token: UPSTASH_REDIS_REST_TOKEN ttl: time to expire keys in seconds if provided, if None keys will never expire namespace: if provided, all keys will be prefixed with this namespace """try:fromupstash_redisimportRedisexceptImportErrorase:raiseImportError("UpstashRedisStore requires the upstash_redis library to be installed. ""pip install upstash_redis")fromeifclientandurl:raiseValueError("Either an Upstash Redis client or a url must be provided, not both.")ifclient:ifnotisinstance(client,Redis):raiseTypeError(f"Expected Upstash Redis client, got {type(client).__name__}.")_client=clientelse:ifnoturlornottoken:raiseValueError("Either an Upstash Redis client or url and token must be provided.")_client=Redis(url=url,token=token)self.client=_clientifnotisinstance(ttl,int)andttlisnotNone:raiseTypeError(f"Expected int or None, got {type(ttl)} instead.")self.ttl=ttlself.namespace=namespacedef_get_prefixed_key(self,key:str)->str:"""Get the key with the namespace prefix. Args: key (str): The original key. Returns: str: The key with the namespace prefix. """delimiter="/"ifself.namespace:returnf"{self.namespace}{delimiter}{key}"returnkeydefmget(self,keys:Sequence[str])->List[Optional[str]]:"""Get the values associated with the given keys."""keys=[self._get_prefixed_key(key)forkeyinkeys]returncast(List[Optional[str]],self.client.mget(*keys),)defmset(self,key_value_pairs:Sequence[Tuple[str,str]])->None:"""Set the given key-value pairs."""forkey,valueinkey_value_pairs:self.client.set(self._get_prefixed_key(key),value,ex=self.ttl)defmdelete(self,keys:Sequence[str])->None:"""Delete the given keys."""_keys=[self._get_prefixed_key(key)forkeyinkeys]self.client.delete(*_keys)defyield_keys(self,*,prefix:Optional[str]=None)->Iterator[str]:"""Yield keys in the store."""ifprefix:pattern=self._get_prefixed_key(prefix)else:pattern=self._get_prefixed_key("*")cursor,keys=self.client.scan(0,match=pattern)forkeyinkeys:ifself.namespace:relative_key=key[len(self.namespace)+1:]yieldrelative_keyelse:yieldkeywhilecursor!=0:cursor,keys=self.client.scan(cursor,match=pattern)forkeyinkeys:ifself.namespace:relative_key=key[len(self.namespace)+1:]yieldrelative_keyelse:yieldkey
[docs]@deprecated("0.0.1",alternative="UpstashRedisByteStore")classUpstashRedisStore(_UpstashRedisStore):""" BaseStore implementation using Upstash Redis as the underlying store to store strings. Deprecated in favor of the more generic UpstashRedisByteStore. """
[docs]classUpstashRedisByteStore(ByteStore):""" BaseStore implementation using Upstash Redis as the underlying store to store raw bytes. """
[docs]defmget(self,keys:Sequence[str])->List[Optional[bytes]]:"""Get the values associated with the given keys."""return[value.encode("utf-8")ifvalueisnotNoneelseNoneforvalueinself.underlying_store.mget(keys)]
[docs]defmset(self,key_value_pairs:Sequence[Tuple[str,bytes]])->None:"""Set the given key-value pairs."""self.underlying_store.mset([(k,v.decode("utf-8"))ifvisnotNoneelseNonefork,vinkey_value_pairs])
[docs]defmdelete(self,keys:Sequence[str])->None:"""Delete the given keys."""self.underlying_store.mdelete(keys)
[docs]defyield_keys(self,*,prefix:Optional[str]=None)->Iterator[str]:"""Yield keys in the store."""yield fromself.underlying_store.yield_keys(prefix=prefix)