Save redis memory state, with the default being to wipe on start still.

This commit is contained in:
BillSchumacher
2023-04-07 00:48:27 -05:00
parent 5a1d9e6d0a
commit cce79695fa
2 changed files with 11 additions and 4 deletions

View File

@@ -64,6 +64,7 @@ class Config(metaclass=Singleton):
self.redis_host = os.getenv("REDIS_HOST") self.redis_host = os.getenv("REDIS_HOST")
self.redis_port = os.getenv("REDIS_PORT") self.redis_port = os.getenv("REDIS_PORT")
self.redis_password = os.getenv("REDIS_PASSWORD") self.redis_password = os.getenv("REDIS_PASSWORD")
self.wipe_redis_on_start = os.getenv("WIPE_REDIS_ON_START", "True") == 'True'
# Note that indexes must be created on db 0 in redis, this is not configureable. # Note that indexes must be created on db 0 in redis, this is not configureable.
self.memory_backend = os.getenv("MEMORY_BACKEND", 'pinecone') self.memory_backend = os.getenv("MEMORY_BACKEND", 'pinecone')

View File

@@ -44,6 +44,7 @@ class RedisMemory(MemoryProviderSingleton):
password=redis_password, password=redis_password,
db=0 # Cannot be changed db=0 # Cannot be changed
) )
if cfg.wipe_redis_on_start:
self.redis.flushall() self.redis.flushall()
try: try:
self.redis.ft("gpt").create_index( self.redis.ft("gpt").create_index(
@@ -55,7 +56,9 @@ class RedisMemory(MemoryProviderSingleton):
) )
except Exception as e: except Exception as e:
print("Error creating Redis search index: ", e) print("Error creating Redis search index: ", e)
self.vec_num = 0 existing_vec_num = self.redis.get('vec_num')
self.vec_num = int(existing_vec_num.decode('utf-8')) if\
existing_vec_num else 0
def add(self, data: str) -> str: def add(self, data: str) -> str:
""" """
@@ -72,10 +75,13 @@ class RedisMemory(MemoryProviderSingleton):
b"data": data, b"data": data,
"embedding": vector "embedding": vector
} }
self.redis.hset(f"gpt:{self.vec_num}", mapping=data_dict) pipe = self.redis.pipeline()
pipe.hset(f"gpt:{self.vec_num}", mapping=data_dict)
_text = f"Inserting data into memory at index: {self.vec_num}:\n"\ _text = f"Inserting data into memory at index: {self.vec_num}:\n"\
f"data: {data}" f"data: {data}"
self.vec_num += 1 self.vec_num += 1
pipe.set('vec_num', self.vec_num)
pipe.execute()
return _text return _text
def get(self, data: str) -> Optional[List[Any]]: def get(self, data: str) -> Optional[List[Any]]: