Skip to content

Commit

Permalink
Showing 2 changed files with 26 additions and 5 deletions.
9 changes: 9 additions & 0 deletions backend/database/redis_db.py
Original file line number Diff line number Diff line change
@@ -267,3 +267,12 @@ def remove_filter_category_item(uid: str, category: str, item: str):

def remove_all_filter_category_items(uid: str, category: str):
r.delete(f'users:{uid}:filters:{category}')


def save_migrated_retrieval_memory_id(memory_id: str):
r.sadd('migrated_retrieval_memory_ids', memory_id)
r.expire('migrated_retrieval_memory_ids', 60 * 60 * 24 * 7)


def has_migrated_retrieval_memory_id(memory_id: str) -> bool:
return r.sismember('migrated_retrieval_memory_ids', memory_id)
22 changes: 17 additions & 5 deletions backend/scripts/users/retrieval.py
Original file line number Diff line number Diff line change
@@ -15,16 +15,28 @@
from database._client import get_users_uid
import database.memories as memories_db
from utils.memories.process_memory import save_structured_vector
from database.redis_db import has_migrated_retrieval_memory_id, save_migrated_retrieval_memory_id

if __name__ == '__main__':
def single(uid, memory, update):
save_structured_vector(uid, memory, update)
save_migrated_retrieval_memory_id(memory.id)


uids = get_users_uid()
for uid in uids:
memories = memories_db.get_memories(uid, limit=2000)
threads = []
for memory in memories:
threads.append(threading.Thread(target=save_structured_vector, args=(uid, Memory(**memory), True)))
if has_migrated_retrieval_memory_id(memory['id']):
print('Skipping', memory['id'])
continue

threads.append(threading.Thread(target=single, args=(uid, Memory(**memory), True)))
if len(threads) == 20:
[t.start() for t in threads]
[t.join() for t in threads]
threads = []

chunks = [threads[i:i + 30] for i in range(0, len(threads), 30)]
for chunk in chunks:
[t.start() for t in chunk]
[t.join() for t in chunk]
[t.start() for t in threads]
[t.join() for t in threads]

0 comments on commit ee89485

Please sign in to comment.