2023-08-29 13:46:41 -06:00
|
|
|
import json
|
2023-09-11 09:51:01 -06:00
|
|
|
import sys
|
2023-08-29 13:46:41 -06:00
|
|
|
|
2023-08-21 21:28:52 -06:00
|
|
|
from flask_caching import Cache
|
2023-08-23 01:14:19 -06:00
|
|
|
from redis import Redis
|
2023-08-27 23:48:10 -06:00
|
|
|
from redis.typing import FieldT
|
2023-09-11 09:51:01 -06:00
|
|
|
import redis as redis_pkg
|
2023-08-21 21:28:52 -06:00
|
|
|
|
2023-08-22 00:01:32 -06:00
|
|
|
cache = Cache(config={'CACHE_TYPE': 'RedisCache', 'CACHE_REDIS_URL': 'redis://localhost:6379/0', 'CACHE_KEY_PREFIX': 'local-llm'})
|
2023-08-23 22:01:06 -06:00
|
|
|
|
|
|
|
|
|
|
|
# redis = Redis()
|
|
|
|
|
|
|
|
class RedisWrapper:
|
|
|
|
"""
|
|
|
|
A wrapper class to set prefixes to keys.
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self, prefix, **kwargs):
|
|
|
|
self.redis = Redis(**kwargs)
|
|
|
|
self.prefix = prefix
|
2023-09-11 09:51:01 -06:00
|
|
|
try:
|
|
|
|
self.set('check_connected', 1)
|
|
|
|
except redis_pkg.exceptions.ConnectionError as e:
|
|
|
|
print('Failed to connect to the Redis server:', e)
|
|
|
|
print('Did you install and start Redis?')
|
|
|
|
sys.exit(1)
|
2023-08-23 22:01:06 -06:00
|
|
|
|
|
|
|
def set(self, key, value):
|
|
|
|
return self.redis.set(f"{self.prefix}:{key}", value)
|
|
|
|
|
|
|
|
def get(self, key):
|
|
|
|
return self.redis.get(f"{self.prefix}:{key}")
|
|
|
|
|
|
|
|
def incr(self, key, amount=1):
|
|
|
|
return self.redis.incr(f"{self.prefix}:{key}", amount)
|
|
|
|
|
|
|
|
def decr(self, key, amount=1):
|
|
|
|
return self.redis.decr(f"{self.prefix}:{key}", amount)
|
|
|
|
|
2023-08-27 23:48:10 -06:00
|
|
|
def sadd(self, key: str, *values: FieldT):
|
|
|
|
return self.redis.sadd(f"{self.prefix}:{key}", *values)
|
|
|
|
|
|
|
|
def srem(self, key: str, *values: FieldT):
|
|
|
|
return self.redis.srem(f"{self.prefix}:{key}", *values)
|
|
|
|
|
|
|
|
def sismember(self, key: str, value: str):
|
|
|
|
return self.redis.sismember(f"{self.prefix}:{key}", value)
|
|
|
|
|
2023-08-29 13:46:41 -06:00
|
|
|
def set_dict(self, key, dict_value):
|
|
|
|
# return self.redis.hset(f"{self.prefix}:{key}", mapping=dict_value)
|
|
|
|
return self.set(f"{self.prefix}:{key}", json.dumps(dict_value))
|
|
|
|
|
|
|
|
def get_dict(self, key):
|
|
|
|
# return self.redis.hgetall(f"{self.prefix}:{key}")
|
|
|
|
r = self.get(f"{self.prefix}:{key}")
|
|
|
|
if not r:
|
|
|
|
return dict()
|
|
|
|
else:
|
|
|
|
return json.loads(r)
|
|
|
|
|
2023-08-23 22:01:06 -06:00
|
|
|
def flush(self):
|
|
|
|
flushed = []
|
|
|
|
for key in self.redis.scan_iter(f'{self.prefix}:*'):
|
|
|
|
flushed.append(key)
|
|
|
|
self.redis.delete(key)
|
|
|
|
return flushed
|
|
|
|
|
|
|
|
|
|
|
|
redis = RedisWrapper('local_llm')
|