32 lines
920 B
Python
32 lines
920 B
Python
|
|
import faiss
|
|
import numpy as np
|
|
import threading
|
|
from app.core.config import settings
|
|
|
|
class FaissIndexWrapper:
|
|
def __init__(self):
|
|
self.index_lock = threading.Lock()
|
|
self.index = None
|
|
self.load_index(settings.INDEX_FILE)
|
|
|
|
def load_index(self, path):
|
|
with self.index_lock:
|
|
self.index = faiss.read_index(path)
|
|
print(f"[FAISS] Index loaded from {path}")
|
|
|
|
def update_index(self, path):
|
|
"""热更新替换当前索引"""
|
|
with self.index_lock:
|
|
new_index = faiss.read_index(path)
|
|
self.index = new_index
|
|
print(f"[FAISS] Index hot-swapped from {path}")
|
|
|
|
def search(self, vector: np.ndarray, top_k: int = 5):
|
|
with self.index_lock:
|
|
D, I = self.index.search(vector.astype(np.float32), top_k)
|
|
return D[0], I[0]
|
|
|
|
# 单例，供 API 层引用
|
|
faiss_index = FaissIndexWrapper()
|