39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
import os
|
|
import faiss
|
|
import threading
|
|
|
|
class UserIndexManager:
|
|
def __init__(self, base_dir="index_data/"):
|
|
self.base_dir = base_dir
|
|
self.index_map = {} # user_id -> faiss.Index
|
|
self.index_lock = threading.Lock()
|
|
self._load_all_indexes()
|
|
|
|
def _load_all_indexes(self):
|
|
for fname in os.listdir(self.base_dir):
|
|
if fname.endswith(".index"):
|
|
user_id = fname.replace(".index", "")
|
|
self._load_single_index(user_id)
|
|
|
|
def _load_single_index(self, user_id):
|
|
path = os.path.join(self.base_dir, f"{user_id}.index")
|
|
if os.path.exists(path):
|
|
index = faiss.read_index(path)
|
|
self.index_map[user_id] = index
|
|
print(f"[INIT] Loaded index for user: {user_id}")
|
|
|
|
def get_index(self, user_id):
|
|
with self.index_lock:
|
|
if user_id not in self.index_map:
|
|
raise FileNotFoundError(f"No FAISS index loaded for user: {user_id}")
|
|
return self.index_map[user_id]
|
|
|
|
def update_index(self, user_id):
|
|
"""重新加载用户索引文件"""
|
|
path = os.path.join(self.base_dir, f"{user_id}.index")
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"No index file found for user: {user_id}")
|
|
with self.index_lock:
|
|
self.index_map[user_id] = faiss.read_index(path)
|
|
print(f"[UPDATE] Index for user {user_id} has been hot-reloaded.")
|