Add configurable index key for redis.

This commit is contained in:
BillSchumacher
2023-04-07 15:27:48 -05:00
parent 5d13fb2546
commit 14e10c9c4d
2 changed files with 9 additions and 7 deletions

View File

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

View File

@@ -44,19 +44,20 @@ class RedisMemory(MemoryProviderSingleton):
password=redis_password,
db=0 # Cannot be changed
)
self.cfg = cfg
if cfg.wipe_redis_on_start:
self.redis.flushall()
try:
self.redis.ft("gpt").create_index(
self.redis.ft(f"{cfg.memory_index}").create_index(
fields=SCHEMA,
definition=IndexDefinition(
prefix=["gpt:"],
prefix=[f"{cfg.memory_index}:"],
index_type=IndexType.HASH
)
)
except Exception as e:
print("Error creating Redis search index: ", e)
existing_vec_num = self.redis.get('vec_num')
existing_vec_num = self.redis.get(f'{cfg.memory_index}-vec_num')
self.vec_num = int(existing_vec_num.decode('utf-8')) if\
existing_vec_num else 0
@@ -76,11 +77,11 @@ class RedisMemory(MemoryProviderSingleton):
"embedding": vector
}
pipe = self.redis.pipeline()
pipe.hset(f"gpt:{self.vec_num}", mapping=data_dict)
pipe.hset(f"{self.cfg.memory_index}:{self.vec_num}", mapping=data_dict)
_text = f"Inserting data into memory at index: {self.vec_num}:\n"\
f"data: {data}"
self.vec_num += 1
pipe.set('vec_num', self.vec_num)
pipe.set(f'{self.cfg.memory_index}-vec_num', self.vec_num)
pipe.execute()
return _text
@@ -126,7 +127,7 @@ class RedisMemory(MemoryProviderSingleton):
query_vector = np.array(query_embedding).astype(np.float32).tobytes()
try:
results = self.redis.ft("gpt").search(
results = self.redis.ft(f"{self.cfg.memory_index}").search(
query, query_params={"vector": query_vector}
)
except Exception as e:
@@ -138,4 +139,4 @@ class RedisMemory(MemoryProviderSingleton):
"""
Returns: The stats of the memory index.
"""
return self.redis.ft("mem").info()
return self.redis.ft(f"{self.cfg.memory_index}").info()