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