File size: 514 Bytes
4fdc679 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
# rag/modules/retriever.py
import numpy as np
from config import TOP_K
_index = None # in-memory FAISS index
def set_index(index_obj):
global _index
_index = index_obj
def has_index() -> bool:
return _index is not None
def retrieve_ids(query_embedding: list[float]) -> list[int]:
if _index is None:
raise RuntimeError("FAISS index is not loaded in memory.")
q = np.array([query_embedding], dtype="float32")
_, idx = _index.search(q, TOP_K)
return [int(i) for i in idx[0]]
|