modules/retriever.py

import faiss
import json
import numpy as np
import os

FAISS_INDEX_PATH = "faiss_index.index"
METADATA_PATH = "faiss_metadata.json"

# Load FAISS index and metadata
def load_faiss_index():
    if not os.path.exists(FAISS_INDEX_PATH) or not os.path.exists(METADATA_PATH):
        print(f"Warning: FAISS index file ({FAISS_INDEX_PATH}) or metadata file ({METADATA_PATH}) not found. Please build the index first.")
        return None, None
    try:
        index = faiss.read_index(FAISS_INDEX_PATH)
        with open(METADATA_PATH, "r") as f:
            metadata = json.load(f)
        print("FAISS index and metadata loaded successfully.")
        return index, metadata
    except Exception as e:
        print(f"Error loading FAISS index or metadata: {e}")
        return None, None

# Retrieve top-k chunks
def retrieve_chunks(query_embedding, index, metadata, k=5):
    if index is None or metadata is None:
        return "Error: FAISS index not loaded. Please upload a knowledge base."
    try:
        D, I = index.search(np.array([query_embedding]), k)
        results = [metadata[str(i)] for i in I[0] if str(i) in metadata and I[0] is not None and len(I[0]) > 0]
        return "\n".join(results)
    except Exception as e:
        print(f"Error during FAISS search: {e}")
        return "Error performing search. Index might be corrupted or empty."